fix

apache · Jun 16, 2023 · 6e7b426 · 6e7b426
1 parent 61b23ff
commit 6e7b426
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 32 deletions.
diff --git a/...ache/seatunnel/connectors/seatunnel/cdc/mongodb/source/fetch/MongodbFetchTaskContext.java b/...ache/seatunnel/connectors/seatunnel/cdc/mongodb/source/fetch/MongodbFetchTaskContext.java
@@ -35,6 +35,7 @@
 import org.bson.BsonValue;
 import org.jetbrains.annotations.NotNull;
 
+import com.mongodb.client.MongoClient;
 import com.mongodb.client.model.changestream.OperationType;
 import io.debezium.connector.base.ChangeEventQueue;
 import io.debezium.pipeline.DataChangeEvent;
@@ -198,6 +199,14 @@ record -> {
 
     @Override
     public void close() {
-        createMongoClient(sourceConfig).close();
+        Runtime.getRuntime()
+                .addShutdownHook(
+                        new Thread(
+                                () -> {
+                                    MongoClient mongoClient = createMongoClient(sourceConfig);
+                                    if (mongoClient != null) {
+                                        mongoClient.close();
+                                    }
+                                }));
     }
 }
diff --git a/.../apache/seatunnel/connectors/seatunnel/cdc/mongodb/source/fetch/MongodbScanFetchTask.java b/.../apache/seatunnel/connectors/seatunnel/cdc/mongodb/source/fetch/MongodbScanFetchTask.java
@@ -75,6 +75,8 @@ public class MongodbScanFetchTask implements FetchTask<SourceSplitBase> {
 
     private volatile boolean taskRunning = false;
 
+    private MongoClient mongoClient;
+
     public MongodbScanFetchTask(SnapshotSplit snapshotSplit) {
         this.snapshotSplit = snapshotSplit;
     }
@@ -158,12 +160,15 @@ public void execute(Context context) throws Exception {
             throw e;
         } finally {
             taskRunning = false;
+            if (mongoClient != null) {
+                mongoClient.close();
+            }
         }
     }
 
     @NotNull private MongoCursor<RawBsonDocument> getSnapshotCursor(
             @NotNull SnapshotSplit snapshotSplit, MongodbSourceConfig sourceConfig) {
-        MongoClient mongoClient = createMongoClient(sourceConfig);
+        mongoClient = createMongoClient(sourceConfig);
         MongoCollection<RawBsonDocument> collection =
                 getMongoCollection(mongoClient, snapshotSplit.getTableId(), RawBsonDocument.class);
         BsonDocument startKey = (BsonDocument) snapshotSplit.getSplitStart()[1];

diff --git a/...pache/seatunnel/connectors/seatunnel/cdc/mongodb/source/fetch/MongodbStreamFetchTask.java b/...pache/seatunnel/connectors/seatunnel/cdc/mongodb/source/fetch/MongodbStreamFetchTask.java
@@ -196,6 +196,9 @@ public void execute(Context context) throws Exception {
             if (changeStreamCursor != null) {
                 changeStreamCursor.close();
             }
+            if (mongoClient != null) {
+                mongoClient.close();
+            }
         }
     }
 
@@ -205,7 +208,9 @@ public boolean isRunning() {
     }
 
     @Override
-    public void shutdown() {}
+    public void shutdown() {
+        taskRunning = false;
+    }
 
     @Override
     public IncrementalSplit getSplit() {

diff --git a/...tunnel-connector-v2-e2e/connector-cdc-mongodb-e2e/src/test/java/mongodb/MongodbCDCIT.java b/...tunnel-connector-v2-e2e/connector-cdc-mongodb-e2e/src/test/java/mongodb/MongodbCDCIT.java
@@ -142,35 +142,41 @@ public void startUp() {
 
     @TestTemplate
     public void testMongodbCdcToMysqlCheckDataE2e(TestContainer container) {
-        CompletableFuture.supplyAsync(
-                () -> {
-                    try {
-                        container.executeJob("/mongodbcdc_to_mysql.conf");
-                    } catch (Exception e) {
-                        log.error("Commit task exception :" + e.getMessage());
-                        throw new RuntimeException(e);
-                    }
-                    return null;
-                });
-        await().atMost(60000, TimeUnit.MILLISECONDS)
-                .untilAsserted(
-                        () -> {
-                            Assertions.assertIterableEquals(
-                                    readMongodbData().stream()
-                                            .peek(e -> e.remove("_id"))
-                                            .map(Document::entrySet)
-                                            .map(Set::stream)
-                                            .map(
-                                                    entryStream ->
-                                                            entryStream
-                                                                    .map(Map.Entry::getValue)
-                                                                    .collect(
-                                                                            Collectors.toCollection(
-                                                                                    ArrayList
-                                                                                            ::new)))
-                                            .collect(Collectors.toList()),
-                                    querySql());
-                        });
+        try {
+            CompletableFuture.supplyAsync(
+                    () -> {
+                        try {
+                            container.executeJob("/mongodbcdc_to_mysql.conf");
+                        } catch (Exception e) {
+                            log.error("Commit task exception :" + e.getMessage());
+                            throw new RuntimeException(e);
+                        }
+                        return null;
+                    });
+            await().atMost(60000, TimeUnit.MILLISECONDS)
+                    .untilAsserted(
+                            () -> {
+                                Assertions.assertIterableEquals(
+                                        readMongodbData().stream()
+                                                .peek(e -> e.remove("_id"))
+                                                .map(Document::entrySet)
+                                                .map(Set::stream)
+                                                .map(
+                                                        entryStream ->
+                                                                entryStream
+                                                                        .map(Map.Entry::getValue)
+                                                                        .collect(
+                                                                                Collectors
+                                                                                        .toCollection(
+                                                                                                ArrayList
+                                                                                                        ::new)))
+                                                .collect(Collectors.toList()),
+                                        querySql());
+                            });
+
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
     }
 
     private Connection getJdbcConnection() throws SQLException {