[Source-Postgres] : Implement WASS algorithm

airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties

@@ -1 +1 @@
-version=0.41.4
+version=0.41.4

airbyte-integrations/connectors/source-mysql/src/main/java/io/airbyte/integrations/source/mysql/initialsync/MySqlInitialLoadRecordIterator.java

@@ -99,8 +99,8 @@ protected AirbyteRecordData computeNext() {
     if (isCdcSync && cdcInitialLoadTimeout.isPresent()
         && Duration.between(startInstant, Instant.now()).compareTo(cdcInitialLoadTimeout.get()) > 0) {
       final String cdcInitialLoadTimeoutMessage = String.format(
-          "Initial load for table %s has taken longer than %s, Canceling sync so that CDC replication can catch-up on subsequent attempt, and then initial snapshotting will resume",
-          getAirbyteStream().get(), cdcInitialLoadTimeout.get());
+          "Initial load has taken longer than %s, Canceling sync so that CDC replication can catch-up on subsequent attempt, and then initial snapshotting will resume",
+          cdcInitialLoadTimeout.get());
       LOGGER.info(cdcInitialLoadTimeoutMessage);
       AirbyteTraceMessageUtility.emitAnalyticsTrace(cdcSnapshotForceShutdownMessage());
       throw new TransientErrorException(cdcInitialLoadTimeoutMessage);

airbyte-integrations/connectors/source-postgres/build.gradle

@@ -12,7 +12,7 @@ java {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.40.10'
+    cdkVersionRequired = '0.41.2'
     features = ['db-sources', 'datastore-postgres']
     useLocalCdk = false
 }

airbyte-integrations/connectors/source-postgres/metadata.yaml

@@ -9,7 +9,7 @@ data:
   connectorSubtype: database
   connectorType: source
   definitionId: decd338e-5647-4c0b-adf4-da0e75f5a750
-  dockerImageTag: 3.4.26
+  dockerImageTag: 3.5.0
   dockerRepository: airbyte/source-postgres
   documentationUrl: https://docs.airbyte.com/integrations/sources/postgres
   githubIssueLabel: source-postgres

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/PostgresSource.java

@@ -556,7 +556,8 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(final
                                                                                                                          * decorateWithStartedStatus=
                                                                                                                          */ true, /*
                                                                                                                                    * decorateWithCompletedStatus=
-                                                                                                                                   */ true));
+                                                                                                                                   */ true,
+          Optional.empty()));
       final List<AutoCloseableIterator<AirbyteMessage>> xminIterators = new ArrayList<>(xminHandler.getIncrementalIterators(
           new ConfiguredAirbyteCatalog().withStreams(xminStreams.streamsForXminSync()), tableNameToTable, emittedAt));
 
