feat: Add CDC sync checkpointing based on time or records

airbyte-config/init/src/main/resources/seed/source_definitions.yaml

@@ -45,7 +45,7 @@
 - name: AlloyDB for PostgreSQL
   sourceDefinitionId: 1fa90628-2b9e-11ed-a261-0242ac120002
   dockerRepository: airbyte/source-alloydb
-  dockerImageTag: 2.0.1
+  dockerImageTag: 2.0.2
   documentationUrl: https://docs.airbyte.com/integrations/sources/alloydb
   icon: alloydb.svg
   sourceType: database
@@ -1501,7 +1501,7 @@
 - name: Postgres
   sourceDefinitionId: decd338e-5647-4c0b-adf4-da0e75f5a750
   dockerRepository: airbyte/source-postgres
-  dockerImageTag: 2.0.0
+  dockerImageTag: 2.0.2
   documentationUrl: https://docs.airbyte.com/integrations/sources/postgres
   icon: postgresql.svg
   sourceType: database

airbyte-config/init/src/main/resources/seed/source_specs.yaml

@@ -370,7 +370,7 @@
     supportsNormalization: false
     supportsDBT: false
     supported_destination_sync_modes: []
-- dockerImage: "airbyte/source-alloydb:2.0.1"
+- dockerImage: "airbyte/source-alloydb:2.0.2"
   spec:
     documentationUrl: "https://docs.airbyte.com/integrations/sources/postgres"
     connectionSpecification:
@@ -11709,7 +11709,7 @@
     supportsNormalization: false
     supportsDBT: false
     supported_destination_sync_modes: []
-- dockerImage: "airbyte/source-postgres:2.0.0"
+- dockerImage: "airbyte/source-postgres:2.0.2"
   spec:
     documentationUrl: "https://docs.airbyte.com/integrations/sources/postgres"
     connectionSpecification:

airbyte-integrations/bases/debezium/src/main/java/io/airbyte/integrations/debezium/AirbyteDebeziumHandler.java

@@ -9,24 +9,17 @@
 import io.airbyte.commons.util.AutoCloseableIterators;
 import io.airbyte.commons.util.MoreIterators;
 import io.airbyte.db.jdbc.JdbcUtils;
-import io.airbyte.integrations.debezium.internals.AirbyteFileOffsetBackingStore;
-import io.airbyte.integrations.debezium.internals.AirbyteSchemaHistoryStorage;
-import io.airbyte.integrations.debezium.internals.DebeziumEventUtils;
-import io.airbyte.integrations.debezium.internals.DebeziumRecordIterator;
-import io.airbyte.integrations.debezium.internals.DebeziumRecordPublisher;
+import io.airbyte.integrations.debezium.internals.*;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
 import io.airbyte.protocol.models.v0.SyncMode;
 import io.debezium.engine.ChangeEvent;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.Iterator;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.concurrent.LinkedBlockingQueue;
-import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -116,30 +109,21 @@ public AutoCloseableIterator<AirbyteMessage> getIncrementalIterators(final Confi
         publisher::close,
         firstRecordWaitTime);
 
