feat(dameng): support dm6 connections via custom jdbc drivers

agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java

@@ -90,7 +90,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
         return unchecked(() -> {
             loadDriver(params);
             try (Connection conn = openTestConnection(params)) {
-                boolean valid = conn != null && conn.isValid(5);
+                boolean valid = isConnectionValid(conn, 5);
                 Map<String, Object> result = new LinkedHashMap<>();
                 result.put("ok", valid);
                 if (valid) {
@@ -275,6 +275,31 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
         return true;
     }
 
+    final boolean isConnectionValid(Connection connection, int timeoutSecs) {
+        if (connection == null) {
+            return false;
+        }
+        try {
+            if (connection.isClosed()) {
+                return false;
+            }
+            String validationQuery = connectionValidationQuery();
+            if (validationQuery == null || validationQuery.isBlank()) {
+                return connection.isValid(timeoutSecs);
+            }
+            try (Statement statement = connection.createStatement()) {
+                try {
+                    statement.setQueryTimeout(timeoutSecs);
+                } catch (Exception | AbstractMethodError ignored) {
+                }
+                statement.execute(validationQuery);
+                return true;
+            }
+        } catch (Exception | AbstractMethodError ignored) {
+            return false;
+        }
+    }
+
     final synchronized boolean usesConnectionPool() {
         return poolRegistry != null;
     }
@@ -485,6 +510,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
     protected void afterDisconnect() throws Exception {
     }
 
+    protected String connectionValidationQuery() {
+        return null;
+    }
+
     protected void beforeQueryExecution(Connection connection, int timeoutSecs) throws Exception {
     }
 
@@ -586,6 +615,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
         return registry.borrow(
             identity,
             JdbcSessionRole.from(params.getSessionRole()),
+            connectionValidationQuery(),
             () -> openInitializedConnection(params)
         );
     }

agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java

@@ -81,10 +81,14 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable {
     }
 
     Lease borrow(String identity, ConnectionFactory connectionFactory) throws Exception {
-        return borrow(identity, JdbcSessionRole.WORKLOAD, connectionFactory);
+        return borrow(identity, JdbcSessionRole.WORKLOAD, null, connectionFactory);
     }
 
     Lease borrow(String identity, JdbcSessionRole role, ConnectionFactory connectionFactory) throws Exception {
+        return borrow(identity, role, null, connectionFactory);
+    }
+
+    Lease borrow(String identity, JdbcSessionRole role, String connectionTestQuery, ConnectionFactory connectionFactory) throws Exception {
         String key = digest(identity);
         while (true) {
             if (closed.get()) {
@@ -96,7 +100,7 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable {
             }
             PoolEntry entry;
             try {
-                entry = pools.computeIfAbsent(key, ignored -> createPoolEntry(key, connectionFactory));
+                entry = pools.computeIfAbsent(key, ignored -> createPoolEntry(key, connectionTestQuery, connectionFactory));
             } catch (PoolCreationException error) {
                 throw error.unwrap();
             }
@@ -130,7 +134,7 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable {
         return entry != null && entry.hasActiveLeases();
     }
 
-    private PoolEntry createPoolEntry(String key, ConnectionFactory connectionFactory) {
+    private PoolEntry createPoolEntry(String key, String connectionTestQuery, ConnectionFactory connectionFactory) {
         try {
             ConnectionFactoryDataSource factoryDataSource = new ConnectionFactoryDataSource(
                 connectionFactory,
@@ -146,6 +150,9 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable {
             config.setMinimumIdle(settings.minimumIdle);
             config.setConnectionTimeout(settings.connectionTimeoutMillis);
             config.setValidationTimeout(settings.validationTimeoutMillis);
+            if (connectionTestQuery != null && !connectionTestQuery.isBlank()) {
+                config.setConnectionTestQuery(connectionTestQuery);
+            }
             config.setIdleTimeout(settings.idleTimeoutMillis);
             config.setMaxLifetime(settings.maxLifetimeMillis);
             config.setInitializationFailTimeout(-1L);

agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java

@@ -171,8 +171,10 @@ public final class JsonRpcServer {
             boolean valid = false;
             if (conn != null) {
                 try {
-                    valid = !conn.isClosed() && conn.isValid(2);
-                } catch (Exception ignored) {
+                    valid = agent instanceof AbstractJdbcAgent jdbcAgent
+                        ? jdbcAgent.isConnectionValid(conn, 2)
+                        : !conn.isClosed() && conn.isValid(2);
+                } catch (Exception | AbstractMethodError ignored) {
                 }
             }
             if (!valid) {
@@ -370,8 +372,10 @@ public final class JsonRpcServer {
         }
         boolean valid = false;
         try {
-            valid = !conn.isClosed() && conn.isValid(2);
-        } catch (Exception ignored) {
+            valid = agent instanceof AbstractJdbcAgent validationAgent
+                ? validationAgent.isConnectionValid(conn, 2)
+                : !conn.isClosed() && conn.isValid(2);
+        } catch (Exception | AbstractMethodError ignored) {
         }
         if (valid) {
             lastConnectionValidationTimeMillis = now;

agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java

@@ -127,6 +127,19 @@ class AbstractJdbcAgentTest {
         assertEquals(1, tracking.closeCount);
     }
 
+    @Test
+    void validatesLegacyJdbcConnectionsWithConfiguredQueryInsteadOfIsValid() {
+        TrackingConnection tracking = new TrackingConnection();
+        tracking.isValidUnsupported = true;
+        TestAgent agent = new TestAgent(tracking);
+        agent.validationQuery = "SELECT 1";
+
+        assertTrue(agent.testConnection(new ConnectParams()));
+
+        assertEquals(0, tracking.isValidCount);
+        assertEquals(Arrays.asList("setQueryTimeout:5", "execute:SELECT 1"), tracking.calls);
+    }
+
     @Test
     void testConnectionCanSkipOpeningAPhysicalConnection() {
         TrackingConnection tracking = new TrackingConnection();
@@ -339,6 +352,7 @@ class AbstractJdbcAgentTest {
         private int afterConnectCount;
         private int afterDisconnectCount;
         private boolean skipTestConnectionOpen;
+        private String validationQuery;
 
         private TestAgent(TrackingConnection tracking) {
             this.tracking = tracking;
@@ -375,6 +389,11 @@ class AbstractJdbcAgentTest {
             afterDisconnectCount += 1;
         }
 
+        @Override
+        protected String connectionValidationQuery() {
+            return validationQuery;
+        }
+
         @Override
         public List<DatabaseInfo> listDatabases() {
             return Collections.emptyList();
@@ -436,6 +455,7 @@ class AbstractJdbcAgentTest {
         private String identifierQuote = "\"";
         private boolean compatibilityQueryFails;
         private int compatibilityQueryCount;
+        private boolean isValidUnsupported;
 
         private Connection connection() {
             return proxy(Connection.class, new MethodHandler() {
@@ -450,6 +470,9 @@ class AbstractJdbcAgentTest {
                     }
                     if ("isValid".equals(name)) {
                         isValidCount += 1;
+                        if (isValidUnsupported) {
+                            throw new AbstractMethodError("legacy JDBC driver");
+                        }
                         return true;
                     }
                     if ("close".equals(name)) {