[Hotfix][Connector] Fixed TDengine connector using jdbc driver to cause loading error (#4598)

hailin0 · web-flow · commit 78f7989b8140 · 2023-04-18T10:40:12.000+08:00
diff --git a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/StableMetadata.java b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/StableMetadata.java
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.tdengine.source;
+
+import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+
+import java.io.Serializable;
+import java.util.List;
+
+@Data
+@AllArgsConstructor
+public class StableMetadata implements Serializable {
+    private final SeaTunnelRowType rowType;
+    private final String timestampFieldName;
+    private final List<String> subTableNames;
+}
diff --git a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSource.java b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSource.java
@@ -38,16 +38,16 @@
 import org.apache.seatunnel.connectors.seatunnel.tdengine.typemapper.TDengineTypeMapper;
 
 import org.apache.commons.lang3.ArrayUtils;
-import org.apache.commons.lang3.StringUtils;
 
 import com.google.auto.service.AutoService;
-import com.google.common.collect.Lists;
 import lombok.SneakyThrows;
 
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.ArrayList;
 import java.util.List;
 
 import static org.apache.seatunnel.connectors.seatunnel.tdengine.config.TDengineSourceConfig.ConfigNames.DATABASE;
@@ -67,7 +67,7 @@
 public class TDengineSource
         implements SeaTunnelSource<SeaTunnelRow, TDengineSourceSplit, TDengineSourceState> {
 
-    private SeaTunnelRowType seaTunnelRowType;
+    private StableMetadata stableMetadata;
     private TDengineSourceConfig tdengineSourceConfig;
 
     @Override
@@ -87,45 +87,7 @@ public void prepare(Config pluginConfig) throws PrepareFailException {
                     "TDengine connection require url/database/stable/username/password. All of these must not be empty.");
         }
         tdengineSourceConfig = buildSourceConfig(pluginConfig);
-
-        // add subtable_name and tags to `seaTunnelRowType`
-        SeaTunnelRowType originRowType = getSTableMetaInfo(tdengineSourceConfig);
-        seaTunnelRowType = addHiddenAttribute(originRowType);
-    }
-
-    @SneakyThrows
-    private SeaTunnelRowType getSTableMetaInfo(TDengineSourceConfig config) {
-        String jdbcUrl =
-                StringUtils.join(
-                        config.getUrl(),
-                        config.getDatabase(),
-                        "?user=",
-                        config.getUsername(),
-                        "&password=",
-                        config.getPassword());
-        Connection conn = DriverManager.getConnection(jdbcUrl);
-        List<String> fieldNames = Lists.newArrayList();
-        List<SeaTunnelDataType<?>> fieldTypes = Lists.newArrayList();
-        try (Statement statement = conn.createStatement()) {
-            final ResultSet metaResultSet =
-                    statement.executeQuery(
-                            "desc " + config.getDatabase() + "." + config.getStable());
-            while (metaResultSet.next()) {
-                fieldNames.add(metaResultSet.getString(1));
-                fieldTypes.add(TDengineTypeMapper.mapping(metaResultSet.getString(2)));
-            }
-        }
-        return new SeaTunnelRowType(
-                fieldNames.toArray(new String[0]), fieldTypes.toArray(new SeaTunnelDataType<?>[0]));
-    }
-
-    private SeaTunnelRowType addHiddenAttribute(SeaTunnelRowType originRowType) {
-        // 0-subtable_name / 1-n field_names /
-        String[] fieldNames = ArrayUtils.add(originRowType.getFieldNames(), 0, "subtable_name");
-        // n+1-> tags
-        SeaTunnelDataType<?>[] fieldTypes =
-                ArrayUtils.add(originRowType.getFieldTypes(), 0, BasicType.STRING_TYPE);
-        return new SeaTunnelRowType(fieldNames, fieldTypes);
+        stableMetadata = getStableMetadata(tdengineSourceConfig);
     }
 
     @Override