-    // convert to airbyte message.
-    final AutoCloseableIterator<AirbyteMessage> messageIterator = AutoCloseableIterators
-        .transform(
-            eventIterator,
-            (event) -> DebeziumEventUtils.toAirbyteMessage(event, cdcMetadataInjector, emittedAt));
-
-    // our goal is to get the state at the time this supplier is called (i.e. after all message records
-    // have been produced)
-    final Supplier<AirbyteMessage> stateMessageSupplier = () -> {
-      final Map<String, String> offset = offsetManager.read();
-      final String dbHistory = trackSchemaHistory ? schemaHistoryManager
-          .orElseThrow(() -> new RuntimeException("Schema History Tracking is true but manager is not initialised")).read() : null;
-
-      return cdcStateHandler.saveState(offset, dbHistory);
-    };
-
-    // wrap the supplier in an iterator so that we can concat it to the message iterator.
-    final Iterator<AirbyteMessage> stateMessageIterator = MoreIterators.singletonIteratorFromSupplier(stateMessageSupplier);
-
-    // this structure guarantees that the debezium engine will be closed, before we attempt to emit the
-    // state file. we want this so that we have a guarantee that the debezium offset file (which we use
-    // to produce the state file) is up-to-date.
-
-    return AutoCloseableIterators.concatWithEagerClose(messageIterator, AutoCloseableIterators.fromIterator(stateMessageIterator));
+    final Duration syncCheckpointSeconds =
+        config.get("sync_checkpoint_seconds") != null ? Duration.ofSeconds(config.get("sync_checkpoint_seconds").asLong())
+            : DebeziumStateDecoratingIterator.SYNC_CHECKPOINT_SECONDS;
+    final Long syncCheckpointRecords = config.get("sync_checkpoint_records") != null ? config.get("sync_checkpoint_records").asLong()
+        : DebeziumStateDecoratingIterator.SYNC_CHECKPOINT_RECORDS;
+    return AutoCloseableIterators.fromIterator(new DebeziumStateDecoratingIterator(
+        eventIterator,
+        cdcStateHandler,
+        cdcMetadataInjector,
+        emittedAt,
+        offsetManager,
+        trackSchemaHistory,
+        schemaHistoryManager.orElse(null),
+        syncCheckpointSeconds,
+        syncCheckpointRecords));
   }
 
   private Optional<AirbyteSchemaHistoryStorage> schemaHistoryManager(final CdcSavedInfoFetcher cdcSavedInfoFetcher) {

airbyte-integrations/bases/debezium/src/main/java/io/airbyte/integrations/debezium/CdcStateHandler.java

@@ -5,16 +5,51 @@
 package io.airbyte.integrations.debezium;
 
 import io.airbyte.protocol.models.v0.AirbyteMessage;
+import io.debezium.engine.ChangeEvent;
 import java.util.Map;
 
 /**
  * This interface is used to allow connectors to save the offset and schema history in the manner
- * which suits them
+ * which suits them. Also, it adds some utils to verify CDC event status.
  */
 public interface CdcStateHandler {
 
   AirbyteMessage saveState(Map<String, String> offset, String dbHistory);
 
   AirbyteMessage saveStateAfterCompletionOfSnapshotOfNewStreams();
 
+  /**
+   * This function indicates if the event is part of the snapshot or not.
+   *
+   * @param event Event from the CDC load
+   * @return Returns `true` when the DB event is part of the snapshot load. Otherwise, returns `false`
+   */
+  default boolean isSnapshotEvent(final ChangeEvent<String, String> event) {
+    return false;
+  }
+
+  /**
+   * This function checks if the event we are processing in the loop is already behind the offset so
+   * the process can safety save the state.
+   *
+   * @param offset DB CDC offset
+   * @param event Event from the CDC load
+   * @return Returns `true` when the record is behind the offset. Otherwise, it returns `false`
+   */
+  default boolean isRecordBehindOffset(final Map<String, String> offset, final ChangeEvent<String, String> event) {
+    return false;
+  }
+
+  /**
+   * This function compares two offsets to make sure both are not pointing to the same position. The
+   * main purpose is to avoid sending same offset multiple times.
+   *
+   * @param offsetA Offset to compare
+   * @param offsetB Offset to compare
+   * @return Returns `true` if both offsets are at the same position. Otherwise, it returns `false`
+   */
+  default boolean isSameOffset(final Map<String, String> offsetA, final Map<String, String> offsetB) {
+    return true;
+  }
+
 }

airbyte-integrations/bases/debezium/src/main/java/io/airbyte/integrations/debezium/internals/DebeziumStateDecoratingIterator.java

@@ -0,0 +1,198 @@
+/*
+ * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.integrations.debezium.internals;
+
+import com.google.common.collect.AbstractIterator;
+import io.airbyte.integrations.debezium.CdcMetadataInjector;
+import io.airbyte.integrations.debezium.CdcStateHandler;
+import io.airbyte.protocol.models.v0.AirbyteMessage;
+import io.airbyte.protocol.models.v0.AirbyteStateMessage;
+import io.debezium.engine.ChangeEvent;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.OffsetDateTime;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This class encapsulates CDC change events and adds the required functionality to create
+ * checkpoints for CDC replications. That way, if the process fails in the middle of a long sync, it
+ * will be able to recover for any acknowledged checkpoint in the next syncs.
+ */
+public class DebeziumStateDecoratingIterator extends AbstractIterator<AirbyteMessage> implements Iterator<AirbyteMessage> {
+
+  public static final Duration SYNC_CHECKPOINT_SECONDS = Duration.ofMinutes(15);
+  public static final Integer SYNC_CHECKPOINT_RECORDS = 10_000;
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(DebeziumStateDecoratingIterator.class);
+
+  private final Iterator<ChangeEvent<String, String>> changeEventIterator;
+  private final CdcStateHandler cdcStateHandler;
+  private final AirbyteFileOffsetBackingStore offsetManager;
+  private final boolean trackSchemaHistory;
+  private final AirbyteSchemaHistoryStorage schemaHistoryManager;
+  private final CdcMetadataInjector cdcMetadataInjector;
+  private final Instant emittedAt;
+
+  private boolean isSyncFinished = false;
+
+  /**
+   * These parameters control when a checkpoint message has to be sent in a CDC integration. We can
+   * emit a checkpoint when any of the following two conditions are met.
+   * <p/>
+   * 1. The amount of records in the current loop ({@code SYNC_CHECKPOINT_RECORDS}) is higher than a
+   * threshold defined by {@code SYNC_CHECKPOINT_RECORDS}.
+   * <p/>
+   * 2. Time between checkpoints ({@code dateTimeLastSync}) is higher than a {@code Duration} defined
+   * at {@code SYNC_CHECKPOINT_SECONDS}.
+   * <p/>
+   */
+  private final Duration syncCheckpointDuration;
+  private final Long syncCheckpointRecords;
+  private OffsetDateTime dateTimeLastSync;
+  private Long recordsLastSync;
+  private boolean sendCheckpointMessage = false;
+
+  /**
+   * `checkpointOffsetToSend` is used as temporal storage for the offset that we want to send as
+   * message. As Debezium is reading records faster that we process them, if we try to send
+   * `offsetManger.read()` offset, it is possible that the state is behind the record we are currently
+   * propagating. To avoid that, we store the offset as soon as we reach the checkpoint threshold
+   * (time or records) and we wait to send it until we are sure that the record we are processing is
+   * behind the offset to be sent.
+   */
+  private final HashMap<String, String> checkpointOffsetToSend = new HashMap<>();
+
+  /**
+   * `previousCheckpointOffset` is used to make sure we don't send duplicated states with the same
+   * offset. Is it possible that the offset Debezium report doesn't move for a period of time, and if
+   * we just rely on the `offsetManger.read()`, there is a chance to sent duplicate states, generating
+   * an unneeded usage of networking and processing.
+   */
+  private final HashMap<String, String> previousCheckpointOffset;
+
+  /**
+   * @param changeEventIterator Base iterator that we want to enrich with checkpoint messages
+   * @param cdcStateHandler Handler to save the offset and schema history
+   * @param offsetManager Handler to read and write debezium offset file
+   * @param trackSchemaHistory Set true if the schema needs to be tracked
+   * @param schemaHistoryManager Handler to write schema. Needs to be initialized if
+   *        trackSchemaHistory is set to true
+   * @param checkpointDuration Duration object with time between syncs
+   * @param checkpointRecords Number of records between syncs
+   */
+  public DebeziumStateDecoratingIterator(final Iterator<ChangeEvent<String, String>> changeEventIterator,
+                                         final CdcStateHandler cdcStateHandler,
+                                         final CdcMetadataInjector cdcMetadataInjector,
+                                         final Instant emittedAt,
+                                         final AirbyteFileOffsetBackingStore offsetManager,
+                                         final boolean trackSchemaHistory,
+                                         final AirbyteSchemaHistoryStorage schemaHistoryManager,
+                                         final Duration checkpointDuration,
+                                         final Long checkpointRecords) {
+    this.changeEventIterator = changeEventIterator;
+    this.cdcStateHandler = cdcStateHandler;
+    this.cdcMetadataInjector = cdcMetadataInjector;
+    this.emittedAt = emittedAt;
+    this.offsetManager = offsetManager;
+    this.trackSchemaHistory = trackSchemaHistory;
+    this.schemaHistoryManager = schemaHistoryManager;
+
+    this.syncCheckpointDuration = checkpointDuration;
+    this.syncCheckpointRecords = checkpointRecords;
+    this.previousCheckpointOffset = (HashMap<String, String>) offsetManager.read();
+    resetCheckpointValues();
+  }
+
+  /**
+   * Computes the next record retrieved from Source stream. Emits state messages as checkpoints based
+   * on number of records or time lapsed.
+   *
+   * <p>
+   * If this method throws an exception, it will propagate outward to the {@code hasNext} or
+   * {@code next} invocation that invoked this method. Any further attempts to use the iterator will
+   * result in an {@link IllegalStateException}.
+   * </p>
+   *
+   * @return {@link AirbyteStateMessage} containing CDC data or state checkpoint message.
+   */
+  @Override
+  protected AirbyteMessage computeNext() {
+    if (isSyncFinished) {
+      return endOfData();
+    }
+
+    if (sendCheckpointMessage) {
+      final AirbyteMessage stateMessage = createStateMessage(checkpointOffsetToSend);
+      previousCheckpointOffset.clear();
+      previousCheckpointOffset.putAll(checkpointOffsetToSend);
+      resetCheckpointValues();
+      return stateMessage;
+    }
+
+    if (changeEventIterator.hasNext()) {
+      try {
+        final ChangeEvent<String, String> event = changeEventIterator.next();
+        recordsLastSync++;
+
+        if (checkpointOffsetToSend.size() == 0 &&
+            (recordsLastSync >= syncCheckpointRecords ||
+                Duration.between(dateTimeLastSync, OffsetDateTime.now()).compareTo(syncCheckpointDuration) > 0)) {
+          // Using temporal variable to avoid reading teh offset twice, one in the condition and another in
+          // the assignation
+          final HashMap<String, String> temporalOffset = (HashMap<String, String>) offsetManager.read();
+          if (!cdcStateHandler.isSameOffset(previousCheckpointOffset, temporalOffset)) {
+            checkpointOffsetToSend.putAll(temporalOffset);
+          }
+        }
+
+        if (checkpointOffsetToSend.size() == 1
+            && changeEventIterator.hasNext()
+            && !cdcStateHandler.isSnapshotEvent(event)
+            && cdcStateHandler.isRecordBehindOffset(checkpointOffsetToSend, event)) {
+          sendCheckpointMessage = true;
+        }
+
+        return DebeziumEventUtils.toAirbyteMessage(event, cdcMetadataInjector, emittedAt);
+      } catch (final Exception e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    isSyncFinished = true;
+    return createStateMessage(offsetManager.read());
+  }
+
+  /**
+   * Initialize or reset the checkpoint variables.
+   */
+  private void resetCheckpointValues() {
+    sendCheckpointMessage = false;
+    checkpointOffsetToSend.clear();
+    recordsLastSync = 0L;
+    dateTimeLastSync = OffsetDateTime.now();
+  }
+
+  /**
+   * Creates {@link AirbyteStateMessage} while updating CDC data, used to checkpoint the state of the
+   * process.
+   *
+   * @return {@link AirbyteStateMessage} which includes offset and schema history if used.
+   */
+  private AirbyteMessage createStateMessage(final Map<String, String> offset) {
+    if (trackSchemaHistory && schemaHistoryManager == null) {
+      throw new RuntimeException("Schema History Tracking is true but manager is not initialised");
+    }
+    if (offsetManager == null) {
+      throw new RuntimeException("Offset can not be null");
+    }
+
+    return cdcStateHandler.saveState(offset, schemaHistoryManager != null ? schemaHistoryManager.read() : null);
+  }
+
+}

airbyte-integrations/connectors/source-alloydb-strict-encrypt/Dockerfile

@@ -16,5 +16,5 @@ ENV APPLICATION source-alloydb-strict-encrypt
 
 COPY --from=build /airbyte /airbyte
 
-LABEL io.airbyte.version=2.0.1
+LABEL io.airbyte.version=2.0.2
 LABEL io.airbyte.name=airbyte/source-alloydb-strict-encrypt

airbyte-integrations/connectors/source-alloydb/Dockerfile

@@ -16,5 +16,5 @@ ENV APPLICATION source-alloydb
 
 COPY --from=build /airbyte /airbyte
 
-LABEL io.airbyte.version=2.0.1
+LABEL io.airbyte.version=2.0.2
 LABEL io.airbyte.name=airbyte/source-alloydb