Skip to content

Commit ba7dbcb

Browse files
feat: Single-pass WAL streaming for LOG_BASED replication (MeltanoLabs#772)
### problem `PostgresLogBasedStream.get_records()` opens its own `LogicalReplicationConnection` per selected stream. With N LOG_BASED streams the tap runs N sequential WAL scans -- each rereads the same segments, with add-tables discarding most records server-side. End-to-end sync time scales ~linearly in N. For pipelines with multiple LOG_BASED streams against a large backlog, this dominates run-time. ### changes A new `SingleConnectionWALReader` opens one logical replication connection with `add-tables` covering all selected LOG_BASED tables, scans the WAL once, and dispatches each parsed wal2json message inline to the owning stream's new `emit_record()` method for immediate Singer RECORD emission. STATE flushes every 30s, and the slot is advanced to the WAL tip on idle/max-run exit. - new modules: `_wal_helpers.py` (FQN/escaping/parsing helpers) and `wal_reader.py` (the reader and read loop). - `client.py` gains `emit_record` method and a config-flag branch in `get_records` - `tap.py` adds the `log_based_single_connection` config setting (default False!) and `_sync_log_based_streams_shared` orchestration - new tests: `tests/test_wal_helpers.py`, `tests/test_wal_reader.py`, `tests/test_consume.py` Full disclosure, I had Claude Code implement those three test modules, and then I iterated a bit. If it's still excessive / not testing usefully -- something Claude is known to do, sigh -- just let me know, and I will take a hatchet to it. This is a very belated follow-up to PR MeltanoLabs#667 and Issue MeltanoLabs#587. ### constraints - `Tap.sync_all` is `@typing.final`, so dispatch can't be restructured at the SDK boundary -- this was a bummer. My next best option was trigger at the first LOG_BASED stream's `get_records()` call, gated by a `_shared_wal_run_completed` flag on the tap so siblings become no-ops. - SCHEMA-before-RECORD across streams: `_sync_log_based_streams_shared` pre-writes every stream's schema before the reader runs. Since the SDK's `Stream.sync()` later calls `_write_schema_message()` again, the override on `PostgresLogBasedStream` is idempotent. (Without that flag every SCHEMA would be emitted twice, which is not great.) - Per-stream LSN filter: Replication opens at `min(start_lsn)` across all streams, so each stream's own bookmark is captured at construction and used to drop messages that it's already past. - I had to dip into private SDK calls -- `_write_record_message()` and `_increment_stream_state()` -- for this to work. I consolidated them in one place -- `emit_record()` -- so SDK renames hit one method. Not sure how stable the API is here... ### questions 1. `_write_schema_message` idempotency: I made the smallest fix I could for the duplicate-SCHEMA bug given the `@final` constraint on `sync_all`. Is there another / cleaner approach? 2. `emit_record()` uses internal SDK calls. Is this going to be an issue? Is there a safer / "public" equivalent? 3. `get_records()` as trigger: The "first stream's get_records fires the shared reader" pattern is a not-great workaround for `sync_all` being final. Is it okay as documented, or is there a cleaner SDK hook? 4. `replication_max_run_seconds` / `replication_idle_exit_seconds` now bound the whole LOG_BASED batch instead of each stream. To me that feels like an improvement, but I don't know the whole system / downstream use caess. For example, does anything downstream assume per-stream bounds? --------- Signed-off-by: Edgar Ramírez Mondragón <edgarrm358@gmail.com> Co-authored-by: Edgar Ramírez Mondragón <16805946+edgarrmondragon@users.noreply.github.com> Co-authored-by: Edgar Ramírez Mondragón <edgarrm358@gmail.com>
1 parent 554ed05 commit ba7dbcb

8 files changed

Lines changed: 1558 additions & 125 deletions

File tree

‎README.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -261,6 +261,19 @@ Log-based replication will modify the schemas output by the tap. Specifically, a
261261
262262
Note that changing what streams are selected after already beginning log-based replication can have unexpected consequences. To ensure consistent output, it is best to keep selected streams the same across invocations of the tap.
263263
264+
#### Adding a new stream to an existing LOG_BASED set
265+
266+
When a stream has no prior bookmark (i.e. it is new to log-based replication), the tap starts reading the replication slot from LSN 0 — the oldest WAL position the slot still retains. This means the tap will walk through all retained WAL for every registered table before reaching the current tip, which can be very slow if the slot has accumulated a large backlog.
267+
268+
Other CDC tools (e.g. [Debezium](https://debezium.io/blog/2021/10/07/incremental-snapshots/), [Airbyte](https://docs.airbyte.com/platform/understanding-airbyte/cdc#adding-new-schemascolumns)) handle this by taking an initial snapshot of the new table at a known LSN and then starting WAL streaming from that point. This tap does not implement that pattern.
269+
270+
The recommended workaround is:
271+
272+
1. Sync the new table once using `FULL_TABLE` or `INCREMENTAL` replication to establish its baseline state.
273+
2. Switch that stream's replication method to `LOG_BASED` for subsequent syncs.
274+
275+
This ensures the new stream has a valid starting LSN bookmark and does not force the other streams to re-read old WAL.
276+
264277
Note also that using log-based replication will cause the replication key for all streams to be set to "_sdc_lsn", which is the Postgres LSN for the record in question.
265278
266279
### How to Set Up Log-Based Replication

‎tap_postgres/_wal_helpers.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
"""Helper functions for LOG_BASED replication."""
2+
3+
from __future__ import annotations
4+
5+
import json
6+
import re
7+
import typing as t
8+
9+
import psycopg2
10+
11+
if t.TYPE_CHECKING:
12+
import logging
13+
14+
from psycopg2 import extras
15+
16+
17+
# wal2json emits enum type names with unescaped double quotes, producing invalid JSON
18+
# like "type":""EnumName"" ... strip the extra quotes via regex and re-parse
19+
_WAL2JSON_ENUM_QUOTE_RE = re.compile(r'"type":""([^"]+)""')
20+
21+
22+
def normalize_fqn(schema: str, table: str) -> str:
23+
"""Generate canonical, fully-qualified name for dispatch.
24+
25+
This is the source of truth for matching a WAL message to a registered stream.
26+
Both sides -- the tap (when registering streams with the reader) and the WAL reader
27+
(when dispatching a parsed payload) -- *must* call this function with the raw schema
28+
and table name strings.
29+
30+
wal2json's format-version=2 output includes ``"schema"`` and ``"table"`` fields
31+
as the raw, unquoted identifier names (wal2json reports whatever Postgres has stored).
32+
Therefore, use the raw names joined by a single dot, with no quoting and no case folding.
33+
34+
Do *not* use ``SQLStream.fully_qualified_name`` for dispatch.
35+
"""
36+
return f"{schema}.{table}"
37+
38+
39+
def escape_for_add_tables(identifier: str) -> str:
40+
"""Escape a schema or table name for use in wal2json's ``add-tables``.
41+
42+
wal2json's ``add-tables`` option takes a comma-separated list of "schema.table" entries.
43+
Backslash is the escape character; comma and dot within an identifier must be escaped,
44+
and backslash itself must be doubled.
45+
46+
References:
47+
- https://github.com/eulerto/wal2json#parameters
48+
"""
49+
return identifier.replace("\\", "\\\\").replace(",", "\\,").replace(".", "\\.")
50+
51+
52+
def build_add_tables_option(fqn_pairs: list[tuple[str, str]]) -> str:
53+
"""Build the wal2json ``add-tables`` option from a list of (schema, table).
54+
55+
Each identifier is escaped with ``escape_for_add_tables`` and joined with
56+
the appropriate separators. For example::
57+
58+
>>> build_add_tables_option([("public", "users"), ("public", "orders")])
59+
'public.users,public.orders'
60+
"""
61+
return ",".join(
62+
f"{escape_for_add_tables(schema)}.{escape_for_add_tables(table)}"
63+
for schema, table in fqn_pairs
64+
)
65+
66+
67+
def parse_wal_message(raw_payload: str, cursor: extras.ReplicationCursor | None) -> dict | None:
68+
"""Parse a raw wal2json JSON payload into a Python dict.
69+
70+
Handles the known wal2json enum-quoting bug via one retry after regex repair.
71+
Returns None if the payload can't be decoded even after repair, in which case
72+
the caller should log and skip.
73+
74+
When ``cursor`` is provided, pre-parses ``text[]`` column values into Python lists
75+
using psycopg2's ``STRINGARRAY`` type caster. This must be done while the cursor is alive,
76+
since ``STRINGARRAY`` reads connection-level encoding info from it.
77+
"""
78+
try:
79+
payload = json.loads(raw_payload)
80+
except json.JSONDecodeError:
81+
try:
82+
payload = json.loads(fix_wal2json_enum_quotes(raw_payload))
83+
except json.JSONDecodeError:
84+
return None
85+
86+
if cursor is not None:
87+
pre_parse_text_arrays(payload, cursor)
88+
89+
return payload
90+
91+
92+
def fix_wal2json_enum_quotes(payload: str) -> str:
93+
"""Repair the wal2json enum-quoting bug in a raw JSON payload.
94+
95+
wal2json outputs enum type names with unescaped double quotes (e.g. "type":""EnumName""),
96+
which is invalid JSON. Normalize this to "type":"EnumName" so a second ``json.loads``
97+
attempt succeeds.
98+
"""
99+
return _WAL2JSON_ENUM_QUOTE_RE.sub(r'"type":"\1"', payload)
100+
101+
102+
def pre_parse_text_arrays(payload: dict, cursor: extras.ReplicationCursor) -> None:
103+
"""Pre-parse ``text[]`` column values in a wal2json payload, in place.
104+
105+
wal2json returns ``text[]`` values as Postgres's array literal string (e.g. '{a,b,c}').
106+
Converting to a Python list requires ``psycopg2.extensions.STRINGARRAY``, which needs
107+
a live cursor for encoding context. Calling this during the WAL read means downstream code
108+
-- ``consume()``, etc. -- can operate on plain Python lists with no cursor dependency.
109+
"""
110+
for key in ("columns", "identity"):
111+
for column in payload.get(key, ()) or ():
112+
if column.get("type") == "text[]" and column.get("value") is not None:
113+
column["value"] = psycopg2.extensions.STRINGARRAY(column["value"], cursor)
114+
115+
116+
def query_current_wal_lsn(dsn: str, logger: logging.Logger) -> int | None:
117+
"""Query ``pg_current_wal_flush_lsn()`` on a non-replication connection."""
118+
try:
119+
conn = psycopg2.connect(dsn)
120+
try:
121+
conn.autocommit = True
122+
with conn.cursor() as cur:
123+
cur.execute("SELECT pg_current_wal_flush_lsn()")
124+
row = cur.fetchone()
125+
if row is None:
126+
return None
127+
return _lsn_str_to_int(row[0])
128+
finally:
129+
conn.close()
130+
except Exception as exc:
131+
logger.warning("Could not query current WAL LSN: %s", exc)
132+
return None
133+
134+
135+
def _lsn_str_to_int(lsn_str: str) -> int:
136+
"""Convert a Postgres LSN string (e.g. '6/4A3B2C10') to an int."""
137+
hi, lo = lsn_str.split("/")
138+
return (int(hi, 16) << 32) + int(lo, 16)

0 commit comments

Comments
 (0)