[FLINK-36236] Fix overriding the connection provider in JdbcSourceBuilder

Author: rajucomp

flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSource.java

@@ -162,20 +162,25 @@ public JdbcSqlSplitEnumeratorBase.Provider<JdbcSourceSplit> getSqlSplitEnumerato
     }
 
     @VisibleForTesting
-    public TypeInformation<OUT> getTypeInformation() {
+    TypeInformation<OUT> getTypeInformation() {
         return typeInformation;
     }
 
     @VisibleForTesting
-    public Configuration getConfiguration() {
+    Configuration getConfiguration() {
         return configuration;
     }
 
     @VisibleForTesting
-    public ResultExtractor<OUT> getResultExtractor() {
+    ResultExtractor<OUT> getResultExtractor() {
         return resultExtractor;
     }
 
+    @VisibleForTesting
+    JdbcConnectionProvider getConnectionProvider() {
+        return connectionProvider;
+    }
+
     @VisibleForTesting
     @Override
     public boolean equals(Object o) {

flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSourceBuilder.java

@@ -267,7 +267,10 @@ public JdbcSourceBuilder<OUT> setOptionalSqlSplitEnumeratorState(
     }
 
     public JdbcSource<OUT> build() {
-        this.connectionProvider = new SimpleJdbcConnectionProvider(connOptionsBuilder.build());
+        if (this.connectionProvider == null) {
+            this.connectionProvider = new SimpleJdbcConnectionProvider(connOptionsBuilder.build());
+        }
+
         if (resultSetFetchSize > 0) {
             this.configuration.set(JdbcSourceOptions.RESULTSET_FETCH_SIZE, resultSetFetchSize);
         }

flink-connector-jdbc-core/src/test/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSourceBuilderTest.java

@@ -20,9 +20,12 @@
 
 import org.apache.flink.api.common.typeinfo.TypeHint;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.connector.jdbc.JdbcConnectionOptions;
 import org.apache.flink.connector.jdbc.core.datastream.source.config.ContinuousUnBoundingSettings;
 import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.SqlTemplateSplitEnumerator;
 import org.apache.flink.connector.jdbc.core.datastream.source.reader.extractor.ResultExtractor;
+import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
+import org.apache.flink.connector.jdbc.datasource.connections.SimpleJdbcConnectionProvider;
 import org.apache.flink.connector.jdbc.split.JdbcGenericParameterValuesProvider;
 import org.apache.flink.connector.jdbc.split.JdbcNumericBetweenParametersProvider;
 import org.apache.flink.connector.jdbc.split.JdbcParameterValuesProvider;
@@ -181,4 +184,23 @@ void testSetStreamingSemantic() {
                 .isInstanceOf(IllegalArgumentException.class)
                 .hasMessage(JdbcSourceBuilder.INVALID_CONTINUOUS_SLIDE_TIMING_HINT);
     }
+
+    @Test
+    void testSetConnectionProvider() {
+        assertThatThrownBy(() -> JdbcSource.builder().setConnectionProvider(null))
+                .isInstanceOf(NullPointerException.class);
+
+        final JdbcConnectionOptions options =
+                new JdbcConnectionOptions.JdbcConnectionOptionsBuilder()
+                        .withUrl(dbUrl)
+                        .withDriverName(driverName)
+                        .build();
+
+        final JdbcConnectionProvider connectionProvider = new SimpleJdbcConnectionProvider(options);
+
+        final JdbcSource<Row> jdbcSource =
+                sourceBuilder.setConnectionProvider(connectionProvider).build();
+
+        assertThat(jdbcSource.getConnectionProvider()).isSameAs(connectionProvider);
+    }
 }