This is an automated email from the ASF dual-hosted git repository.
jerryshao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 9014793d4c [#13591] fix(jdbc): Avoid nested connection borrows in JDBC
catalog operations (#13729)
9014793d4c is described below
commit 9014793d4c7393483a6afc3e0c35b840a6a94268
Author: Qi Yu <[email protected]>
AuthorDate: Sat Oct 10 18:45:07 2026 +0800
[#13591] fix(jdbc): Avoid nested connection borrows in JDBC catalog
operations (#13729)
### What changes were proposed in this pull request?
Stop JDBC catalog operations from holding one pooled connection while
borrowing another from the same data source.
- Table create/alter and database drop generate SQL before borrowing the
execution connection. Generators may load the table, query Doris
backends/version, or check that a database is empty.
- ClickHouse validates alter settings and delegates to the base method.
OceanBase still generates and applies each change in order, borrowing
separately for each non-empty statement.
- Base and ClickHouse `load` prefetch the driver version on their held
connection. PostgreSQL and Hologres inherit base `load`. Successful
reads, including null, are cached; failed reads use the existing
null-version precision fallback and retry on the next load's held
connection. Initialization clears the cache for the new data source.
- SQL-generator Javadoc documents the connection-borrowing contract.
### Why are the changes needed?
With the default pool size of 10, concurrent alters on distinct MySQL
tables could fill the pool with outer connections and then wait for
another inside SQL generation, failing with `Timeout waiting for idle
object` after 30 seconds. Loading datetime columns had the same nested
borrow, including when driver-version metadata returned null or failed.
Fix: #13591
### Does this PR introduce _any_ user-facing change?
No API or configuration changes. The changed operation paths avoid
nested borrows. Side effects:
- Generator validation and Doris backend/version queries precede a
missing-database error from the execution connection.
- Empty generated SQL skips the execution connection, so that step no
longer reports a missing database.
- OceanBase borrows per non-empty statement. Each statement remains
standalone DDL with the catalog set on every borrow; sequential
application and partial-failure behavior are preserved.
- Driver-version metadata is read once per data source after success,
including a null result. A failed read is logged at debug level and
retried by a later load; it does not borrow again during column parsing.
### How was this patch tested?
- All non-Docker unit tests in JDBC common, ClickHouse, OceanBase,
MySQL, Doris, PostgreSQL, Hologres and StarRocks: 538 tests passed.
- JDBC common connection-usage coverage now has 12 cases, including two
real single-connection SQLite/DBCP reproductions and counted base-load
coverage for null metadata, failed-read recovery, uncached getter
fallback, and reinitialization. ClickHouse's actual load test verifies
driver-version prefetch; OceanBase verifies generation after the
preceding DDL.
- Four targeted mutations independently remove the failed-read guard,
base-load prefetch, initialization reset, and ClickHouse-load prefetch;
each is detected by its regression test.
- Spotless and Javadoc pass for the changed modules. Docker catalog
tests and a real 32-client MySQL reproduction were not run locally.
---
.../operations/ClickHouseTableOperations.java | 14 +-
...stClickHouseTableOperationsConnectionUsage.java | 67 +++
.../TestClickHouseTableOperationsUnit.java | 1 +
.../operation/OceanBaseTableOperations.java | 23 +-
...estOceanBaseTableOperationsConnectionUsage.java | 80 ++++
.../jdbc/operation/JdbcDatabaseOperations.java | 8 +-
.../jdbc/operation/JdbcTableOperations.java | 134 +++++-
.../TestJdbcOperationsConnectionUsage.java | 480 +++++++++++++++++++++
.../jdbc/utils/ConnectionCountingDataSource.java | 168 ++++++++
9 files changed, 929 insertions(+), 46 deletions(-)
diff --git
a/catalogs-contrib/catalog-jdbc-clickhouse/src/main/java/org/apache/gravitino/catalog/clickhouse/operations/ClickHouseTableOperations.java
b/catalogs-contrib/catalog-jdbc-clickhouse/src/main/java/org/apache/gravitino/catalog/clickhouse/operations/ClickHouseTableOperations.java
index dac8cd44ed..bb07965db0 100644
---
a/catalogs-contrib/catalog-jdbc-clickhouse/src/main/java/org/apache/gravitino/catalog/clickhouse/operations/ClickHouseTableOperations.java
+++
b/catalogs-contrib/catalog-jdbc-clickhouse/src/main/java/org/apache/gravitino/catalog/clickhouse/operations/ClickHouseTableOperations.java
@@ -761,18 +761,7 @@ public class ClickHouseTableOperations extends
JdbcTableOperations {
public void alterTable(String databaseName, String tableName, TableChange...
changes)
throws NoSuchTableException {
validateTableSettingChanges(changes);
- LOG.info("Attempting to alter table {} from database {}", tableName,
databaseName);
- try (Connection connection = getConnection(databaseName)) {
- String sql = generateAlterTableSql(databaseName, tableName, changes);
- if (StringUtils.isEmpty(sql)) {
- LOG.info("No changes to alter table {} from database {}", tableName,
databaseName);
- return;
- }
- JdbcConnectorUtils.executeUpdate(connection, sql);
- LOG.info("Alter table {} from database {}", tableName, databaseName);
- } catch (final SQLException se) {
- throw this.exceptionMapper.toGravitinoException(se);
- }
+ super.alterTable(databaseName, tableName, changes);
}
@Override
@@ -892,6 +881,7 @@ public class ClickHouseTableOperations extends
JdbcTableOperations {
@Override
public JdbcTable load(String databaseName, String tableName) throws
NoSuchTableException {
try (Connection connection = getConnection(databaseName)) {
+ cacheDriverVersion(connection);
ResultSet tables = getTable(connection, databaseName, tableName);
JdbcTable.Builder jdbcTableBuilder = getTableBuilder(tables,
databaseName, tableName);
diff --git
a/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsConnectionUsage.java
b/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsConnectionUsage.java
new file mode 100644
index 0000000000..d9ce7a3188
--- /dev/null
+++
b/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsConnectionUsage.java
@@ -0,0 +1,67 @@
+/*
+ * 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.gravitino.catalog.clickhouse.operations;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import
org.apache.gravitino.catalog.clickhouse.converter.ClickHouseExceptionConverter;
+import org.apache.gravitino.catalog.jdbc.utils.ConnectionCountingDataSource;
+import org.apache.gravitino.rel.TableChange;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+/** Tests that altering a table never holds one pooled connection while
borrowing another. */
+public class TestClickHouseTableOperationsConnectionUsage {
+
+ private static final String ALTER_SQL = "ALTER TABLE `orders` MODIFY COMMENT
'comment'";
+
+ @Test
+ public void testAlterTableGeneratesSqlWithoutHoldingAConnection() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<Integer> borrowedDuringGeneration = new ArrayList<>();
+ ClickHouseTableOperations operations =
+ new ClickHouseTableOperations() {
+ {
+ dataSource = connections.dataSource();
+ exceptionMapper = new ClickHouseExceptionConverter();
+ }
+
+ @Override
+ protected String generateAlterTableSql(
+ String databaseName, String tableName, TableChange... changes) {
+ borrowedDuringGeneration.add(connections.borrowed());
+ // Like loading the original table, borrow a connection while
generating.
+ try (Connection ignored = dataSource.getConnection()) {
+ return ALTER_SQL;
+ } catch (SQLException e) {
+ throw new IllegalStateException(e);
+ }
+ }
+ };
+
+ operations.alterTable("db", "orders",
TableChange.updateComment("comment"));
+
+ Assertions.assertEquals(List.of(0), borrowedDuringGeneration);
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(List.of(ALTER_SQL), connections.executedSql());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+}
diff --git
a/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsUnit.java
b/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsUnit.java
index cc8c1f4c85..1a77ef0a1a 100644
---
a/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsUnit.java
+++
b/catalogs-contrib/catalog-jdbc-clickhouse/src/test/java/org/apache/gravitino/catalog/clickhouse/operations/TestClickHouseTableOperationsUnit.java
@@ -822,6 +822,7 @@ public class TestClickHouseTableOperationsUnit {
JdbcTable table = ops.load("db", "t");
+ Mockito.verify(metadata).getDriverVersion();
Assertions.assertEquals(1, table.columns().length);
Assertions.assertEquals("MergeTree",
table.properties().get(TableConstants.ENGINE));
Assertions.assertEquals("",
table.properties().get(TableConstants.PARTITION_KEY));
diff --git
a/catalogs-contrib/catalog-jdbc-oceanbase/src/main/java/org/apache/gravitino/catalog/oceanbase/operation/OceanBaseTableOperations.java
b/catalogs-contrib/catalog-jdbc-oceanbase/src/main/java/org/apache/gravitino/catalog/oceanbase/operation/OceanBaseTableOperations.java
index bd09770adc..fd1aceeec5 100644
---
a/catalogs-contrib/catalog-jdbc-oceanbase/src/main/java/org/apache/gravitino/catalog/oceanbase/operation/OceanBaseTableOperations.java
+++
b/catalogs-contrib/catalog-jdbc-oceanbase/src/main/java/org/apache/gravitino/catalog/oceanbase/operation/OceanBaseTableOperations.java
@@ -211,19 +211,22 @@ public class OceanBaseTableOperations extends
JdbcTableOperations {
public void alterTable(String databaseName, String tableName, TableChange...
changes)
throws NoSuchTableException {
LOG.info("Attempting to alter table {} from database {}", tableName,
databaseName);
- try (Connection connection = getConnection(databaseName)) {
- for (TableChange change : changes) {
- String sql = generateAlterTableSql(databaseName, tableName, change);
- if (StringUtils.isEmpty(sql)) {
- LOG.info("No changes to alter table {} from database {}", tableName,
databaseName);
- continue;
- }
+ for (TableChange change : changes) {
+ // Generate each statement before borrowing the connection that executes
it, as the
+ // generator may load the table through another connection. Each change
is generated after
+ // the previous one is applied, so it sees the updated table.
+ String sql = generateAlterTableSql(databaseName, tableName, change);
+ if (StringUtils.isEmpty(sql)) {
+ LOG.info("No changes to alter table {} from database {}", tableName,
databaseName);
+ continue;
+ }
+ try (Connection connection = getConnection(databaseName)) {
JdbcConnectorUtils.executeUpdate(connection, sql);
+ } catch (final SQLException se) {
+ throw this.exceptionMapper.toGravitinoException(se);
}
- LOG.info("Alter table {} from database {}", tableName, databaseName);
- } catch (final SQLException se) {
- throw this.exceptionMapper.toGravitinoException(se);
}
+ LOG.info("Alter table {} from database {}", tableName, databaseName);
}
@Override
diff --git
a/catalogs-contrib/catalog-jdbc-oceanbase/src/test/java/org/apache/gravitino/catalog/oceanbase/operation/TestOceanBaseTableOperationsConnectionUsage.java
b/catalogs-contrib/catalog-jdbc-oceanbase/src/test/java/org/apache/gravitino/catalog/oceanbase/operation/TestOceanBaseTableOperationsConnectionUsage.java
new file mode 100644
index 0000000000..839422ba42
--- /dev/null
+++
b/catalogs-contrib/catalog-jdbc-oceanbase/src/test/java/org/apache/gravitino/catalog/oceanbase/operation/TestOceanBaseTableOperationsConnectionUsage.java
@@ -0,0 +1,80 @@
+/*
+ * 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.gravitino.catalog.oceanbase.operation;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.gravitino.catalog.jdbc.converter.JdbcExceptionConverter;
+import org.apache.gravitino.catalog.jdbc.utils.ConnectionCountingDataSource;
+import org.apache.gravitino.rel.TableChange;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+/** Tests that altering a table never holds one pooled connection while
borrowing another. */
+public class TestOceanBaseTableOperationsConnectionUsage {
+
+ @Test
+ public void
testEachChangeIsGeneratedWithoutHoldingAConnectionAfterThePreviousOneIsApplied()
{
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<Integer> borrowedDuringGeneration = new ArrayList<>();
+ List<Integer> executedBeforeGeneration = new ArrayList<>();
+ OceanBaseTableOperations operations =
+ new OceanBaseTableOperations() {
+ {
+ dataSource = connections.dataSource();
+ exceptionMapper = new JdbcExceptionConverter();
+ }
+
+ @Override
+ protected String generateAlterTableSql(
+ String databaseName, String tableName, TableChange... changes) {
+ borrowedDuringGeneration.add(connections.borrowed());
+ executedBeforeGeneration.add(connections.executedSql().size());
+ if (changes[0] instanceof TableChange.SetProperty) {
+ return "";
+ }
+ // Like loading the original table, borrow a connection while
generating.
+ try (Connection ignored = dataSource.getConnection()) {
+ return "ALTER TABLE `orders` "
+ + ((TableChange.UpdateComment) changes[0]).getNewComment();
+ } catch (SQLException e) {
+ throw new IllegalStateException(e);
+ }
+ }
+ };
+
+ operations.alterTable(
+ "db",
+ "orders",
+ TableChange.updateComment("first"),
+ TableChange.setProperty("key", "value"),
+ TableChange.updateComment("second"));
+
+ Assertions.assertEquals(List.of(0, 0, 0), borrowedDuringGeneration);
+ // Each change is generated after the previous statement ran, so it sees
the updated table.
+ Assertions.assertEquals(List.of(0, 1, 1), executedBeforeGeneration);
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(
+ List.of("ALTER TABLE `orders` first", "ALTER TABLE `orders` second"),
+ connections.executedSql());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+}
diff --git
a/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java
b/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java
index 7dd96aa444..cf6f89d95a 100644
---
a/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java
+++
b/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java
@@ -131,8 +131,11 @@ public abstract class JdbcDatabaseOperations implements
DatabaseOperation {
}
protected void dropDatabase(String databaseName, boolean cascade) {
+ // Generate the SQL before borrowing the connection that executes it, as a
non-cascading drop
+ // checks for tables through another connection.
+ String sql = generateDropDatabaseSql(databaseName, cascade);
try (final Connection connection = getConnection()) {
- JdbcConnectorUtils.executeUpdate(connection,
generateDropDatabaseSql(databaseName, cascade));
+ JdbcConnectorUtils.executeUpdate(connection, sql);
} catch (final SQLException se) {
throw this.exceptionMapper.toGravitinoException(se);
}
@@ -161,6 +164,9 @@ public abstract class JdbcDatabaseOperations implements
DatabaseOperation {
* The default implementation of this method is based on MySQL syntax, and
if the catalog does not
* support MySQL syntax, this method needs to be rewritten.
*
+ * <p>A non-cascading drop checks for tables through a connection from
{@link #dataSource}, so
+ * callers must not hold one while calling this method.
+ *
* @param databaseName The name of the database.
* @param cascade cascade If set to true, drops all the tables in the schema
as well.
* @return the SQL statement to drop a database with the given name.
diff --git
a/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcTableOperations.java
b/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcTableOperations.java
index 516b3c01b9..00cf4a9111 100644
---
a/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcTableOperations.java
+++
b/catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcTableOperations.java
@@ -85,6 +85,16 @@ public abstract class JdbcTableOperations implements
TableOperation {
protected JdbcColumnDefaultValueConverter columnDefaultValueConverter;
+ // The driver version is fixed for a data source. It is cached so that
column parsing in load,
+ // which holds a connection, does not borrow another one to read it. A
driver may report a null
+ // version, so a separate flag records that it was read; written after the
version, it publishes
+ // it.
+ private String driverVersion;
+ private volatile boolean driverVersionCached;
+ // Even a failed attempt must prevent a fallback borrow while load holds its
connection.
+ // cacheDriverVersion retries failures using the connection held by a later
load.
+ private volatile boolean driverVersionAttempted;
+
@Override
public void initialize(
DataSource dataSource,
@@ -93,6 +103,9 @@ public abstract class JdbcTableOperations implements
TableOperation {
JdbcColumnDefaultValueConverter jdbcColumnDefaultValueConverter,
Map<String, String> conf) {
this.dataSource = dataSource;
+ driverVersion = null;
+ driverVersionCached = false;
+ driverVersionAttempted = false;
this.exceptionMapper = exceptionMapper;
this.typeConverter = jdbcTypeConverter;
this.columnDefaultValueConverter = jdbcColumnDefaultValueConverter;
@@ -142,18 +155,19 @@ public abstract class JdbcTableOperations implements
TableOperation {
SortOrder[] sortOrders)
throws TableAlreadyExistsException {
LOG.info("Attempting to create table {} in database {}", tableName,
databaseName);
+ // Generate the SQL before borrowing the connection that executes it, as
generators may query
+ // the server through another connection (for example, Doris checks its
backends and version).
+ String sql =
+ generateCreateTableSql(
+ tableName,
+ columns,
+ comment,
+ properties,
+ partitioning,
+ distribution,
+ indexes,
+ sortOrders);
try (Connection connection = getConnection(databaseName)) {
- String sql =
- generateCreateTableSql(
- tableName,
- columns,
- comment,
- properties,
- partitioning,
- distribution,
- indexes,
- sortOrders);
-
JdbcConnectorUtils.executeUpdate(connection, sql);
LOG.info("Created table {} in database {} with SQL:\n{}", tableName,
databaseName, sql);
} catch (final SQLException se) {
@@ -248,6 +262,7 @@ public abstract class JdbcTableOperations implements
TableOperation {
// 2. MySQL treats 'a_b' as a wildcard, matching any table name that
begins with 'a', followed
// by any character, and ending with 'b'.
try (Connection connection = getConnection(databaseName)) {
+ cacheDriverVersion(connection);
// 1. Get table information, The result of tables may be more than one
due to the reason
// above, so we need to check the result.
ResultSet tables = getTable(connection, databaseName, tableName);
@@ -353,12 +368,15 @@ public abstract class JdbcTableOperations implements
TableOperation {
public void alterTable(String databaseName, String tableName, TableChange...
changes)
throws NoSuchTableException {
LOG.info("Attempting to alter table {} from database {}", tableName,
databaseName);
+ // Generate the SQL before borrowing the connection that executes it.
Generators may load the
+ // table, which borrows another connection; holding both lets concurrent
alters on distinct
+ // tables exhaust the pool and wait for each other until the borrow
timeout.
+ String sql = generateAlterTableSql(databaseName, tableName, changes);
+ if (StringUtils.isEmpty(sql)) {
+ LOG.info("No changes to alter table {} from database {}", tableName,
databaseName);
+ return;
+ }
try (Connection connection = getConnection(databaseName)) {
- String sql = generateAlterTableSql(databaseName, tableName, changes);
- if (StringUtils.isEmpty(sql)) {
- LOG.info("No changes to alter table {} from database {}", tableName,
databaseName);
- return;
- }
JdbcConnectorUtils.executeUpdate(connection, sql);
LOG.info("Alter table {} from database {}", tableName, databaseName);
} catch (final SQLException se) {
@@ -511,6 +529,21 @@ public abstract class JdbcTableOperations implements
TableOperation {
return fieldNames.stream().map(colName -> new String[]
{colName}).toArray(String[][]::new);
}
+ /**
+ * Generates the SQL statement that creates a table without sort orders.
+ *
+ * <p>Implementations may borrow connections from {@link #dataSource}, so
callers must not hold
+ * one while calling this method.
+ *
+ * @param tableName the name of the table
+ * @param columns the columns of the table
+ * @param comment the comment of the table
+ * @param properties the properties of the table
+ * @param partitioning the partitioning of the table
+ * @param distribution the distribution of the table
+ * @param indexes the indexes of the table
+ * @return the SQL statement to create the table
+ */
protected abstract String generateCreateTableSql(
String tableName,
JdbcColumn[] columns,
@@ -520,6 +553,22 @@ public abstract class JdbcTableOperations implements
TableOperation {
Distribution distribution,
Index[] indexes);
+ /**
+ * Generates the SQL statement that creates a table.
+ *
+ * <p>Implementations may borrow connections from {@link #dataSource}, so
callers must not hold
+ * one while calling this method.
+ *
+ * @param tableName the name of the table
+ * @param columns the columns of the table
+ * @param comment the comment of the table
+ * @param properties the properties of the table
+ * @param partitioning the partitioning of the table
+ * @param distribution the distribution of the table
+ * @param indexes the indexes of the table
+ * @param sortOrders the sort orders of the table
+ * @return the SQL statement to create the table
+ */
protected String generateCreateTableSql(
String tableName,
JdbcColumn[] columns,
@@ -567,6 +616,17 @@ public abstract class JdbcTableOperations implements
TableOperation {
return generatePurgeTableSql(tableName);
}
+ /**
+ * Generates the SQL statement that applies the given changes to a table.
+ *
+ * <p>Implementations may load the table, which borrows a connection from
{@link #dataSource}, so
+ * callers must not hold one while calling this method.
+ *
+ * @param databaseName the name of the database
+ * @param tableName the name of the table
+ * @param changes the changes to apply
+ * @return the SQL statement, or null or an empty string if nothing needs to
be executed
+ */
protected abstract String generateAlterTableSql(
String databaseName, String tableName, TableChange... changes);
@@ -754,21 +814,49 @@ public abstract class JdbcTableOperations implements
TableOperation {
}
/**
- * Get MySQL driver version from DatabaseMetaData
+ * Get MySQL driver version from DatabaseMetaData.
+ *
+ * <p>The version is cached once read. {@link #load} caches it from the
connection it holds before
+ * parsing columns (see {@link #cacheDriverVersion}), so this method borrows
a connection only
+ * before the first version-read attempt. Failed reads return null until a
later load retries
+ * using its own connection, avoiding a nested borrow during column parsing.
*
* @return the driver version string, or null if not available
*/
protected String getMySQLDriverVersion() {
+ if (driverVersionCached || driverVersionAttempted || dataSource == null) {
+ return driverVersion;
+ }
+ try (Connection connection = dataSource.getConnection()) {
+ cacheDriverVersion(connection);
+ } catch (SQLException e) {
+ LOG.debug("Failed to get driver version", e);
+ }
+ return driverVersion;
+ }
+
+ /**
+ * Caches the driver version from a connection the caller already holds.
+ *
+ * <p>Overrides of {@link #load} that parse columns must call this first, so
that {@link
+ * #getMySQLDriverVersion()} does not borrow another connection while theirs
is held.
+ *
+ * @param connection the connection held by the caller
+ */
+ protected void cacheDriverVersion(Connection connection) {
+ if (driverVersionCached) {
+ return;
+ }
try {
- if (dataSource != null) {
- try (Connection connection = dataSource.getConnection()) {
- return connection.getMetaData().getDriverVersion();
- }
- }
+ driverVersion = connection.getMetaData().getDriverVersion();
+ driverVersionCached = true;
} catch (SQLException e) {
+ // A later load retries on its held connection; column parsing must not
retry by borrowing
+ // another connection. Until recovery, preserve the existing
null-version precision fallback.
LOG.debug("Failed to get driver version", e);
+ } finally {
+ driverVersionAttempted = true;
}
- return null;
}
/**
diff --git
a/catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/operation/TestJdbcOperationsConnectionUsage.java
b/catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/operation/TestJdbcOperationsConnectionUsage.java
new file mode 100644
index 0000000000..f8d70176c0
--- /dev/null
+++
b/catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/operation/TestJdbcOperationsConnectionUsage.java
@@ -0,0 +1,480 @@
+/*
+ * 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.gravitino.catalog.jdbc.operation;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Proxy;
+import java.nio.file.Files;
+import java.sql.Connection;
+import java.sql.DatabaseMetaData;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import javax.sql.DataSource;
+import org.apache.commons.dbcp2.BasicDataSource;
+import org.apache.commons.io.FileUtils;
+import org.apache.gravitino.catalog.jdbc.JdbcColumn;
+import org.apache.gravitino.catalog.jdbc.JdbcTable;
+import org.apache.gravitino.catalog.jdbc.config.JdbcConfig;
+import
org.apache.gravitino.catalog.jdbc.converter.SqliteColumnDefaultValueConverter;
+import org.apache.gravitino.catalog.jdbc.converter.SqliteExceptionConverter;
+import org.apache.gravitino.catalog.jdbc.converter.SqliteTypeConverter;
+import org.apache.gravitino.catalog.jdbc.utils.ConnectionCountingDataSource;
+import org.apache.gravitino.catalog.jdbc.utils.DataSourceUtils;
+import org.apache.gravitino.exceptions.NoSuchTableException;
+import org.apache.gravitino.rel.Column;
+import org.apache.gravitino.rel.TableChange;
+import org.apache.gravitino.rel.expressions.distributions.Distribution;
+import org.apache.gravitino.rel.expressions.distributions.Distributions;
+import org.apache.gravitino.rel.expressions.sorts.SortOrder;
+import org.apache.gravitino.rel.expressions.transforms.Transform;
+import org.apache.gravitino.rel.expressions.transforms.Transforms;
+import org.apache.gravitino.rel.indexes.Index;
+import org.apache.gravitino.rel.indexes.Indexes;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+/** Tests that JDBC operations never hold one pooled connection while
borrowing another. */
+public class TestJdbcOperationsConnectionUsage {
+
+ private static final String DATABASE = "test";
+ private static final String TABLE = "orders";
+ private static final String ALTER_SQL = "ALTER TABLE orders ADD COLUMN note
INTEGER";
+
+ @Test
+ public void testAlterTableGeneratesSqlWithoutHoldingAConnection() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<Integer> borrowedDuringGeneration = new ArrayList<>();
+ JdbcTableOperations operations =
+ borrowingGenerator(connections, borrowedDuringGeneration, ALTER_SQL);
+
+ operations.alterTable(DATABASE, TABLE,
TableChange.updateComment("comment"));
+
+ Assertions.assertEquals(List.of(0), borrowedDuringGeneration);
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(List.of(ALTER_SQL), connections.executedSql());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void testAlterTableWithoutChangesExecutesNothing() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<Integer> borrowedDuringGeneration = new ArrayList<>();
+ JdbcTableOperations operations = borrowingGenerator(connections,
borrowedDuringGeneration, "");
+
+ operations.alterTable(DATABASE, TABLE,
TableChange.updateComment("comment"));
+
+ Assertions.assertEquals(List.of(0), borrowedDuringGeneration);
+ Assertions.assertTrue(connections.executedSql().isEmpty());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void testAlterTableGeneratorFailureBorrowsNoExecutionConnection() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ JdbcTableOperations operations =
+ new SqliteTableOperations() {
+ @Override
+ protected String generateAlterTableSql(
+ String databaseName, String tableName, TableChange... changes) {
+ throw new NoSuchTableException("Table %s does not exist",
tableName);
+ }
+ };
+ initialize(operations, connections.dataSource());
+
+ NoSuchTableException exception =
+ Assertions.assertThrows(
+ NoSuchTableException.class,
+ () -> operations.alterTable(DATABASE, TABLE,
TableChange.updateComment("comment")));
+
+ Assertions.assertTrue(exception.getMessage().contains(TABLE));
+ Assertions.assertEquals(0, connections.peakBorrowed());
+ Assertions.assertTrue(connections.executedSql().isEmpty());
+ }
+
+ @Test
+ public void testCreateTableGeneratesSqlWithoutHoldingAConnection() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<Integer> borrowedDuringGeneration = new ArrayList<>();
+ String createSql = "CREATE TABLE orders (id INTEGER)";
+ JdbcTableOperations operations =
+ new SqliteTableOperations() {
+ @Override
+ protected String generateCreateTableSql(
+ String tableName,
+ JdbcColumn[] columns,
+ String comment,
+ Map<String, String> properties,
+ Transform[] partitioning,
+ Distribution distribution,
+ Index[] indexes,
+ SortOrder[] sortOrders) {
+ borrowedDuringGeneration.add(connections.borrowed());
+ // Like Doris checking its backends and version, query the server
while generating.
+ borrowAndRelease(this.dataSource);
+ return createSql;
+ }
+ };
+ initialize(operations, connections.dataSource());
+
+ operations.create(
+ DATABASE,
+ TABLE,
+ new JdbcColumn[0],
+ null,
+ Collections.emptyMap(),
+ Transforms.EMPTY_TRANSFORM,
+ Distributions.NONE,
+ Indexes.EMPTY_INDEXES);
+
+ Assertions.assertEquals(List.of(0), borrowedDuringGeneration);
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(List.of(createSql), connections.executedSql());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void testDropDatabaseGeneratesSqlWithoutHoldingAConnection() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<Integer> borrowedDuringGeneration = new ArrayList<>();
+ String dropSql = "DROP DATABASE `test`";
+ JdbcDatabaseOperations operations =
+ new SqliteDatabaseOperations("unused") {
+ @Override
+ public String generateDropDatabaseSql(String databaseName, boolean
cascade) {
+ borrowedDuringGeneration.add(connections.borrowed());
+ // Like the default non-cascading drop checking for tables, borrow
while generating.
+ borrowAndRelease(this.dataSource);
+ return dropSql;
+ }
+ };
+ operations.initialize(
+ connections.dataSource(), new SqliteExceptionConverter(),
Collections.emptyMap());
+
+ operations.dropDatabase(DATABASE, false);
+
+ Assertions.assertEquals(List.of(0), borrowedDuringGeneration);
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(List.of(dropSql), connections.executedSql());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void
testAlterTableSucceedsWhenGeneratorLoadsTableFromSingleConnectionPool()
+ throws IOException, SQLException {
+ // The pool size reproduces the starvation deterministically: with the
execution connection
+ // borrowed first, the generator's load would wait for a connection that
can never be returned.
+ File directory =
Files.createTempDirectory("gravitino-jdbc-alter").toFile();
+ BasicDataSource dataSource = null;
+ try {
+ dataSource = singleConnectionPool(new File(directory,
DATABASE).getPath());
+ try (Connection connection = dataSource.getConnection();
+ Statement statement = connection.createStatement()) {
+ statement.executeUpdate("CREATE TABLE orders (id INTEGER)");
+ }
+ JdbcTableOperations operations = loadingGenerator(dataSource);
+
+ operations.alterTable(DATABASE, TABLE,
TableChange.updateComment("comment"));
+
+ JdbcTable table = operations.load(DATABASE, TABLE);
+ Assertions.assertArrayEquals(
+ new String[] {"id", "note"},
+
Arrays.stream(table.columns()).map(Column::name).toArray(String[]::new));
+ } finally {
+ if (dataSource != null) {
+ dataSource.close();
+ }
+ FileUtils.deleteQuietly(directory);
+ }
+ }
+
+ @Test
+ public void testLoadReadsDriverVersionWithoutBorrowingAnotherConnection()
+ throws IOException, SQLException {
+ // MySQL, OceanBase and Doris read the driver version while parsing each
datetime column, which
+ // happens while load holds its connection.
+ File directory = Files.createTempDirectory("gravitino-jdbc-load").toFile();
+ BasicDataSource dataSource = null;
+ try {
+ dataSource = singleConnectionPool(new File(directory,
DATABASE).getPath());
+ try (Connection connection = dataSource.getConnection();
+ Statement statement = connection.createStatement()) {
+ statement.executeUpdate("CREATE TABLE orders (id INTEGER, created
TEXT)");
+ }
+ List<String> driverVersions = new ArrayList<>();
+ JdbcTableOperations operations =
+ new SqliteTableOperations() {
+ @Override
+ public Integer calculateDatetimePrecision(String typeName, int
columnSize, int scale) {
+ driverVersions.add(getMySQLDriverVersion());
+ return null;
+ }
+ };
+ initialize(operations, dataSource);
+
+ operations.load(DATABASE, TABLE);
+
+ Assertions.assertEquals(2, driverVersions.size());
+ Assertions.assertNotNull(driverVersions.get(0));
+ Assertions.assertEquals(driverVersions.get(0), driverVersions.get(1));
+ } finally {
+ if (dataSource != null) {
+ dataSource.close();
+ }
+ FileUtils.deleteQuietly(directory);
+ }
+ }
+
+ @Test
+ public void testNullDriverVersionIsCachedFromTheHeldConnection() throws
SQLException {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ JdbcTableOperations operations = new SqliteTableOperations();
+ initialize(operations, connections.dataSource());
+
+ try (Connection held = connections.dataSource().getConnection()) {
+ operations.cacheDriverVersion(held);
+ // A null version is still cached, so later lookups never borrow while
one is held.
+ Assertions.assertNull(operations.getMySQLDriverVersion());
+ Assertions.assertNull(operations.getMySQLDriverVersion());
+ }
+
+ Assertions.assertEquals(1, connections.totalBorrows());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void testLoadWithNullVersionUsesOneCountedConnection() throws
SQLException {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ List<String> versions = new ArrayList<>();
+ JdbcTableOperations operations = countedLoadOperations(connections,
versions);
+
+ JdbcTable table = operations.load(DATABASE, TABLE);
+
+ Assertions.assertEquals(2, table.columns().length);
+ Assertions.assertEquals(Arrays.asList(null, null), versions);
+ Assertions.assertEquals(1, connections.totalBorrows());
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void
testFailedVersionReadDoesNotBorrowDuringLoadAndRetriesOnNextLoad()
+ throws SQLException {
+ AtomicInteger reads = new AtomicInteger();
+ DatabaseMetaData metadata =
+ metadata(
+ () -> {
+ if (reads.incrementAndGet() == 1) {
+ throw new SQLException("version unavailable");
+ }
+ return "driver-1";
+ });
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource(metadata);
+ List<String> versions = new ArrayList<>();
+ JdbcTableOperations operations = countedLoadOperations(connections,
versions);
+
+ Assertions.assertEquals(2, operations.load(DATABASE,
TABLE).columns().length);
+ Assertions.assertEquals(Arrays.asList(null, null), versions);
+ Assertions.assertEquals(1, connections.totalBorrows());
+ Assertions.assertEquals(2, operations.load(DATABASE,
TABLE).columns().length);
+
+ Assertions.assertEquals(Arrays.asList(null, null, "driver-1", "driver-1"),
versions);
+ Assertions.assertEquals(2, reads.get());
+ Assertions.assertEquals(2, connections.totalBorrows());
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void testDriverVersionFallbackBorrowsOnceAndCachesNull() {
+ ConnectionCountingDataSource connections = new
ConnectionCountingDataSource();
+ JdbcTableOperations operations = new SqliteTableOperations();
+ initialize(operations, connections.dataSource());
+
+ Assertions.assertNull(operations.getMySQLDriverVersion());
+ Assertions.assertNull(operations.getMySQLDriverVersion());
+
+ Assertions.assertEquals(1, connections.totalBorrows());
+ Assertions.assertEquals(1, connections.peakBorrowed());
+ Assertions.assertEquals(0, connections.borrowed());
+ }
+
+ @Test
+ public void testInitializeClearsPreviouslyCachedVersion() throws
SQLException {
+ JdbcTableOperations operations = new SqliteTableOperations();
+ initialize(operations, new ConnectionCountingDataSource().dataSource());
+ Assertions.assertNull(operations.getMySQLDriverVersion());
+ DatabaseMetaData metadata = metadata(() -> "new-driver");
+ ConnectionCountingDataSource replacement = new
ConnectionCountingDataSource(metadata);
+
+ initialize(operations, replacement.dataSource());
+
+ Assertions.assertEquals("new-driver", operations.getMySQLDriverVersion());
+ Assertions.assertEquals(1, replacement.totalBorrows());
+ Assertions.assertEquals(0, replacement.borrowed());
+ }
+
+ private static JdbcTableOperations countedLoadOperations(
+ ConnectionCountingDataSource connections, List<String> versions) {
+ JdbcTableOperations operations =
+ new SqliteTableOperations() {
+ @Override
+ protected ResultSet getTable(Connection connection, String
databaseName, String tableName)
+ throws SQLException {
+ return resultSet(Collections.singletonList(Map.of("TABLE_NAME",
tableName)));
+ }
+
+ @Override
+ protected ResultSet getColumns(
+ Connection connection, String databaseName, String tableName)
throws SQLException {
+ return resultSet(
+ Arrays.asList(
+ Map.of("TABLE_NAME", tableName, "COLUMN_NAME", "id",
"TYPE_NAME", "TEXT"),
+ Map.of(
+ "TABLE_NAME", tableName, "COLUMN_NAME", "created",
"TYPE_NAME", "TEXT")));
+ }
+
+ @Override
+ protected List<Index> getIndexes(
+ Connection connection, String databaseName, String tableName) {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public Integer calculateDatetimePrecision(String typeName, int
columnSize, int scale) {
+ Assertions.assertEquals(1, connections.borrowed());
+ versions.add(getMySQLDriverVersion());
+ return null;
+ }
+ };
+ initialize(operations, connections.dataSource());
+ return operations;
+ }
+
+ private interface VersionReader {
+ String read() throws SQLException;
+ }
+
+ private static DatabaseMetaData metadata(VersionReader reader) {
+ return (DatabaseMetaData)
+ Proxy.newProxyInstance(
+ DatabaseMetaData.class.getClassLoader(),
+ new Class<?>[] {DatabaseMetaData.class},
+ (proxy, method, args) -> {
+ if ("getDriverVersion".equals(method.getName())) {
+ return reader.read();
+ }
+ throw new UnsupportedOperationException(method.getName());
+ });
+ }
+
+ private static ResultSet resultSet(List<Map<String, String>> rows) {
+ AtomicInteger position = new AtomicInteger(-1);
+ return (ResultSet)
+ Proxy.newProxyInstance(
+ ResultSet.class.getClassLoader(),
+ new Class<?>[] {ResultSet.class},
+ (proxy, method, args) -> {
+ switch (method.getName()) {
+ case "next":
+ return position.incrementAndGet() < rows.size();
+ case "close":
+ return null;
+ case "getString":
+ return rows.get(position.get()).get(args[0]);
+ case "getInt":
+ return 0;
+ case "getBoolean":
+ return false;
+ default:
+ throw new UnsupportedOperationException(method.getName());
+ }
+ });
+ }
+
+ // Like the MySQL generator loading the original table, borrows a connection
while generating.
+ private static JdbcTableOperations borrowingGenerator(
+ ConnectionCountingDataSource connections,
+ List<Integer> borrowedDuringGeneration,
+ String sql) {
+ JdbcTableOperations operations =
+ new SqliteTableOperations() {
+ @Override
+ protected String generateAlterTableSql(
+ String databaseName, String tableName, TableChange... changes) {
+ borrowedDuringGeneration.add(connections.borrowed());
+ borrowAndRelease(this.dataSource);
+ return sql;
+ }
+ };
+ initialize(operations, connections.dataSource());
+ return operations;
+ }
+
+ private static JdbcTableOperations loadingGenerator(DataSource dataSource) {
+ JdbcTableOperations operations =
+ new SqliteTableOperations() {
+ @Override
+ protected String generateAlterTableSql(
+ String databaseName, String tableName, TableChange... changes) {
+ load(databaseName, tableName);
+ return ALTER_SQL;
+ }
+ };
+ initialize(operations, dataSource);
+ return operations;
+ }
+
+ private static void borrowAndRelease(DataSource dataSource) {
+ try (Connection ignored = dataSource.getConnection()) {
+ // Only the borrow matters.
+ } catch (SQLException e) {
+ throw new IllegalStateException(e);
+ }
+ }
+
+ private static void initialize(JdbcTableOperations operations, DataSource
dataSource) {
+ operations.initialize(
+ dataSource,
+ new SqliteExceptionConverter(),
+ new SqliteTypeConverter(),
+ new SqliteColumnDefaultValueConverter(),
+ Collections.emptyMap());
+ }
+
+ private static BasicDataSource singleConnectionPool(String path) {
+ Map<String, String> properties = new HashMap<>();
+ properties.put(JdbcConfig.JDBC_DRIVER.getKey(), "org.sqlite.JDBC");
+ properties.put(JdbcConfig.JDBC_URL.getKey(), "jdbc:sqlite:" + path);
+ properties.put(JdbcConfig.USERNAME.getKey(), "test");
+ properties.put(JdbcConfig.PASSWORD.getKey(), "test");
+ properties.put(JdbcConfig.POOL_MIN_SIZE.getKey(), "1");
+ properties.put(JdbcConfig.POOL_MAX_SIZE.getKey(), "1");
+ properties.put(JdbcConfig.POOL_MAX_WAIT_MS.getKey(), "1000");
+ return (BasicDataSource) DataSourceUtils.createDataSource(properties);
+ }
+}
diff --git
a/catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/utils/ConnectionCountingDataSource.java
b/catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/utils/ConnectionCountingDataSource.java
new file mode 100644
index 0000000000..e2a947745b
--- /dev/null
+++
b/catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/utils/ConnectionCountingDataSource.java
@@ -0,0 +1,168 @@
+/*
+ * 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.gravitino.catalog.jdbc.utils;
+
+import java.lang.reflect.Proxy;
+import java.sql.Connection;
+import java.sql.DatabaseMetaData;
+import java.sql.Statement;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import javax.sql.DataSource;
+
+/**
+ * A test {@link DataSource} whose connections only record the SQL they
execute. It tracks how many
+ * connections are borrowed at once, so tests can assert that an operation
never holds one
+ * connection while borrowing another.
+ */
+public class ConnectionCountingDataSource {
+
+ private final DatabaseMetaData metadata;
+
+ private final AtomicInteger borrowed = new AtomicInteger();
+ private final AtomicInteger peakBorrowed = new AtomicInteger();
+ private final AtomicInteger totalBorrows = new AtomicInteger();
+ private final List<String> executedSql = new CopyOnWriteArrayList<>();
+ private final DataSource dataSource =
+ (DataSource)
+ Proxy.newProxyInstance(
+ DataSource.class.getClassLoader(),
+ new Class<?>[] {DataSource.class},
+ (proxy, method, args) -> {
+ if ("getConnection".equals(method.getName())) {
+ return borrow();
+ }
+ return objectMethod(proxy, method.getName(), args,
"DataSource");
+ });
+
+ /** Creates a counting data source whose driver reports a null version. */
+ public ConnectionCountingDataSource() {
+ this(metaData());
+ }
+
+ /**
+ * Creates a counting data source with the supplied JDBC metadata.
+ *
+ * @param metadata the metadata returned by each borrowed connection
+ */
+ public ConnectionCountingDataSource(DatabaseMetaData metadata) {
+ this.metadata = metadata;
+ }
+
+ /** Returns the data source to pass to the operations under test. */
+ public DataSource dataSource() {
+ return dataSource;
+ }
+
+ /** Returns the number of connections currently borrowed and not yet closed.
*/
+ public int borrowed() {
+ return borrowed.get();
+ }
+
+ /** Returns the largest number of connections that were borrowed at the same
time. */
+ public int peakBorrowed() {
+ return peakBorrowed.get();
+ }
+
+ /** Returns how many connections were borrowed in total. */
+ public int totalBorrows() {
+ return totalBorrows.get();
+ }
+
+ /** Returns the update statements executed so far, in order. */
+ public List<String> executedSql() {
+ return executedSql;
+ }
+
+ private Connection borrow() {
+ totalBorrows.incrementAndGet();
+ peakBorrowed.accumulateAndGet(borrowed.incrementAndGet(), Math::max);
+ AtomicBoolean closed = new AtomicBoolean();
+ return (Connection)
+ Proxy.newProxyInstance(
+ Connection.class.getClassLoader(),
+ new Class<?>[] {Connection.class},
+ (proxy, method, args) -> {
+ switch (method.getName()) {
+ case "close":
+ if (closed.compareAndSet(false, true)) {
+ borrowed.decrementAndGet();
+ }
+ return null;
+ case "setCatalog":
+ case "setSchema":
+ return null;
+ case "createStatement":
+ return statement();
+ case "getMetaData":
+ return metadata;
+ default:
+ return objectMethod(proxy, method.getName(), args,
"Connection");
+ }
+ });
+ }
+
+ // Answers Object methods so the proxies can be printed or compared, and
rejects anything else.
+ private static Object objectMethod(Object proxy, String name, Object[] args,
String type) {
+ switch (name) {
+ case "toString":
+ return "Counting" + type + "@" +
Integer.toHexString(System.identityHashCode(proxy));
+ case "hashCode":
+ return System.identityHashCode(proxy);
+ case "equals":
+ return proxy == args[0];
+ default:
+ throw new UnsupportedOperationException(type + "." + name);
+ }
+ }
+
+ // Reports a null driver version, which the JDBC contract allows.
+ private static DatabaseMetaData metaData() {
+ return (DatabaseMetaData)
+ Proxy.newProxyInstance(
+ DatabaseMetaData.class.getClassLoader(),
+ new Class<?>[] {DatabaseMetaData.class},
+ (proxy, method, args) -> {
+ if ("getDriverVersion".equals(method.getName())) {
+ return null;
+ }
+ return objectMethod(proxy, method.getName(), args,
"DatabaseMetaData");
+ });
+ }
+
+ private Statement statement() {
+ return (Statement)
+ Proxy.newProxyInstance(
+ Statement.class.getClassLoader(),
+ new Class<?>[] {Statement.class},
+ (proxy, method, args) -> {
+ switch (method.getName()) {
+ case "executeUpdate":
+ executedSql.add((String) args[0]);
+ return 0;
+ case "close":
+ return null;
+ default:
+ return objectMethod(proxy, method.getName(), args,
"Statement");
+ }
+ });
+ }
+}
