From 4e5cb621026fae09d3ccdb3cafce3776483c59c8 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 26 Jun 2026 15:48:18 +0800 Subject: [PATCH] fix(kingbase): list views and routines --- .../com/dbx/agent/kingbase/KingbaseAgent.java | 100 ++++++++++++++++-- .../dbx/agent/kingbase/KingbaseAgentTest.java | 85 ++++++++++++++- 2 files changed, 176 insertions(+), 9 deletions(-) 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 4949460c1..6ec41a3b3 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 @@ -90,21 +90,84 @@ public final class KingbaseAgent extends PostgresLikeAgent { @Override public List listTables(String schema) { - return listTables(schema, "table_type = 'BASE TABLE'"); + if (isMysqlCompatMode()) { + return listTables(schema, "table_type IN ('BASE TABLE', 'VIEW')"); + } + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT c.relname AS table_name, " + + "CASE c.relkind " + + "WHEN 'r' THEN 'TABLE' " + + "WHEN 'p' THEN 'TABLE' " + + "WHEN 'v' THEN 'VIEW' " + + "WHEN 'm' THEN 'MATERIALIZED_VIEW' " + + "WHEN 'f' THEN 'FOREIGN_TABLE' " + + "ELSE 'TABLE' END AS table_type, " + + "d.description AS table_comment " + + "FROM sys_catalog.sys_class c " + + "JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace " + + "LEFT JOIN sys_catalog.sys_description d ON d.objoid = c.oid AND d.objsubid = 0 " + + "WHERE n.nspname = ? AND c.relkind IN ('r','p','v','m','f') " + + "ORDER BY c.relname"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, effectiveSchema(schema)); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new TableInfo( + rs.getString("table_name"), + normalizeTableType(rs.getString("table_type")), + rs.getString("table_comment") + )); + } + } + } + return result; + }); } @Override public List listObjects(String schema) { - List result = new ArrayList<>(); - for (TableInfo table : listTables(schema)) { - result.add(new ObjectInfo(table.getName(), table.getTable_type(), effectiveSchema(schema), table.getComment())); - } - return result; + return unchecked(() -> { + String effectiveSchema = effectiveSchema(schema); + List result = new ArrayList<>(); + for (TableInfo table : listTables(effectiveSchema)) { + result.add(new ObjectInfo(table.getName(), table.getTable_type(), effectiveSchema, table.getComment())); + } + if (isMysqlCompatMode()) { + return result; + } + + String sql = "SELECT p.proname AS routine_name, " + + "CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS routine_type, " + + "d.description AS routine_comment " + + "FROM sys_catalog.sys_proc p " + + "JOIN sys_catalog.sys_namespace n ON n.oid = p.pronamespace " + + "LEFT JOIN sys_catalog.sys_description d ON d.objoid = p.oid AND d.objsubid = 0 " + + "WHERE n.nspname = ? AND p.prokind IN ('p','f') " + + "ORDER BY p.proname"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, effectiveSchema); + try (ResultSet rs = stmt.executeQuery()) { + while (rs.next()) { + result.add(new ObjectInfo( + rs.getString("routine_name"), + rs.getString("routine_type"), + effectiveSchema, + rs.getString("routine_comment") + )); + } + } + } + return result; + }); } @Override public ObjectSource getObjectSource(String schema, String name, String objectType) { - if (!"VIEW".equalsIgnoreCase(objectType)) { + if ("FUNCTION".equalsIgnoreCase(objectType) || "PROCEDURE".equalsIgnoreCase(objectType)) { + return routineSource(schema, name, objectType); + } + if (!"VIEW".equalsIgnoreCase(objectType) && !"MATERIALIZED_VIEW".equalsIgnoreCase(objectType)) { return new ObjectSource(name, objectType, effectiveSchema(schema), ""); } return unchecked(() -> { @@ -124,6 +187,29 @@ public final class KingbaseAgent extends PostgresLikeAgent { }); } + private ObjectSource routineSource(String schema, String name, String objectType) { + return unchecked(() -> { + String source = ""; + String prokind = "PROCEDURE".equalsIgnoreCase(objectType) ? "p" : "f"; + String sql = "SELECT sys_get_functiondef(p.oid) AS source " + + "FROM sys_catalog.sys_proc p " + + "JOIN sys_catalog.sys_namespace n ON n.oid = p.pronamespace " + + "WHERE n.nspname = ? AND p.proname = ? AND p.prokind = ? " + + "ORDER BY p.oid LIMIT 1"; + try (PreparedStatement stmt = requireConnected().prepareStatement(sql)) { + stmt.setString(1, effectiveSchema(schema)); + stmt.setString(2, name); + stmt.setString(3, prokind); + try (ResultSet rs = stmt.executeQuery()) { + if (rs.next()) { + source = coalesce(rs.getString("source")); + } + } + } + return new ObjectSource(name, objectType, effectiveSchema(schema), source); + }); + } + @Override public List getColumns(String schema, String table) { 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 05ddda31d..3ef64f7ea 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 @@ -2,6 +2,9 @@ package com.dbx.agent.kingbase; import com.dbx.agent.DatabaseAgent; import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.ObjectInfo; +import com.dbx.agent.ObjectSource; +import com.dbx.agent.TableInfo; import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest; import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Assertions; @@ -98,6 +101,72 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertFalse(sql.get(0).contains("SHOW")); } + @Test + void regularListTablesUsesKingbaseCatalogAndIncludesViews() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"table_name", "table_type", "table_comment"}, + new Object[][]{{"app_table", "TABLE", "table comment"}, {"app_view", "VIEW", "view comment"}} + ))); + + List tables = agent.listTables("public"); + + Assertions.assertEquals(2, tables.size()); + Assertions.assertEquals("app_table", tables.get(0).getName()); + Assertions.assertEquals("TABLE", tables.get(0).getTable_type()); + Assertions.assertEquals("app_view", tables.get(1).getName()); + Assertions.assertEquals("VIEW", tables.get(1).getTable_type()); + Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_class"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("c.relkind IN ('r','p','v','m','f')"), sql.get(0)); + } + + @Test + void regularListObjectsIncludesKingbaseViewsProceduresAndFunctions() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, + resultSet( + new String[]{"table_name", "table_type", "table_comment"}, + new Object[][]{{"app_table", "TABLE", null}, {"app_view", "VIEW", "view comment"}} + ), + resultSet( + new String[]{"routine_name", "routine_type", "routine_comment"}, + new Object[][]{{"refresh_stats", "PROCEDURE", "proc comment"}, {"format_name", "FUNCTION", "fn comment"}} + ) + )); + + List objects = agent.listObjects("public"); + + Assertions.assertEquals(4, objects.size()); + Assertions.assertEquals("app_table", objects.get(0).getName()); + Assertions.assertEquals("TABLE", objects.get(0).getObject_type()); + Assertions.assertEquals("app_view", objects.get(1).getName()); + Assertions.assertEquals("VIEW", objects.get(1).getObject_type()); + Assertions.assertEquals("refresh_stats", objects.get(2).getName()); + Assertions.assertEquals("PROCEDURE", objects.get(2).getObject_type()); + Assertions.assertEquals("format_name", objects.get(3).getName()); + Assertions.assertEquals("FUNCTION", objects.get(3).getObject_type()); + Assertions.assertTrue(sql.get(1).contains("FROM sys_catalog.sys_proc"), sql.get(1)); + Assertions.assertTrue(sql.get(1).contains("p.prokind IN ('p','f')"), sql.get(1)); + } + + @Test + void regularRoutineSourceUsesKingbaseFunctionDefinition() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, resultSet( + new String[]{"source"}, + new Object[][]{{"CREATE FUNCTION public.format_name() RETURNS text AS $$ SELECT 'x'; $$"}} + ))); + + ObjectSource source = agent.getObjectSource("public", "format_name", "FUNCTION"); + + Assertions.assertTrue(source.getSource().startsWith("CREATE FUNCTION public.format_name()"), source.getSource()); + Assertions.assertTrue(sql.get(0).contains("SELECT sys_get_functiondef(p.oid) AS source"), sql.get(0)); + Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_proc"), sql.get(0)); + } + @Test void mysqlCompatTimestampTypeNameIsReadAsTimestampText() throws Exception { Timestamp timestamp = Timestamp.valueOf("2026-06-22 11:29:00"); @@ -110,9 +179,19 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { } private static Connection preparedConnection(List sql, ResultSet rs) { + return preparedConnection(sql, new ResultSet[]{rs}); + } + + private static Connection preparedConnection(List sql, ResultSet... resultSets) { + int[] resultSetIndex = {0}; PreparedStatement statement = proxy(PreparedStatement.class, (method, args) -> { if ("executeQuery".equals(method.getName())) { - return rs; + int current = Math.min(resultSetIndex[0], resultSets.length - 1); + resultSetIndex[0] += 1; + return resultSets[current]; + } + if ("setString".equals(method.getName())) { + return null; } if ("close".equals(method.getName())) { return null; @@ -122,7 +201,9 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Statement plainStatement = proxy(Statement.class, (method, args) -> { if ("executeQuery".equals(method.getName())) { sql.add(String.valueOf(args[0])); - return rs; + int current = Math.min(resultSetIndex[0], resultSets.length - 1); + resultSetIndex[0] += 1; + return resultSets[current]; } if ("close".equals(method.getName())) { return null;