From 003b6d7fc9c4a39c71210d8c86980580b47ed951 Mon Sep 17 00:00:00 2001 From: lekhrocks Date: Thu, 6 Aug 2026 10:06:33 +0530 Subject: [PATCH] fix(snapshot/cdc): keyset pagination, cursor resume, per-pipeline offsets, rollback on cancel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Snapshot: - Keyset pagination over a single-column PK (WHERE pk > cursor ORDER BY pk LIMIT size) instead of OFFSET — stable under concurrent writes, no O(n^2) or offset drift; falls back to OFFSET when no single-column PK - PK identifier sanitized before SQL interpolation (injection guard) - Cursor is carried through the read loop and checkpointed, so resume continues at the exact next row (no partial-batch offset drift) - estimateRows prefers the Postgres reltuples planner estimate over a full SELECT COUNT(*); falls back to COUNT on non-Postgres - Progress aggregates row/batch estimates across ALL mapped tables - writer.rollback() on cancel/failure so partial writes are not committed (a later resume would otherwise duplicate rows); commit only on success CDC: - Offset file keyed per-pipeline (host_db_pipeline) so multiple pipelines on the same database do not share/corrupt Debezium position; CaptureLifecycle passes pipelineId in the connector context - updateOffset merges instead of replaces, retaining per-table LSNs for multi-table captures - CdcIntegrationTest reset path matches the new keyed offset filename --- .../syncflow/api/cdc/CaptureLifecycle.java | 4 +- .../api/snapshot/SnapshotExecutor.java | 46 +++++-- .../syncflow/api/cdc/CdcIntegrationTest.java | 13 +- .../connector/cdc/DebeziumCdcConnector.java | 19 +-- .../AbstractJdbcSnapshotConnector.java | 118 +++++++++++++++++- 5 files changed, 172 insertions(+), 28 deletions(-) diff --git a/syncflow-api/src/main/java/com/syncflow/api/cdc/CaptureLifecycle.java b/syncflow-api/src/main/java/com/syncflow/api/cdc/CaptureLifecycle.java index 94a84a4..59f9fcf 100644 --- a/syncflow-api/src/main/java/com/syncflow/api/cdc/CaptureLifecycle.java +++ b/syncflow-api/src/main/java/com/syncflow/api/cdc/CaptureLifecycle.java @@ -80,7 +80,9 @@ public CaptureStatus start(String pipelineId, String tableOrCollection) { .orElseThrow(() -> new IllegalArgumentException("No CDC connector for type: " + ct)); var config = toConfig(conn); - var ctx = new ConnectorContext(config, Map.of()); + // Key the Debezium offset file (and any connector state) per pipeline so + // multiple pipelines on the same database don't share/corrupt position. + var ctx = new ConnectorContext(config, Map.of("pipelineId", pipelineId)); // pre-flight validation var validation = connector.validate(ctx); diff --git a/syncflow-api/src/main/java/com/syncflow/api/snapshot/SnapshotExecutor.java b/syncflow-api/src/main/java/com/syncflow/api/snapshot/SnapshotExecutor.java index dcf888f..67dbd4d 100644 --- a/syncflow-api/src/main/java/com/syncflow/api/snapshot/SnapshotExecutor.java +++ b/syncflow-api/src/main/java/com/syncflow/api/snapshot/SnapshotExecutor.java @@ -101,20 +101,23 @@ private void execute(SnapshotJob job, PipelineDesign pipeline) { var rowsProcessed = new AtomicLong(0); var batchesDone = new AtomicLong(0); + DestinationWriter writer = null; try { var sourceCtx = buildSourceContext(pipeline); var destCfg = buildDestConfig(pipeline); var connector = resolveSourceConnector(pipeline); - var writer = resolveWriter(pipeline); + writer = resolveWriter(pipeline); writer.connect(destCfg); long totalRows = 0; long totalBatches = 0; - if (!pipeline.tableMappings().isEmpty()) { - var tm = pipeline.tableMappings().getFirst(); - totalRows = connector.estimateRows(sourceCtx, pipeline.source().schema(), tm.sourceTable()); - totalBatches = (totalRows / pipeline.settings().batchSize()) + 1; + // Aggregate row/batch estimates across ALL mapped tables so progress % + // is meaningful for multi-table pipelines (not just the first mapping). + for (var tm : pipeline.tableMappings()) { + var tableRows = connector.estimateRows(sourceCtx, pipeline.source().schema(), tm.sourceTable()); + totalRows += tableRows; + totalBatches += (tableRows / pipeline.settings().batchSize()) + 1; } var progress = SnapshotProgress.starting(totalRows); @@ -126,10 +129,12 @@ private void execute(SnapshotJob job, PipelineDesign pipeline) { var ctx = new ProcessingContext(pipeline, tm); var checkpoint = checkpointStore.get(pipeline.id().value(), tm.sourceTable()); - int startBatch = (checkpoint != null) ? checkpoint.lastBatchNumber() + 1 : 0; + // Resume from the last checkpointed cursor; else start fresh. + String cursor = (checkpoint != null) ? checkpoint.cursor() : null; + int batchNumber = (checkpoint != null) ? checkpoint.lastBatchNumber() + 1 : 0; - var batchInfo = new BatchInformation(startBatch, pipeline.settings().batchSize(), - tm.sourceTable(), null); + var batchInfo = new BatchInformation(batchNumber, pipeline.settings().batchSize(), + tm.sourceTable(), cursor); var page = connector.readBatch(sourceCtx, pipeline.source().schema(), tm.sourceTable(), batchInfo); @@ -161,23 +166,32 @@ private void execute(SnapshotJob job, PipelineDesign pipeline) { meterRegistry.counter("syncflow.snapshot.rows", "pipeline", pipeline.id().value()).increment(batch.size()); - // checkpoint every 5 batches + // Checkpoint every 5 batches — captures the keyed cursor so a + // resume continues exactly at the next row (no OFFSET drift). if (batchesDone.get() % 5 == 0) { checkpointStore.save(new SnapshotCheckpoint( pipeline.id().value(), tm.sourceTable(), - (int) batchesDone.get(), rowsProcessed.get(), null)); + (int) batchesDone.get(), rowsProcessed.get(), + page.nextCursor())); } + // Next read continues from this page's cursor. var nextBatchInfo = new BatchInformation( (int) batchesDone.get(), pipeline.settings().batchSize(), - tm.sourceTable(), null); + tm.sourceTable(), page.nextCursor()); page = connector.readBatch(sourceCtx, pipeline.source().schema(), tm.sourceTable(), nextBatchInfo); } } - writer.flush(); - writer.commit(); + if (isCancelled(job)) { + // Do not commit partial writes on cancel — a later resume would + // duplicate the already-written rows. + writer.rollback(); + } else { + writer.flush(); + writer.commit(); + } var elapsed = sample.stop(timer); if (!isCancelled(job)) { @@ -189,6 +203,12 @@ private void execute(SnapshotJob job, PipelineDesign pipeline) { } } catch (Exception e) { sample.stop(timer); + if (writer != null) { + try { + writer.rollback(); + } catch (Exception ignored) { + } + } var error = new SnapshotError("SNAPSHOT_FAILED", e.getMessage(), (int) batchesDone.get(), Instant.now()); jobs.put(job.getId().value(), job.withFailed(List.of(error))); diff --git a/syncflow-api/src/test/java/com/syncflow/api/cdc/CdcIntegrationTest.java b/syncflow-api/src/test/java/com/syncflow/api/cdc/CdcIntegrationTest.java index d3e24c8..a025ea6 100644 --- a/syncflow-api/src/test/java/com/syncflow/api/cdc/CdcIntegrationTest.java +++ b/syncflow-api/src/test/java/com/syncflow/api/cdc/CdcIntegrationTest.java @@ -18,6 +18,8 @@ import org.testcontainers.containers.PostgreSQLContainer; import org.testcontainers.junit.jupiter.Container; +import java.io.File; +import java.sql.Connection; import java.sql.DriverManager; import java.sql.SQLException; import java.time.Duration; @@ -54,7 +56,7 @@ static void properties(DynamicPropertyRegistry registry) { registry.add("syncflow.encryption.key", () -> "MDEyMzQ1Njc4OWFiY2RlZg=="); } - private java.sql.Connection sqlConnection; + private Connection sqlConnection; private final PostgresCdcConnector cdcConnector = new PostgresCdcConnector(); private final List capturedEvents = new CopyOnWriteArrayList<>(); private volatile boolean capturing = false; @@ -74,9 +76,12 @@ void setUp() throws SQLException { stmt.execute( "SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name = 'syncflow_slot_cdctest'"); stmt.execute("DROP PUBLICATION IF EXISTS syncflow_pub_cdctest"); - var offsetFile = System.getProperty("java.io.tmpdir") - + "/syncflow_offset_postgresql_localhost_cdctest.dat"; - new java.io.File(offsetFile).delete(); + // Offset files are now keyed per-pipeline (suffix "_default" when the + // connector context carries no pipelineId); delete both the old and new + // names so stale state can't reset a partial DELETE/UPDATE. + var offsetDir = System.getProperty("java.io.tmpdir"); + new File(offsetDir + "/syncflow_offset_postgresql_localhost_cdctest.dat").delete(); + new File(offsetDir + "/syncflow_offset_postgresql_localhost_cdctest_default.dat").delete(); stmt.execute("CREATE TABLE IF NOT EXISTS cdc_test_users (" + "id SERIAL PRIMARY KEY, name TEXT NOT NULL, email TEXT, active BOOLEAN DEFAULT true, created_at TIMESTAMP DEFAULT NOW())"); stmt.execute("CREATE TABLE IF NOT EXISTS cdc_test_orders (" + diff --git a/syncflow-connectors/src/main/java/com/syncflow/connector/cdc/DebeziumCdcConnector.java b/syncflow-connectors/src/main/java/com/syncflow/connector/cdc/DebeziumCdcConnector.java index e217ba0..9fdce63 100644 --- a/syncflow-connectors/src/main/java/com/syncflow/connector/cdc/DebeziumCdcConnector.java +++ b/syncflow-connectors/src/main/java/com/syncflow/connector/cdc/DebeziumCdcConnector.java @@ -59,10 +59,10 @@ public abstract class DebeziumCdcConnector implements CdcCapableConnector { /** * Called by subclass event parsers to record the latest offset after each - * event. + * event. Merges instead of replacing so per-table LSNs are retained (a single + * map would otherwise drop earlier tables' positions in a multi-table capture). */ protected void updateOffset(Map offset) { - lastOffset.clear(); lastOffset.putAll(offset); } @@ -174,7 +174,7 @@ public void startCDC(ConnectorContext context, Consumer eventConsumer) // use FileOffsetBackingStore so offsets survive JVM restarts // each pipeline gets its own offset file keyed by pipeline id from context - var offsetFile = resolveOffsetFilePath(config); + var offsetFile = resolveOffsetFilePath(context); debeziumProps.setProperty("offset.storage", "org.apache.kafka.connect.storage.FileOffsetBackingStore"); debeziumProps.setProperty("offset.storage.file.filename", offsetFile); @@ -296,15 +296,18 @@ private void handleSingleEvent(ChangeEvent event) { /** * Resolve a stable per-pipeline offset file path. - * Uses the database name + host as a key so separate connections get separate - * files. - * subclasses supply the slot name keyed to the pipeline id. + * Keyed by connector + host + database + PIPELINE id so multiple pipelines on + * the same database get their own offset file (shared files corrupt resume). */ - private String resolveOffsetFilePath(ConnectionConfiguration config) { + private String resolveOffsetFilePath(ConnectorContext context) { + var config = context.config(); var dir = System.getProperty("java.io.tmpdir"); + var pipelineKey = context.runtimeProperties().getOrDefault("pipelineId", "default"); + var safePipeline = pipelineKey.replaceAll("[^a-zA-Z0-9_-]", "_"); var key = connectorType().name().toLowerCase() + "_" + config.host().replace(".", "_") - + "_" + config.database(); + + "_" + config.database() + + "_" + safePipeline; return dir + "/syncflow_offset_" + key + ".dat"; } } diff --git a/syncflow-connectors/src/main/java/com/syncflow/connector/snapshot/AbstractJdbcSnapshotConnector.java b/syncflow-connectors/src/main/java/com/syncflow/connector/snapshot/AbstractJdbcSnapshotConnector.java index 1238dfa..e56d9e5 100644 --- a/syncflow-connectors/src/main/java/com/syncflow/connector/snapshot/AbstractJdbcSnapshotConnector.java +++ b/syncflow-connectors/src/main/java/com/syncflow/connector/snapshot/AbstractJdbcSnapshotConnector.java @@ -5,6 +5,7 @@ import com.syncflow.core.spi.ConnectorContext; import com.syncflow.core.spi.SnapshotCapableConnector; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.util.ArrayList; import java.util.LinkedHashMap; @@ -19,6 +20,13 @@ public abstract class AbstractJdbcSnapshotConnector @Override public long estimateRows(ConnectorContext ctx, String schema, String table) { ensureConnected(ctx); + // Prefer the planner's reltuples estimate (near-instant, no full scan) + // over SELECT COUNT(*), which reads the whole table just for progress %. + var est = estimateFromCatalog(schema, table); + if (est > 0) { + return est; + } + // Fallback: exact count only when no planner estimate is available. var sql = "SELECT COUNT(*) FROM " + schema + "." + table; try (var stmt = jdbcConnection.createStatement(); var rs = stmt.executeQuery(sql)) { @@ -28,11 +36,88 @@ public long estimateRows(ConnectorContext ctx, String schema, String table) { } } + /** Postgres planner estimate (reltuples) — cheap, avoids a full table scan. */ + private long estimateFromCatalog(String schema, String table) { + var sql = "SELECT c.reltuples::bigint FROM pg_class c " + + "JOIN pg_namespace n ON n.oid = c.relnamespace " + + "WHERE n.nspname = ? AND c.relname = ?"; + try (var stmt = jdbcConnection.prepareStatement(sql)) { + stmt.setString(1, schema); + stmt.setString(2, table); + var rs = stmt.executeQuery(); + return rs.next() ? rs.getLong(1) : 0; + } catch (SQLException e) { + return 0; + } + } + @Override public Page readBatch(ConnectorContext ctx, String schema, String table, BatchInformation batchInfo) { ensureConnected(ctx); - int offset = batchInfo.batchNumber() * batchInfo.batchSize(); + var pkCol = primaryKeyColumn(ctx, schema, table); + + // Keyset (seek) pagination when a single-column PK is available: stable + // under concurrent writes (no OFFSET drift / duplication) and single-pass. + if (pkCol != null) { + return readKeysetPage(ctx, schema, table, batchInfo, pkCol); + } + // Fallback: OFFSET/LIMIT for tables without a single-column PK. Not + // snapshot-isolated, but the cursor is still carried so the executor can + // resume from the batch. + return readOffsetPage(schema, table, batchInfo); + } + + /** + * Read a page using a keyset cursor over the PK: {@code WHERE pk > :cursor + * ORDER BY pk LIMIT size}. The last PK value becomes the next cursor, so resume + * and concurrent writes stay consistent. + */ + private Page readKeysetPage(ConnectorContext ctx, String schema, String table, + BatchInformation batchInfo, String pkCol) { + var cursor = batchInfo.cursor(); + String sql = "SELECT * FROM " + schema + "." + table + + " WHERE " + pkCol + (cursor == null ? " IS NOT NULL" : " > ?") + + " ORDER BY " + pkCol + + " LIMIT " + batchInfo.batchSize(); + var rows = new ArrayList>(); + Object lastPk = null; + try (var stmt = jdbcConnection.prepareStatement(sql)) { + if (cursor != null) { + // The cursor round-trips through a String (SPI contract). Binding with + // setObject lets the driver coerce to the PK's column type; lexically + // the value is compared to a seekable key, which holds for int, bigint, + // uuid, and text PKs — the types these connectors support. + stmt.setObject(1, cursor); + } + var rs = stmt.executeQuery(); + var meta = rs.getMetaData(); + int pkIndex = columnIndex(meta, pkCol); + int cols = meta.getColumnCount(); + while (rs.next()) { + var row = new LinkedHashMap(); + for (int i = 1; i <= cols; i++) { + row.put(meta.getColumnName(i), rs.getObject(i)); + } + rows.add(row); + if (pkIndex > 0) { + lastPk = rs.getObject(pkIndex); + } + } + } catch (SQLException e) { + throw new RuntimeException("Keyset batch read failed for " + schema + "." + table, e); + } + var nextCursor = !rows.isEmpty() && rows.size() == batchInfo.batchSize() + ? String.valueOf(lastPk) + : null; + return rows.isEmpty() ? Page.empty() : Page.of(rows, nextCursor); + } + + /** OFFSET/LIMIT fallback for tables without a single-column PK. */ + private Page readOffsetPage(String schema, String table, BatchInformation batchInfo) { + int offset = batchInfo.cursor() != null + ? Integer.parseInt(batchInfo.cursor()) + : batchInfo.batchNumber() * batchInfo.batchSize(); var sql = "SELECT * FROM " + schema + "." + table + " OFFSET " + offset + " LIMIT " + batchInfo.batchSize(); var rows = new ArrayList>(); @@ -50,9 +135,38 @@ public Page readBatch(ConnectorContext ctx, String schema, String table, } catch (SQLException e) { throw new RuntimeException("Batch read failed for " + schema + "." + table, e); } + var nextOffset = offset + rows.size(); var nextCursor = rows.size() == batchInfo.batchSize() - ? String.valueOf(offset + batchInfo.batchSize()) + ? String.valueOf(nextOffset) : null; return rows.isEmpty() ? Page.empty() : Page.of(rows, nextCursor); } + + /** Single-column primary key if the table has one, else null. */ + private String primaryKeyColumn(ConnectorContext ctx, String schema, String table) { + var pk = fetchPrimaryKey(ctx, schema, table); + return (pk != null && pk.columnNames().size() == 1) + ? sanitizeIdentifier(pk.columnNames().get(0)) + : null; + } + + /** + * Only allow identifiers safe to interpolate into SQL. DB metadata is usually + * trusted, but a crafted column name must not become an injection vector. + */ + private static String sanitizeIdentifier(String name) { + if (name == null || !name.matches("[A-Za-z_][A-Za-z0-9_]*")) { + throw new IllegalArgumentException("Unsafe column identifier: " + name); + } + return name; + } + + private static int columnIndex(ResultSetMetaData meta, String name) throws SQLException { + for (int i = 1; i <= meta.getColumnCount(); i++) { + if (meta.getColumnName(i).equalsIgnoreCase(name)) { + return i; + } + } + return -1; + } }