diff --git a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java index 87b12f90e0..9729df18a0 100644 --- a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java +++ b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java @@ -49,21 +49,36 @@ public final class KingbaseAgent extends PostgresLikeAgent { @Override public List listDatabases() { return unchecked(() -> { - String sql = isMysqlCompatMode() - ? "SELECT current_database() AS database_name" - : "SELECT datname AS database_name FROM sys_database WHERE datistemplate = false ORDER BY datname"; - try (PreparedStatement stmt = requireConnected().prepareStatement(sql); - ResultSet rs = stmt.executeQuery()) { - List result = new ArrayList<>(); - while (rs.next()) { - result.add(new DatabaseInfo(rs.getString("database_name"))); - } + if (isMysqlCompatMode()) { + List result = queryDatabases("SELECT current_database() AS database_name"); if (!result.isEmpty()) return result; } + for (String sql : List.of( + "SELECT datname AS database_name FROM sys_catalog.sys_database WHERE datistemplate = false ORDER BY datname", + "SELECT datname AS database_name FROM pg_database WHERE datistemplate = false ORDER BY datname" + )) { + try { + List result = queryDatabases(sql); + if (!result.isEmpty()) return result; + } catch (Exception ignored) { + // Kingbase catalog names differ across compatibility modes and versions. + } + } return Collections.singletonList(new DatabaseInfo(getConfiguredDatabase())); }); } + private List queryDatabases(String sql) throws Exception { + try (PreparedStatement stmt = requireConnected().prepareStatement(sql); + ResultSet rs = stmt.executeQuery()) { + List result = new ArrayList<>(); + while (rs.next()) { + result.add(new DatabaseInfo(rs.getString("database_name"))); + } + return result; + } + } + @Override public List listSchemas() { return unchecked(() -> { diff --git a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java index a29c7e36ef..80fdd5842f 100644 --- a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java +++ b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java @@ -19,6 +19,7 @@ import java.lang.reflect.Proxy; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Statement; import java.sql.Timestamp; import java.sql.Types; @@ -80,7 +81,24 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertEquals(2, databases.size()); Assertions.assertEquals("app", databases.get(0).getName()); Assertions.assertEquals("analytics", databases.get(1).getName()); - Assertions.assertTrue(sql.get(0).contains("FROM sys_database"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_database"), sql.get(0)); + } + + @Test + void regularListDatabasesFallsBackToPostgresCatalog() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnectionWithFailure(sql, "sys_catalog.sys_database", resultSet( + new String[]{"database_name"}, + new Object[][]{{"test"}} + ))); + + List databases = agent.listDatabases(); + + Assertions.assertEquals(1, databases.size()); + Assertions.assertEquals("test", databases.get(0).getName()); + Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_database"), sql.get(0)); + Assertions.assertTrue(sql.get(1).contains("FROM pg_database"), sql.get(1)); } @Test @@ -390,6 +408,31 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { }); } + private static Connection preparedConnectionWithFailure(List sql, String failingSqlFragment, ResultSet fallback) { + return proxy(Connection.class, (method, args) -> { + if ("prepareStatement".equals(method.getName())) { + String preparedSql = String.valueOf(args[0]); + sql.add(preparedSql); + return proxy(PreparedStatement.class, (statementMethod, statementArgs) -> { + if ("executeQuery".equals(statementMethod.getName())) { + if (preparedSql.contains(failingSqlFragment)) { + throw new SQLException("relation does not exist: " + failingSqlFragment); + } + return fallback; + } + if ("close".equals(statementMethod.getName())) { + return null; + } + return defaultValue(statementMethod.getReturnType()); + }); + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + private static ResultSet resultSet(String[] columns, Object[][] rows) { int[] index = {-1}; return proxy(ResultSet.class, (method, args) -> {