patrashu commented on code in PR #234: URL: https://github.com/apache/flink-connector-jdbc/pull/234#discussion_r4139055819
########## 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: Switched to the JDBC 4.2 `getObject(int, Class)` overload for every temporal SQL type (DATE/TIME/TIMESTAMP, with/without timezone), reading them explicitly as LocalDate/LocalTime/LocalDateTime/OffsetTime/OffsetDateTime. Added a TIMESTAMP-column test against H2 as a regression check. -- 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]
