ck-8803

Author: FrommyMind

seatunnel-connectors-v2/connector-clickhouse/src/main/java/org/apache/seatunnel/connectors/seatunnel/clickhouse/source/ClickhouseSource.java

@@ -18,33 +18,30 @@
 package org.apache.seatunnel.connectors.seatunnel.clickhouse.source;
 
 import org.apache.seatunnel.api.source.Boundedness;
-import org.apache.seatunnel.api.source.SeaTunnelSource;
-import org.apache.seatunnel.api.source.SourceReader;
-import org.apache.seatunnel.api.source.SourceSplitEnumerator;
-import org.apache.seatunnel.api.source.SupportColumnProjection;
-import org.apache.seatunnel.api.source.SupportParallelism;
 import org.apache.seatunnel.api.table.catalog.CatalogTable;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
-import org.apache.seatunnel.connectors.seatunnel.clickhouse.state.ClickhouseSourceState;
+import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 
 import com.clickhouse.client.ClickHouseNode;
+import org.apache.seatunnel.connectors.seatunnel.common.source.AbstractSingleSplitReader;
+import org.apache.seatunnel.connectors.seatunnel.common.source.AbstractSingleSplitSource;
+import org.apache.seatunnel.connectors.seatunnel.common.source.SingleSplitReaderContext;
 
 import java.util.Collections;
 import java.util.List;
 
