Support subtable and fieldNames in tdengine source. support fieldNames in sink.

Author: chncaesar

docs/en/connector-v2/sink/TDengine.md

@@ -15,14 +15,15 @@ Used to write data to TDengine. You need to create stable before running seatunn
 
 ## Options
 
-|   name   |  type  | required | default value |
-|----------|--------|----------|---------------|
-| url      | string | yes      | -             |
-| username | string | yes      | -             |
-| password | string | yes      | -             |
-| database | string | yes      |               |
-| stable   | string | yes      | -             |
-| timezone | string | no       | UTC           |
+| name        |  type  | required | default value |
+|-------------|--------|----------|---------------|
+| url         | string | yes      | -             |
+| username    | string | yes      | -             |
+| password    | string | yes      | -             |
+| database    | string | yes      |               |
+| stable      | string | yes      | -             |
+| timezone    | string | no       | UTC           |
+| field_names | string | no       | -             |
 
 ### url [string]
 
@@ -54,6 +55,9 @@ the stable of the TDengine when you select
 
 the timeznoe of the TDengine sever, it's important to the ts field
 
+### field_names [string]
+the field names of the TDengine when you insert, if not set, all fields will be written
+
 ## Example
 
 ### sink
@@ -67,6 +71,7 @@ sink {
           database : "power2"
           stable : "meters2"
           timezone: UTC
+          field_names: "ts, voltage, current, power"
         }
 }
 ```

docs/en/connector-v2/source/TDengine.md

@@ -22,15 +22,17 @@ supports query SQL and can achieve projection effect.
 
 ## Options
 
-|    name     |  type  | required | default value |
-|-------------|--------|----------|---------------|
-| url         | string | yes      | -             |
-| username    | string | yes      | -             |
-| password    | string | yes      | -             |
-| database    | string | yes      |               |
-| stable      | string | yes      | -             |
-| lower_bound | long   | yes      | -             |
-| upper_bound | long   | yes      | -             |
+| name          |  type  | required | default value |
+|---------------|--------|----------|---------------|
+| url           | string | yes      | -             |
+| username      | string | yes      | -             |
+| password      | string | yes      | -             |
+| database      | string | yes      |               |
+| stable        | string | yes      | -             |
+| sub_table     | string | no       | -             |
+| lower_bound   | long   | yes      | -             |
+| upper_bound   | long   | yes      | -             |
+| field_names   | string | no       | -             |
 
 ### url [string]
 
@@ -58,6 +60,10 @@ the database of the TDengine when you select
 
 the stable of the TDengine when you select
 
+### sub_table [string]
+the sub_table of the TDengine, separated by commas. If not specified, all sub-tables will be selected. 
+If specified, only the specified sub-tables will be selected.
+
 ### lower_bound [long]
 
 the lower_bound of the migration period
@@ -66,6 +72,10 @@ the lower_bound of the migration period
 
 the upper_bound of the migration period
 
+### field_names [string]
+The field_names of the TDengine when you select, separated by commas. If not specified, all fields will be selected.
+If specified, only the specified fields will be selected.
+
 ## Example
 
 ### source
@@ -78,9 +88,11 @@ source {
           password : "taosdata"
           database : "power"
           stable : "meters"
+          sub_table : "meter_1,meter_2"
           lower_bound : "2018-10-03 14:38:05.000"
           upper_bound : "2018-10-03 14:38:16.800"
           plugin_output = "tdengine_result"
+          field_names : "ts,voltage,current,power"
         }
 }
 ```

seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/config/TDengineCommonOptions.java

@@ -54,4 +54,13 @@ public abstract class TDengineCommonOptions {
                     .stringType()
                     .noDefaultValue()
                     .withDescription("The TDengine super table name");
+
+    public static final Option<String> FIELD_NAMES =
+            Options.key("field_names")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription(
+                            "The field names to be written to TDengine, separated by commas. "
+                                    + "If not specified, all fields will be written. "
+                                    + "This option is useful when the source schema does not match the TDengine table schema.");
 }

seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/config/TDengineSinkConfig.java