@@ -600,7 +601,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(final
       final List<AutoCloseableIterator<AirbyteMessage>> initialSyncCtidIterators = new ArrayList<>(
           cursorBasedCtidHandler.getInitialSyncCtidIterator(new ConfiguredAirbyteCatalog().withStreams(finalListOfStreamsToBeSyncedViaCtid),
               tableNameToTable,
-              emittedAt, /* decorateWithStartedStatus= */ true, /* decorateWithCompletedStatus= */ true));
+              emittedAt, /* decorateWithStartedStatus= */ true, /* decorateWithCompletedStatus= */ true, Optional.empty()));
       final List<AutoCloseableIterator<AirbyteMessage>> cursorBasedIterators = new ArrayList<>(super.getIncrementalIterators(database,
           new ConfiguredAirbyteCatalog().withStreams(
               cursorBasedStreamsCategorised.remainingStreams()

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/InitialSyncCtidIterator.java

@@ -4,6 +4,7 @@
 
 package io.airbyte.integrations.source.postgres.ctid;
 
+import static io.airbyte.cdk.db.DbAnalyticsUtils.cdcSnapshotForceShutdownMessage;
 import static io.airbyte.cdk.integrations.source.relationaldb.RelationalDbQueryUtils.getFullyQualifiedTableNameWithQuoting;
 import static io.airbyte.integrations.source.postgres.ctid.InitialSyncCtidIteratorConstants.EIGHT_KB;
 import static io.airbyte.integrations.source.postgres.ctid.InitialSyncCtidIteratorConstants.GIGABYTE;
@@ -14,7 +15,9 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
 import io.airbyte.cdk.db.jdbc.JdbcDatabase;
+import io.airbyte.cdk.integrations.base.AirbyteTraceMessageUtility;
 import io.airbyte.cdk.integrations.source.relationaldb.RelationalDbQueryUtils;
+import io.airbyte.commons.exceptions.TransientErrorException;
 import io.airbyte.commons.stream.AirbyteStreamUtils;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.commons.util.AutoCloseableIterators;
@@ -25,6 +28,8 @@
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
@@ -67,6 +72,10 @@ public class InitialSyncCtidIterator extends AbstractIterator<RowDataWithCtid> i
   private boolean subQueriesInitialized = false;
   private final boolean tidRangeScanCapableDBServer;
 
+  private final Instant startInstant;
+  private Optional<Duration> cdcInitialLoadTimeout;
+  private boolean isCdcSync;
+
   public InitialSyncCtidIterator(final CtidStateManager ctidStateManager,
                                  final JdbcDatabase database,
                                  final CtidPostgresSourceOperations sourceOperations,
@@ -79,7 +88,9 @@ public InitialSyncCtidIterator(final CtidStateManager ctidStateManager,
                                  final int maxTuple,
                                  final FileNodeHandler fileNodeHandler,
                                  final boolean tidRangeScanCapableDBServer,
-                                 final boolean useTestPageSize) {
+                                 final boolean useTestPageSize,
+                                 final Instant startInstant,
+                                 final Optional<Duration> cdcInitialLoadTimeout) {
     this.airbyteStream = AirbyteStreamUtils.convertFromNameAndNamespace(tableName, schemaName);
     this.blockSize = blockSize;
     this.maxTuple = maxTuple;
@@ -95,11 +106,23 @@ public InitialSyncCtidIterator(final CtidStateManager ctidStateManager,
     this.tableSize = tableSize;
     this.tidRangeScanCapableDBServer = tidRangeScanCapableDBServer;
     this.useTestPageSize = useTestPageSize;
+    this.startInstant = startInstant;
+    this.cdcInitialLoadTimeout = cdcInitialLoadTimeout;
+    this.isCdcSync = isCdcSync(ctidStateManager);
   }
 
   @CheckForNull
   @Override
   protected RowDataWithCtid computeNext() {
+    if (isCdcSync && cdcInitialLoadTimeout.isPresent()
+        && Duration.between(startInstant, Instant.now()).compareTo(cdcInitialLoadTimeout.get()) > 0) {
+      final String cdcInitialLoadTimeoutMessage = String.format(
+          "Initial load for table %s has taken longer than %s, Canceling sync so that CDC replication can catch-up on subsequent attempt, and then initial snapshotting will resume",
+          getAirbyteStream().get(), cdcInitialLoadTimeout.get());
+      LOGGER.info(cdcInitialLoadTimeoutMessage);
+      AirbyteTraceMessageUtility.emitAnalyticsTrace(cdcSnapshotForceShutdownMessage());
+      throw new TransientErrorException(cdcInitialLoadTimeoutMessage);
+    }
     try {
       if (!subQueriesInitialized) {
         initSubQueries();
@@ -330,4 +353,14 @@ public void close() throws Exception {
     }
   }
 
+  private boolean isCdcSync(CtidStateManager initialLoadStateManager) {
+    if (initialLoadStateManager instanceof CtidGlobalStateManager) {
+      LOGGER.info("Running a cdc sync");
+      return true;
+    } else {
+      LOGGER.info("Not running a cdc sync");
+      return false;
+    }
+  }
+
 }

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/PostgresCtidHandler.java

@@ -80,10 +80,12 @@ public PostgresCtidHandler(final JsonNode config,
     this.tidRangeScanCapableDBServer = CtidUtils.isTidRangeScanCapableDBServer(database);
   }
 
+  @NotNull
   @Override
   public AutoCloseableIterator<AirbyteMessage> getIteratorForStream(@NotNull ConfiguredAirbyteStream airbyteStream,
                                                                     @NotNull TableInfo<CommonField<PostgresType>> table,
-                                                                    @NotNull Instant emittedAt) {
+                                                                    @NotNull Instant emittedAt,
+                                                                    @NotNull Optional<Duration> cdcInitialLoadTimeout) {
     final AirbyteStream stream = airbyteStream.getStream();
     final String streamName = stream.getName();
     final String namespace = stream.getNamespace();
@@ -100,7 +102,7 @@ public AutoCloseableIterator<AirbyteMessage> getIteratorForStream(@NotNull Confi
         table.getName(),
         tableBlockSizes.get(pair).tableSize(),
         tableBlockSizes.get(pair).blockSize(),
-        tablesMaxTuple.orElseGet(() -> Map.of(pair, -1)).get(pair));
+        tablesMaxTuple.orElseGet(() -> Map.of(pair, -1)).get(pair), emittedAt, cdcInitialLoadTimeout);
     final AutoCloseableIterator<AirbyteMessageWithCtid> recordIterator =
         getRecordIterator(queryStream, streamName, namespace, emittedAt.toEpochMilli());
     final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, airbyteStream);
@@ -110,9 +112,10 @@ public AutoCloseableIterator<AirbyteMessage> getIteratorForStream(@NotNull Confi
   public List<AutoCloseableIterator<AirbyteMessage>> getInitialSyncCtidIterator(
                                                                                 final ConfiguredAirbyteCatalog catalog,
                                                                                 final Map<String, TableInfo<CommonField<PostgresType>>> tableNameToTable,
-                                                                                final Instant emmitedAt,
+                                                                                final Instant emittedAt,
                                                                                 final boolean decorateWithStartedStatus,
-                                                                                final boolean decorateWithCompletedStatus) {
+                                                                                final boolean decorateWithCompletedStatus,
+                                                                                final Optional<Duration> cdcInitialLoadTimeout) {
     final List<AutoCloseableIterator<AirbyteMessage>> iteratorList = new ArrayList<>();
     for (final ConfiguredAirbyteStream airbyteStream : catalog.getStreams()) {
       final AirbyteStream stream = airbyteStream.getStream();
@@ -134,7 +137,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getInitialSyncCtidIterator(
         // Grab the selected fields to sync
         final TableInfo<CommonField<PostgresType>> table = tableNameToTable
             .get(fullyQualifiedTableName);
-        final var iterator = getIteratorForStream(airbyteStream, table, emmitedAt);
+        final var iterator = getIteratorForStream(airbyteStream, table, emittedAt, cdcInitialLoadTimeout);
         iteratorList.add(iterator);
 
         if (decorateWithCompletedStatus) {
@@ -152,12 +155,15 @@ private AutoCloseableIterator<RowDataWithCtid> queryTableCtid(
                                                                 final String tableName,
                                                                 final long tableSize,
                                                                 final long blockSize,
-                                                                final int maxTuple) {
+                                                                final int maxTuple,
+                                                                final Instant emittedAt,
+                                                                @NotNull final Optional<Duration> cdcInitialLoadTimeout) {
 
     LOGGER.info("Queueing query for table: {}", tableName);
     return new InitialSyncCtidIterator(ctidStateManager, database, sourceOperations, quoteString, columnNames, schemaName, tableName, tableSize,
         blockSize, maxTuple, fileNodeHandler, tidRangeScanCapableDBServer,
-        config.has(USE_TEST_CHUNK_SIZE) && config.get(USE_TEST_CHUNK_SIZE).asBoolean());
+        config.has(USE_TEST_CHUNK_SIZE) && config.get(USE_TEST_CHUNK_SIZE).asBoolean(), emittedAt,
+        cdcInitialLoadTimeout);
   }
 
   // Transforms the given iterator to create an {@link AirbyteRecordMessage}

airbyte-integrations/connectors/source-postgres/src/main/resources/spec.json

@@ -304,6 +304,16 @@
                 "enum": ["Fail sync", "Re-sync data"],
                 "default": "Fail sync",
                 "order": 9
+              },
+              "initial_load_timeout_hours": {
+                "type": "integer",
+                "title": "Initial Load Timeout in Hours (Advanced)",
+                "description": "The amount of time an initial load is allowed to continue for before catching up on CDC logs.",
+                "default": 8,
+                "min": 4,
+                "max": 24,
+                "order": 10,
+                "always_show": true
               }
             }
           },

airbyte-integrations/connectors/source-postgres/src/test-integration/resources/expected_cloud_deployment_spec.json

@@ -305,6 +305,16 @@
                 "enum": ["Fail sync", "Re-sync data"],
                 "default": "Fail sync",
                 "order": 9
+              },
+              "initial_load_timeout_hours": {
+                "type": "integer",
+                "title": "Initial Load Timeout in Hours (Advanced)",
+                "description": "The amount of time an initial load is allowed to continue for before catching up on CDC logs.",
+                "default": 8,
+                "min": 4,
+                "max": 24,
+                "order": 10,
+                "always_show": true
               }
             }
           },

airbyte-integrations/connectors/source-postgres/src/test-integration/resources/expected_spec.json

@@ -304,6 +304,16 @@
                 "enum": ["Fail sync", "Re-sync data"],
                 "default": "Fail sync",
                 "order": 9
+              },
+              "initial_load_timeout_hours": {
+                "type": "integer",
+                "title": "Initial Load Timeout in Hours (Advanced)",
+                "description": "The amount of time an initial load is allowed to continue for before catching up on CDC logs.",
+                "default": 8,
+                "min": 4,
+                "max": 24,
+                "order": 10,
+                "always_show": true
               }
             }
           },