diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java b/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java index bf64bc995..57f7ca3d2 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java @@ -1475,7 +1475,7 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable { call("physical_set_network_timeout", () -> { connection.setNetworkTimeout(networkTimeoutExecutor, networkTimeoutMillis); return null; - }, factoryDataSource, timeoutMillis); + }, factoryDataSource, timeoutMillis, true); } private T call( @@ -1483,6 +1483,16 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable { PhysicalConnectionCall call, ConnectionFactoryDataSource factoryDataSource, long timeoutMillis + ) throws SQLException { + return call(operation, call, factoryDataSource, timeoutMillis, false); + } + + private T call( + String operation, + PhysicalConnectionCall call, + ConnectionFactoryDataSource factoryDataSource, + long timeoutMillis, + boolean preserveCompletedFailure ) throws SQLException { CompletableFuture outcome = new CompletableFuture<>(); try { @@ -1511,6 +1521,13 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable { factoryDataSource.poison(failure); throw failure; } catch (ExecutionException error) { + if (preserveCompletedFailure) { + Throwable cause = error.getCause(); + if (cause instanceof SQLException sqlError) { + throw sqlError; + } + throw new SQLException("JDBC physical operation failed: " + operation, cause); + } SQLException failure = new PhysicalConnectionStateUnknownException(error.getCause()); factoryDataSource.poison(failure); throw failure; diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java index a4d16761f..dcc47ef8b 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java @@ -850,6 +850,27 @@ class JdbcConnectionPoolingTest { } } + @RepeatedTest(5) + void unsupportedNetworkTimeoutDoesNotPoisonIdentity() throws Exception { + AtomicInteger physicalOpens = new AtomicInteger(); + String url = h2Url("unsupported_network_timeout"); + try (Connection ignored = openH2(url, physicalOpens)) { + // Keep H2 bootstrap outside the setup classification watchdog. + } + physicalOpens.set(0); + try (JdbcConnectionPoolRegistry registry = new JdbcConnectionPoolRegistry(shortTimeoutPoolSettings(1, 32))) { + for (int attempt = 0; attempt < 2; attempt++) { + try (JdbcConnectionPoolRegistry.Lease lease = registry.borrow( + "unsupported-network-timeout", + () -> unsupportedNetworkTimeoutConnection(openH2(url, physicalOpens)) + )) { + assertTrue(lease.connection().isValid(1)); + } + } + assertEquals(1, physicalOpens.get()); + } + } + @Test void blockedSetupAfterKnownFailurePoisonsCurrentAttemptGeneration() throws Exception { AtomicInteger connectionAttempts = new AtomicInteger(); @@ -2136,6 +2157,23 @@ class JdbcConnectionPoolingTest { ); } + private static Connection unsupportedNetworkTimeoutConnection(Connection delegate) { + return (Connection) Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + (proxy, method, args) -> { + if ("setNetworkTimeout".equals(method.getName())) { + throw new SQLException("Does not support setNetworkTimeout"); + } + try { + return method.invoke(delegate, args); + } catch (InvocationTargetException error) { + throw error.getCause(); + } + } + ); + } + private static Connection asynchronousAbortConnection( Connection delegate, CountDownLatch abortScheduled,