DBZ-6416 Add initial snapshot notifications

stevemann2705 · Jun 22, 2023 · 34e28ac · 34e28ac
1 parent d7e4063
commit 34e28ac
Show file tree

Hide file tree

Showing 26 changed files with 540 additions and 86 deletions.
diff --git a/...-mongodb/src/main/java/io/debezium/connector/mongodb/MongoDbChangeEventSourceFactory.java b/...-mongodb/src/main/java/io/debezium/connector/mongodb/MongoDbChangeEventSourceFactory.java
@@ -61,7 +61,8 @@ public MongoDbChangeEventSourceFactory(MongoDbConnectorConfig configuration, Err
     }
 
     @Override
-    public SnapshotChangeEventSource<MongoDbPartition, MongoDbOffsetContext> getSnapshotChangeEventSource(SnapshotProgressListener<MongoDbPartition> snapshotProgressListener) {
+    public SnapshotChangeEventSource<MongoDbPartition, MongoDbOffsetContext> getSnapshotChangeEventSource(SnapshotProgressListener<MongoDbPartition> snapshotProgressListener,
+                                                                                                          NotificationService<MongoDbPartition, MongoDbOffsetContext> notificationService) {
         return new MongoDbSnapshotChangeEventSource(
                 configuration,
                 taskContext,
@@ -70,7 +71,8 @@ public SnapshotChangeEventSource<MongoDbPartition, MongoDbOffsetContext> getSnap
                 dispatcher,
                 clock,
                 snapshotProgressListener,
-                errorHandler);
+                errorHandler,
+                notificationService);
     }
 
     @Override

