[source-mongodb-v2] : Implement WASS algorithm (#42561)

Fixes airbytehq/airbyte-internal-issues#8730

Port of #38240 for MongoDB

Author: theyueli

airbyte-integrations/connectors/source-mongodb-v2/build.gradle

@@ -3,7 +3,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.40.7'
+    cdkVersionRequired = '0.43.4'
     features = ['db-sources', 'datastore-mongo']
     useLocalCdk = false
 }

airbyte-integrations/connectors/source-mongodb-v2/integration_tests/expected_spec.json

@@ -185,6 +185,16 @@
         "default": "Lookup",
         "order": 12,
         "group": "advanced"
+      },
+      "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": 13,
+        "group": "advanced"
       }
     },
     "groups": [

airbyte-integrations/connectors/source-mongodb-v2/metadata.yaml

@@ -8,7 +8,7 @@ data:
   connectorSubtype: database
   connectorType: source
   definitionId: b2e713cd-cc36-4c0a-b5bd-b47cb8a0561e
-  dockerImageTag: 1.4.3
+  dockerImageTag: 1.5.0
   dockerRepository: airbyte/source-mongodb-v2
   documentationUrl: https://docs.airbyte.com/integrations/sources/mongodb-v2
   githubIssueLabel: source-mongodb-v2

airbyte-integrations/connectors/source-mongodb-v2/src/main/java/io/airbyte/integrations/source/mongodb/InitialSnapshotHandler.java

@@ -24,6 +24,8 @@
 import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage;
 import io.airbyte.protocol.models.v0.CatalogHelpers;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -51,7 +53,9 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIterators(
                                                                   final MongoDatabase database,
                                                                   final MongoDbSourceConfig config,
                                                                   final boolean decorateWithStartedStatus,
-                                                                  final boolean decorateWithCompletedStatus) {
+                                                                  final boolean decorateWithCompletedStatus,
+                                                                  final Instant emittedAt,
+                                                                  final Optional<Duration> cdcInitialLoadTimeout) {
     final boolean isEnforceSchema = config.getEnforceSchema();
     final var checkpointInterval = config.getCheckpointInterval();
     final String MULTIPLE_ID_TYPES_ANALYTICS_MESSAGE_KEY = "db-sources-mongo-multiple-id-types";
@@ -85,7 +89,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIterators(
 
           final Optional<CollectionStatistics> collectionStatistics = MongoUtil.getCollectionStatistics(database, airbyteStream);
           final var recordIterator = new MongoDbInitialLoadRecordIterator(collection, fields, existingState, isEnforceSchema,
-              MongoUtil.getChunkSizeForCollection(collectionStatistics, airbyteStream));
+              MongoUtil.getChunkSizeForCollection(collectionStatistics, airbyteStream), emittedAt, cdcInitialLoadTimeout);
           final var stateIterator =
               new SourceStateIterator<>(recordIterator, airbyteStream, stateManager, new StateEmitFrequency(checkpointInterval,
                   MongoConstants.CHECKPOINT_DURATION));

airbyte-integrations/connectors/source-mongodb-v2/src/main/java/io/airbyte/integrations/source/mongodb/MongoDbInitialLoadRecordIterator.java

@@ -4,6 +4,7 @@
 
 package io.airbyte.integrations.source.mongodb;
 
+import static io.airbyte.cdk.db.DbAnalyticsUtils.cdcSnapshotForceShutdownMessage;
 import static io.airbyte.integrations.source.mongodb.state.IdType.idToStringRepresenation;
 import static io.airbyte.integrations.source.mongodb.state.IdType.parseBinaryIdString;
 import static io.airbyte.integrations.source.mongodb.state.InitialSnapshotStatus.IN_PROGRESS;
@@ -13,10 +14,14 @@
 import com.mongodb.client.MongoCursor;
 import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.Sorts;
+import io.airbyte.cdk.integrations.base.AirbyteTraceMessageUtility;
 import io.airbyte.commons.exceptions.ConfigErrorException;
+import io.airbyte.commons.exceptions.TransientErrorException;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.integrations.source.mongodb.state.IdType;
 import io.airbyte.integrations.source.mongodb.state.MongoDbStreamState;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Optional;
 import org.bson.*;
 import org.bson.conversions.Bson;
@@ -44,21 +49,37 @@ public class MongoDbInitialLoadRecordIterator extends AbstractIterator<Document>
 
   private int numSubqueries = 0;
 
+  private final Instant startInstant;
+  private Optional<Duration> cdcInitialLoadTimeout;
+
   MongoDbInitialLoadRecordIterator(final MongoCollection<Document> collection,
                                    final Bson fields,
                                    final Optional<MongoDbStreamState> existingState,
                                    final boolean isEnforceSchema,
-                                   final int chunkSize) {
+                                   final int chunkSize,
+                                   final Instant startInstant,
+                                   final Optional<Duration> cdcInitialLoadTimeout) {
     this.collection = collection;
     this.fields = fields;
     this.currentState = existingState;
     this.isEnforceSchema = isEnforceSchema;
     this.chunkSize = chunkSize;
     this.currentIterator = buildNewQueryIterator();
+    this.startInstant = startInstant;
+    this.cdcInitialLoadTimeout = cdcInitialLoadTimeout;
   }
 
   @Override
   protected Document computeNext() {
+    if (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);
+    }
     if (shouldBuildNextQuery()) {
       try {
         LOGGER.info("Finishing subquery number : {}, processing at id : {}", numSubqueries,

airbyte-integrations/connectors/source-mongodb-v2/src/main/java/io/airbyte/integrations/source/mongodb/MongoDbSource.java

@@ -27,6 +27,7 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -184,7 +185,9 @@ List<AutoCloseableIterator<AirbyteMessage>> createFullRefreshIterators(final Mon
         mongoClient.getDatabase(sourceConfig.getDatabaseName()),
         sourceConfig,
         true,
-        true);
+        true,
+        emmitedAt,
+        Optional.empty());
 
     return fullRefreshIterators;
   }

airbyte-integrations/connectors/source-mongodb-v2/src/main/java/io/airbyte/integrations/source/mongodb/cdc/MongoDbCdcInitializer.java

@@ -12,6 +12,7 @@
 import com.mongodb.client.MongoDatabase;
 import io.airbyte.cdk.integrations.base.AirbyteTraceMessageUtility;
 import io.airbyte.cdk.integrations.debezium.AirbyteDebeziumHandler;
+import io.airbyte.cdk.integrations.source.relationaldb.InitialLoadTimeoutUtil;
 import io.airbyte.cdk.integrations.source.relationaldb.streamstatus.StreamStatusTraceEmitterIterator;
 import io.airbyte.commons.exceptions.ConfigErrorException;
 import io.airbyte.commons.json.Jsons;
@@ -21,6 +22,7 @@
 import io.airbyte.integrations.source.mongodb.InitialSnapshotHandler;
 import io.airbyte.integrations.source.mongodb.MongoDbSourceConfig;
 import io.airbyte.integrations.source.mongodb.MongoUtil;
+import io.airbyte.integrations.source.mongodb.state.InitialSnapshotStatus;
 import io.airbyte.integrations.source.mongodb.state.MongoDbStateManager;
 import io.airbyte.protocol.models.v0.*;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
@@ -30,6 +32,7 @@
 import java.time.Instant;
 import java.util.*;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.bson.BsonDocument;
 import org.bson.BsonTimestamp;
@@ -91,19 +94,25 @@ public List<AutoCloseableIterator<AirbyteMessage>> createCdcIterators(
     // a workaround to allow making subsequent wait time configurable.
     final Duration subsequentRecordWaitTime = firstRecordWaitTime;
     LOGGER.info("Subsequent cdc record wait time: {} seconds", subsequentRecordWaitTime);
+    final Duration initialLoadTimeout = InitialLoadTimeoutUtil.getInitialLoadTimeout(config.rawConfig());
+
     final int queueSize = MongoUtil.getDebeziumEventQueueSize(config);
     final String databaseName = config.getDatabaseName();
     final boolean isEnforceSchema = config.getEnforceSchema();
+
     final Properties defaultDebeziumProperties = MongoDbCdcProperties.getDebeziumProperties();
     logOplogInfo(mongoClient);
+
     final BsonDocument initialResumeToken =
         MongoDbResumeTokenHelper.getMostRecentResumeToken(mongoClient, databaseName, incrementalOnlyStreamsCatalog);
     final JsonNode initialDebeziumState =
         mongoDbDebeziumStateUtil.constructInitialDebeziumState(initialResumeToken, databaseName);
+
     final MongoDbCdcState cdcState =
         (stateManager.getCdcState() == null || stateManager.getCdcState().state() == null || stateManager.getCdcState().state().isNull())
             ? new MongoDbCdcState(initialDebeziumState, isEnforceSchema)
             : new MongoDbCdcState(Jsons.clone(stateManager.getCdcState().state()), stateManager.getCdcState().schema_enforced());
+
     final Optional<BsonDocument> optSavedOffset = mongoDbDebeziumStateUtil.savedOffset(
         Jsons.clone(defaultDebeziumProperties),
         incrementalOnlyStreamsCatalog,
@@ -127,6 +136,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> createCdcIterators(
         throw new ConfigErrorException(
             "Saved offset is not valid. Please reset the connection, and then increase oplog retention and/or increase sync frequency to prevent his from happening in the future. See https://docs.airbyte.com/integrations/sources/mongodb-v2#mongodb-oplog-and-change-streams for more details");
       }
+
       LOGGER.info("Saved offset is not valid. Airbyte will trigger a full refresh.");
       // If the offset in the state is invalid, reset the state to the initial STATE
       stateManager.resetState(new MongoDbCdcState(initialDebeziumState, config.getEnforceSchema()));
@@ -145,33 +155,42 @@ public List<AutoCloseableIterator<AirbyteMessage>> createCdcIterators(
     final List<ConfiguredAirbyteStream> initialSnapshotStreams =
         MongoDbCdcInitialSnapshotUtils.getStreamsForInitialSnapshot(mongoClient, stateManager, incrementalOnlyStreamsCatalog, savedOffsetIsValid);
     final InitialSnapshotHandler initialSnapshotHandler = new InitialSnapshotHandler();
+
+    final Set<AirbyteStreamNameNamespacePair> streamsStillInInitialSnapshot = stateManager.getStreamStates().entrySet().stream()
+        .filter(e -> InitialSnapshotStatus.IN_PROGRESS.equals(e.getValue().status()))
+        .map(Map.Entry::getKey)
+        .collect(Collectors.toSet());
+
+    // Fetch the streams from the catalog that still need to complete the initial snapshot sync
+    List<ConfiguredAirbyteStream> inProgressSnapshotStreams = new ArrayList<>(incrementalOnlyStreamsCatalog.getStreams().stream()
+        .filter(stream -> streamsStillInInitialSnapshot.contains(AirbyteStreamNameNamespacePair.fromAirbyteStream(stream.getStream())))
+        .map(Jsons::clone)
+        .toList());
+    final var startedCdcStreamList = incrementalOnlyStreamsCatalog.getStreams().stream()
+        .filter(stream -> (!initialSnapshotStreams.contains(stream) || inProgressSnapshotStreams.contains(stream)))
+        .map(stream -> stream.getStream().getNamespace() + "." + stream.getStream().getName()).toList();
+
     final List<AutoCloseableIterator<AirbyteMessage>> initialSnapshotIterators =
         initialSnapshotHandler.getIterators(initialSnapshotStreams, stateManager, mongoClient.getDatabase(databaseName),
-            config, true, false);
+            config, false, false, emittedAt, Optional.of(initialLoadTimeout));
 
     final AirbyteDebeziumHandler<BsonTimestamp> handler = new AirbyteDebeziumHandler<>(config.getDatabaseConfig(),
         new MongoDbCdcTargetPosition(initialResumeToken), false, firstRecordWaitTime, queueSize, false);
+
     final MongoDbCdcStateHandler mongoDbCdcStateHandler = new MongoDbCdcStateHandler(stateManager);
     final MongoDbCdcSavedInfoFetcher cdcSavedInfoFetcher = new MongoDbCdcSavedInfoFetcher(stateToBeUsed);
+
     final var cdcStreamList = incrementalOnlyStreamsCatalog.getStreams().stream()
         .filter(stream -> stream.getSyncMode() == SyncMode.INCREMENTAL)
         .map(s -> s.getStream().getNamespace() + "\\." + s.getStream().getName())
         .toList();
-    final var propertiesManager =
-        new MongoDbDebeziumPropertiesManager(defaultDebeziumProperties, config.getDatabaseConfig(), incrementalOnlyStreamsCatalog, cdcStreamList);
-    final var eventConverter =
-        new MongoDbDebeziumEventConverter(cdcMetadataInjector, incrementalOnlyStreamsCatalog, emittedAt, config.getDatabaseConfig());
-
-    final Supplier<AutoCloseableIterator<AirbyteMessage>> incrementalIteratorSupplier = () -> handler.getIncrementalIterators(
-        propertiesManager, eventConverter, cdcSavedInfoFetcher, mongoDbCdcStateHandler);
 
     // We can close the client after the initial snapshot is complete, incremental
     // iterator does not make use of the client.
     final AutoCloseableIterator<AirbyteMessage> initialSnapshotIterator = AutoCloseableIterators.appendOnClose(
         AutoCloseableIterators.concatWithEagerClose(initialSnapshotIterators), mongoClient::close);
 
     final List<AutoCloseableIterator<AirbyteMessage>> cdcStreamsStartStatusEmitters = incrementalOnlyStreamsCatalog.getStreams().stream()
-        .filter(stream -> !initialSnapshotStreams.contains(stream))
         .map(stream -> (AutoCloseableIterator<AirbyteMessage>) new StreamStatusTraceEmitterIterator(new AirbyteStreamStatusHolder(
             new io.airbyte.protocol.models.AirbyteStreamNameNamespacePair(stream.getStream().getName(), stream.getStream().getNamespace()),
             AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.STARTED)))
@@ -183,9 +202,69 @@ public List<AutoCloseableIterator<AirbyteMessage>> createCdcIterators(
             AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.COMPLETE)))
         .toList();
 
-    return Stream.of(Collections.singletonList(initialSnapshotIterator), cdcStreamsStartStatusEmitters,
-        Collections.singletonList(AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)),
-        cdcStreamsCompleteStatusEmitters).flatMap(Collection::stream).toList();
+    if (startedCdcStreamList.isEmpty()) {
+      LOGGER.info("First sync - no cdc streams have been completed or started");
+      /*
+       * This is the first run case - no initial loads have been started. In this case, we want to run the
+       * iterators in the following order: 1. Run the initial load iterators. This step will timeout and
+       * throw a transient error if run for too long (> 8hrs by default). 2. Run the debezium iterators
+       * with ALL of the incremental streams configured. This is because if step 1 completes, the initial
+       * load can be considered finished.
+       */
+      final var propertiesManager =
+          new MongoDbDebeziumPropertiesManager(defaultDebeziumProperties, config.getDatabaseConfig(), incrementalOnlyStreamsCatalog, cdcStreamList);
+      final var eventConverter =
+          new MongoDbDebeziumEventConverter(cdcMetadataInjector, incrementalOnlyStreamsCatalog, emittedAt, config.getDatabaseConfig());
+      final Supplier<AutoCloseableIterator<AirbyteMessage>> incrementalIteratorSupplier = () -> handler.getIncrementalIterators(
+          propertiesManager, eventConverter, cdcSavedInfoFetcher, mongoDbCdcStateHandler);
+
+      return Stream.of(
+          cdcStreamsStartStatusEmitters,
+          Collections.singletonList(initialSnapshotIterator),
+          Collections.singletonList(AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)),
+          cdcStreamsCompleteStatusEmitters).flatMap(Collection::stream).toList();
+    } else if (initialSnapshotIterators.isEmpty()) {
+      LOGGER.info("Initial load has finished completely - only reading the oplog");
+      /*
+       * In this case, the initial load has completed and only debezium should be run. The iterators
+       * should be run in the following order: 1. Run the debezium iterators with ALL of the incremental
+       * streams configured.
+       */
+      final var propertiesManager =
+          new MongoDbDebeziumPropertiesManager(defaultDebeziumProperties, config.getDatabaseConfig(), incrementalOnlyStreamsCatalog, cdcStreamList);
+      final var eventConverter =
+          new MongoDbDebeziumEventConverter(cdcMetadataInjector, incrementalOnlyStreamsCatalog, emittedAt, config.getDatabaseConfig());
+      final Supplier<AutoCloseableIterator<AirbyteMessage>> incrementalIteratorSupplier = () -> handler.getIncrementalIterators(
+          propertiesManager, eventConverter, cdcSavedInfoFetcher, mongoDbCdcStateHandler);
+      return Stream.of(
+          cdcStreamsStartStatusEmitters,
+          Collections.singletonList(AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)),
+          cdcStreamsCompleteStatusEmitters).flatMap(Collection::stream).toList();
+    } else {
+      LOGGER.info("Initial load is in progress - reading oplog first and then resuming with initial load.");
+      /*
+       * In this case, the initial load has partially completed (WASS case). The iterators should be run
+       * in the following order: 1. Run the debezium iterators with only the incremental streams which
+       * have been fully or partially completed configured. 2. Resume initial load for partially completed
+       * and not started streams. This step will timeout and throw a transient error if run for too long
+       * (> 8hrs by default). 3. Emit a transient error. This is to signal to the platform to restart the
+       * sync to clear the oplog. We cannot simply add the same cdc iterators as their target end position
+       * is fixed to the tip of the oplog at the start of the sync.
+       */
+      final var propertiesManager =
+          new MongoDbDebeziumPropertiesManager(defaultDebeziumProperties, config.getDatabaseConfig(), incrementalOnlyStreamsCatalog,
+              startedCdcStreamList);
+      final var eventConverter =
+          new MongoDbDebeziumEventConverter(cdcMetadataInjector, incrementalOnlyStreamsCatalog, emittedAt, config.getDatabaseConfig());
+      final Supplier<AutoCloseableIterator<AirbyteMessage>> incrementalIteratorSupplier = () -> handler.getIncrementalIterators(
+          propertiesManager, eventConverter, cdcSavedInfoFetcher, mongoDbCdcStateHandler);
+      return Stream.of(
+          cdcStreamsStartStatusEmitters,
+          Collections.singletonList(AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)),
+          Collections.singletonList(initialSnapshotIterator),
+          cdcStreamsCompleteStatusEmitters)
+          .flatMap(Collection::stream).toList();
+    }
   }
 
   private void logOplogInfo(final MongoClient mongoClient) {

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

@@ -185,6 +185,16 @@
         "default": "Lookup",
         "order": 12,
         "group": "advanced"
+      },
+      "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": 13,
+        "group": "advanced"
       }
     },
     "groups": [