-public class ClickhouseSource
-        implements SeaTunnelSource<SeaTunnelRow, ClickhouseSourceSplit, ClickhouseSourceState>,
-                SupportParallelism,
-                SupportColumnProjection {
+public class ClickhouseSource extends AbstractSingleSplitSource<SeaTunnelRow> {
 
-    private List<ClickHouseNode> servers;
-    private CatalogTable catalogTable;
-    private String sql;
+    private final List<ClickHouseNode> servers;
+    private final CatalogTable catalogTable;
+    private final String sql;
+    private final SeaTunnelRowType rowTypeInfo;
 
-    public ClickhouseSource(List<ClickHouseNode> servers, CatalogTable catalogTable, String sql) {
+    public ClickhouseSource(List<ClickHouseNode> servers, CatalogTable catalogTable, String sql, SeaTunnelRowType rowTypeInfo) {
         this.servers = servers;
         this.catalogTable = catalogTable;
         this.sql = sql;
+        this.rowTypeInfo = rowTypeInfo;
     }
 
     @Override
@@ -62,25 +59,10 @@ public List<CatalogTable> getProducedCatalogTables() {
         return Collections.singletonList(catalogTable);
     }
 
-    @Override
-    public SourceReader<SeaTunnelRow, ClickhouseSourceSplit> createReader(
-            SourceReader.Context readerContext) throws Exception {
-        return new ClickhouseSourceReader(
-                servers, readerContext, this.catalogTable.getSeaTunnelRowType(), sql);
-    }
-
-    @Override
-    public SourceSplitEnumerator<ClickhouseSourceSplit, ClickhouseSourceState> createEnumerator(
-            SourceSplitEnumerator.Context<ClickhouseSourceSplit> enumeratorContext)
-            throws Exception {
-        return new ClickhouseSourceSplitEnumerator(enumeratorContext);
-    }
 
     @Override
-    public SourceSplitEnumerator<ClickhouseSourceSplit, ClickhouseSourceState> restoreEnumerator(
-            SourceSplitEnumerator.Context<ClickhouseSourceSplit> enumeratorContext,
-            ClickhouseSourceState checkpointState)
-            throws Exception {
-        return new ClickhouseSourceSplitEnumerator(enumeratorContext);
+    public AbstractSingleSplitReader<SeaTunnelRow> createReader(
+            SingleSplitReaderContext readerContext) {
+        return new ClickhouseSourceReader(servers, readerContext, sql, rowTypeInfo);
     }
 }

seatunnel-connectors-v2/connector-clickhouse/src/main/java/org/apache/seatunnel/connectors/seatunnel/clickhouse/source/ClickhouseSourceFactory.java

@@ -22,14 +22,12 @@
 import org.apache.seatunnel.api.configuration.util.OptionRule;
 import org.apache.seatunnel.api.source.SeaTunnelSource;
 import org.apache.seatunnel.api.source.SourceSplit;
-import org.apache.seatunnel.api.table.catalog.CatalogTable;
-import org.apache.seatunnel.api.table.catalog.PhysicalColumn;
-import org.apache.seatunnel.api.table.catalog.TableIdentifier;
-import org.apache.seatunnel.api.table.catalog.TableSchema;
+import org.apache.seatunnel.api.table.catalog.*;
 import org.apache.seatunnel.api.table.connector.TableSource;
 import org.apache.seatunnel.api.table.factory.Factory;
 import org.apache.seatunnel.api.table.factory.TableSourceFactory;
 import org.apache.seatunnel.api.table.factory.TableSourceFactoryContext;
+import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.common.constants.PluginType;
 import org.apache.seatunnel.connectors.seatunnel.clickhouse.exception.ClickhouseConnectorException;
 import org.apache.seatunnel.connectors.seatunnel.clickhouse.util.ClickhouseUtil;
@@ -67,7 +65,7 @@ public String factoryIdentifier() {
             TableSource<T, SplitT, StateT> createSource(TableSourceFactoryContext context) {
         ReadonlyConfig readonlyConfig = context.getOptions();
         List<ClickHouseNode> nodes = ClickhouseUtil.createNodes(readonlyConfig);
-
+        SeaTunnelRowType rowTypeInfo = CatalogTableUtil.buildWithConfig(readonlyConfig).getSeaTunnelRowType();
         String sql = readonlyConfig.get(SQL);
         ClickHouseNode currentServer = nodes.get(ThreadLocalRandom.current().nextInt(nodes.size()));
         try (ClickHouseClient client = ClickHouseClient.newInstance(currentServer.getProtocol());
@@ -103,7 +101,7 @@ TableSource<T, SplitT, StateT> createSource(TableSourceFactoryContext context) {
                             catalogName);
             return () ->
                     (SeaTunnelSource<T, SplitT, StateT>)
-                            new ClickhouseSource(nodes, catalogTable, sql);
+                            new ClickhouseSource(nodes, catalogTable, sql, rowTypeInfo);
         } catch (ClickHouseException e) {
             throw new ClickhouseConnectorException(
                     SeaTunnelAPIErrorCode.CONFIG_VALIDATION_FAILED,

seatunnel-connectors-v2/connector-clickhouse/src/main/java/org/apache/seatunnel/connectors/seatunnel/clickhouse/source/ClickhouseSourceReader.java

@@ -19,7 +19,6 @@
 
 import org.apache.seatunnel.api.source.Boundedness;
 import org.apache.seatunnel.api.source.Collector;
-import org.apache.seatunnel.api.source.SourceReader;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.connectors.seatunnel.clickhouse.util.TypeConvertUtil;
@@ -30,36 +29,33 @@
 import com.clickhouse.client.ClickHouseRequest;
 import com.clickhouse.client.ClickHouseResponse;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.seatunnel.connectors.seatunnel.common.source.AbstractSingleSplitReader;
+import org.apache.seatunnel.connectors.seatunnel.common.source.SingleSplitReaderContext;
 
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Random;
 
 @Slf4j
-public class ClickhouseSourceReader implements SourceReader<SeaTunnelRow, ClickhouseSourceSplit> {
+public class ClickhouseSourceReader extends AbstractSingleSplitReader<SeaTunnelRow> {
 
     private final List<ClickHouseNode> servers;
     private ClickHouseClient client;
     private final SeaTunnelRowType rowTypeInfo;
-    private final SourceReader.Context readerContext;
+    private final SingleSplitReaderContext readerContext;
     private ClickHouseRequest<?> request;
     private final String sql;
-    private volatile boolean noMoreSplit;
 
-    private final List<ClickhouseSourceSplit> splits;
 
     ClickhouseSourceReader(
             List<ClickHouseNode> servers,
-            SourceReader.Context readerContext,
-            SeaTunnelRowType rowTypeInfo,
-            String sql) {
+            SingleSplitReaderContext readerContext,
+            String sql,
+            SeaTunnelRowType rowTypeInfo) {
         this.servers = servers;
         this.readerContext = readerContext;
-        this.rowTypeInfo = rowTypeInfo;
         this.sql = sql;
-        this.splits = new ArrayList<>();
+        this.rowTypeInfo = rowTypeInfo;
     }
 
     @Override
@@ -80,31 +76,28 @@ public void close() throws IOException {
     @Override
     public void pollNext(Collector<SeaTunnelRow> output) throws Exception {
         synchronized (output.getCheckpointLock()) {
-            if (!splits.isEmpty()) {
-                try (ClickHouseResponse response = this.request.query(sql).executeAndWait()) {
-                    response.stream()
-                            .forEach(
-                                    record -> {
-                                        Object[] values =
-                                                new Object[this.rowTypeInfo.getFieldNames().length];
-                                        for (int i = 0; i < record.size(); i++) {
-                                            if (record.getValue(i).isNullOrEmpty()) {
-                                                values[i] = null;
-                                            } else {
-                                                values[i] =
-                                                        TypeConvertUtil.valueUnwrap(
-                                                                this.rowTypeInfo.getFieldType(i),
-                                                                record.getValue(i));
-                                            }
+
+            try (ClickHouseResponse response = this.request.query(sql).executeAndWait()) {
+                response.stream()
+                        .forEach(
+                                record -> {
+                                    Object[] values =
+                                            new Object[this.rowTypeInfo.getFieldNames().length];
+                                    for (int i = 0; i < record.size(); i++) {
+                                        if (record.getValue(i).isNullOrEmpty()) {
+                                            values[i] = null;
+                                        } else {
+                                            values[i] =
+                                                    TypeConvertUtil.valueUnwrap(
+                                                            this.rowTypeInfo.getFieldType(i),
+                                                            record.getValue(i));
                                         }
-                                        output.collect(new SeaTunnelRow(values));
-                                    });
-                }
-                signalNoMoreElement();
+                                    }
+                                    output.collect(new SeaTunnelRow(values));
+                                });
             }
-            if (noMoreSplit
-                    && splits.isEmpty()
-                    && Boundedness.BOUNDED.equals(readerContext.getBoundedness())) {
+            signalNoMoreElement();
+            if (Boundedness.BOUNDED.equals(readerContext.getBoundedness())) {
                 signalNoMoreElement();
             }
         }
@@ -113,24 +106,5 @@ record -> {
     private void signalNoMoreElement() {
         log.info("Closed the bounded ClickHouse source");
         this.readerContext.signalNoMoreElement();
-        this.splits.clear();
     }
-
-    @Override
-    public List<ClickhouseSourceSplit> snapshotState(long checkpointId) throws Exception {
-        return Collections.emptyList();
-    }
-
-    @Override
-    public void addSplits(List<ClickhouseSourceSplit> splits) {
-        this.splits.addAll(splits);
-    }
-
-    @Override
-    public void handleNoMoreSplits() {
-        noMoreSplit = true;
-    }
-
-    @Override
-    public void notifyCheckpointComplete(long checkpointId) throws Exception {}
 }