[FLINK-38606][Connectors/MongoDB] Support for Flink 2.1.0 API

Author: neoneo008

flink-connector-mongodb/src/main/java/org/apache/flink/connector/mongodb/sink/MongoSink.java

@@ -21,6 +21,7 @@
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.connector.sink2.Sink;
 import org.apache.flink.api.connector.sink2.SinkWriter;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
 import org.apache.flink.api.java.ClosureCleaner;
 import org.apache.flink.connector.base.DeliveryGuarantee;
 import org.apache.flink.connector.mongodb.common.config.MongoConnectionOptions;
@@ -78,7 +79,7 @@ public static <IN> MongoSinkBuilder<IN> builder() {
     }
 
     @Override
-    public SinkWriter<IN> createWriter(InitContext context) {
+    public SinkWriter<IN> createWriter(WriterInitContext context) {
         return new MongoWriter<>(
                 connectionOptions,
                 writeOptions,

flink-connector-mongodb/src/main/java/org/apache/flink/connector/mongodb/sink/writer/MongoWriter.java

@@ -22,8 +22,8 @@
 import org.apache.flink.api.common.functions.util.ListCollector;
 import org.apache.flink.api.common.operators.MailboxExecutor;
 import org.apache.flink.api.common.serialization.SerializationSchema;
-import org.apache.flink.api.connector.sink2.Sink;
 import org.apache.flink.api.connector.sink2.SinkWriter;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
 import org.apache.flink.connector.mongodb.common.config.MongoConnectionOptions;
 import org.apache.flink.connector.mongodb.sink.config.MongoWriteOptions;
 import org.apache.flink.connector.mongodb.sink.writer.context.DefaultMongoSinkContext;
@@ -89,7 +89,7 @@ public MongoWriter(
             MongoConnectionOptions connectionOptions,
             MongoWriteOptions writeOptions,
             boolean flushOnCheckpoint,
-            Sink.InitContext initContext,
+            WriterInitContext initContext,
             MongoSerializationSchema<IN> serializationSchema) {
         this.connectionOptions = checkNotNull(connectionOptions);
         this.writeOptions = checkNotNull(writeOptions);

flink-connector-mongodb/src/main/java/org/apache/flink/connector/mongodb/sink/writer/context/DefaultMongoSinkContext.java

@@ -18,23 +18,23 @@
 package org.apache.flink.connector.mongodb.sink.writer.context;
 
 import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.connector.sink2.Sink;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
 import org.apache.flink.connector.mongodb.sink.config.MongoWriteOptions;
 
 /** Default {@link MongoSinkContext} implementation. */
 @Internal
 public class DefaultMongoSinkContext implements MongoSinkContext {
 
-    private final Sink.InitContext initContext;
+    private final WriterInitContext initContext;
     private final MongoWriteOptions writeOptions;
 
-    public DefaultMongoSinkContext(Sink.InitContext initContext, MongoWriteOptions writeOptions) {
+    public DefaultMongoSinkContext(WriterInitContext initContext, MongoWriteOptions writeOptions) {
         this.initContext = initContext;
         this.writeOptions = writeOptions;
     }
 
     @Override
-    public Sink.InitContext getInitContext() {
+    public WriterInitContext getInitContext() {
         return initContext;
     }
 

flink-connector-mongodb/src/main/java/org/apache/flink/connector/mongodb/sink/writer/context/MongoSinkContext.java

@@ -18,7 +18,7 @@
 package org.apache.flink.connector.mongodb.sink.writer.context;
 
 import org.apache.flink.annotation.PublicEvolving;
-import org.apache.flink.api.connector.sink2.Sink;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
 import org.apache.flink.connector.mongodb.sink.config.MongoWriteOptions;
 import org.apache.flink.connector.mongodb.sink.writer.serializer.MongoSerializationSchema;
 
@@ -27,7 +27,7 @@
 public interface MongoSinkContext {
 
     /** Returns the current sink's init context. */
-    Sink.InitContext getInitContext();
+    WriterInitContext getInitContext();
 
     /** Returns the current process time in flink. */
     long processTime();

flink-connector-mongodb/src/main/java/org/apache/flink/connector/mongodb/source/MongoSource.java

@@ -26,9 +26,7 @@
 import org.apache.flink.api.connector.source.SplitEnumerator;
 import org.apache.flink.api.connector.source.SplitEnumeratorContext;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
-import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
-import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
 import org.apache.flink.connector.mongodb.common.config.MongoConnectionOptions;
 import org.apache.flink.connector.mongodb.source.config.MongoReadOptions;
 import org.apache.flink.connector.mongodb.source.enumerator.MongoSourceEnumState;
@@ -134,8 +132,6 @@ public Boundedness getBoundedness() {
 
     @Override
     public SourceReader<OUT, MongoSourceSplit> createReader(SourceReaderContext readerContext) {
-        FutureCompletingBlockingQueue<RecordsWithSplitIds<BsonDocument>> elementsQueue =
-                new FutureCompletingBlockingQueue<>();
 
         MongoSourceReaderContext mongoReaderContext =
                 new MongoSourceReaderContext(readerContext, limit);
@@ -150,7 +146,6 @@ public SourceReader<OUT, MongoSourceSplit> createReader(SourceReaderContext read
                                 mongoReaderContext);
 
         return new MongoSourceReader<>(
-                elementsQueue,
                 splitReaderSupplier,
                 new MongoRecordEmitter<>(deserializationSchema),
                 mongoReaderContext);

flink-connector-mongodb/src/main/java/org/apache/flink/connector/mongodb/source/reader/MongoSourceReader.java

@@ -19,11 +19,9 @@
 
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.connector.base.source.reader.RecordEmitter;
-import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
 import org.apache.flink.connector.base.source.reader.SingleThreadMultiplexSourceReaderBase;
 import org.apache.flink.connector.base.source.reader.fetcher.SingleThreadFetcherManager;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
-import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
 import org.apache.flink.connector.mongodb.source.split.MongoScanSourceSplit;
 import org.apache.flink.connector.mongodb.source.split.MongoScanSourceSplitState;
 import org.apache.flink.connector.mongodb.source.split.MongoSourceSplit;
@@ -49,13 +47,11 @@ public class MongoSourceReader<OUT>
     private static final Logger LOG = LoggerFactory.getLogger(MongoSourceReader.class);
 
     public MongoSourceReader(
-            FutureCompletingBlockingQueue<RecordsWithSplitIds<BsonDocument>> elementQueue,
             Supplier<SplitReader<BsonDocument, MongoSourceSplit>> splitReaderSupplier,
             RecordEmitter<BsonDocument, OUT, MongoSourceSplitState> recordEmitter,
             MongoSourceReaderContext readerContext) {
         super(
-                elementQueue,
-                new SingleThreadFetcherManager<>(elementQueue, splitReaderSupplier),
+                new SingleThreadFetcherManager<>(splitReaderSupplier, readerContext.getConfiguration()),
                 recordEmitter,
                 readerContext.getConfiguration(),
                 readerContext);

flink-connector-mongodb/src/test/java/org/apache/flink/connector/mongodb/sink/MongoSinkITCase.java

@@ -18,8 +18,9 @@
 package org.apache.flink.connector.mongodb.sink;
 
 import org.apache.flink.api.common.functions.MapFunction;
-import org.apache.flink.api.common.restartstrategy.RestartStrategies;
 import org.apache.flink.api.common.state.CheckpointListener;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.RestartStrategyOptions;
 import org.apache.flink.connector.base.DeliveryGuarantee;
 import org.apache.flink.connector.mongodb.sink.writer.context.MongoSinkContext;
 import org.apache.flink.connector.mongodb.sink.writer.serializer.MongoSerializationSchema;
@@ -101,9 +102,10 @@ void testWriteToMongoWithDeliveryGuarantee(DeliveryGuarantee deliveryGuarantee)
             throws Exception {
         final String collection = "test-sink-with-delivery-" + deliveryGuarantee;
         final MongoSink<Document> sink = createSink(collection, deliveryGuarantee);
-        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+        Configuration config = new Configuration();
+        config.set(RestartStrategyOptions.RESTART_STRATEGY, "disable");
+        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config);
         env.enableCheckpointing(100L);
-        env.setRestartStrategy(RestartStrategies.noRestart());
 
         env.fromSequence(1, 5).map(new TestMapFunction()).sinkTo(sink);
         env.execute();

flink-connector-mongodb/src/test/java/org/apache/flink/connector/mongodb/sink/writer/MongoWriterITCase.java

@@ -243,7 +243,8 @@ void testSinkContext() throws Exception {
 
         MongoSerializationSchema<Document> testSerializationSchema =
                 (element, context) -> {
-                    assertThat(context.getInitContext().getSubtaskId()).isEqualTo(0);
+                    assertThat(context.getInitContext().getTaskInfo().getIndexOfThisSubtask())
+                            .isEqualTo(0);
                     assertThat(context.getWriteOptions()).isEqualTo(expectOptions);
                     assertThat(context.processTime())
                             .isEqualTo(

flink-connector-mongodb/src/test/java/org/apache/flink/connector/mongodb/table/MongoFilterPushDownVisitorTest.java

@@ -286,7 +286,7 @@ private List<ResolvedExpression> resolveSQLFilterToExpression(
         RexNodeExpression rexExp =
                 (RexNodeExpression) tbImpl.getParser().parseSqlExpression(sqlExp, sourceType, null);
 
-        RexNode cnf = FlinkRexUtil.toCnf(rexBuilder, -1, rexExp.getRexNode());
+        RexNode cnf = FlinkRexUtil.toCnf(rexBuilder, rexExp.getRexNode());
         // converts the cnf condition to a list of AND conditions
         List<RexNode> conjunctions = RelOptUtil.conjunctions(cnf);
 

pom.xml

@@ -57,7 +57,7 @@ under the License.
 		<mongodb7.version>7.0.12</mongodb7.version>
 		<mongodb.version>${mongodb4.version}</mongodb.version>
 
-		<flink.version>1.20.0</flink.version>
+		<flink.version>2.1.0</flink.version>
 		<scala.binary.version>2.12</scala.binary.version>
 		<scala-library.version>2.12.7</scala-library.version>
 		<junit5.version>5.8.1</junit5.version>