@@ -135,7 +97,7 @@ public Boundedness getBoundedness() {
 
     @Override
     public SeaTunnelDataType<SeaTunnelRow> getProducedType() {
-        return seaTunnelRowType;
+        return stableMetadata.getRowType();
     }
 
     @Override
@@ -147,14 +109,74 @@ public SourceReader<SeaTunnelRow, TDengineSourceSplit> createReader(Context read
     public SourceSplitEnumerator<TDengineSourceSplit, TDengineSourceState> createEnumerator(
             SourceSplitEnumerator.Context<TDengineSourceSplit> enumeratorContext) {
         return new TDengineSourceSplitEnumerator(
-                seaTunnelRowType, tdengineSourceConfig, enumeratorContext);
+                stableMetadata, tdengineSourceConfig, enumeratorContext);
     }
 
     @Override
     public SourceSplitEnumerator<TDengineSourceSplit, TDengineSourceState> restoreEnumerator(
             SourceSplitEnumerator.Context<TDengineSourceSplit> enumeratorContext,
             TDengineSourceState checkpointState) {
         return new TDengineSourceSplitEnumerator(
-                seaTunnelRowType, tdengineSourceConfig, checkpointState, enumeratorContext);
+                stableMetadata, tdengineSourceConfig, checkpointState, enumeratorContext);
+    }
+
+    private StableMetadata getStableMetadata(TDengineSourceConfig config) throws SQLException {
+        String timestampFieldName = null;
+        List<String> subTableNames = new ArrayList<>();
+        List<String> fieldNames = new ArrayList<>();
+        List<SeaTunnelDataType<?>> fieldTypes = new ArrayList<>();
+
+        String jdbcUrl =
+                String.join(
+                        "",
+                        config.getUrl(),
+                        config.getDatabase(),
+                        "?user=",
+                        config.getUsername(),
+                        "&password=",
+                        config.getPassword());
+        try (Connection conn = DriverManager.getConnection(jdbcUrl)) {
+            try (Statement statement = conn.createStatement()) {
+                ResultSet metaResultSet =
+                        statement.executeQuery(
+                                "desc " + config.getDatabase() + "." + config.getStable());
+                while (metaResultSet.next()) {
+                    if (timestampFieldName == null) {
+                        timestampFieldName = metaResultSet.getString(1);
+                    }
+                    fieldNames.add(metaResultSet.getString(1));
+                    fieldTypes.add(TDengineTypeMapper.mapping(metaResultSet.getString(2)));
+                }
+            }
+            try (Statement statement = conn.createStatement()) {
+                String metaSQL =
+                        "select table_name from information_schema.ins_tables where db_name = '"
+                                + config.getDatabase()
+                                + "' and stable_name='"
+                                + config.getStable()
+                                + "';";
+                ResultSet subTableNameResultSet = statement.executeQuery(metaSQL);
+                while (subTableNameResultSet.next()) {
+                    String subTableName = subTableNameResultSet.getString(1);
+                    subTableNames.add(subTableName);
+                }
+            }
+        }
+
+        SeaTunnelRowType rowType = addHiddenAttribute(fieldNames, fieldTypes);
+        return new StableMetadata(rowType, timestampFieldName, subTableNames);
+    }
+
+    private SeaTunnelRowType addHiddenAttribute(
+            List<String> fieldNames, List<SeaTunnelDataType<?>> fieldTypes) {
+        // add subtable_name and tags to `seaTunnelRowType`
+        // 0-subtable_name / 1-n field_names /
+        String[] newFieldNames =
+                ArrayUtils.add(fieldNames.toArray(new String[0]), 0, "subtable_name");
+        // n+1-> tags
+        SeaTunnelDataType<?>[] newFieldTypes =
+                ArrayUtils.add(
+                        fieldTypes.toArray(new SeaTunnelDataType[0]), 0, BasicType.STRING_TYPE);
+        return new SeaTunnelRowType(newFieldNames, newFieldTypes);
     }
 }
diff --git a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumerator.java b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumerator.java
@@ -19,22 +19,9 @@
 
 import org.apache.seatunnel.api.source.SourceEvent;
 import org.apache.seatunnel.api.source.SourceSplitEnumerator;
