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 6ec41a3b3..966640a40 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 @@ -214,6 +214,62 @@ public final class KingbaseAgent extends PostgresLikeAgent { public List getColumns(String schema, String table) { return unchecked(() -> { Set primaryKeys = primaryKeys(schema, table); + if (!isMysqlCompatMode()) { + return getRegularColumns(schema, table, primaryKeys); + } + return getInformationSchemaColumns(schema, table, primaryKeys); + }); + } + + private List getRegularColumns(String schema, String table, Set primaryKeys) { + return unchecked(() -> { + List result = new ArrayList<>(); + String sql = "SELECT a.attname AS column_name, " + + "format_type(a.atttypid, a.atttypmod) AS data_type, " + + "NOT a.attnotnull AS is_nullable, " + + "sys_get_expr(ad.adbin, ad.adrelid) AS column_default, " + + "d.description AS column_comment, " + + "CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 " + + "THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, " + + "CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 " + + "THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, " + + "CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 " + + "THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length " + + "FROM sys_catalog.sys_attribute a " + + "JOIN sys_catalog.sys_type t ON t.oid = a.atttypid " + + "JOIN sys_catalog.sys_class c ON c.oid = a.attrelid " + + "JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace " + + "LEFT JOIN sys_catalog.sys_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum " + + "LEFT JOIN sys_catalog.sys_description d ON d.objoid = a.attrelid AND d.objsubid = a.attnum " + + "WHERE n.nspname = " + sqlString(effectiveSchema(schema)) + + " AND c.relname = " + sqlString(table) + " " + + "AND a.attnum > 0 AND NOT a.attisdropped " + + "ORDER BY a.attnum"; + try (Statement stmt = requireConnected().createStatement()) { + try (ResultSet rs = stmt.executeQuery(sql)) { + while (rs.next()) { + String columnName = rs.getString("column_name"); + result.add(new ColumnInfo( + columnName, + rs.getString("data_type"), + rs.getBoolean("is_nullable"), + rs.getString("column_default"), + primaryKeys.contains(columnName), + null, + rs.getString("column_comment"), + intObject(rs, "numeric_precision"), + intObject(rs, "numeric_scale"), + intObject(rs, "character_maximum_length") + )); + } + } + } + return result; + }); + } + + private List getInformationSchemaColumns(String schema, String table, Set primaryKeys) { + return unchecked(() -> { List result = new ArrayList<>(); String sql = "SELECT column_name, data_type, is_nullable, column_default, " + "numeric_precision, numeric_scale, character_maximum_length " + 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 3ef64f7ea..6ad79b1bb 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 @@ -1,5 +1,6 @@ package com.dbx.agent.kingbase; +import com.dbx.agent.ColumnInfo; import com.dbx.agent.DatabaseAgent; import com.dbx.agent.DatabaseInfo; import com.dbx.agent.ObjectInfo; @@ -167,6 +168,78 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_proc"), sql.get(0)); } + @Test + void regularGetColumnsUsesFormattedCatalogTypes() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, + resultSet( + new String[]{"column_name"}, + new Object[][]{{"id"}} + ), + resultSet( + new String[]{ + "column_name", + "data_type", + "is_nullable", + "column_default", + "column_comment", + "numeric_precision", + "numeric_scale", + "character_maximum_length" + }, + new Object[][]{ + {"id", "integer", false, "nextval('orders_id_seq'::regclass)", "identifier", 32, 0, null}, + {"create_time", "timestamp with time zone", true, null, null, null, null, null}, + {"name", "character varying(64)", true, null, "display name", null, null, 64} + } + ) + )); + + List columns = agent.getColumns("public", "orders"); + + Assertions.assertEquals(3, columns.size()); + Assertions.assertEquals("integer", columns.get(0).getData_type()); + Assertions.assertTrue(columns.get(0).getIs_primary_key()); + Assertions.assertFalse(columns.get(0).getIs_nullable()); + Assertions.assertEquals("timestamp with time zone", columns.get(1).getData_type()); + Assertions.assertNotEquals("USER-DEFINED", columns.get(1).getData_type()); + Assertions.assertEquals(Integer.valueOf(64), columns.get(2).getCharacter_maximum_length()); + Assertions.assertTrue(sql.get(1).contains("format_type(a.atttypid, a.atttypmod) AS data_type"), sql.get(1)); + Assertions.assertTrue(sql.get(1).contains("FROM sys_catalog.sys_attribute"), sql.get(1)); + Assertions.assertFalse(sql.get(1).contains("information_schema.columns"), sql.get(1)); + } + + @Test + void mysqlCompatGetColumnsKeepsInformationSchemaPath() { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + agent.setMysqlCompatMode(true); + TestSupport.setPrivateConnection(agent, preparedConnection(sql, + resultSet( + new String[]{"column_name"}, + new Object[][]{{"id"}} + ), + resultSet( + new String[]{ + "column_name", + "data_type", + "is_nullable", + "column_default", + "numeric_precision", + "numeric_scale", + "character_maximum_length" + }, + new Object[][]{{"id", "int", "NO", null, 32, 0, null}} + ) + )); + + List columns = agent.getColumns("PUBLIC", "orders"); + + Assertions.assertEquals("int", columns.get(0).getData_type()); + Assertions.assertTrue(sql.get(1).contains("FROM information_schema.columns"), sql.get(1)); + } + @Test void mysqlCompatTimestampTypeNameIsReadAsTimestampText() throws Exception { Timestamp timestamp = Timestamp.valueOf("2026-06-22 11:29:00"); @@ -243,6 +316,13 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { } } return null; + case "getBoolean": + Object booleanValue = columnValue(columns, rows[index[0]], args[0]); + if (booleanValue instanceof Boolean) return booleanValue; + if (booleanValue instanceof Number) return ((Number) booleanValue).intValue() != 0; + return Boolean.parseBoolean(String.valueOf(booleanValue)); + case "getObject": + return columnValue(columns, rows[index[0]], args[0]); case "wasNull": return false; case "close": @@ -253,6 +333,18 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { }); } + private static Object columnValue(String[] columns, Object[] row, Object key) { + if (key instanceof Number) { + return row[((Number) key).intValue() - 1]; + } + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(String.valueOf(key))) { + return row[i]; + } + } + return null; + } + private static ResultSet timestampResultSet(Timestamp timestamp) { return proxy(ResultSet.class, (method, args) -> { switch (method.getName()) { diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 285a99628..1a0f5ccba 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -1197,6 +1197,9 @@ public final class DbxJdbcPlugin { if (driverQuirks(connection).useOracleMetadata()) { return oracleGetColumns(conn, oracleEffectiveSchema(conn, schema), table); } + if (isKingbaseUrl(optionalText(connection, "connection_string"))) { + return kingbaseGetColumns(conn, schema, table); + } DatabaseMetaData meta = conn.getMetaData(); JdbcDriverQuirks quirks = driverQuirks(connection); String catalog = metadataCatalog(database, quirks); @@ -1420,10 +1423,87 @@ public final class DbxJdbcPlugin { }; } + private static JsonNode kingbaseGetColumns(Connection conn, String schema, String table) throws SQLException { + ArrayNode result = MAPPER.createArrayNode(); + String effectiveSchema = emptyToNull(schema) == null ? "PUBLIC" : schema; + Set primaryKeys = kingbasePrimaryKeys(conn, effectiveSchema, table); + String sql = "SELECT a.attname AS column_name, " + + "format_type(a.atttypid, a.atttypmod) AS data_type, " + + "NOT a.attnotnull AS is_nullable, " + + "sys_get_expr(ad.adbin, ad.adrelid) AS column_default, " + + "d.description AS column_comment, " + + "CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 " + + "THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, " + + "CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 " + + "THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, " + + "CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 " + + "THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length " + + "FROM sys_catalog.sys_attribute a " + + "JOIN sys_catalog.sys_type t ON t.oid = a.atttypid " + + "JOIN sys_catalog.sys_class c ON c.oid = a.attrelid " + + "JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace " + + "LEFT JOIN sys_catalog.sys_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum " + + "LEFT JOIN sys_catalog.sys_description d ON d.objoid = a.attrelid AND d.objsubid = a.attnum " + + "WHERE n.nspname = " + sqlString(effectiveSchema) + + " AND c.relname = " + sqlString(table) + " " + + "AND a.attnum > 0 AND NOT a.attisdropped " + + "ORDER BY a.attnum"; + try (Statement statement = conn.createStatement()) { + try (ResultSet rs = statement.executeQuery(sql)) { + while (rs.next()) { + String name = rs.getString("column_name"); + ObjectNode item = columnNode(result, name); + item.put("data_type", rs.getString("data_type")); + item.put("is_nullable", rs.getBoolean("is_nullable")); + putNullablePreferValue(item, "column_default", rs.getString("column_default")); + item.put("is_primary_key", primaryKeys.contains(name)); + item.putNull("extra"); + putNullablePreferValue(item, "comment", rs.getString("column_comment")); + putNullableInt(item, "numeric_precision", rs.getObject("numeric_precision")); + putNullableInt(item, "numeric_scale", rs.getObject("numeric_scale")); + putNullableInt(item, "character_maximum_length", rs.getObject("character_maximum_length")); + } + } + } + return result; + } + + private static Set kingbasePrimaryKeys(Connection conn, String schema, String table) { + Set primaryKeys = new HashSet<>(); + String sql = "SELECT a.attname AS column_name " + + "FROM sys_catalog.sys_constraint co " + + "JOIN sys_catalog.sys_class c ON c.oid = co.conrelid " + + "JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace " + + "JOIN LATERAL (SELECT unnest(co.conkey) AS attnum, generate_series(1, array_length(co.conkey, 1)) AS ord) AS pk_cols ON true " + + "JOIN sys_catalog.sys_attribute a ON a.attrelid = c.oid AND a.attnum = pk_cols.attnum " + + "WHERE co.contype = 'p' " + + "AND n.nspname = " + sqlString(schema) + " " + + "AND c.relname = " + sqlString(table) + " " + + "ORDER BY pk_cols.ord"; + try (Statement statement = conn.createStatement()) { + try (ResultSet rs = statement.executeQuery(sql)) { + while (rs.next()) { + primaryKeys.add(rs.getString("column_name")); + } + } + } catch (SQLException ignored) { + return Collections.emptySet(); + } + return primaryKeys; + } + + private static boolean isKingbaseUrl(String url) { + return urlMatchesPrefix(url, "jdbc:kingbase"); + } + private static String quoteAnsiIdentifier(String identifier) { return "\"" + identifier.replace("\"", "\"\"") + "\""; } + private static String sqlString(String value) { + return "'" + (value == null ? "" : value).replace("'", "''") + "'"; + } + private static void appendColumns( ArrayNode result, DatabaseMetaData meta, diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index 1af882c54..d56dcfd26 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -660,6 +660,26 @@ final class DbxJdbcPluginTest { assertEquals(true, response.path("result").path(0).path("is_primary_key").asBoolean()); } + @Test + void kingbaseGetColumnsUsesFormattedCatalogTypes() throws Exception { + Method method = DbxJdbcPlugin.class.getDeclaredMethod("kingbaseGetColumns", Connection.class, String.class, String.class); + method.setAccessible(true); + List sql = new ArrayList<>(); + + JsonNode result = (JsonNode) method.invoke(null, kingbaseColumnsConnection(sql), "dbx_issue_1942", "t_timestamp_type"); + + assertEquals("id", result.path(0).path("name").asText()); + assertEquals("INTEGER", result.path(0).path("data_type").asText()); + assertEquals(true, result.path(0).path("is_primary_key").asBoolean()); + assertEquals("create_time", result.path(1).path("name").asText()); + assertEquals("TIMESTAMP WITH TIME ZONE", result.path(1).path("data_type").asText()); + assertEquals("create_by", result.path(2).path("name").asText()); + assertEquals("CHARACTER VARYING(64 byte)", result.path(2).path("data_type").asText()); + assertEquals(64, result.path(2).path("character_maximum_length").asInt()); + assertEquals(true, sql.get(1).contains("format_type(a.atttypid, a.atttypmod) AS data_type")); + assertEquals(true, sql.get(1).contains("FROM sys_catalog.sys_attribute")); + } + @Test void columnIsNullablePrefersIsNullableStringWhenNullableCodeIsWrong() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod("columnIsNullable", ResultSet.class); @@ -939,6 +959,116 @@ final class DbxJdbcPluginTest { ); } + private static Connection kingbaseColumnsConnection(List sql) { + ResultSet primaryKeys = rowsResultSet( + new String[] { "column_name" }, + new Object[][] { { "id" } } + ); + ResultSet columns = rowsResultSet( + new String[] { + "column_name", + "data_type", + "is_nullable", + "column_default", + "column_comment", + "numeric_precision", + "numeric_scale", + "character_maximum_length" + }, + new Object[][] { + { "id", "INTEGER", false, null, null, 32, 0, null }, + { "create_time", "TIMESTAMP WITH TIME ZONE", true, null, null, null, null, null }, + { "create_by", "CHARACTER VARYING(64 byte)", true, null, null, null, null, 64 } + } + ); + int[] index = { 0 }; + return (Connection) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Connection.class }, + (proxy, method, args) -> switch (method.getName()) { + case "createStatement" -> { + yield statement(sql, index[0]++ == 0 ? primaryKeys : columns); + } + case "isClosed" -> false; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static Statement statement(List sql, ResultSet rs) { + return (Statement) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Statement.class }, + (proxy, method, args) -> switch (method.getName()) { + case "executeQuery" -> { + sql.add(String.valueOf(args[0])); + yield rs; + } + case "close" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static PreparedStatement preparedStatement(ResultSet rs) { + return (PreparedStatement) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { PreparedStatement.class }, + (proxy, method, args) -> switch (method.getName()) { + case "executeQuery" -> rs; + case "setString", "close" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static ResultSet rowsResultSet(String[] columns, Object[][] rows) { + return (ResultSet) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { ResultSet.class }, + new java.lang.reflect.InvocationHandler() { + private int index = -1; + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return switch (method.getName()) { + case "next" -> ++index < rows.length; + case "getString" -> stringValue(columns, rows[index], args[0]); + case "getBoolean" -> booleanValue(columns, rows[index], args[0]); + case "getObject" -> columnValue(columns, rows[index], args[0]); + case "close" -> null; + default -> defaultValue(method.getReturnType()); + }; + } + } + ); + } + + private static String stringValue(String[] columns, Object[] row, Object key) { + Object value = columnValue(columns, row, key); + return value == null ? null : String.valueOf(value); + } + + private static boolean booleanValue(String[] columns, Object[] row, Object key) { + Object value = columnValue(columns, row, key); + if (value instanceof Boolean bool) return bool; + if (value instanceof Number number) return number.intValue() != 0; + return Boolean.parseBoolean(String.valueOf(value)); + } + + private static Object columnValue(String[] columns, Object[] row, Object key) { + if (key instanceof Number number) { + return row[number.intValue() - 1]; + } + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(String.valueOf(key))) { + return row[i]; + } + } + return null; + } + private static ResultSet columnNullableResultSet(String isNullable, int nullableCode) { return (ResultSet) Proxy.newProxyInstance( DbxJdbcPluginTest.class.getClassLoader(),