MartijnVisser commented on code in PR #234: URL: https://github.com/apache/flink-connector-jdbc/pull/234#discussion_r4081427911
########## flink-connector-jdbc-core/src/test/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/AbstractDynamicSplitterEnumeratorTest.java: ########## @@ -0,0 +1,165 @@ +/* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.connector.jdbc.core.datastream.source.split.CheckpointedOffset; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; Review Comment: Mockito is gone from main since FLINK-40729, so this won't compile after a rebase. A provider like `StaticConnectionProvider` in `BoundaryQuerySplitterEnumeratorTest` would work here. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/AbstractDynamicSplitterEnumerator.java: ########## @@ -0,0 +1,124 @@ + /* Review Comment: `spotless:check` fails on this leading space and on the formatting in `BoundaryQuerySplitterEnumeratorTest`. `mvn spotless:apply` fixes both. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/AbstractDynamicSplitterEnumerator.java: ########## @@ -0,0 +1,124 @@ + /* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import java.io.Serializable; +import java.sql.Connection; +import java.util.Collections; +import java.util.List; + +/** + * Base class for {@link SplitterEnumerator} implementations that must open a real database + * connection at job runtime to discover their splits (e.g. by executing a user-provided query), as + * opposed to computing splits purely from static configuration. + * + * <p>Split discovery is deliberately deferred to the first call of {@link #enumerateSplits()} + * rather than performed in {@link #start(JdbcConnectionProvider)}: the enclosing {@code + * JdbcSourceEnumerator} calls {@code start()} synchronously on the JobManager coordinator thread, + * but invokes {@code enumerateSplits()} through Flink's own {@code + * SplitEnumeratorContext#callAsync}, which already runs off that thread. Deferring the actual I/O + * this way avoids blocking job startup and avoids having to manage any dedicated thread of our own. + * + * <p>Because {@code enumerateSplits()} can be invoked concurrently by the coordinator before the + * first call returns, discovery is guarded so it only ever runs once. + */ +@Internal +public abstract class AbstractDynamicSplitterEnumerator implements SplitterEnumerator { + + private transient JdbcConnectionProvider connectionProvider; + + private List<JdbcSourceSplit> discoveredSplits; + + private boolean finished = false; + + @Override + public Boundedness getBoundedness() { + return Boundedness.BOUNDED; + } + + @Override + public void start(JdbcConnectionProvider connectionProvider) { + // Intentionally no I/O here - see class javadoc. + this.connectionProvider = connectionProvider; + } + + @Override + public synchronized List<JdbcSourceSplit> enumerateSplits() { + if (finished) { + return Collections.emptyList(); + } + if (discoveredSplits == null) { + try { + Connection connection = connectionProvider.getOrEstablishConnection(); + discoveredSplits = discoverSplits(connection); + } catch (Exception e) { + throw new RuntimeException("Failed to discover splits for " + describeSource(), e); + } + } + finished = true; + return discoveredSplits; + } + + /** + * Discover the full set of splits using the given, already-established connection. Called at + * most once per enumerator instance. + */ + protected abstract List<JdbcSourceSplit> discoverSplits(Connection connection) throws Exception; + + /** Short human-readable description of what this enumerator discovers splits for. */ + protected abstract String describeSource(); + + @Override + public synchronized boolean isAllSplitsFinished() { Review Comment: `synchronized` here makes the coordinator thread wait for the whole boundary query. In my test `start()` blocked on it. I think a volatile `finished` is enough. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/AbstractDynamicSplitterEnumerator.java: ########## @@ -0,0 +1,124 @@ + /* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import java.io.Serializable; +import java.sql.Connection; +import java.util.Collections; +import java.util.List; + +/** + * Base class for {@link SplitterEnumerator} implementations that must open a real database + * connection at job runtime to discover their splits (e.g. by executing a user-provided query), as + * opposed to computing splits purely from static configuration. + * + * <p>Split discovery is deliberately deferred to the first call of {@link #enumerateSplits()} + * rather than performed in {@link #start(JdbcConnectionProvider)}: the enclosing {@code + * JdbcSourceEnumerator} calls {@code start()} synchronously on the JobManager coordinator thread, + * but invokes {@code enumerateSplits()} through Flink's own {@code + * SplitEnumeratorContext#callAsync}, which already runs off that thread. Deferring the actual I/O + * this way avoids blocking job startup and avoids having to manage any dedicated thread of our own. + * + * <p>Because {@code enumerateSplits()} can be invoked concurrently by the coordinator before the + * first call returns, discovery is guarded so it only ever runs once. + */ +@Internal +public abstract class AbstractDynamicSplitterEnumerator implements SplitterEnumerator { + + private transient JdbcConnectionProvider connectionProvider; + + private List<JdbcSourceSplit> discoveredSplits; + + private boolean finished = false; + + @Override + public Boundedness getBoundedness() { + return Boundedness.BOUNDED; + } + + @Override + public void start(JdbcConnectionProvider connectionProvider) { + // Intentionally no I/O here - see class javadoc. + this.connectionProvider = connectionProvider; + } + + @Override + public synchronized List<JdbcSourceSplit> enumerateSplits() { + if (finished) { + return Collections.emptyList(); + } + if (discoveredSplits == null) { + try { + Connection connection = connectionProvider.getOrEstablishConnection(); + discoveredSplits = discoverSplits(connection); + } catch (Exception e) { + throw new RuntimeException("Failed to discover splits for " + describeSource(), e); + } + } + finished = true; + return discoveredSplits; + } + + /** + * Discover the full set of splits using the given, already-established connection. Called at + * most once per enumerator instance. + */ + protected abstract List<JdbcSourceSplit> discoverSplits(Connection connection) throws Exception; + + /** Short human-readable description of what this enumerator discovers splits for. */ + protected abstract String describeSource(); + + @Override + public synchronized boolean isAllSplitsFinished() { + return finished; + } + + @Override + public void close() { + if (connectionProvider != null) { + connectionProvider.closeConnection(); + } + } + + @Override + public List<String> lineageQueries() { + return Collections.singletonList(describeSource()); + } + + @Override + public Serializable serializableState() { + // All splits are discovered and handed out in a single enumerateSplits() batch, so by + // the time any checkpoint happens either nothing was discovered yet (safe to redo from + // scratch on restore) or discovery already fully completed (any not-yet-assigned splits + // are tracked separately by JdbcSourceEnumeratorState). The only thing that must survive + // a restore is "don't run discovery again" - a single boolean is sufficient for that. + return finished; Review Comment: A checkpoint taken after `finished` flips on the worker thread, but before the handler adds the splits, stores `true` without them. After restore my test read nothing. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/BoundaryQuerySplitterEnumerator.java: ########## @@ -0,0 +1,259 @@ +/* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.jdbc.core.database.dialect.JdbcDialect; +import org.apache.flink.connector.jdbc.core.datastream.source.split.CheckpointedOffset; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeFamily; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.util.Preconditions; + +import java.io.Serializable; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * A {@link AbstractDynamicSplitterEnumerator} that discovers split boundaries by executing a + * user-provided {@code scan.partition.boundary-query} against the source database, instead of + * requiring a statically-known numeric {@code [lower, upper]} range. + * + * <p>The boundary query is expected to return a single, comparable column; each returned value + * (after being sorted ascending) becomes a split point. {@code N} returned values produce {@code N + * + 1} partitions: + * + * <ul> + * <li>first: {@code col < v(0) OR col IS NULL} + * <li>middle {@code i} (0 < i < N): {@code col >= v(i-1) AND col < v(i)} + * <li>last: {@code col >= v(N-1)} + * </ul> + * + * <p>Unlike every other splitter in this connector, the SQL predicate text itself differs between + * splits rather than only the bound parameter values, since a single {@code BETWEEN ? AND ?} + * template cannot express the "first"/"last" open-ended ranges above. + */ +@Internal +public class BoundaryQuerySplitterEnumerator extends AbstractDynamicSplitterEnumerator { + + private static final Set<Integer> NUMERIC_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.TINYINT, + Types.SMALLINT, + Types.INTEGER, + Types.BIGINT, + Types.DECIMAL, + Types.NUMERIC, + Types.FLOAT, + Types.REAL, + Types.DOUBLE)); + + private static final Set<Integer> STRING_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.CHAR, + Types.VARCHAR, + Types.LONGVARCHAR, + Types.NCHAR, + Types.NVARCHAR, + Types.LONGNVARCHAR)); + + private static final Set<Integer> DATETIME_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.DATE, + Types.TIME, + Types.TIMESTAMP, + Types.TIME_WITH_TIMEZONE, + Types.TIMESTAMP_WITH_TIMEZONE)); + + private final String baseSqlTemplate; + private final String boundaryQuery; + private final String partitionColumn; + private final LogicalType partitionColumnType; + private final JdbcDialect dialect; + private final int numPartitions; + + public BoundaryQuerySplitterEnumerator( + String baseSqlTemplate, + String boundaryQuery, + String partitionColumn, + LogicalType partitionColumnType, + JdbcDialect dialect, + int numPartitions) { + this.baseSqlTemplate = Preconditions.checkNotNull(baseSqlTemplate); + this.boundaryQuery = Preconditions.checkNotNull(boundaryQuery); + this.partitionColumn = Preconditions.checkNotNull(partitionColumn); + this.partitionColumnType = Preconditions.checkNotNull(partitionColumnType); + this.dialect = Preconditions.checkNotNull(dialect); + Preconditions.checkArgument(numPartitions > 0, "numPartitions must be positive"); + this.numPartitions = numPartitions; + } + + @Override + protected List<JdbcSourceSplit> discoverSplits(Connection connection) throws SQLException { + List<Serializable> boundaryValues = fetchBoundaryValues(connection); + return buildSplits(boundaryValues); + } + + @Override + protected String describeSource() { + // Used only for error messages and AbstractDynamicSplitterEnumerator's discovery + // description - deliberately not the same query lineageQueries() reports below. + return boundaryQuery; + } + + @Override + public List<String> lineageQueries() { + // Report the actual data query for lineage purposes, not the boundary-query used only to + // discover split points - the latter often targets a different table (e.g. a stats table). + return Collections.singletonList(baseSqlTemplate); + } + + private List<Serializable> fetchBoundaryValues(Connection connection) throws SQLException { + // Wrap with ORDER BY 1 so split points are always ascending, even if the user's query + // doesn't specify an order itself. + String wrappedQuery = + "SELECT * FROM (" + boundaryQuery + ") AS boundary_query_result ORDER BY 1"; + + // A misconfigured boundary-query (e.g. missing an aggregation/sampling step) could + // otherwise return an unbounded number of rows and exhaust JobManager memory before the + // count is ever checked, so the max-allowed check happens per-row instead of after + // buffering the full result. + int maxAllowed = numPartitions - 1; + List<Serializable> values = new ArrayList<>(); + try (PreparedStatement statement = connection.prepareStatement(wrappedQuery); + ResultSet resultSet = statement.executeQuery()) { + ResultSetMetaData metaData = resultSet.getMetaData(); + validateSingleColumn(metaData); + validateColumnTypeComparable(metaData.getColumnType(1)); + + while (resultSet.next()) { + if (values.size() >= maxAllowed) { + throw new IllegalStateException( + "The boundary-query [" + + boundaryQuery + + "] returned more than " + + maxAllowed + + " values, exceeding the limit allowed for scan.partition.num=" + + numPartitions + + "."); + } + Object value = resultSet.getObject(1); + if (value == null) { + throw new IllegalStateException( + "The boundary-query [" + + boundaryQuery + + "] returned a NULL boundary value, which is not allowed."); + } + values.add((Serializable) value); + } + } + + return values; + } + + // Validation happens here, against a live connection at split-discovery time, rather than + // statically when the job is submitted - unlike every other scan.partition.* option, the + // legality of a boundary-query can only be known by actually running it. + private void validateSingleColumn(ResultSetMetaData metaData) throws SQLException { + if (metaData.getColumnCount() != 1) { + throw new IllegalStateException( + "The boundary-query [" + + boundaryQuery + + "] must return exactly one column, but returned " + + metaData.getColumnCount() + + "."); + } + } + + private void validateColumnTypeComparable(int sqlType) { + if (!isTypeComparable(sqlType, partitionColumnType)) { + throw new IllegalStateException( + "The boundary-query's returned column type (java.sql.Types code " + + sqlType + + ") is not comparable with the type of '" + + partitionColumn + + "' (" + + partitionColumnType + + ")."); + } + } + + private static boolean isTypeComparable(int sqlType, LogicalType partitionColumnType) { + LogicalTypeRoot root = partitionColumnType.getTypeRoot(); + if (NUMERIC_SQL_TYPES.contains(sqlType)) { + return root.getFamilies().contains(LogicalTypeFamily.NUMERIC); + } + if (STRING_SQL_TYPES.contains(sqlType)) { + return root.getFamilies().contains(LogicalTypeFamily.CHARACTER_STRING); + } + if (DATETIME_SQL_TYPES.contains(sqlType)) { + return root.getFamilies().contains(LogicalTypeFamily.DATETIME); + } + return false; + } + + private List<JdbcSourceSplit> buildSplits(List<Serializable> boundaryValues) { + int partitionCount = boundaryValues.size() + 1; + List<JdbcSourceSplit> splits = new ArrayList<>(partitionCount); + + if (partitionCount == 1) { + splits.add(newSplit(0, baseSqlTemplate, null)); + return splits; + } + + String quotedColumn = dialect.quoteIdentifier(partitionColumn); + for (int i = 0; i < partitionCount; i++) { + String predicate; + Serializable[] params; + if (i == 0) { + predicate = quotedColumn + " < ? OR " + quotedColumn + " IS NULL"; + params = new Serializable[] {boundaryValues.get(0)}; + } else if (i == partitionCount - 1) { + predicate = quotedColumn + " >= ?"; + params = new Serializable[] {boundaryValues.get(i - 1)}; + } else { + predicate = quotedColumn + " >= ? AND " + quotedColumn + " < ?"; + params = new Serializable[] {boundaryValues.get(i - 1), boundaryValues.get(i)}; + } + String sql = baseSqlTemplate + " WHERE (" + predicate + ")"; Review Comment: This assumes `baseSqlTemplate` has no WHERE or LIMIT, but `JdbcDynamicTableSource` adds both for filter and limit push-down. How will FLINK-40394 combine them? ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/BoundaryQuerySplitterEnumerator.java: ########## @@ -0,0 +1,259 @@ +/* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.jdbc.core.database.dialect.JdbcDialect; +import org.apache.flink.connector.jdbc.core.datastream.source.split.CheckpointedOffset; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeFamily; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.util.Preconditions; + +import java.io.Serializable; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * A {@link AbstractDynamicSplitterEnumerator} that discovers split boundaries by executing a + * user-provided {@code scan.partition.boundary-query} against the source database, instead of + * requiring a statically-known numeric {@code [lower, upper]} range. + * + * <p>The boundary query is expected to return a single, comparable column; each returned value + * (after being sorted ascending) becomes a split point. {@code N} returned values produce {@code N + * + 1} partitions: + * + * <ul> + * <li>first: {@code col < v(0) OR col IS NULL} + * <li>middle {@code i} (0 < i < N): {@code col >= v(i-1) AND col < v(i)} + * <li>last: {@code col >= v(N-1)} + * </ul> + * + * <p>Unlike every other splitter in this connector, the SQL predicate text itself differs between + * splits rather than only the bound parameter values, since a single {@code BETWEEN ? AND ?} + * template cannot express the "first"/"last" open-ended ranges above. + */ +@Internal +public class BoundaryQuerySplitterEnumerator extends AbstractDynamicSplitterEnumerator { + + private static final Set<Integer> NUMERIC_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.TINYINT, + Types.SMALLINT, + Types.INTEGER, + Types.BIGINT, + Types.DECIMAL, + Types.NUMERIC, + Types.FLOAT, + Types.REAL, + Types.DOUBLE)); + + private static final Set<Integer> STRING_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.CHAR, + Types.VARCHAR, + Types.LONGVARCHAR, + Types.NCHAR, + Types.NVARCHAR, + Types.LONGNVARCHAR)); + + private static final Set<Integer> DATETIME_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.DATE, + Types.TIME, + Types.TIMESTAMP, + Types.TIME_WITH_TIMEZONE, + Types.TIMESTAMP_WITH_TIMEZONE)); + + private final String baseSqlTemplate; + private final String boundaryQuery; + private final String partitionColumn; + private final LogicalType partitionColumnType; + private final JdbcDialect dialect; + private final int numPartitions; + + public BoundaryQuerySplitterEnumerator( + String baseSqlTemplate, + String boundaryQuery, + String partitionColumn, + LogicalType partitionColumnType, + JdbcDialect dialect, + int numPartitions) { + this.baseSqlTemplate = Preconditions.checkNotNull(baseSqlTemplate); + this.boundaryQuery = Preconditions.checkNotNull(boundaryQuery); + this.partitionColumn = Preconditions.checkNotNull(partitionColumn); + this.partitionColumnType = Preconditions.checkNotNull(partitionColumnType); + this.dialect = Preconditions.checkNotNull(dialect); + Preconditions.checkArgument(numPartitions > 0, "numPartitions must be positive"); + this.numPartitions = numPartitions; + } + + @Override + protected List<JdbcSourceSplit> discoverSplits(Connection connection) throws SQLException { + List<Serializable> boundaryValues = fetchBoundaryValues(connection); + return buildSplits(boundaryValues); + } + + @Override + protected String describeSource() { + // Used only for error messages and AbstractDynamicSplitterEnumerator's discovery + // description - deliberately not the same query lineageQueries() reports below. + return boundaryQuery; + } + + @Override + public List<String> lineageQueries() { + // Report the actual data query for lineage purposes, not the boundary-query used only to + // discover split points - the latter often targets a different table (e.g. a stats table). + return Collections.singletonList(baseSqlTemplate); + } + + private List<Serializable> fetchBoundaryValues(Connection connection) throws SQLException { + // Wrap with ORDER BY 1 so split points are always ascending, even if the user's query + // doesn't specify an order itself. + String wrappedQuery = + "SELECT * FROM (" + boundaryQuery + ") AS boundary_query_result ORDER BY 1"; + + // A misconfigured boundary-query (e.g. missing an aggregation/sampling step) could + // otherwise return an unbounded number of rows and exhaust JobManager memory before the + // count is ever checked, so the max-allowed check happens per-row instead of after + // buffering the full result. + int maxAllowed = numPartitions - 1; + List<Serializable> values = new ArrayList<>(); + try (PreparedStatement statement = connection.prepareStatement(wrappedQuery); + ResultSet resultSet = statement.executeQuery()) { + ResultSetMetaData metaData = resultSet.getMetaData(); + validateSingleColumn(metaData); + validateColumnTypeComparable(metaData.getColumnType(1)); + + while (resultSet.next()) { + if (values.size() >= maxAllowed) { + throw new IllegalStateException( + "The boundary-query [" + + boundaryQuery + + "] returned more than " + + maxAllowed + + " values, exceeding the limit allowed for scan.partition.num=" + + numPartitions + + "."); + } + Object value = resultSet.getObject(1); Review Comment: On Oracle 21c `getObject` returns `oracle.sql.TIMESTAMP` for a TIMESTAMP column, so a driver class ends up in the split. Reading it as `LocalDateTime` might make sense here. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/AbstractDynamicSplitterEnumerator.java: ########## @@ -0,0 +1,124 @@ + /* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import java.io.Serializable; +import java.sql.Connection; +import java.util.Collections; +import java.util.List; + +/** + * Base class for {@link SplitterEnumerator} implementations that must open a real database + * connection at job runtime to discover their splits (e.g. by executing a user-provided query), as + * opposed to computing splits purely from static configuration. + * + * <p>Split discovery is deliberately deferred to the first call of {@link #enumerateSplits()} + * rather than performed in {@link #start(JdbcConnectionProvider)}: the enclosing {@code + * JdbcSourceEnumerator} calls {@code start()} synchronously on the JobManager coordinator thread, + * but invokes {@code enumerateSplits()} through Flink's own {@code + * SplitEnumeratorContext#callAsync}, which already runs off that thread. Deferring the actual I/O + * this way avoids blocking job startup and avoids having to manage any dedicated thread of our own. + * + * <p>Because {@code enumerateSplits()} can be invoked concurrently by the coordinator before the + * first call returns, discovery is guarded so it only ever runs once. + */ +@Internal +public abstract class AbstractDynamicSplitterEnumerator implements SplitterEnumerator { + + private transient JdbcConnectionProvider connectionProvider; + + private List<JdbcSourceSplit> discoveredSplits; + + private boolean finished = false; + + @Override + public Boundedness getBoundedness() { + return Boundedness.BOUNDED; + } + + @Override + public void start(JdbcConnectionProvider connectionProvider) { + // Intentionally no I/O here - see class javadoc. + this.connectionProvider = connectionProvider; + } + + @Override + public synchronized List<JdbcSourceSplit> enumerateSplits() { + if (finished) { + return Collections.emptyList(); + } + if (discoveredSplits == null) { + try { + Connection connection = connectionProvider.getOrEstablishConnection(); + discoveredSplits = discoverSplits(connection); + } catch (Exception e) { + throw new RuntimeException("Failed to discover splits for " + describeSource(), e); Review Comment: This doesn't fail the job: `JdbcSourceEnumerator.onSplitsDiscovered` logs it and retries, 21 times in 20 cycles in my test. Could you make it fail, as FLIP-602 describes? ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/BoundaryQuerySplitterEnumerator.java: ########## @@ -0,0 +1,259 @@ +/* + * 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.flink.connector.jdbc.core.datastream.source.enumerator.splitter; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.jdbc.core.database.dialect.JdbcDialect; +import org.apache.flink.connector.jdbc.core.datastream.source.split.CheckpointedOffset; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeFamily; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.util.Preconditions; + +import java.io.Serializable; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * A {@link AbstractDynamicSplitterEnumerator} that discovers split boundaries by executing a + * user-provided {@code scan.partition.boundary-query} against the source database, instead of + * requiring a statically-known numeric {@code [lower, upper]} range. + * + * <p>The boundary query is expected to return a single, comparable column; each returned value + * (after being sorted ascending) becomes a split point. {@code N} returned values produce {@code N + * + 1} partitions: + * + * <ul> + * <li>first: {@code col < v(0) OR col IS NULL} + * <li>middle {@code i} (0 < i < N): {@code col >= v(i-1) AND col < v(i)} + * <li>last: {@code col >= v(N-1)} + * </ul> + * + * <p>Unlike every other splitter in this connector, the SQL predicate text itself differs between + * splits rather than only the bound parameter values, since a single {@code BETWEEN ? AND ?} + * template cannot express the "first"/"last" open-ended ranges above. + */ +@Internal +public class BoundaryQuerySplitterEnumerator extends AbstractDynamicSplitterEnumerator { + + private static final Set<Integer> NUMERIC_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.TINYINT, + Types.SMALLINT, + Types.INTEGER, + Types.BIGINT, + Types.DECIMAL, + Types.NUMERIC, + Types.FLOAT, + Types.REAL, + Types.DOUBLE)); + + private static final Set<Integer> STRING_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.CHAR, + Types.VARCHAR, + Types.LONGVARCHAR, + Types.NCHAR, + Types.NVARCHAR, + Types.LONGNVARCHAR)); + + private static final Set<Integer> DATETIME_SQL_TYPES = + new HashSet<>( + Arrays.asList( + Types.DATE, + Types.TIME, + Types.TIMESTAMP, + Types.TIME_WITH_TIMEZONE, + Types.TIMESTAMP_WITH_TIMEZONE)); + + private final String baseSqlTemplate; + private final String boundaryQuery; + private final String partitionColumn; + private final LogicalType partitionColumnType; + private final JdbcDialect dialect; + private final int numPartitions; + + public BoundaryQuerySplitterEnumerator( + String baseSqlTemplate, + String boundaryQuery, + String partitionColumn, + LogicalType partitionColumnType, + JdbcDialect dialect, + int numPartitions) { + this.baseSqlTemplate = Preconditions.checkNotNull(baseSqlTemplate); + this.boundaryQuery = Preconditions.checkNotNull(boundaryQuery); + this.partitionColumn = Preconditions.checkNotNull(partitionColumn); + this.partitionColumnType = Preconditions.checkNotNull(partitionColumnType); + this.dialect = Preconditions.checkNotNull(dialect); + Preconditions.checkArgument(numPartitions > 0, "numPartitions must be positive"); + this.numPartitions = numPartitions; + } + + @Override + protected List<JdbcSourceSplit> discoverSplits(Connection connection) throws SQLException { + List<Serializable> boundaryValues = fetchBoundaryValues(connection); + return buildSplits(boundaryValues); + } + + @Override + protected String describeSource() { + // Used only for error messages and AbstractDynamicSplitterEnumerator's discovery + // description - deliberately not the same query lineageQueries() reports below. + return boundaryQuery; + } + + @Override + public List<String> lineageQueries() { + // Report the actual data query for lineage purposes, not the boundary-query used only to + // discover split points - the latter often targets a different table (e.g. a stats table). + return Collections.singletonList(baseSqlTemplate); + } + + private List<Serializable> fetchBoundaryValues(Connection connection) throws SQLException { + // Wrap with ORDER BY 1 so split points are always ascending, even if the user's query + // doesn't specify an order itself. + String wrappedQuery = + "SELECT * FROM (" + boundaryQuery + ") AS boundary_query_result ORDER BY 1"; Review Comment: On Oracle 21c this fails with ORA-00933, since Oracle doesn't accept `AS` before a table alias. Dropping it works there and is standard SQL. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
