MartijnVisser commented on code in PR #209: URL: https://github.com/apache/flink-connector-jdbc/pull/209#discussion_r4057281265
########## flink-connector-jdbc-postgres/src/main/java/org/apache/flink/connector/jdbc/postgres/datastream/connection/PostgresConnectionProvider.java: ########## @@ -0,0 +1,307 @@ +/* + * 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.postgres.datastream.connection; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.connector.jdbc.core.datastream.connection.AbstractConnectionProvider; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionException; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionOptions; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionProvider; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.Table; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableBounds; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; + +import com.zaxxer.hikari.HikariDataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +/** + * Postgres {@link ConnectionProvider} implementation backed by a pooled JDBC connection. Only the + * Postgres-specific table/partition discovery and bound-query building live here; the connection + * lifecycle, pooling, and statement plumbing are inherited from {@link AbstractConnectionProvider}. + */ +@PublicEvolving +public class PostgresConnectionProvider extends AbstractConnectionProvider { + + private static final Logger LOG = LoggerFactory.getLogger(PostgresConnectionProvider.class); + + private static final String TABLE_PARTITIONED_TYPE = "PARTITIONED TABLE"; + private static final Set<String> TABLE_TYPES = + new HashSet<>(Arrays.asList("TABLE", TABLE_PARTITIONED_TYPE)); + private static final Set<String> VIEW_TYPES = + new HashSet<>(Arrays.asList("VIEW", "MATERIALIZED VIEW")); + private static final String[] SUPPORTED_TYPES = + Stream.concat(TABLE_TYPES.stream(), VIEW_TYPES.stream()).toArray(String[]::new); + + private static final int POOL_SIZE = 4; + + private String snapshotId; + private String snapshotTxId; + + public PostgresConnectionProvider(ConnectionOptions jdbcOptions) { + super(jdbcOptions); + } + + private PostgresConnectionProvider(ConnectionOptions jdbcOptions, HikariDataSource pool) { + super(jdbcOptions, pool); + } + + @Override + public ConnectionProvider newInstance() { + return new PostgresConnectionProvider(jdbcOptions, getOrCreatePool()); + } + + @Override + protected int maxPoolSize() { + return POOL_SIZE; + } + + @Override + protected String poolName() { + return "postgres-splitter-pool"; + } + + @Override + protected void onConnectionEstablished() { + checkTransactionSnapshotTxId(); + } + + public void createGlobalSnapshotId() throws SQLException, ClassNotFoundException { + Connection currentConn = getOrEstablishConnection(); + currentConn.setAutoCommit(false); + currentConn.setTransactionIsolation(Connection.TRANSACTION_REPEATABLE_READ); + try (Statement statement = currentConn.createStatement(); + ResultSet resultSet = statement.executeQuery("SELECT pg_export_snapshot()")) { + if (resultSet.next()) { + snapshotId = resultSet.getString(1); + LOG.info("Created global snapshot id: {}", snapshotId); + } + snapshotTxId = currentSnapshotTxId(); + } + } + + private void checkTransactionSnapshotTxId() { + try { + if (snapshotId != null + && isConnectionValid() + && !snapshotTxId.equalsIgnoreCase(currentSnapshotTxId())) { + Connection currentConn = getConnection(); + assert currentConn != null; + if (!currentConn.getAutoCommit()) { + currentConn.rollback(); + } + currentConn.setAutoCommit(false); + currentConn.setTransactionIsolation(Connection.TRANSACTION_REPEATABLE_READ); + LOG.info("Setting connection with snapshot id: {}", snapshotId); + try (PreparedStatement statement = + currentConn.prepareStatement("SET TRANSACTION SNAPSHOT ?")) { Review Comment: Postgres takes no bind parameter here; this throws `syntax error at or near "$1"` on 16. Needs the id as a literal, and `execute()`. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/snapshot/DatabaseSplitterEnumerator.java: ########## @@ -0,0 +1,217 @@ +/* + * 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.snapshot; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.Table; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.stream.Collectors; + +/** Splitter enumerator that fans out over every table in a database/schema. */ +@PublicEvolving +public class DatabaseSplitterEnumerator extends AsyncSnapshotSplitterEnumerator<JdbcSourceSplit> { + + private final String catalog; + private final String schema; + private final Set<String> tables; + private final Set<String> lineageQueries; + private final int chunkSize; + + // Bound concurrent table splitters so we don't exhaust the connection pool. + // Must be <= pool max size in ConnectionProvider. + private static final int MAX_CONCURRENT_TABLE_SPLITTERS = 4; + + private transient Queue<TableSplitterEnumerator> pendingTableSplitters; + private transient List<TableSplitterEnumerator> activeTableSplitters; + + public DatabaseSplitterEnumerator( + String catalog, String schema, Set<String> tables, int chunkSize) { + super(schema); + this.catalog = catalog; + this.schema = schema; + this.tables = tables; + this.lineageQueries = new LinkedHashSet<>(); + this.chunkSize = chunkSize; + } + + public static DatabaseSplitterEnumeratorBuilder builder() { + return new DatabaseSplitterEnumeratorBuilder(); + } + + @Override + public void start(JdbcConnectionProvider connectionProvider) { + initConnection(connectionProvider); + this.pendingTableSplitters = new ConcurrentLinkedQueue<>(); + this.activeTableSplitters = new CopyOnWriteArrayList<>(); + prepareTableSplitters(); Review Comment: `start()` is on the SourceCoordinator thread and this queries per table inline, stalling split assignment and checkpoints. Move it to the background thread below. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/connection/AbstractConnectionProvider.java: ########## @@ -0,0 +1,451 @@ +/* + * 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.connection; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; +import org.apache.flink.util.Preconditions; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; + +import java.sql.Connection; +import java.sql.Driver; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.time.Duration; +import java.util.Enumeration; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +/** + * Base {@link ConnectionProvider} that implements the connection lifecycle, pooling, and + * prepared-statement plumbing shared by every dialect-specific provider. Subclasses only need to + * implement the dialect-specific table discovery and bound-query building methods declared by + * {@link ConnectionProvider} (e.g. {@code getTables}, {@code queryMinMax}, {@code + * queryNextChunkMax}, {@code createQueryWithBounds}, {@code newInstance}). + */ +@Internal +public abstract class AbstractConnectionProvider implements ConnectionProvider { + + private static final Logger LOG = LoggerFactory.getLogger(AbstractConnectionProvider.class); + + private static final int DEFAULT_POOL_SIZE = 4; + private static final int MINIMUM_POOL_SIZE = 1; + + protected final ConnectionOptions jdbcOptions; + private int queryTimeoutSeconds; + private transient Driver loadedDriver; + private transient Connection connection; + private final Map<String, PreparedStatement> statementCache; + private transient HikariDataSource connectionPool; + private final boolean poolOwner; + + static { + // Load DriverManager first to avoid deadlock between DriverManager's + // static initialization block and specific driver class's static + // initialization block when two different driver classes are loading + // concurrently using Class.forName while DriverManager is uninitialized + // before. + // + // This could happen in JDK 8 but not above as driver loading has been + // moved out of DriverManager's static initialization block since JDK 9. + DriverManager.getDrivers(); + } + + protected AbstractConnectionProvider(ConnectionOptions jdbcOptions) { + this.jdbcOptions = jdbcOptions; + this.statementCache = new ConcurrentHashMap<>(); + this.queryTimeoutSeconds = jdbcOptions.getConnectionQueryTimeoutSeconds(); + this.poolOwner = true; + try { + // Establish the raw connection directly rather than via getOrEstablishConnection(): + // that method invokes the overridable onConnectionEstablished() hook, which a + // subclass may depend on its own fields for — fields that aren't initialized yet + // while this superclass constructor is still running. + establishConnection(); + } catch (Exception e) { + throw new ConnectionException( + "Failed to establish initial connection during connection provider construction.", + e); + } + } + + /** + * Creates a pooled instance that borrows its connection from the given pool. This instance does + * NOT own the pool and will not shut it down on close. + */ + protected AbstractConnectionProvider(ConnectionOptions jdbcOptions, HikariDataSource pool) { + this.jdbcOptions = jdbcOptions; + this.statementCache = new ConcurrentHashMap<>(); + this.queryTimeoutSeconds = jdbcOptions.getConnectionQueryTimeoutSeconds(); + this.connectionPool = pool; + this.poolOwner = false; + try { + this.connection = pool.getConnection(); + } catch (SQLException e) { + throw new ConnectionException("Failed to borrow connection from pool.", e); + } + } + + protected synchronized HikariDataSource getOrCreatePool() { + if (connectionPool == null) { + connectionPool = createConnectionPool(); + } + return connectionPool; + } + + private HikariDataSource createConnectionPool() { + HikariConfig config = new HikariConfig(); + // Provide a DataSource directly so HikariCP doesn't try to load the driver + // via its own classloader. In Flink, the JDBC driver lives in the user + // classloader and is not visible to HikariCP's threads. Using this driver + // delegate guarantees connections are created with the same code path as + // getOrEstablishConnection(). + config.setDataSource(new ConnectionDataSource(jdbcOptions, this::getLoadedDriver)); + config.setMinimumIdle(MINIMUM_POOL_SIZE); + config.setMaximumPoolSize(maxPoolSize()); + config.setConnectionTimeout( + Duration.ofSeconds(jdbcOptions.getConnectionCheckTimeoutSeconds()).toMillis()); + config.setPoolName(poolName()); + LOG.info("Creating HikariCP connection pool with maxPoolSize={}", maxPoolSize()); + return new HikariDataSource(config); + } + + /** Maximum number of pooled connections. Override to change the pool size. */ + protected int maxPoolSize() { + return DEFAULT_POOL_SIZE; + } + + /** Name used for the HikariCP pool, shown in logs/metrics. */ + protected String poolName() { + return getClass().getSimpleName() + "-pool"; + } + + /** + * Hook invoked every time a connection is (re-)established, before it's handed back to the + * caller. No-op by default; dialect-specific subclasses can override to re-apply + * connection-scoped state (e.g. re-syncing a shared snapshot transaction id). + */ + protected void onConnectionEstablished() throws SQLException {} + + @Override + public Connection getConnection() { + return connection; + } + + @Nonnull + @Override + public Properties getProperties() { + return jdbcOptions.getProperties(); + } + + @Override + public boolean isConnectionValid() throws SQLException { + return connection != null + && !connection.isClosed() + && connection.isValid(jdbcOptions.getConnectionCheckTimeoutSeconds()); + } + + private Driver loadDriver(String driverName) throws SQLException, ClassNotFoundException { + Preconditions.checkNotNull(driverName); + Enumeration<Driver> drivers = DriverManager.getDrivers(); + while (drivers.hasMoreElements()) { + Driver driver = drivers.nextElement(); + if (driver.getClass().getName().equals(driverName)) { + return driver; + } + } + // We could reach here for reasons: + // * Class loader hell of DriverManager(see JDK-8146872). + // * driver is not installed as a service provider. + Class<?> clazz = + Class.forName(driverName, true, Thread.currentThread().getContextClassLoader()); + try { + return (Driver) clazz.getDeclaredConstructor().newInstance(); + } catch (Exception ex) { + throw new SQLException("Fail to create driver of class " + driverName, ex); + } + } + + private Driver getLoadedDriver() throws SQLException, ClassNotFoundException { + if (loadedDriver == null) { + loadedDriver = loadDriver(jdbcOptions.getDriverName()); + } + return loadedDriver; + } + + @Override + public Connection getOrEstablishConnection() throws SQLException, ClassNotFoundException { + if (isConnectionValid()) { + onConnectionEstablished(); + return connection; + } + establishConnection(); + onConnectionEstablished(); + return connection; + } + + /** + * Establishes a fresh {@link #connection}. Does not invoke {@link #onConnectionEstablished()}. + */ + private void establishConnection() throws SQLException, ClassNotFoundException { Review Comment: This replaces `connection` but leaves `statementCache` bound to the old one, so a dropped connection never recovers. Verified against Postgres 16. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/snapshot/DatabaseSplitterEnumerator.java: ########## @@ -0,0 +1,217 @@ +/* + * 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.snapshot; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.Table; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.stream.Collectors; + +/** Splitter enumerator that fans out over every table in a database/schema. */ +@PublicEvolving +public class DatabaseSplitterEnumerator extends AsyncSnapshotSplitterEnumerator<JdbcSourceSplit> { + + private final String catalog; + private final String schema; + private final Set<String> tables; + private final Set<String> lineageQueries; + private final int chunkSize; + + // Bound concurrent table splitters so we don't exhaust the connection pool. + // Must be <= pool max size in ConnectionProvider. + private static final int MAX_CONCURRENT_TABLE_SPLITTERS = 4; + + private transient Queue<TableSplitterEnumerator> pendingTableSplitters; + private transient List<TableSplitterEnumerator> activeTableSplitters; + + public DatabaseSplitterEnumerator( + String catalog, String schema, Set<String> tables, int chunkSize) { + super(schema); + this.catalog = catalog; + this.schema = schema; + this.tables = tables; + this.lineageQueries = new LinkedHashSet<>(); + this.chunkSize = chunkSize; + } + + public static DatabaseSplitterEnumeratorBuilder builder() { + return new DatabaseSplitterEnumeratorBuilder(); + } + + @Override + public void start(JdbcConnectionProvider connectionProvider) { + initConnection(connectionProvider); + this.pendingTableSplitters = new ConcurrentLinkedQueue<>(); + this.activeTableSplitters = new CopyOnWriteArrayList<>(); + prepareTableSplitters(); + startBackgroundWork(); + } + + @Override + public List<String> lineageQueries() { + return new ArrayList<>(this.lineageQueries); + } + + @Override + protected void runBackgroundWork() throws InterruptedException { + // Start an initial batch of table splitters — each one borrows its + // own pooled connection. We refill below as splitters finish so we + // never exceed the pool capacity. + fillActiveTableSplitters(); + + while (!Thread.currentThread().isInterrupted() + && (!activeTableSplitters.isEmpty() || !pendingTableSplitters.isEmpty())) { + boolean producedSplits = false; + List<TableSplitterEnumerator> finished = new ArrayList<>(); + + for (TableSplitterEnumerator tableSplitter : activeTableSplitters) { + if (tableSplitter.isAllSplitsFinished()) { + lineageQueries.addAll(tableSplitter.lineageQueries()); Review Comment: Plain `LinkedHashSet`, written here on the background thread and copied from the job manager thread. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/snapshot/AsyncSnapshotSplitterEnumerator.java: ########## @@ -0,0 +1,189 @@ +/* + * 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.snapshot; + +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionProvider; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.SplitterEnumerator; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Shared async background-computation machinery for the snapshot splitters ({@link + * TableSplitterEnumerator}, {@link DatabaseSplitterEnumerator}): a single background daemon thread + * computes {@code T} items and offers them into a queue, while {@link #enumerateSplits()} drains + * and converts whatever's ready — non-blocking apart from a short wait for the very first item. + * + * @param <T> the type of item the background thread produces, converted to a {@link + * JdbcSourceSplit} at drain time via {@link #toSplit} + */ +abstract class AsyncSnapshotSplitterEnumerator<T> implements SplitterEnumerator { + + private final Logger log = LoggerFactory.getLogger(getClass()); + private final String name; + private final Queue<T> outputQueue = new ConcurrentLinkedQueue<>(); + + protected transient ConnectionProvider connection; + + private transient ExecutorService executor; + private transient AtomicBoolean workDone; + private transient CountDownLatch firstReady; + private transient volatile Throwable backgroundFailure; + + protected AsyncSnapshotSplitterEnumerator(String name) { + this.name = name; + } + + /** Validates and stores the connection provider for use by subclasses. */ + protected final void initConnection(JdbcConnectionProvider connectionProvider) { + if (!(connectionProvider instanceof ConnectionProvider)) { + throw new IllegalArgumentException( + "Connection provider must be an instance of " + + ConnectionProvider.class.getSimpleName()); + } + this.connection = (ConnectionProvider) connectionProvider; + } + + @Override + public final Boundedness getBoundedness() { + return Boundedness.CONTINUOUS_UNBOUNDED; + } + + /** Starts the background thread that runs {@link #runBackgroundWork()}. */ + protected final void startBackgroundWork() { + this.workDone = new AtomicBoolean(false); + this.firstReady = new CountDownLatch(1); + this.executor = + Executors.newSingleThreadExecutor( + r -> { + Thread t = new Thread(r, "snapshot-compute-" + name); + t.setDaemon(true); + return t; + }); + executor.submit( + () -> { + try { + runBackgroundWork(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Throwable e) { + log.error("Background computation failed for {}", name, e); + backgroundFailure = e; + } finally { + workDone.set(true); + firstReady.countDown(); + } + }); + } + + /** Subclass-specific unit of work; push results via {@link #offer}/{@link #offerAll}. */ + protected abstract void runBackgroundWork() throws Exception; + + /** Converts a queued item into an emittable split. */ + protected abstract JdbcSourceSplit toSplit(T item); + + /** Subclass-specific resource cleanup, called after the background thread has stopped. */ + protected abstract void closeResources(); + + /** Offers a single computed item and wakes up anyone waiting on the first-ready signal. */ + protected final void offer(T item) { + outputQueue.add(item); + firstReady.countDown(); + } + + /** Offers a batch of computed items and wakes up anyone waiting on the first-ready signal. */ + protected final void offerAll(Collection<T> items) { + if (!items.isEmpty()) { + outputQueue.addAll(items); + firstReady.countDown(); + } + } + + @Override + public final boolean isAllSplitsFinished() { + return workDone != null && workDone.get() && outputQueue.isEmpty(); + } + + @Override + public final List<JdbcSourceSplit> enumerateSplits() { + // Wait briefly for the background thread to produce at least one item. + if (firstReady != null) { + try { + firstReady.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + if (backgroundFailure != null) { + throw new IllegalStateException( + "Split computation failed for " + name + " — refusing to emit partial splits", + backgroundFailure); + } + + List<JdbcSourceSplit> splits = new ArrayList<>(); + T item; + while ((item = outputQueue.poll()) != null) { + splits.add(toSplit(item)); + } + return splits; + } + + @Override + public final void close() { + if (executor != null) { + executor.shutdownNow(); + try { + if (!executor.awaitTermination(5, TimeUnit.SECONDS)) { + log.warn( + "Background computation for {} did not stop within the shutdown grace period — a" + + " query may still be running against its connection.", + name); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + closeResources(); + } + + @Override + public final Serializable serializableState() { Review Comment: Returns null and `restoreState()` ignores its argument, so a restore re-emits every split. Both `final`, and the test asserts it as the contract. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/splitter/snapshot/AsyncSnapshotSplitterEnumerator.java: ########## @@ -0,0 +1,189 @@ +/* + * 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.snapshot; + +import org.apache.flink.api.connector.source.Boundedness; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionProvider; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.SplitterEnumerator; +import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit; +import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Shared async background-computation machinery for the snapshot splitters ({@link + * TableSplitterEnumerator}, {@link DatabaseSplitterEnumerator}): a single background daemon thread + * computes {@code T} items and offers them into a queue, while {@link #enumerateSplits()} drains + * and converts whatever's ready — non-blocking apart from a short wait for the very first item. + * + * @param <T> the type of item the background thread produces, converted to a {@link + * JdbcSourceSplit} at drain time via {@link #toSplit} + */ +abstract class AsyncSnapshotSplitterEnumerator<T> implements SplitterEnumerator { + + private final Logger log = LoggerFactory.getLogger(getClass()); + private final String name; + private final Queue<T> outputQueue = new ConcurrentLinkedQueue<>(); + + protected transient ConnectionProvider connection; + + private transient ExecutorService executor; + private transient AtomicBoolean workDone; + private transient CountDownLatch firstReady; + private transient volatile Throwable backgroundFailure; + + protected AsyncSnapshotSplitterEnumerator(String name) { + this.name = name; + } + + /** Validates and stores the connection provider for use by subclasses. */ + protected final void initConnection(JdbcConnectionProvider connectionProvider) { + if (!(connectionProvider instanceof ConnectionProvider)) { + throw new IllegalArgumentException( + "Connection provider must be an instance of " + + ConnectionProvider.class.getSimpleName()); + } + this.connection = (ConnectionProvider) connectionProvider; + } + + @Override + public final Boundedness getBoundedness() { Review Comment: Back to the open thread above: boundedness does not control split production, so table size is not the reason. BATCH mode is impossible as this stands. ########## flink-connector-jdbc-postgres/src/main/java/org/apache/flink/connector/jdbc/postgres/datastream/connection/PostgresConnectionProvider.java: ########## @@ -0,0 +1,307 @@ +/* + * 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.postgres.datastream.connection; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.connector.jdbc.core.datastream.connection.AbstractConnectionProvider; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionException; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionOptions; +import org.apache.flink.connector.jdbc.core.datastream.connection.ConnectionProvider; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.Table; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableBounds; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; + +import com.zaxxer.hikari.HikariDataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +/** + * Postgres {@link ConnectionProvider} implementation backed by a pooled JDBC connection. Only the + * Postgres-specific table/partition discovery and bound-query building live here; the connection + * lifecycle, pooling, and statement plumbing are inherited from {@link AbstractConnectionProvider}. + */ +@PublicEvolving +public class PostgresConnectionProvider extends AbstractConnectionProvider { + + private static final Logger LOG = LoggerFactory.getLogger(PostgresConnectionProvider.class); + + private static final String TABLE_PARTITIONED_TYPE = "PARTITIONED TABLE"; + private static final Set<String> TABLE_TYPES = + new HashSet<>(Arrays.asList("TABLE", TABLE_PARTITIONED_TYPE)); + private static final Set<String> VIEW_TYPES = + new HashSet<>(Arrays.asList("VIEW", "MATERIALIZED VIEW")); + private static final String[] SUPPORTED_TYPES = + Stream.concat(TABLE_TYPES.stream(), VIEW_TYPES.stream()).toArray(String[]::new); + + private static final int POOL_SIZE = 4; + + private String snapshotId; + private String snapshotTxId; + + public PostgresConnectionProvider(ConnectionOptions jdbcOptions) { + super(jdbcOptions); + } + + private PostgresConnectionProvider(ConnectionOptions jdbcOptions, HikariDataSource pool) { + super(jdbcOptions, pool); + } + + @Override + public ConnectionProvider newInstance() { + return new PostgresConnectionProvider(jdbcOptions, getOrCreatePool()); + } + + @Override + protected int maxPoolSize() { + return POOL_SIZE; + } + + @Override + protected String poolName() { + return "postgres-splitter-pool"; + } + + @Override + protected void onConnectionEstablished() { Review Comment: This runs on every `getOrEstablishConnection()`, not only on establish, so each call adds an `isValid()` round trip and a `pg_current_snapshot()` query. ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/connection/AbstractConnectionProvider.java: ########## @@ -0,0 +1,451 @@ +/* + * 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.connection; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; +import org.apache.flink.util.Preconditions; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; + +import java.sql.Connection; +import java.sql.Driver; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.time.Duration; +import java.util.Enumeration; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +/** + * Base {@link ConnectionProvider} that implements the connection lifecycle, pooling, and + * prepared-statement plumbing shared by every dialect-specific provider. Subclasses only need to + * implement the dialect-specific table discovery and bound-query building methods declared by + * {@link ConnectionProvider} (e.g. {@code getTables}, {@code queryMinMax}, {@code + * queryNextChunkMax}, {@code createQueryWithBounds}, {@code newInstance}). + */ +@Internal +public abstract class AbstractConnectionProvider implements ConnectionProvider { + + private static final Logger LOG = LoggerFactory.getLogger(AbstractConnectionProvider.class); + + private static final int DEFAULT_POOL_SIZE = 4; + private static final int MINIMUM_POOL_SIZE = 1; + + protected final ConnectionOptions jdbcOptions; + private int queryTimeoutSeconds; + private transient Driver loadedDriver; + private transient Connection connection; + private final Map<String, PreparedStatement> statementCache; Review Comment: `PreparedStatement` is not serialisable and this is not `transient`, unlike the three below. Intentional? ########## flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/connection/AbstractConnectionProvider.java: ########## @@ -0,0 +1,451 @@ +/* + * 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.connection; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableColumn; +import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.snapshot.domain.TableId; +import org.apache.flink.util.Preconditions; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; + +import java.sql.Connection; +import java.sql.Driver; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.time.Duration; +import java.util.Enumeration; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +/** + * Base {@link ConnectionProvider} that implements the connection lifecycle, pooling, and + * prepared-statement plumbing shared by every dialect-specific provider. Subclasses only need to + * implement the dialect-specific table discovery and bound-query building methods declared by + * {@link ConnectionProvider} (e.g. {@code getTables}, {@code queryMinMax}, {@code + * queryNextChunkMax}, {@code createQueryWithBounds}, {@code newInstance}). + */ +@Internal +public abstract class AbstractConnectionProvider implements ConnectionProvider { + + private static final Logger LOG = LoggerFactory.getLogger(AbstractConnectionProvider.class); + + private static final int DEFAULT_POOL_SIZE = 4; + private static final int MINIMUM_POOL_SIZE = 1; + + protected final ConnectionOptions jdbcOptions; + private int queryTimeoutSeconds; + private transient Driver loadedDriver; + private transient Connection connection; + private final Map<String, PreparedStatement> statementCache; + private transient HikariDataSource connectionPool; + private final boolean poolOwner; + + static { + // Load DriverManager first to avoid deadlock between DriverManager's + // static initialization block and specific driver class's static + // initialization block when two different driver classes are loading + // concurrently using Class.forName while DriverManager is uninitialized + // before. + // + // This could happen in JDK 8 but not above as driver loading has been + // moved out of DriverManager's static initialization block since JDK 9. + DriverManager.getDrivers(); + } + + protected AbstractConnectionProvider(ConnectionOptions jdbcOptions) { + this.jdbcOptions = jdbcOptions; + this.statementCache = new ConcurrentHashMap<>(); + this.queryTimeoutSeconds = jdbcOptions.getConnectionQueryTimeoutSeconds(); + this.poolOwner = true; + try { + // Establish the raw connection directly rather than via getOrEstablishConnection(): + // that method invokes the overridable onConnectionEstablished() hook, which a + // subclass may depend on its own fields for — fields that aren't initialized yet + // while this superclass constructor is still running. + establishConnection(); + } catch (Exception e) { + throw new ConnectionException( + "Failed to establish initial connection during connection provider construction.", + e); + } + } + + /** + * Creates a pooled instance that borrows its connection from the given pool. This instance does + * NOT own the pool and will not shut it down on close. + */ + protected AbstractConnectionProvider(ConnectionOptions jdbcOptions, HikariDataSource pool) { + this.jdbcOptions = jdbcOptions; + this.statementCache = new ConcurrentHashMap<>(); + this.queryTimeoutSeconds = jdbcOptions.getConnectionQueryTimeoutSeconds(); + this.connectionPool = pool; + this.poolOwner = false; + try { + this.connection = pool.getConnection(); + } catch (SQLException e) { + throw new ConnectionException("Failed to borrow connection from pool.", e); + } + } + + protected synchronized HikariDataSource getOrCreatePool() { + if (connectionPool == null) { + connectionPool = createConnectionPool(); + } + return connectionPool; + } + + private HikariDataSource createConnectionPool() { + HikariConfig config = new HikariConfig(); + // Provide a DataSource directly so HikariCP doesn't try to load the driver + // via its own classloader. In Flink, the JDBC driver lives in the user + // classloader and is not visible to HikariCP's threads. Using this driver + // delegate guarantees connections are created with the same code path as + // getOrEstablishConnection(). + config.setDataSource(new ConnectionDataSource(jdbcOptions, this::getLoadedDriver)); + config.setMinimumIdle(MINIMUM_POOL_SIZE); + config.setMaximumPoolSize(maxPoolSize()); + config.setConnectionTimeout( + Duration.ofSeconds(jdbcOptions.getConnectionCheckTimeoutSeconds()).toMillis()); + config.setPoolName(poolName()); + LOG.info("Creating HikariCP connection pool with maxPoolSize={}", maxPoolSize()); + return new HikariDataSource(config); + } + + /** Maximum number of pooled connections. Override to change the pool size. */ + protected int maxPoolSize() { + return DEFAULT_POOL_SIZE; + } + + /** Name used for the HikariCP pool, shown in logs/metrics. */ + protected String poolName() { + return getClass().getSimpleName() + "-pool"; + } + + /** + * Hook invoked every time a connection is (re-)established, before it's handed back to the + * caller. No-op by default; dialect-specific subclasses can override to re-apply + * connection-scoped state (e.g. re-syncing a shared snapshot transaction id). + */ + protected void onConnectionEstablished() throws SQLException {} + + @Override + public Connection getConnection() { + return connection; + } + + @Nonnull + @Override + public Properties getProperties() { + return jdbcOptions.getProperties(); + } + + @Override + public boolean isConnectionValid() throws SQLException { + return connection != null + && !connection.isClosed() + && connection.isValid(jdbcOptions.getConnectionCheckTimeoutSeconds()); + } + + private Driver loadDriver(String driverName) throws SQLException, ClassNotFoundException { + Preconditions.checkNotNull(driverName); + Enumeration<Driver> drivers = DriverManager.getDrivers(); + while (drivers.hasMoreElements()) { + Driver driver = drivers.nextElement(); + if (driver.getClass().getName().equals(driverName)) { + return driver; + } + } + // We could reach here for reasons: + // * Class loader hell of DriverManager(see JDK-8146872). + // * driver is not installed as a service provider. + Class<?> clazz = + Class.forName(driverName, true, Thread.currentThread().getContextClassLoader()); + try { + return (Driver) clazz.getDeclaredConstructor().newInstance(); + } catch (Exception ex) { + throw new SQLException("Fail to create driver of class " + driverName, ex); + } + } + + private Driver getLoadedDriver() throws SQLException, ClassNotFoundException { + if (loadedDriver == null) { + loadedDriver = loadDriver(jdbcOptions.getDriverName()); + } + return loadedDriver; + } + + @Override + public Connection getOrEstablishConnection() throws SQLException, ClassNotFoundException { + if (isConnectionValid()) { + onConnectionEstablished(); + return connection; + } + establishConnection(); + onConnectionEstablished(); + return connection; + } + + /** + * Establishes a fresh {@link #connection}. Does not invoke {@link #onConnectionEstablished()}. + */ + private void establishConnection() throws SQLException, ClassNotFoundException { + if (!poolOwner) { + // connectionPool is guaranteed non-null here: the pool-borrowing constructor requires it. + connection = connectionPool.getConnection(); + return; + } + String connectionUrl = jdbcOptions.getDbURL(); + if (jdbcOptions.getDriverName() == null) { + connection = DriverManager.getConnection(connectionUrl, getProperties()); + } else { + Driver driver = getLoadedDriver(); + connection = driver.connect(connectionUrl, getProperties()); + if (connection == null) { + // Throw same exception as DriverManager.getConnection when no driver found to match + // caller expectation. + throw new SQLException("No suitable driver found for " + connectionUrl, "08001"); + } + } + } + + @Override + public Connection reestablishConnection() throws SQLException, ClassNotFoundException { + closeConnection(); + return getOrEstablishConnection(); + } + + public void withQueryTimeout(Duration queryTimeout) { + Objects.requireNonNull(queryTimeout, "queryTimeout must be provided"); + if (queryTimeout.isZero() || queryTimeout.isNegative()) { + throw new IllegalArgumentException("queryTimeout must be positive"); + } + this.queryTimeoutSeconds = (int) queryTimeout.getSeconds(); + } + + private Set<String> getPrimaryKeys(TableId tableId) { + Set<String> primaryKeys = new HashSet<>(); + try (ResultSet rs = + getOrEstablishConnection() + .getMetaData() + .getPrimaryKeys( + tableId.catalogName(), tableId.schemaName(), tableId.tableName())) { + while (rs.next()) { + primaryKeys.add(rs.getString(4)); + } + } catch (SQLException | ClassNotFoundException e) { + throw new ConnectionException("Failed to get primary keys for table " + tableId, e); + } + return primaryKeys; + } + + @Override + public Set<TableColumn> getTableColumns(TableId tableId) { + Set<String> tablePrimaryKeys = getPrimaryKeys(tableId); + Set<TableColumn> tableColumns = new LinkedHashSet<>(); + try (ResultSet rs = + getOrEstablishConnection() + .getMetaData() + .getColumns( + tableId.catalogName(), + tableId.schemaName(), + tableId.tableName(), + (String) null)) { + while (rs.next()) { + String columnName = rs.getString(4); + TableColumn column = + TableColumn.builder() + .withColumnName(columnName) + .withColumnType(rs.getString(6)) + .withColumnPosition(rs.getInt(17)) + .withColumnNullable(isNullable(rs.getInt(11))) + .withColumnPk(tablePrimaryKeys.contains(columnName)) + .build(); + tableColumns.add(column); + } + } catch (SQLException | ClassNotFoundException e) { + throw new ConnectionException( + String.format("Failed to get columns for table %s", tableId), e); + } + return tableColumns; + } + + protected static boolean isNullable(int jdbcNullable) { + return jdbcNullable == ResultSetMetaData.columnNullable + || jdbcNullable == ResultSetMetaData.columnNullableUnknown; + } + + protected <T> T queryAndMap(String query, ResultSetMapper<T> mapper) { + Objects.requireNonNull(mapper, "Mapper must be provided"); + try (Statement statement = createStatement()) { + if (LOG.isTraceEnabled()) { + LOG.trace("running '{}' with {}s timeout", query, this.queryTimeoutSeconds); + } + + try (ResultSet resultSet = statement.executeQuery(query)) { + return mapper.apply(resultSet); + } + } catch (Exception e) { + throw new ConnectionException(String.format("Failed executing query %s", query), e); + } + } + + private Statement createStatement() { + try { + final Statement statement = getOrEstablishConnection().createStatement(); + initializeStatement(statement); + return statement; + } catch (SQLException | ClassNotFoundException e) { + throw new ConnectionException("Failed to create statement from factory", e); + } + } + + protected <T> T prepareQueryAndMap( + String preparedQuery, StatementPreparer preparer, ResultSetMapper<T> mapper) { + Objects.requireNonNull(mapper, "Mapper must be provided"); + try { + PreparedStatement statement = prepareQuery(preparedQuery, preparer); + try (ResultSet resultSet = statement.executeQuery()) { + return mapper.apply(resultSet); + } + } catch (Exception e) { + throw new ConnectionException( + String.format("Failed executing query %s", preparedQuery), e); + } + } + + private PreparedStatement prepareQuery(String preparedQuery, StatementPreparer preparer) { + try { + PreparedStatement statement = this.createPreparedStatement(preparedQuery); + preparer.accept(statement); + return statement; + } catch (SQLException e) { + throw new ConnectionException("Failed to prepare query", e); + } + } + + private PreparedStatement createPreparedStatement(String preparedQueryString) { + return this.statementCache.computeIfAbsent( + preparedQueryString, + (query) -> { + try { + LOG.trace( + "Inserting prepared statement '{}' that does not exist in the cache", + query); + PreparedStatement preparedStatement = + getOrEstablishConnection().prepareStatement(query); + initializeStatement(preparedStatement); + if (LOG.isTraceEnabled()) { + LOG.trace( + "PreparedStatement '{}' with {}s timeout", + preparedQueryString, + this.queryTimeoutSeconds); + } + return preparedStatement; + } catch (SQLException | ClassNotFoundException e) { + throw new ConnectionException(e); + } + }); + } + + private void initializeStatement(Statement statement) { + try { + statement.setQueryTimeout(queryTimeoutSeconds); + } catch (SQLException e) { + throw new ConnectionException("Failed to add timeout to statement", e); + } + } + + private void closePreparedStatement(PreparedStatement statement) { + LOG.trace("Closing prepared statement '{}' removed from cache", statement); + try { + statement.close(); + } catch (Exception e) { + LOG.info("Exception while closing a prepared statement removed from cache", e); + } + } + + @Override + public void closeConnection() { + if (connection == null) { + return; + } + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future<Object> futureClose = + executor.submit( + () -> { + this.connection.close(); + LOG.info("Connection gracefully closed"); + return null; + }); + + try { + futureClose.get(10L, TimeUnit.SECONDS); + } catch (ExecutionException e) { + throw new ConnectionException(e.getCause()); + } catch (InterruptedException | TimeoutException e) { Review Comment: The interrupt flag is dropped. Please restore it before falling through to `abort()`. -- 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]