diff --git a/.../main/java/io/debezium/connector/mongodb/MongoDbIncrementalSnapshotChangeEventSource.java b/.../main/java/io/debezium/connector/mongodb/MongoDbIncrementalSnapshotChangeEventSource.java
@@ -160,7 +160,7 @@ protected void sendEvent(EventDispatcher<MongoDbPartition, CollectionId> dispatc
         context.sendEvent(keyFromRow(row));
 
         MongoDbOffsetContext mongoDbOffsetContext = getMongoDbOffsetContext(offsetContext);
-        ReplicaSet replicaSet = replicaSets.getIncrementalSnapshotReplicaSet();
+        ReplicaSet replicaSet = replicaSets.getSnapshotReplicaSet();
         ReplicaSetOffsetContext replicaSetOffsetContext = mongoDbOffsetContext.getReplicaSetOffsetContext(replicaSet);
         ReplicaSetPartition replicaSetPartition = mongoDbOffsetContext.getReplicaSetPartition(replicaSet);
 

diff --git a/...mongodb/src/main/java/io/debezium/connector/mongodb/MongoDbSnapshotChangeEventSource.java b/...mongodb/src/main/java/io/debezium/connector/mongodb/MongoDbSnapshotChangeEventSource.java
@@ -33,25 +33,27 @@
 
 import io.debezium.DebeziumException;
 import io.debezium.connector.SnapshotRecord;
-import io.debezium.connector.mongodb.connection.ConnectionContext;
 import io.debezium.connector.mongodb.connection.MongoDbConnection;
 import io.debezium.connector.mongodb.connection.ReplicaSet;
 import io.debezium.connector.mongodb.recordemitter.MongoDbSnapshotRecordEmitter;
 import io.debezium.pipeline.ErrorHandler;
 import io.debezium.pipeline.EventDispatcher;
 import io.debezium.pipeline.EventDispatcher.SnapshotReceiver;
+import io.debezium.pipeline.notification.NotificationService;
 import io.debezium.pipeline.source.AbstractSnapshotChangeEventSource;
 import io.debezium.pipeline.source.spi.SnapshotChangeEventSource;
 import io.debezium.pipeline.source.spi.SnapshotProgressListener;
 import io.debezium.pipeline.spi.ChangeRecordEmitter;
+import io.debezium.pipeline.spi.OffsetContext;
+import io.debezium.pipeline.spi.Offsets;
 import io.debezium.pipeline.spi.SnapshotResult;
 import io.debezium.pipeline.txmetadata.TransactionContext;
 import io.debezium.util.Clock;
 import io.debezium.util.Strings;
 import io.debezium.util.Threads;
 
 /**
- * A {@link SnapshotChangeEventSource} that performs multi-threaded snapshots of replica sets.
+ * A {@link SnapshotChangeEventSource} that performs multithreaded snapshots of replica sets.
  *
  * @author Chris Cranford
  */
@@ -62,30 +64,60 @@ public class MongoDbSnapshotChangeEventSource extends AbstractSnapshotChangeEven
     private final MongoDbConnectorConfig connectorConfig;
     private final MongoDbTaskContext taskContext;
     private final MongoDbConnection.ChangeEventSourceConnectionFactory connections;
-    private final ConnectionContext connectionContext;
     private final ReplicaSets replicaSets;
     private final EventDispatcher<MongoDbPartition, CollectionId> dispatcher;
     private final Clock clock;
     private final SnapshotProgressListener<MongoDbPartition> snapshotProgressListener;
     private final ErrorHandler errorHandler;
-    private AtomicBoolean aborted = new AtomicBoolean(false);
+    private final AtomicBoolean aborted = new AtomicBoolean(false);
 
     public MongoDbSnapshotChangeEventSource(MongoDbConnectorConfig connectorConfig, MongoDbTaskContext taskContext,
                                             MongoDbConnection.ChangeEventSourceConnectionFactory connections, ReplicaSets replicaSets,
                                             EventDispatcher<MongoDbPartition, CollectionId> dispatcher, Clock clock,
-                                            SnapshotProgressListener<MongoDbPartition> snapshotProgressListener, ErrorHandler errorHandler) {
-        super(connectorConfig, snapshotProgressListener);
+                                            SnapshotProgressListener<MongoDbPartition> snapshotProgressListener, ErrorHandler errorHandler,
+                                            NotificationService<MongoDbPartition, MongoDbOffsetContext> notificationService) {
+        super(connectorConfig, snapshotProgressListener, notificationService);
         this.connectorConfig = connectorConfig;
         this.taskContext = taskContext;
         this.connections = connections;
-        this.connectionContext = taskContext.getConnectionContext();
         this.replicaSets = replicaSets;
         this.dispatcher = dispatcher;
         this.clock = clock;
         this.snapshotProgressListener = snapshotProgressListener;
         this.errorHandler = errorHandler;
     }
 
+    /*
+     * This is required because MongoDbPartition and MongoDbOffsetContext are not managed well for MongoDB. They are only correctly initialized just before starting CDC streaming
+     * In the future only ReplicaSetPartition and ReplicaSetOffset should be present and initialized in the MongoDbConnectorTask
+     */
+    @Override
+    protected Offsets<MongoDbPartition, OffsetContext> getOffsets(SnapshotContext<MongoDbPartition, MongoDbOffsetContext> snapshotContext,
+                                                                  MongoDbOffsetContext mongoDbOffsetContext, SnapshottingTask snapshottingTask) {
+
+        final MongoDbSnapshottingTask mongoDbSnapshottingTask = (MongoDbSnapshottingTask) snapshottingTask;
+        final MongoDbSnapshotContext mongoDbSnapshotContext = (MongoDbSnapshotContext) snapshotContext;
+
+        ReplicaSet replicaSet = getReplicaSet(mongoDbSnapshottingTask);
+        if (mongoDbOffsetContext == null) {
+            initSnapshotStartOffsets(mongoDbSnapshotContext);
+            return Offsets.of(snapshotContext.offset.getReplicaSetPartition(replicaSet),
+                    snapshotContext.offset.getReplicaSetOffsetContext(replicaSet));
+        }
+
+        return Offsets.of(mongoDbOffsetContext.getReplicaSetPartition(replicaSet),
+                mongoDbOffsetContext.getReplicaSetOffsetContext(replicaSet));
+    }
+
+    private ReplicaSet getReplicaSet(MongoDbSnapshottingTask mongoDbSnapshottingTask) {
+
+        // In case of a Sharded Cluster, for snapshot, only the connection.mode=sharded is supported. In this case only one ReplicaSet is present.
+        if (mongoDbSnapshottingTask.getReplicaSetsToSnapshot().isEmpty()) { // When snapshot mode is never
+            return replicaSets.getSnapshotReplicaSet();
+        }
+        return mongoDbSnapshottingTask.getReplicaSetsToSnapshot().get(0);
+    }
+
     @Override
     protected SnapshotResult<MongoDbOffsetContext> doExecute(ChangeEventSourceContext context,
                                                              MongoDbOffsetContext prevOffsetCtx,

diff --git a/debezium-connector-mongodb/src/main/java/io/debezium/connector/mongodb/ReplicaSets.java b/debezium-connector-mongodb/src/main/java/io/debezium/connector/mongodb/ReplicaSets.java
@@ -96,12 +96,12 @@ public List<ReplicaSet> all() {
     }
 
     /**
-     * Get the ReplicaSet for the incremental snapshot
+     * Get the ReplicaSet for the snapshot
      *
      * @return in case of a ReplicaSet deployments return the only ReplicaSet available.
      * In case of a Sharded Cluster, for incremental snapshot, only the connection.mode=sharded is supported. In this case only one ReplicaSet is present.
      */
-    public ReplicaSet getIncrementalSnapshotReplicaSet() {
+    public ReplicaSet getSnapshotReplicaSet() {
         return all().get(0);
     }