@@ -37,6 +37,7 @@ public class TDengineSinkConfig implements Serializable {
     private String database;
     private String stable;
     private String timezone;
+    private String fieldNames;
 
     public static TDengineSinkConfig of(ReadonlyConfig config) {
         Builder builder = TDengineSinkConfig.builder();
@@ -50,6 +51,8 @@ public static TDengineSinkConfig of(ReadonlyConfig config) {
         Optional<String> optionalTimezone = config.getOptional(TDengineSinkOptions.TIMEZONE);
 
         builder.timezone(optionalTimezone.orElseGet(TDengineSinkOptions.TIMEZONE::defaultValue));
+        Optional<String> optionalFieldNames = config.getOptional(TDengineCommonOptions.FIELD_NAMES);
+        builder.fieldNames(optionalFieldNames.orElse(""));
 
         return builder.build();
     }

seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/config/TDengineSourceConfig.java

@@ -22,7 +22,10 @@
 import lombok.Data;
 
 import java.io.Serializable;
+import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 @Data
 public class TDengineSourceConfig implements Serializable {
@@ -36,8 +39,9 @@ public class TDengineSourceConfig implements Serializable {
     private String stable;
     private String lowerBound;
     private String upperBound;
-    private List<String> fields;
     private List<String> tags;
+    private Set<String> subTables;
+    private Set<String> fieldNames;
 
     public static TDengineSourceConfig buildSourceConfig(ReadonlyConfig pluginConfig) {
         TDengineSourceConfig tdengineSourceConfig = new TDengineSourceConfig();
@@ -48,6 +52,19 @@ public static TDengineSourceConfig buildSourceConfig(ReadonlyConfig pluginConfig
         tdengineSourceConfig.setPassword(pluginConfig.get(TDengineSourceOptions.PASSWORD));
         tdengineSourceConfig.setUpperBound(pluginConfig.get(TDengineSourceOptions.UPPER_BOUND));
         tdengineSourceConfig.setLowerBound(pluginConfig.get(TDengineSourceOptions.LOWER_BOUND));
+        if (pluginConfig.getOptional(TDengineSourceOptions.SUB_TABLE).isPresent()) {
+            String plugin = pluginConfig.get(TDengineSourceOptions.SUB_TABLE);
+            String[] subTableArray = plugin.split(",");
+            tdengineSourceConfig.setSubTables(
+                    Arrays.stream(subTableArray).collect(Collectors.toSet()));
+        }
+        if (pluginConfig.getOptional(TDengineCommonOptions.FIELD_NAMES).isPresent()) {
+            String fieldNames = pluginConfig.get(TDengineCommonOptions.FIELD_NAMES);
+            tdengineSourceConfig.setFieldNames(
+                    Arrays.stream(fieldNames.split(",")).collect(Collectors.toSet()));
+        } else {
+            tdengineSourceConfig.setFieldNames(null);
+        }
         return tdengineSourceConfig;
     }
 }

seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/config/TDengineSourceOptions.java

@@ -33,4 +33,12 @@ public class TDengineSourceOptions extends TDengineCommonOptions {
                     .stringType()
                     .noDefaultValue()
                     .withDescription("The upper bound for data query range");
+
+    public static final Option<String> SUB_TABLE =
+            Options.key("sub_table")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription(
+                            "The sub table names to query data from, separated by comma , "
+                                    + "if not specified, all sub tables will be queried");
 }

seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/sink/TDengineSinkWriter.java

@@ -90,6 +90,7 @@ public TDengineSinkWriter(TDengineSinkConfig config, SeaTunnelRowType seaTunnelR
     @SneakyThrows
     @Override
     public void write(SeaTunnelRow element) {
+
         final ArrayList<Object> tags = Lists.newArrayList();
         for (int i = element.getArity() - tagsNum; i < element.getArity(); i++) {
             tags.add(element.getField(i));
@@ -103,11 +104,16 @@ public void write(SeaTunnelRow element) {
                 conn.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
             String sql =
                     String.format(
-                            "INSERT INTO %s using %s tags ( %s ) VALUES ( %s );",
+                            "INSERT INTO %s using %s tags ( %s ) %s VALUES ( %s );",
                             element.getField(0),
                             config.getStable(),
                             tagValues,
+                            config.getFieldNames().isEmpty()
+                                    ? ""
+                                    : "( " + config.getFieldNames() + " )",
                             StringUtils.join(convertDataType(metrics), ","));
+            statement.addBatch(sql);
+            statement.executeBatch();
             final int rowCount = statement.executeUpdate(sql);
             if (rowCount == 0) {
                 Throwables.propagateIfPossible(
@@ -140,6 +146,7 @@ Object[] convertDataType(Object[] objects) {
                             if (object == null) {
                                 return null;
                             }
+
                             if (LocalDateTime.class.equals(object.getClass())) {
                                 // transform timezone according to the config
                                 return "'"

seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSource.java

@@ -36,6 +36,7 @@
 import org.apache.commons.lang3.ArrayUtils;
 
 import com.taosdata.jdbc.TSDBDriver;
+import lombok.Getter;
 import lombok.SneakyThrows;
 
 import java.sql.Connection;
@@ -59,8 +60,7 @@
  */
 public class TDengineSource
         implements SeaTunnelSource<SeaTunnelRow, TDengineSourceSplit, TDengineSourceState> {
-
-    private final StableMetadata stableMetadata;
+    @Getter private final StableMetadata stableMetadata;
     private final TDengineSourceConfig tdengineSourceConfig;
     private final CatalogTable catalogTable;
 
@@ -137,12 +137,22 @@ private StableMetadata getStableMetadata(TDengineSourceConfig config) throws SQL
                 if (timestampFieldName == null) {
                     timestampFieldName = metaResultSet.getString(1);
                 }
+                if (config.getFieldNames() != null
+                        && !config.getFieldNames().isEmpty()
+                        && !config.getFieldNames().contains(metaResultSet.getString(1))) {
+                    continue;
+                }
                 fieldNames.add(metaResultSet.getString(1));
                 fieldTypes.add(TDengineTypeMapper.mapping(metaResultSet.getString(2)));
             }
 
             while (subTableNameResultSet.next()) {
                 String subTableName = subTableNameResultSet.getString(1);
+                if (config.getSubTables() != null
+                        && !config.getSubTables().isEmpty()
+                        && !config.getSubTables().contains(subTableName)) {
+                    continue;
+                }
                 subTableNames.add(subTableName);
             }
         }

seatunnel-connectors-v2/connector-tdengine/src/test/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceReaderTest.java

@@ -18,22 +18,42 @@
 
 package org.apache.seatunnel.connectors.seatunnel.tdengine.source;
 
+import org.apache.seatunnel.api.configuration.ReadonlyConfig;
 import org.apache.seatunnel.api.source.Collector;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.connectors.seatunnel.tdengine.exception.TDengineConnectorException;
 
+import org.apache.commons.lang3.StringUtils;
+
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
 
+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.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Properties;
 import java.util.Random;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.when;
+
 class TDengineSourceReaderTest {
     Logger logger;
     TDengineSourceReader tDengineSourceReader;
@@ -103,6 +123,62 @@ void testPoll() throws InterruptedException {
         pool.awaitTermination(3, TimeUnit.SECONDS);
     }
 
+    @Test
+    public void testGetStableMetadata() throws SQLException {
+
+        try (MockedStatic<DriverManager> dm = mockStatic(DriverManager.class)) {
+
+            Connection mockConn = mock(Connection.class);
+            Statement mockStatement = mock(Statement.class);
+            ResultSet metadataResultSet = mock(ResultSet.class);
+            ResultSet tableResultSet = mock(ResultSet.class);
+
+            dm.when(() -> DriverManager.getConnection(anyString(), any(Properties.class)))
+                    .thenReturn(mockConn);
+
+            when(mockConn.createStatement()).thenReturn(mockStatement);
+
+            when(mockStatement.executeQuery(
+                            argThat(
+                                    sql ->
+                                            StringUtils.isNotEmpty(sql)
+                                                    && sql.trim()
+                                                            .toLowerCase()
+                                                            .startsWith("desc"))))
+                    .thenReturn(metadataResultSet);
+            when(metadataResultSet.next()).thenReturn(true, true, false);
+            when(metadataResultSet.getString(1)).thenReturn("ts", "col1", "col1", "col2");
+            when(metadataResultSet.getString(2)).thenReturn("INT", "VARCHAR(20)");
+
+            when(mockStatement.executeQuery(
+                            argThat(
+                                    sql ->
+                                            sql.trim()
+                                                    .toLowerCase()
+                                                    .startsWith(
+                                                            "select table_name from information_schema.ins_tables"))))
+                    .thenReturn(tableResultSet);
+            when(tableResultSet.next()).thenReturn(true, true, false);
+            when(tableResultSet.getString(1)).thenReturn("sub_table_1", "sub_table_2");
+            Map<String, Object> map = new HashMap<>();
+            map.put("url", "jdbc:TAOS-RS://localhost:6041/");
+            map.put("database", "test_db");
+            map.put("username", "root");
+            map.put("password", "taosdata");
+            map.put("stable", "stable");
+            map.put("sub_table", "sub_table_1");
+            map.put("field_names", "col1");
+
+            ReadonlyConfig config = ReadonlyConfig.fromMap(map);
+            TDengineSource source = new TDengineSource(config);
+            StableMetadata stableMetadata = source.getStableMetadata();
+            Assertions.assertEquals(1, stableMetadata.getSubTableNames().size());
+            Assertions.assertEquals("sub_table_1", stableMetadata.getSubTableNames().get(0));
+            Assertions.assertEquals(2, stableMetadata.getRowType().getFieldNames().length);
+            Assertions.assertEquals("col1", stableMetadata.getRowType().getFieldNames()[1]);
+        }
+    }
+
     private static class TestCollector implements Collector<SeaTunnelRow> {
 
         private final List<SeaTunnelRow> rows = new ArrayList<>();