diff --git a/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/main/java/ai/chat2db/plugin/mysql/enums/type/MysqlColumnTypeEnum.java b/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/main/java/ai/chat2db/plugin/mysql/enums/type/MysqlColumnTypeEnum.java index e79753f74..fa9f3efb9 100644 --- a/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/main/java/ai/chat2db/plugin/mysql/enums/type/MysqlColumnTypeEnum.java +++ b/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/main/java/ai/chat2db/plugin/mysql/enums/type/MysqlColumnTypeEnum.java @@ -299,7 +299,7 @@ private String buildExt(TableColumn column, MysqlColumnTypeEnum type) { if (!type.columnType.isSupportExtent() || StringUtils.isEmpty(column.getExtent())) { return ""; } - return column.getComment(); + return column.getExtent(); } private String buildDefaultValue(TableColumn column, MysqlColumnTypeEnum type) { diff --git a/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/test/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilderTest.java b/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/test/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilderTest.java index f275597a1..e75d3ad9d 100644 --- a/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/test/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilderTest.java +++ b/chat2db-community-server/chat2db-community-plugins/chat2db-community-mysql/src/test/java/ai/chat2db/plugin/mysql/builder/MysqlSqlBuilderTest.java @@ -14,6 +14,7 @@ import ai.chat2db.community.domain.api.model.result.ResultOperation; import ai.chat2db.spi.model.datasource.ConnectInfo; import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.plugin.mysql.enums.type.MysqlColumnTypeEnum; import org.junit.jupiter.api.Test; import java.util.List; @@ -24,6 +25,20 @@ class MysqlSqlBuilderTest { + @Test + void shouldUseEnumExtentInsteadOfColumnComment() { + TableColumn column = TableColumn.builder() + .name("status") + .columnType("ENUM") + .extent("('draft','published')") + .comment("workflow status") + .build(); + + String sql = MysqlColumnTypeEnum.ENUM.buildCreateColumnSql(column); + + assertTrue(sql.contains("('draft','published')"), sql); + } + @Test void shouldBuildCreateDatabaseWithCharsetAndCollation() { MysqlSqlBuilder builder = new MysqlSqlBuilder(); diff --git a/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilder.java b/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilder.java index 4354d33ea..a04048f6a 100644 --- a/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilder.java +++ b/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilder.java @@ -294,7 +294,7 @@ protected void buildTableName(String databaseName, String schemaName, String tab if (StringUtils.isNotBlank(databaseName)) { script.append(SQLConstants.OPEN_SQUARE_BRACKET).append(databaseName).append(SQLConstants.CLOSE_SQUARE_BRACKET).append('.'); } - if (StringUtils.isNotBlank(databaseName)) { + if (StringUtils.isNotBlank(schemaName)) { script.append(SQLConstants.OPEN_SQUARE_BRACKET).append(schemaName).append(SQLConstants.CLOSE_SQUARE_BRACKET).append('.'); } if (!(tableName.startsWith(SQLConstants.OPEN_SQUARE_BRACKET) && tableName.endsWith(SQLConstants.CLOSE_SQUARE_BRACKET))) { diff --git a/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/test/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilderTest.java b/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/test/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilderTest.java index 5e151af20..aff415ef3 100644 --- a/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/test/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilderTest.java +++ b/chat2db-community-server/chat2db-community-plugins/chat2db-community-sqlserver/src/test/java/ai/chat2db/plugin/sqlserver/builder/SqlServerSqlBuilderTest.java @@ -29,6 +29,14 @@ void shouldUseTopWhenLimitingSingleRowDeleteAndUpdate() { builder.appendSingleRowLimit("UPDATE", "[t]", where, "UPDATE [t] set [a] = 1" + where)); } + @Test + void shouldIncludeSchemaWhenDatabaseNameIsBlank() { + ExposedSqlServerSqlBuilder builder = new ExposedSqlServerSqlBuilder(); + + assertEquals("[dbo].[orders]", builder.tableName(null, "dbo", "orders")); + assertEquals("[analytics].[dbo].[orders]", builder.tableName("analytics", "dbo", "orders")); + } + @Test void shouldQuoteAndEscapeQualifiedViewName() { SqlServerSqlBuilder builder = new SqlServerSqlBuilder(); @@ -45,4 +53,12 @@ void shouldQuoteAndEscapeQualifiedViewName() { + "'order] schema', 'VIEW', 'select] view'", builder.buildCreateView(view)); } + + private static final class ExposedSqlServerSqlBuilder extends SqlServerSqlBuilder { + private String tableName(String databaseName, String schemaName, String tableName) { + StringBuilder script = new StringBuilder(); + buildTableName(databaseName, schemaName, tableName, script); + return script.toString(); + } + } } diff --git a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultDBManager.java b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultDBManager.java index 5d01fba66..4f992232c 100644 --- a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultDBManager.java +++ b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultDBManager.java @@ -126,12 +126,14 @@ public Connection getConnection(ConnectInfo connectInfo) { try { connection.setCatalog(connectInfo.getDatabaseName()); } catch (SQLException e) { + log.warn("Failed to set catalog to '{}': {}", connectInfo.getDatabaseName(), e.getMessage()); } } if (StringUtils.isNotBlank(connectInfo.getSchemaName())) { try { connection.setSchema(connectInfo.getSchemaName()); } catch (SQLException e) { + log.warn("Failed to set schema to '{}': {}", connectInfo.getSchemaName(), e.getMessage()); } } } diff --git a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultSQLExecutor.java b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultSQLExecutor.java index 85e916bff..92a55939e 100644 --- a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultSQLExecutor.java +++ b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/DefaultSQLExecutor.java @@ -46,6 +46,7 @@ import java.sql.*; import java.util.*; +import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Function; @@ -57,6 +58,8 @@ public class DefaultSQLExecutor implements ICommandExecutor { private static final int STREAMING_ROW_BATCH_SIZE = 200; + private static final Executor DIRECT_ABORT_EXECUTOR = Runnable::run; + private static final DefaultSQLExecutor INSTANCE = new DefaultSQLExecutor(); @@ -280,9 +283,10 @@ public Long count(String sql, Connection connection) throws SQLException { boolean query = stmt.execute(); if (query) { long n = 0; - ResultSet rs = stmt.getResultSet(); - while (rs.next()) { - n++; + try (ResultSet rs = stmt.getResultSet()) { + while (rs.next()) { + n++; + } } return n; } else { @@ -1476,19 +1480,83 @@ public void fetchAllTableRecords(FetchAllTableRecordsRequest request) { String sql = request.getSql(); int batchSize = request.getBatchSize(); IResultSetConsumer consumer = request.getConsumer(); + final boolean manageTransaction; + try { + manageTransaction = connection.getAutoCommit(); + } catch (SQLException ex) { + throw new RuntimeException("Failed to read original autoCommit", ex); + } + boolean transactionStarted = false; + boolean transactionResolved = false; + boolean discardRequired = false; + Exception failure = null; try (PreparedStatement stmt = connection.prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY )) { - connection.setAutoCommit(false); + if (manageTransaction) { + transactionStarted = true; + connection.setAutoCommit(false); + } stmt.setFetchSize(batchSize); - ResultSet rs = stmt.executeQuery(); - consumer.accept(rs); - connection.commit(); - connection.setAutoCommit(true); + try (ResultSet rs = stmt.executeQuery()) { + consumer.accept(rs); + } + if (transactionStarted) { + connection.commit(); + transactionResolved = true; + } } catch (Exception e) { - log.error("Failed to fetch table records. Query: {}", sql, e); - throw new RuntimeException(e); + failure = e; + } + + if (failure != null && transactionStarted && !transactionResolved) { + try { + connection.rollback(); + transactionResolved = true; + } catch (Exception rollbackEx) { + failure.addSuppressed(rollbackEx); + discardRequired = true; + } + } + + if (transactionStarted && !discardRequired) { + try { + connection.setAutoCommit(true); + } catch (Exception restoreEx) { + if (failure == null) { + failure = restoreEx; + } else { + failure.addSuppressed(restoreEx); + } + discardRequired = true; + } + } + + if (discardRequired) { + discardConnection(connection, failure); + } + if (failure != null) { + log.error("Failed to fetch table records. Query: {}", sql, failure); + throw new RuntimeException(failure); + } + } + + private void discardConnection(Connection connection, Throwable failure) { + ConnectInfo connectInfo = Chat2DBContext.getConnectInfo(); + if (connectInfo != null && connectInfo.getConnection() == connection) { + connectInfo.setConnection(null); + } + try { + connection.abort(DIRECT_ABORT_EXECUTOR); + return; + } catch (Exception abortEx) { + failure.addSuppressed(abortEx); + } + try { + connection.close(); + } catch (Exception closeEx) { + failure.addSuppressed(closeEx); } } diff --git a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/model/datasource/ConnectInfo.java b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/model/datasource/ConnectInfo.java index 5d74633ce..871ba3476 100644 --- a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/model/datasource/ConnectInfo.java +++ b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/model/datasource/ConnectInfo.java @@ -139,6 +139,9 @@ public void setProject(String project) { private Date lastAccessTime; + // Runtime-only marker used to reject leases created before a pool invalidation. + private transient volatile Long poolGeneration; + public String getDbVersion() { return dbVersion; @@ -213,7 +216,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(dataSourceId, consoleId, databaseName); + return Objects.hash(dataSourceId, gmtModified); } public Long getDataSourceId() { @@ -490,6 +493,14 @@ public void setLastAccessTime(Date lastAccessTime) { this.lastAccessTime = lastAccessTime; } + public Long poolGeneration() { + return poolGeneration; + } + + public void updatePoolGeneration(Long poolGeneration) { + this.poolGeneration = poolGeneration; + } + private AtomicBoolean inUse = new AtomicBoolean(false); public boolean trySetInUse() { diff --git a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java index 72a0040cd..73e9af953 100644 --- a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java +++ b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java @@ -48,12 +48,23 @@ public class Chat2DBContext { } } + private static IPlugin getPlugin(String dbType) { + if (StringUtils.isBlank(dbType)) { + throw new IllegalArgumentException("Database type must not be blank. Registered types: " + PLUGIN_MAP.keySet()); + } + IPlugin plugin = PLUGIN_MAP.get(dbType); + if (plugin == null) { + throw new IllegalArgumentException("Unsupported database type: " + dbType + ". Registered types: " + PLUGIN_MAP.keySet()); + } + return plugin; + } + public static DriverConfig getDefaultDriverConfig(String dbType) { - return PLUGIN_MAP.get(dbType).getDBConfig().getDefaultDriverConfig(); + return getPlugin(dbType).getDBConfig().getDefaultDriverConfig(); } public static ISqlBuilder getSqlBuilder() { - return PLUGIN_MAP.get(getConnectInfo().getDbType()).getDbMetaData().getSqlBuilder(); + return getPlugin(getConnectInfo().getDbType()).getDbMetaData().getSqlBuilder(); } @@ -62,18 +73,18 @@ public static ConnectInfo getConnectInfo() { } public static IDbMetaData getDbMetaData() { - return PLUGIN_MAP.get(getConnectInfo().getDbType()).getDbMetaData(); + return getPlugin(getConnectInfo().getDbType()).getDbMetaData(); } public static IDbMetaData getDbMetaData(String dbType) { if (StringUtils.isBlank(dbType)) { return getDbMetaData(); } - return PLUGIN_MAP.get(dbType).getDbMetaData(); + return getPlugin(dbType).getDbMetaData(); } public static DBConfig getDBConfig(String dbType) { - return PLUGIN_MAP.get(dbType).getDBConfig(); + return getPlugin(dbType).getDBConfig(); } public static DBConfig getDBConfig() { @@ -81,15 +92,15 @@ public static DBConfig getDBConfig() { if (connectInfo == null) { return null; } - return PLUGIN_MAP.get(connectInfo.getDbType()).getDBConfig(); + return getPlugin(connectInfo.getDbType()).getDBConfig(); } public static IDbManager getDbManager() { - return PLUGIN_MAP.get(getConnectInfo().getDbType()).getDbManager(); + return getPlugin(getConnectInfo().getDbType()).getDbManager(); } public static IDbManager getDbManager(String dbType) { - return PLUGIN_MAP.get(dbType).getDbManager(); + return getPlugin(dbType).getDbManager(); } public static IAccountManager getAccountManager() { diff --git a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/ConnectionPool.java b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/ConnectionPool.java index c2701d5c9..98105b905 100644 --- a/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/ConnectionPool.java +++ b/chat2db-community-server/chat2db-community-spi/src/main/java/ai/chat2db/spi/sql/ConnectionPool.java @@ -28,42 +28,15 @@ public class ConnectionPool { private static ConcurrentHashMap>> CONNECTION_MAP = new ConcurrentHashMap<>(); + private static final ConcurrentHashMap CONNECTION_GENERATIONS = new ConcurrentHashMap<>(); + static { new Thread(() -> { while (true) { try { Thread.sleep(1000 * 60 * 1); - log.info("CONNECTION_MAP size:{}", CONNECTION_MAP.size()); - for (Map.Entry>> entry : CONNECTION_MAP.entrySet()) { - ConcurrentHashMap> map = entry.getValue(); - if (map == null) { - continue; - } - for (Map.Entry> e : map.entrySet()) { - LinkedBlockingQueue queue = e.getValue(); - if (queue == null) { - continue; - } - log.info(e.getKey() + " queue size:{}", queue.size()); - Iterator iterator = queue.iterator(); - while (iterator.hasNext()) { - ConnectInfo connectInfo = iterator.next(); - log.info("check connection:{},usage:{}", connectInfo.getKey(), connectInfo.isInUse()); - if (connectInfo.trySetInUse()) { - try { - if (connectInfo.getLastAccessTime().getTime() + 1000 * 60 * 30 < System.currentTimeMillis()) { - closeQuietly(connectInfo); - } else if (!checkConnectionIsActive(connectInfo)) { - closeQuietly(connectInfo); - } - } finally { - connectInfo.releaseInUse(); - } - } - } - } - } + cleanupConnections(); } catch (Exception e) { log.error("close connection error", e); } @@ -71,6 +44,105 @@ public class ConnectionPool { }).start(); } + static void cleanupConnections() { + log.info("CONNECTION_MAP size:{}", CONNECTION_MAP.size()); + for (Map.Entry>> entry : CONNECTION_MAP.entrySet()) { + ConcurrentHashMap> map = entry.getValue(); + if (map == null) { + continue; + } + for (Map.Entry> queueEntry : map.entrySet()) { + LinkedBlockingQueue queue = queueEntry.getValue(); + if (queue == null) { + continue; + } + log.info(queueEntry.getKey() + " queue size:{}", queue.size()); + cleanupQueue(entry.getKey(), queueEntry.getKey(), queue); + } + } + } + + static void cleanupQueue(LinkedBlockingQueue queue) { + cleanupQueue(null, null, queue); + } + + private static void cleanupQueue(Long datasourceId, String connectionKey, + LinkedBlockingQueue queue) { + int connectionsToCheck = queue.size(); + for (int i = 0; i < connectionsToCheck; i++) { + ConnectInfo connectInfo = queue.poll(); + if (connectInfo == null) { + return; + } + log.info("check connection:{},usage:{}", connectInfo.getKey(), connectInfo.isInUse()); + Date lastAccessTime = connectInfo.getLastAccessTime(); + if (!connectInfo.trySetInUse()) { + returnToQueue(datasourceId, connectionKey, queue, connectInfo, false); + continue; + } + boolean reusable = true; + try { + boolean expired = lastAccessTime != null + && lastAccessTime.getTime() + 1000 * 60 * 30 < System.currentTimeMillis(); + if (expired || !checkConnectionIsActive(connectInfo)) { + closeQuietly(connectInfo); + reusable = false; + } + } finally { + connectInfo.releaseInUse(); + if (reusable) { + returnToQueue(datasourceId, connectionKey, queue, connectInfo, true); + } + } + } + } + + static LinkedBlockingQueue newConnectionQueue() { + return new LinkedBlockingQueue<>(MAX_CONNECTIONS); + } + + static LinkedBlockingQueue getOrCreateConnectionQueue(Long datasourceId, + String connectionKey) { + ConcurrentHashMap> map = + CONNECTION_MAP.computeIfAbsent(datasourceId, key -> new ConcurrentHashMap<>()); + return map.computeIfAbsent(connectionKey, key -> newConnectionQueue()); + } + + static void offerOrClose(LinkedBlockingQueue queue, ConnectInfo connectInfo) { + if (!queue.offer(connectInfo)) { + closeQuietly(connectInfo); + } + } + + private static void returnToQueue(Long datasourceId, String connectionKey, + LinkedBlockingQueue queue, + ConnectInfo connectInfo, boolean closeIfRejected) { + if (datasourceId == null) { + if (!queue.offer(connectInfo)) { + handleRejectedReturn(connectInfo, closeIfRejected); + } + return; + } + boolean[] returnedToCurrentQueue = {false}; + CONNECTION_MAP.computeIfPresent(datasourceId, (key, keyMap) -> { + if (keyMap.get(connectionKey) == queue) { + returnedToCurrentQueue[0] = queue.offer(connectInfo); + } + return keyMap; + }); + if (!returnedToCurrentQueue[0]) { + handleRejectedReturn(connectInfo, closeIfRejected); + } + } + + private static void handleRejectedReturn(ConnectInfo connectInfo, boolean closeIfRejected) { + if (closeIfRejected) { + closeQuietly(connectInfo); + } else { + log.warn("Dropped duplicate pooled connection reference for {}", connectInfo.getKey()); + } + } + private static boolean checkConnectionIsActive(ConnectInfo connectInfo) { try { Connection connection = connectInfo.getConnection(); @@ -117,8 +189,7 @@ private static boolean validateByProbeSql(ConnectInfo connectInfo, Connection co } - private static boolean isRecentlyUsed(ConnectInfo connectInfo) { - Date lastAccessTime = connectInfo.getLastAccessTime(); + private static boolean isRecentlyUsed(Date lastAccessTime) { return lastAccessTime != null && System.currentTimeMillis() - lastAccessTime.getTime() < SKIP_VALIDATION_IF_RECENTLY_USED_MS; } @@ -136,10 +207,19 @@ private static Connection tryGetExistingConnection(ConnectInfo connectInfo) { return null; } + static long currentGeneration(Long datasourceId) { + return datasourceId == null ? 0L : CONNECTION_GENERATIONS.getOrDefault(datasourceId, 0L); + } + public static Connection createNewConnection(ConnectInfo connectInfo) { + return createNewConnection(connectInfo, currentGeneration(connectInfo.getDataSourceId())); + } + + private static Connection createNewConnection(ConnectInfo connectInfo, long generation) { log.info("Creating new individual connection"); Connection connection = Chat2DBContext.getDbManager(connectInfo.getDbType()).getConnection(connectInfo); connectInfo.setConnection(connection); + connectInfo.updatePoolGeneration(generation); return connection; } @@ -148,32 +228,59 @@ public static Connection getConnection(ConnectInfo connectInfo) { if (connection != null) { return connection; } - ConcurrentHashMap> map = CONNECTION_MAP.get(connectInfo.getDataSourceId()); - if (map == null) { + Long datasourceId = connectInfo.getDataSourceId(); + if (datasourceId == null) { return createNewConnection(connectInfo); } - LinkedBlockingQueue queue = map.get(connectInfo.getKey()); - if (queue != null) { - ConnectInfo c = queue.poll(); - if (c != null && c.trySetInUse()) { - Connection conn = tryGetExistingConnection(c); - if (conn != null && (isRecentlyUsed(c) || checkConnectionIsActive(c))) { - connectInfo.setConnection(conn); - log.info("Got connection from pool"); - return conn; - } else { - closeQuietly(c); - return createNewConnection(connectInfo); - } - } else { - return createNewConnection(connectInfo); - } - } else { - return createNewConnection(connectInfo); + long acquisitionGeneration = currentGeneration(datasourceId); + LinkedBlockingQueue queue = + getOrCreateConnectionQueue(datasourceId, connectInfo.getKey()); + Connection pooledConnection = tryBorrowConnection(connectInfo, datasourceId, + connectInfo.getKey(), queue, acquisitionGeneration); + if (pooledConnection != null) { + return pooledConnection; + } + return createNewConnection(connectInfo, acquisitionGeneration); + } + + static Connection tryBorrowConnection(ConnectInfo connectInfo, LinkedBlockingQueue queue) { + return tryBorrowConnection(connectInfo, null, null, queue, 0L); + } + + private static Connection tryBorrowConnection(ConnectInfo connectInfo, Long datasourceId, String connectionKey, + LinkedBlockingQueue queue, + long acquisitionGeneration) { + ConnectInfo pooledConnectInfo = queue.poll(); + if (pooledConnectInfo == null) { + return null; + } + Date lastAccessTime = pooledConnectInfo.getLastAccessTime(); + if (!pooledConnectInfo.trySetInUse()) { + returnToQueue(datasourceId, connectionKey, queue, pooledConnectInfo, false); + return null; + } + if (datasourceId != null + && !Objects.equals(pooledConnectInfo.poolGeneration(), acquisitionGeneration)) { + closeQuietly(pooledConnectInfo); + return null; + } + Connection connection = tryGetExistingConnection(pooledConnectInfo); + if (connection != null + && (isRecentlyUsed(lastAccessTime) || checkConnectionIsActive(pooledConnectInfo))) { + connectInfo.setConnection(connection); + connectInfo.updatePoolGeneration(pooledConnectInfo.poolGeneration()); + log.info("Got connection from pool"); + return connection; } + closeQuietly(pooledConnectInfo); + return null; } public static void removeConnection(Long datasourceId) { + if (datasourceId == null) { + return; + } + CONNECTION_GENERATIONS.merge(datasourceId, 1L, Long::sum); CONNECTION_MAP.computeIfPresent(datasourceId, (key, keyMap) -> { for (Map.Entry> entry : keyMap.entrySet()) { LinkedBlockingQueue queue = entry.getValue(); @@ -225,14 +332,31 @@ public static void close(ConnectInfo connectInfo) { return; } - String key = connectInfo.getKey(); - ConcurrentHashMap> map = CONNECTION_MAP.computeIfAbsent(connectInfo.getDataSourceId(), k -> new ConcurrentHashMap<>()); - LinkedBlockingQueue queue = map.computeIfAbsent(key, k -> new LinkedBlockingQueue<>()); - if (queue.size() < MAX_CONNECTIONS) { - queue.offer(connectInfo); - } else { + Long datasourceId = connectInfo.getDataSourceId(); + if (datasourceId == null) { closeQuietly(connectInfo); + return; } + Long leaseGeneration = connectInfo.poolGeneration(); + if (leaseGeneration == null) { + closeQuietly(connectInfo); + return; + } + long expectedGeneration = leaseGeneration; + String connectionKey = connectInfo.getKey(); + CONNECTION_MAP.compute(datasourceId, (key, keyMap) -> { + if (expectedGeneration != currentGeneration(datasourceId) || keyMap == null) { + closeQuietly(connectInfo); + return keyMap; + } + LinkedBlockingQueue queue = keyMap.get(connectionKey); + if (queue == null) { + closeQuietly(connectInfo); + } else { + offerOrClose(queue, connectInfo); + } + return keyMap; + }); } } diff --git a/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/community/test/spi/sql/DefaultSQLExecutorTransactionTest.java b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/community/test/spi/sql/DefaultSQLExecutorTransactionTest.java new file mode 100644 index 000000000..6e6790668 --- /dev/null +++ b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/community/test/spi/sql/DefaultSQLExecutorTransactionTest.java @@ -0,0 +1,314 @@ +package ai.chat2db.community.test.spi.sql; + +import ai.chat2db.community.domain.api.config.DriverConfig; +import ai.chat2db.spi.DefaultSQLExecutor; +import ai.chat2db.spi.model.datasource.ConnectInfo; +import ai.chat2db.spi.model.request.FetchAllTableRecordsRequest; +import ai.chat2db.spi.sql.Chat2DBContext; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class DefaultSQLExecutorTransactionTest { + + @Test + void shouldNotCommitCallerManagedTransaction() throws Exception { + String url = "jdbc:h2:mem:fetch_records_caller_commit;DB_CLOSE_DELAY=-1"; + try (Connection connection = DriverManager.getConnection(url); + Connection observer = DriverManager.getConnection(url)) { + createTable(connection); + connection.setAutoCommit(false); + executeUpdate(connection, "INSERT INTO records VALUES (1)"); + AtomicReference resultSet = new AtomicReference<>(); + + DefaultSQLExecutor.getInstance().fetchAllTableRecords(request(connection, rs -> { + resultSet.set(rs); + while (rs.next()) { + } + })); + + assertFalse(connection.getAutoCommit()); + assertTrue(resultSet.get().isClosed()); + assertEquals(0, countRows(observer)); + connection.rollback(); + } + } + + @Test + void shouldNotRollbackCallerManagedTransaction() throws Exception { + String url = "jdbc:h2:mem:fetch_records_caller_rollback;DB_CLOSE_DELAY=-1"; + try (Connection connection = DriverManager.getConnection(url); + Connection observer = DriverManager.getConnection(url)) { + createTable(connection); + connection.setAutoCommit(false); + executeUpdate(connection, "INSERT INTO records VALUES (1)"); + + assertThrows(RuntimeException.class, () -> DefaultSQLExecutor.getInstance() + .fetchAllTableRecords(request(connection, rs -> { + throw new IllegalStateException("consumer failed"); + }))); + + assertFalse(connection.getAutoCommit()); + connection.commit(); + assertEquals(1, countRows(observer)); + } + } + + @Test + void shouldRollbackAndRestoreTransactionStartedByExecutor() throws Exception { + String url = "jdbc:h2:mem:fetch_records_executor_rollback;DB_CLOSE_DELAY=-1"; + try (Connection connection = DriverManager.getConnection(url); + Connection observer = DriverManager.getConnection(url)) { + createTable(connection); + + assertThrows(RuntimeException.class, () -> DefaultSQLExecutor.getInstance() + .fetchAllTableRecords(request(connection, rs -> { + executeUpdate(connection, "INSERT INTO records VALUES (1)"); + throw new IllegalStateException("consumer failed"); + }))); + + assertTrue(connection.getAutoCommit()); + assertEquals(0, countRows(observer)); + } + } + + @Test + void shouldDiscardConnectionAndPreserveRollbackFailure() throws Exception { + String url = "jdbc:h2:mem:fetch_records_rollback_failure;DB_CLOSE_DELAY=-1"; + try (Connection delegate = DriverManager.getConnection(url); + Connection observer = DriverManager.getConnection(url)) { + createTable(delegate); + List lifecycleCalls = new ArrayList<>(); + AtomicBoolean abortExecutorCompleted = new AtomicBoolean(); + Connection connection = proxyConnection(delegate, true, false, false, false, + lifecycleCalls, abortExecutorCompleted); + ConnectInfo connectInfo = putContext(connection); + IllegalStateException consumerFailure = new IllegalStateException("consumer failed"); + try { + RuntimeException thrown = assertThrows(RuntimeException.class, + () -> DefaultSQLExecutor.getInstance().fetchAllTableRecords(request(connection, rs -> { + executeUpdate(connection, "INSERT INTO records VALUES (1)"); + throw consumerFailure; + }))); + + assertSame(consumerFailure, thrown.getCause()); + assertEquals(1, consumerFailure.getSuppressed().length); + assertEquals("rollback failed", consumerFailure.getSuppressed()[0].getMessage()); + assertEquals(List.of("rollback", "abort"), lifecycleCalls); + assertTrue(abortExecutorCompleted.get()); + assertTrue(connection.isClosed()); + assertNull(connectInfo.getConnection()); + assertEquals(0, countRows(observer)); + } finally { + Chat2DBContext.removeContext(); + } + } + } + + @Test + void shouldDiscardConnectionWhenAutoCommitRestorationFails() throws Exception { + String url = "jdbc:h2:mem:fetch_records_restore_failure;DB_CLOSE_DELAY=-1"; + try (Connection delegate = DriverManager.getConnection(url)) { + createTable(delegate); + List lifecycleCalls = new ArrayList<>(); + AtomicBoolean abortExecutorCompleted = new AtomicBoolean(); + Connection connection = proxyConnection(delegate, false, true, false, false, + lifecycleCalls, abortExecutorCompleted); + ConnectInfo connectInfo = putContext(connection); + try { + RuntimeException thrown = assertThrows(RuntimeException.class, + () -> DefaultSQLExecutor.getInstance().fetchAllTableRecords(request(connection, rs -> { + while (rs.next()) { + } + }))); + + assertTrue(thrown.getCause() instanceof SQLException); + assertEquals("restore autoCommit failed", thrown.getCause().getMessage()); + assertEquals(List.of("abort"), lifecycleCalls); + assertTrue(abortExecutorCompleted.get()); + assertTrue(connection.isClosed()); + assertNull(connectInfo.getConnection()); + } finally { + Chat2DBContext.removeContext(); + } + } + } + + @Test + void shouldFallBackToCloseWhenAbortFails() throws Exception { + String url = "jdbc:h2:mem:fetch_records_abort_failure;DB_CLOSE_DELAY=-1"; + try (Connection delegate = DriverManager.getConnection(url)) { + createTable(delegate); + List lifecycleCalls = new ArrayList<>(); + AtomicBoolean abortExecutorCompleted = new AtomicBoolean(); + Connection connection = proxyConnection(delegate, true, false, true, false, + lifecycleCalls, abortExecutorCompleted); + ConnectInfo connectInfo = putContext(connection); + IllegalStateException consumerFailure = new IllegalStateException("consumer failed"); + try { + RuntimeException thrown = assertThrows(RuntimeException.class, + () -> DefaultSQLExecutor.getInstance().fetchAllTableRecords(request(connection, rs -> { + throw consumerFailure; + }))); + + assertSame(consumerFailure, thrown.getCause()); + assertEquals(List.of("rollback", "abort", "close"), lifecycleCalls); + assertEquals(2, consumerFailure.getSuppressed().length); + assertEquals("rollback failed", consumerFailure.getSuppressed()[0].getMessage()); + assertEquals("abort failed", consumerFailure.getSuppressed()[1].getMessage()); + assertTrue(abortExecutorCompleted.get()); + assertTrue(connection.isClosed()); + assertNull(connectInfo.getConnection()); + } finally { + Chat2DBContext.removeContext(); + } + } + } + + @Test + void shouldPreserveAbortAndCloseFailuresWhenDiscardCannotComplete() throws Exception { + String url = "jdbc:h2:mem:fetch_records_abort_close_failure;DB_CLOSE_DELAY=-1"; + try (Connection delegate = DriverManager.getConnection(url)) { + createTable(delegate); + List lifecycleCalls = new ArrayList<>(); + AtomicBoolean abortExecutorCompleted = new AtomicBoolean(); + Connection connection = proxyConnection(delegate, true, false, true, true, + lifecycleCalls, abortExecutorCompleted); + ConnectInfo connectInfo = putContext(connection); + IllegalStateException consumerFailure = new IllegalStateException("consumer failed"); + try { + RuntimeException thrown = assertThrows(RuntimeException.class, + () -> DefaultSQLExecutor.getInstance().fetchAllTableRecords(request(connection, rs -> { + throw consumerFailure; + }))); + + assertSame(consumerFailure, thrown.getCause()); + assertEquals(List.of("rollback", "abort", "close"), lifecycleCalls); + assertEquals(3, consumerFailure.getSuppressed().length); + assertEquals("rollback failed", consumerFailure.getSuppressed()[0].getMessage()); + assertEquals("abort failed", consumerFailure.getSuppressed()[1].getMessage()); + assertEquals("close failed", consumerFailure.getSuppressed()[2].getMessage()); + assertTrue(abortExecutorCompleted.get()); + assertNull(connectInfo.getConnection()); + } finally { + Chat2DBContext.removeContext(); + } + } + } + + private static FetchAllTableRecordsRequest request( + Connection connection, ai.chat2db.spi.IResultSetConsumer consumer) { + return FetchAllTableRecordsRequest.builder() + .connection(connection) + .sql("SELECT * FROM records") + .batchSize(10) + .consumer(consumer) + .build(); + } + + private static void createTable(Connection connection) throws java.sql.SQLException { + executeUpdate(connection, "CREATE TABLE records (id INT PRIMARY KEY)"); + } + + private static void executeUpdate(Connection connection, String sql) throws java.sql.SQLException { + try (Statement statement = connection.createStatement()) { + statement.executeUpdate(sql); + } + } + + private static int countRows(Connection connection) throws java.sql.SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery("SELECT COUNT(*) FROM records")) { + resultSet.next(); + return resultSet.getInt(1); + } + } + + private static ConnectInfo putContext(Connection connection) { + ConnectInfo connectInfo = new ConnectInfo(); + connectInfo.setConnection(connection); + connectInfo.setDriverConfig(new DriverConfig()); + Chat2DBContext.putContext(connectInfo); + return connectInfo; + } + + private static Connection proxyConnection( + Connection delegate, boolean failRollback, boolean failAutoCommitRestore) { + return proxyConnection(delegate, failRollback, failAutoCommitRestore, false, false, + new ArrayList<>(), new AtomicBoolean()); + } + + private static Connection proxyConnection( + Connection delegate, boolean failRollback, boolean failAutoCommitRestore, + boolean failAbort, boolean failClose, List lifecycleCalls, + AtomicBoolean abortExecutorCompleted) { + AtomicBoolean autoCommitDisabled = new AtomicBoolean(); + return (Connection) Proxy.newProxyInstance( + DefaultSQLExecutorTransactionTest.class.getClassLoader(), + new Class[]{Connection.class}, + (proxy, method, args) -> { + if ("rollback".equals(method.getName())) { + lifecycleCalls.add("rollback"); + if (failRollback) { + throw new SQLException("rollback failed"); + } + } + if ("abort".equals(method.getName())) { + lifecycleCalls.add("abort"); + AtomicBoolean taskCompleted = new AtomicBoolean(); + ((Executor) args[0]).execute(() -> taskCompleted.set(true)); + abortExecutorCompleted.set(taskCompleted.get()); + if (failAbort) { + throw new SQLException("abort failed"); + } + delegate.close(); + return null; + } + if ("close".equals(method.getName())) { + lifecycleCalls.add("close"); + if (failClose) { + throw new SQLException("close failed"); + } + } + if (failAutoCommitRestore && "setAutoCommit".equals(method.getName())) { + boolean autoCommit = (Boolean) args[0]; + if (autoCommit && autoCommitDisabled.get()) { + throw new SQLException("restore autoCommit failed"); + } + Object result = invoke(delegate, method, args); + if (!autoCommit) { + autoCommitDisabled.set(true); + } + return result; + } + return invoke(delegate, method, args); + }); + } + + private static Object invoke(Connection delegate, Method method, Object[] args) throws Throwable { + try { + return method.invoke(delegate, args); + } catch (InvocationTargetException ex) { + throw ex.getCause(); + } + } +} diff --git a/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/model/datasource/ConnectInfoTest.java b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/model/datasource/ConnectInfoTest.java new file mode 100644 index 000000000..1c3052fae --- /dev/null +++ b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/model/datasource/ConnectInfoTest.java @@ -0,0 +1,28 @@ +package ai.chat2db.spi.model.datasource; + +import org.junit.jupiter.api.Test; + +import java.time.LocalDateTime; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class ConnectInfoTest { + + @Test + void equalInstancesShouldHaveEqualHashCodes() { + LocalDateTime modifiedAt = LocalDateTime.of(2026, 7, 25, 12, 0); + ConnectInfo first = new ConnectInfo(); + first.setDataSourceId(42L); + first.setGmtModified(modifiedAt); + first.setConsoleId(1L); + first.setDatabaseName("first_database"); + ConnectInfo second = new ConnectInfo(); + second.setDataSourceId(42L); + second.setGmtModified(modifiedAt); + second.setConsoleId(2L); + second.setDatabaseName("second_database"); + + assertEquals(first, second); + assertEquals(first.hashCode(), second.hashCode()); + } +} diff --git a/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/sql/Chat2DBContextTest.java b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/sql/Chat2DBContextTest.java new file mode 100644 index 000000000..b936c9ad6 --- /dev/null +++ b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/sql/Chat2DBContextTest.java @@ -0,0 +1,25 @@ +package ai.chat2db.spi.sql; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class Chat2DBContextTest { + + @Test + void shouldDescribeUnsupportedDatabaseType() { + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> Chat2DBContext.getDBConfig("NOT_REGISTERED_FOR_TEST")); + + assertTrue(exception.getMessage().contains("Unsupported database type: NOT_REGISTERED_FOR_TEST")); + } + + @Test + void shouldRejectBlankDatabaseType() { + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> Chat2DBContext.getDBConfig(" ")); + + assertTrue(exception.getMessage().contains("Database type must not be blank")); + } +} diff --git a/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/sql/ConnectionPoolTest.java b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/sql/ConnectionPoolTest.java new file mode 100644 index 000000000..2a9adae4a --- /dev/null +++ b/chat2db-community-server/chat2db-community-spi/src/test/java/ai/chat2db/spi/sql/ConnectionPoolTest.java @@ -0,0 +1,311 @@ +package ai.chat2db.spi.sql; + +import ai.chat2db.spi.model.datasource.ConnectInfo; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Date; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class ConnectionPoolTest { + + @Test + void cleanupShouldCloseAndRemoveInvalidIdleConnection() { + AtomicBoolean closed = new AtomicBoolean(); + ConnectInfo connectInfo = connectInfo(connection(false, closed)); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + queue.offer(connectInfo); + + ConnectionPool.cleanupQueue(queue); + + assertTrue(queue.isEmpty()); + assertTrue(closed.get()); + assertNull(connectInfo.getConnection()); + } + + @Test + void cleanupShouldCloseExpiredConnectionEvenWhenItIsValid() { + AtomicBoolean closed = new AtomicBoolean(); + ConnectInfo connectInfo = connectInfo(connection(true, closed)); + connectInfo.setLastAccessTime(new Date(System.currentTimeMillis() - TimeUnit.MINUTES.toMillis(31))); + LinkedBlockingQueue queue = ConnectionPool.newConnectionQueue(); + queue.offer(connectInfo); + + ConnectionPool.cleanupQueue(queue); + + assertTrue(queue.isEmpty()); + assertTrue(closed.get()); + assertNull(connectInfo.getConnection()); + } + + @Test + void cleanupShouldLeaveBorrowedConnectionInQueue() { + ConnectInfo connectInfo = connectInfo(connection(true, new AtomicBoolean())); + assertTrue(connectInfo.trySetInUse()); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + queue.offer(connectInfo); + + ConnectionPool.cleanupQueue(queue); + + assertSame(connectInfo, queue.peek()); + connectInfo.releaseInUse(); + } + + @Test + void failedBorrowShouldReturnConnectionInfoToQueue() { + ConnectInfo pooled = connectInfo(connection(true, new AtomicBoolean())); + assertTrue(pooled.trySetInUse()); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + queue.offer(pooled); + + assertNull(ConnectionPool.tryBorrowConnection(new ConnectInfo(), queue)); + assertSame(pooled, queue.peek()); + pooled.releaseInUse(); + } + + @Test + void failedBorrowShouldNotCloseConnectionOwnedByAnotherThreadWhenQueueRefills() { + AtomicBoolean pooledClosed = new AtomicBoolean(); + ConnectInfo pooled = connectInfo(connection(true, pooledClosed)); + ConnectInfo replacement = connectInfo(connection(true, new AtomicBoolean())); + assertTrue(pooled.trySetInUse()); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(1) { + private boolean refillAfterPoll = true; + + @Override + public ConnectInfo poll() { + ConnectInfo result = super.poll(); + if (refillAfterPoll) { + refillAfterPoll = false; + assertTrue(super.offer(replacement)); + } + return result; + } + }; + queue.offer(pooled); + + assertNull(ConnectionPool.tryBorrowConnection(new ConnectInfo(), queue)); + + assertSame(replacement, queue.peek()); + assertFalse(pooledClosed.get()); + assertNotNull(pooled.getConnection()); + pooled.releaseInUse(); + } + + @Test + void staleConnectionShouldBeValidatedBeforeBorrow() { + AtomicBoolean closed = new AtomicBoolean(); + AtomicInteger validationCalls = new AtomicInteger(); + ConnectInfo pooled = connectInfo(connection(false, closed, validationCalls)); + pooled.setLastAccessTime(new Date(System.currentTimeMillis() - TimeUnit.MINUTES.toMillis(1))); + LinkedBlockingQueue queue = ConnectionPool.newConnectionQueue(); + queue.offer(pooled); + + assertNull(ConnectionPool.tryBorrowConnection(new ConnectInfo(), queue)); + + assertEquals(1, validationCalls.get()); + assertTrue(closed.get()); + assertNull(pooled.getConnection()); + } + + @Test + void concurrentReturnsShouldKeepQueueBoundedAndCloseOverflowConnections() throws Exception { + int connectionCount = 32; + LinkedBlockingQueue queue = ConnectionPool.newConnectionQueue(); + List connectInfos = new ArrayList<>(connectionCount); + List closedConnections = new ArrayList<>(connectionCount); + for (int i = 0; i < connectionCount; i++) { + AtomicBoolean closed = new AtomicBoolean(); + closedConnections.add(closed); + connectInfos.add(connectInfo(connection(true, closed))); + } + + ExecutorService executor = Executors.newFixedThreadPool(connectionCount); + CountDownLatch ready = new CountDownLatch(connectionCount); + CountDownLatch start = new CountDownLatch(1); + try { + List> futures = new ArrayList<>(connectionCount); + for (ConnectInfo connectInfo : connectInfos) { + futures.add(executor.submit(() -> { + ready.countDown(); + assertTrue(start.await(5, TimeUnit.SECONDS)); + ConnectionPool.offerOrClose(queue, connectInfo); + return null; + })); + } + assertTrue(ready.await(5, TimeUnit.SECONDS)); + start.countDown(); + for (Future future : futures) { + future.get(5, TimeUnit.SECONDS); + } + + assertEquals(2, queue.size()); + Set retained = Collections.newSetFromMap(new IdentityHashMap<>()); + retained.addAll(queue); + for (int i = 0; i < connectionCount; i++) { + assertEquals(!retained.contains(connectInfos.get(i)), closedConnections.get(i).get()); + } + } finally { + start.countDown(); + executor.shutdownNow(); + } + } + + @Test + void cleanupShouldCloseConnectionWhenDatasourceIsRemovedDuringValidation() throws Exception { + long datasourceId = -1924L; + AtomicBoolean closed = new AtomicBoolean(); + CountDownLatch validationStarted = new CountDownLatch(1); + CountDownLatch finishValidation = new CountDownLatch(1); + ConnectInfo connectInfo = connectInfo(blockingValidationConnection( + closed, validationStarted, finishValidation)); + connectInfo.setDataSourceId(datasourceId); + connectInfo.updatePoolGeneration(ConnectionPool.currentGeneration(datasourceId)); + LinkedBlockingQueue queue = + ConnectionPool.getOrCreateConnectionQueue(datasourceId, connectInfo.getKey()); + assertTrue(queue.offer(connectInfo)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future cleanup = executor.submit(ConnectionPool::cleanupConnections); + try { + assertTrue(validationStarted.await(5, TimeUnit.SECONDS)); + ConnectionPool.removeConnection(datasourceId); + finishValidation.countDown(); + cleanup.get(5, TimeUnit.SECONDS); + + assertTrue(closed.get()); + assertNull(connectInfo.getConnection()); + } finally { + finishValidation.countDown(); + ConnectionPool.removeConnection(datasourceId); + executor.shutdownNow(); + } + } + + @Test + void returnedLeaseShouldCloseWhenDatasourceWasRemovedAndPoolWasRecreated() { + long datasourceId = -1925L; + AtomicBoolean oldClosed = new AtomicBoolean(); + ConnectInfo pooled = connectInfo(connection(true, oldClosed)); + pooled.setDataSourceId(datasourceId); + pooled.updatePoolGeneration(ConnectionPool.currentGeneration(datasourceId)); + LinkedBlockingQueue oldQueue = + ConnectionPool.getOrCreateConnectionQueue(datasourceId, pooled.getKey()); + assertTrue(oldQueue.offer(pooled)); + + ConnectInfo oldLease = new ConnectInfo(); + oldLease.setDataSourceId(datasourceId); + assertNotNull(ConnectionPool.getConnection(oldLease)); + + ConnectionPool.removeConnection(datasourceId); + + AtomicBoolean replacementClosed = new AtomicBoolean(); + ConnectInfo replacement = connectInfo(connection(true, replacementClosed)); + replacement.setDataSourceId(datasourceId); + LinkedBlockingQueue replacementQueue = + ConnectionPool.getOrCreateConnectionQueue(datasourceId, replacement.getKey()); + assertTrue(replacementQueue.offer(replacement)); + ConnectionPool.close(oldLease); + + assertTrue(oldClosed.get()); + assertNull(oldLease.getConnection()); + assertFalse(replacementClosed.get()); + assertSame(replacement, replacementQueue.peek()); + ConnectionPool.removeConnection(datasourceId); + assertTrue(replacementClosed.get()); + } + + private static ConnectInfo connectInfo(Connection connection) { + ConnectInfo connectInfo = new ConnectInfo(); + connectInfo.setConnection(connection); + connectInfo.setLastAccessTime(new Date()); + return connectInfo; + } + + private static Connection connection(boolean valid, AtomicBoolean closed) { + return connection(valid, closed, new AtomicInteger()); + } + + private static Connection connection(boolean valid, AtomicBoolean closed, AtomicInteger validationCalls) { + return (Connection) Proxy.newProxyInstance( + ConnectionPoolTest.class.getClassLoader(), + new Class[]{Connection.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "isClosed": + return closed.get(); + case "isValid": + validationCalls.incrementAndGet(); + return valid; + case "close": + closed.set(true); + return null; + case "toString": + return "TestConnection"; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static Connection blockingValidationConnection(AtomicBoolean closed, + CountDownLatch validationStarted, + CountDownLatch finishValidation) { + return (Connection) Proxy.newProxyInstance( + ConnectionPoolTest.class.getClassLoader(), + new Class[]{Connection.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "isClosed": + return closed.get(); + case "isValid": + validationStarted.countDown(); + if (!finishValidation.await(5, TimeUnit.SECONDS)) { + throw new SQLException("Timed out waiting to finish validation"); + } + return true; + case "close": + closed.set(true); + return null; + case "toString": + return "BlockingValidationConnection"; + default: + return defaultValue(method.getReturnType()); + } + }); + } + + private static Object defaultValue(Class type) { + if (!type.isPrimitive()) { + return null; + } + if (type == boolean.class) { + return false; + } + if (type == char.class) { + return '\0'; + } + return 0; + } +}