-import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
-import org.apache.seatunnel.common.exception.CommonErrorCode;
 import org.apache.seatunnel.connectors.seatunnel.tdengine.config.TDengineSourceConfig;
-import org.apache.seatunnel.connectors.seatunnel.tdengine.exception.TDengineConnectorException;
 import org.apache.seatunnel.connectors.seatunnel.tdengine.state.TDengineSourceState;
 
-import org.apache.commons.lang3.StringUtils;
-
-import com.google.common.collect.Sets;
-import lombok.SneakyThrows;
-
-import java.sql.Connection;
-import java.sql.DriverManager;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.Statement;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -49,26 +36,25 @@ public class TDengineSourceSplitEnumerator
 
     private final SourceSplitEnumerator.Context<TDengineSourceSplit> context;
     private final TDengineSourceConfig config;
+    private final StableMetadata stableMetadata;
     private Set<TDengineSourceSplit> pendingSplit = new HashSet<>();
     private Set<TDengineSourceSplit> assignedSplit = new HashSet<>();
-    private Connection conn;
-    private SeaTunnelRowType seaTunnelRowType;
 
     public TDengineSourceSplitEnumerator(
-            SeaTunnelRowType seaTunnelRowType,
+            StableMetadata stableMetadata,
             TDengineSourceConfig config,
             SourceSplitEnumerator.Context<TDengineSourceSplit> context) {
-        this(seaTunnelRowType, config, null, context);
+        this(stableMetadata, config, null, context);
     }
 
     public TDengineSourceSplitEnumerator(
-            SeaTunnelRowType seaTunnelRowType,
+            StableMetadata stableMetadata,
             TDengineSourceConfig config,
             TDengineSourceState sourceState,
             SourceSplitEnumerator.Context<TDengineSourceSplit> context) {
         this.config = config;
         this.context = context;
-        this.seaTunnelRowType = seaTunnelRowType;
+        this.stableMetadata = stableMetadata;
         if (sourceState != null) {
             this.assignedSplit = sourceState.getAssignedSplit();
         }
@@ -78,64 +64,33 @@ private static int getSplitOwner(String tp, int numReaders) {
         return (tp.hashCode() & Integer.MAX_VALUE) % numReaders;
     }
 
-    @SneakyThrows
     @Override
-    public void open() {
-        String jdbcUrl =
-                StringUtils.join(
-                        config.getUrl(),
-                        config.getDatabase(),
-                        "?user=",
-                        config.getUsername(),
-                        "&password=",
-                        config.getPassword());
-        conn = DriverManager.getConnection(jdbcUrl);
-    }
+    public void open() {}
 
     @Override
-    public void run() throws SQLException {
+    public void run() {
         pendingSplit = getAllSplits();
         assignSplit(context.registeredReaders());
     }
 
     /*
-     * 1. get timestampField
-     * 2. get all sub tables of configured super table
-     * 3. each split has one sub table
+     * each split has one sub table
      */
-    private Set<TDengineSourceSplit> getAllSplits() throws SQLException {
-        final String timestampFieldName;
-        try (Statement statement = conn.createStatement()) {
-            final ResultSet fieldNameResultSet =
-                    statement.executeQuery(
-                            "desc " + config.getDatabase() + "." + config.getStable());
-            fieldNameResultSet.next();
-            timestampFieldName = fieldNameResultSet.getString(1);
-        }
-
-        final Set<TDengineSourceSplit> splits = Sets.newHashSet();
-        try (Statement statement = conn.createStatement()) {
-            String metaSQL =
-                    "select table_name from information_schema.ins_tables where db_name = '"
-                            + config.getDatabase()
-                            + "' and stable_name='"
-                            + config.getStable()
-                            + "';";
-            ResultSet subTableNameResultSet = statement.executeQuery(metaSQL);
-            while (subTableNameResultSet.next()) {
-                final String subTableName = subTableNameResultSet.getString(1);
-                final TDengineSourceSplit splitBySubTable =
-                        createSplitBySubTable(subTableName, timestampFieldName);
-                splits.add(splitBySubTable);
-            }
+    private Set<TDengineSourceSplit> getAllSplits() {
+        final String timestampFieldName = stableMetadata.getTimestampFieldName();
+        final Set<TDengineSourceSplit> splits = new HashSet<>();
+        for (String subTableName : stableMetadata.getSubTableNames()) {
+            TDengineSourceSplit splitBySubTable =
+                    createSplitBySubTable(subTableName, timestampFieldName);
+            splits.add(splitBySubTable);
         }
         return splits;
     }
 
     private TDengineSourceSplit createSplitBySubTable(
             String subTableName, String timestampFieldName) {
         String selectFields =
-                Arrays.stream(seaTunnelRowType.getFieldNames())
+                Arrays.stream(stableMetadata.getRowType().getFieldNames())
                         .skip(1)
                         .collect(Collectors.joining(","));
         String subTableSQL =
@@ -152,7 +107,7 @@ private TDengineSourceSplit createSplitBySubTable(
             if (end != null) {
                 endCondition = timestampFieldName + " < '" + end + "'";
             }
-            String query = StringUtils.join(new String[] {startCondition, endCondition}, " and ");
+            String query = String.join(" and ", startCondition, endCondition);
             subTableSQL = subTableSQL + " where " + query;
         }
 
@@ -220,18 +175,7 @@ public void notifyCheckpointAborted(long checkpointId) throws Exception {
     }
 
     @Override
-    public void close() {
-        try {
-            if (!Objects.isNull(conn)) {
-                conn.close();
-            }
-        } catch (SQLException e) {
-            throw new TDengineConnectorException(
-                    CommonErrorCode.READER_OPERATION_FAILED,
-                    "TDengine split_enumerator connection close failed",
-                    e);
-        }
-    }
+    public void close() {}
 
     @Override
     public void handleSplitRequest(int subtaskId) {
diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-tdengine-e2e/src/test/java/org/apache/seatunnel/e2e/connector/tdengine/TDengineIT.java b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-tdengine-e2e/src/test/java/org/apache/seatunnel/e2e/connector/tdengine/TDengineIT.java
@@ -20,6 +20,7 @@
 import org.apache.seatunnel.e2e.common.TestResource;
 import org.apache.seatunnel.e2e.common.TestSuiteBase;
 import org.apache.seatunnel.e2e.common.container.TestContainer;
+import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
@@ -48,6 +49,10 @@
 import static org.awaitility.Awaitility.given;
 
 @Slf4j
+@DisabledOnContainer(
+        value = {},
+        type = {},
+        disabledReason = "Override TestSuiteBase @DisabledOnContainer")
 public class TDengineIT extends TestSuiteBase implements TestResource {
     private static final String DOCKER_IMAGE = "tdengine/tdengine:3.0.2.1";
     private static final String NETWORK_ALIASES1 = "flink_e2e_tdengine_src";
diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-tdengine-e2e/src/test/resources/tdengine/tdengine_source_to_sink.conf b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-tdengine-e2e/src/test/resources/tdengine/tdengine_source_to_sink.conf
@@ -21,7 +21,7 @@
 env {
   # You can set flink configuration here
   execution.parallelism = 2
-  job.mode = "STREAMING"
+  job.mode = "BATCH"
   #execution.checkpoint.interval = 10000
   #execution.checkpoint.data-uri = "hdfs://localhost:9000/checkpoint"
 }

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@`
`21`	`21`	`env {`
`22`	`22`	`# You can set flink configuration here`
`23`	`23`	`execution.parallelism = 2`
`24`		`- job.mode = "STREAMING"`
	`24`	`+ job.mode = "BATCH"`
`25`	`25`	`#execution.checkpoint.interval = 10000`
`26`	`26`	`#execution.checkpoint.data-uri = "hdfs://localhost:9000/checkpoint"`
`27`	`27`	`}`