Skip to content

Commit

Permalink
fix: use psycopg text array decoding in wal2json messages
Browse files Browse the repository at this point in the history
  • Loading branch information
hudovisk committed Aug 27, 2024
1 parent ba7ad87 commit e3f56a0
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 5 deletions.
16 changes: 12 additions & 4 deletions tap_postgres/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -429,7 +429,7 @@ def get_records(self, context: Context | None) -> Iterable[dict[str, Any]]:
while True:
message = logical_replication_cursor.read_message()
if message:
row = self.consume(message)
row = self.consume(message, logical_replication_cursor)
if row:
yield row
else:
Expand All @@ -456,7 +456,7 @@ def get_records(self, context: Context | None) -> Iterable[dict[str, Any]]:
logical_replication_cursor.close()
logical_replication_connection.close()

def consume(self, message) -> dict | None:
def consume(self, message, cursor) -> dict | None:
"""Ingest WAL message."""
try:
message_payload = json.loads(message.payload)
Expand All @@ -476,12 +476,14 @@ def consume(self, message) -> dict | None:

if message_payload["action"] in upsert_actions:
for column in message_payload["columns"]:
row.update({column["name"]: column["value"]})
value = self._parse_column_value(column, cursor)
row.update({column["name"]: value})
row.update({"_sdc_deleted_at": None})
row.update({"_sdc_lsn": message.data_start})
elif message_payload["action"] in delete_actions:
for column in message_payload["identity"]:
row.update({column["name"]: column["value"]})
value = self._parse_column_value(column, cursor)
row.update({column["name"]: value})
row.update(
{
"_sdc_deleted_at": datetime.datetime.utcnow().strftime(
Expand Down Expand Up @@ -517,6 +519,12 @@ def consume(self, message) -> dict | None:

return row

def _parse_column_value(self, column, cursor):
if column["type"] == "text[]":
return psycopg2.extensions.STRINGARRAY(column["value"], cursor)

return column["value"]

def logical_replication_connection(self):
"""A logical replication connection to the database.
Expand Down
52 changes: 51 additions & 1 deletion tests/test_log_based.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import json

import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import BIGINT, TEXT
from sqlalchemy.dialects.postgresql import ARRAY, BIGINT, TEXT

from tap_postgres.tap import TapPostgres
from tests.test_core import PostgresTestRunner
Expand Down Expand Up @@ -57,3 +57,53 @@ def test_null_append():
tap_class=TapPostgres, config=LOG_BASED_CONFIG, catalog=tap_catalog
)
test_runner.sync_all()


def test_string_array_column():
"""LOG_BASED syncs failed with array columns.
This test checks that even when a catalog contains properties with types represented
as arrays (ex: "text[]") LOG_BASED replication can properly decode their value.
"""
table_name = "test_array_column"
engine = sa.create_engine(
"postgresql://postgres:postgres@localhost:5434/postgres", future=True
)

metadata_obj = sa.MetaData()
table = sa.Table(
table_name,
metadata_obj,
sa.Column("id", BIGINT, primary_key=True),
sa.Column("data", ARRAY(TEXT), nullable=True),
)
with engine.begin() as conn:
table.drop(conn, checkfirst=True)
metadata_obj.create_all(conn)
insert = table.insert().values(id=123, data=["1", "2"])
conn.execute(insert)
insert = table.insert().values(id=321, data=['This is a "test"', "2"])
conn.execute(insert)

tap = TapPostgres(config=LOG_BASED_CONFIG)
tap_catalog = json.loads(tap.catalog_json_text)
altered_table_name = f"public-{table_name}"

for stream in tap_catalog["streams"]:
if stream.get("stream") and altered_table_name not in stream["stream"]:
for metadata in stream["metadata"]:
metadata["metadata"]["selected"] = False
else:
stream["replication_method"] = "LOG_BASED"
stream["replication_key"] = "_sdc_lsn"
for metadata in stream["metadata"]:
metadata["metadata"]["selected"] = True
if metadata["breadcrumb"] == []:
metadata["metadata"]["replication-method"] = "LOG_BASED"

test_runner = PostgresTestRunner(
tap_class=TapPostgres, config=LOG_BASED_CONFIG, catalog=tap_catalog
)
test_runner.sync_all()
assert test_runner.record_messages[0]["record"]["data"] == ["1", "2"]
assert test_runner.record_messages[1]["record"]["data"] == ['This is a "test"', "2"]

0 comments on commit e3f56a0

Please sign in to comment.