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 b03161252b..bc1b520e60 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 @@ -495,19 +495,13 @@ public final class KingbaseAgent extends PostgresLikeAgent { "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, " + - (sqlServerIdentityCatalogMode - ? "ic.seed_value AS identity_seed, ic.increment_value AS identity_increment " - : "NULL AS identity_seed, NULL AS identity_increment ") + + "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 " + - (sqlServerIdentityCatalogMode - ? "LEFT JOIN sys.identity_columns ic ON ic.object_id = c.oid AND ic.column_id = a.attnum " - : "") + "WHERE n.nspname = " + sqlString(effectiveSchema(schema)) + " AND c.relname = " + sqlString(table) + " " + "AND a.attnum > 0 AND NOT a.attisdropped " + @@ -522,7 +516,7 @@ public final class KingbaseAgent extends PostgresLikeAgent { rs.getBoolean("is_nullable"), rs.getString("column_default"), primaryKeys.contains(columnName), - identityExtra(rs), + null, rs.getString("column_comment"), intObject(rs, "numeric_precision"), intObject(rs, "numeric_scale"), @@ -531,11 +525,40 @@ public final class KingbaseAgent extends PostgresLikeAgent { } } } + applySqlServerIdentityMetadata(schema, table, result); return result; }); } - private static String identityExtra(ResultSet rs) throws Exception { + private void applySqlServerIdentityMetadata(String schema, String table, List columns) { + if (!sqlServerIdentityCatalogMode || columns.isEmpty()) return; + String sql = "SELECT a.attname AS column_name, ic.seed_value AS identity_seed, " + + "ic.increment_value AS identity_increment " + + "FROM sys.identity_columns ic " + + "JOIN sys_catalog.sys_class c ON c.oid = ic.object_id " + + "JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace " + + "JOIN sys_catalog.sys_attribute a ON a.attrelid = c.oid AND a.attnum = ic.column_id " + + "WHERE n.nspname = " + sqlString(effectiveSchema(schema)) + + " AND c.relname = " + sqlString(table); + try (Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery(sql)) { + Map columnsByName = new LinkedHashMap<>(); + for (ColumnInfo column : columns) { + columnsByName.put(column.getName(), column); + } + while (rs.next()) { + ColumnInfo column = columnsByName.get(rs.getString("column_name")); + if (column != null) { + column.setExtra(identityExtra(rs)); + } + } + } catch (SQLException ignored) { + // Identity metadata is optional and some Kingbase versions expose a broken compatibility view. + sqlServerIdentityCatalogMode = false; + } + } + + private static String identityExtra(ResultSet rs) throws SQLException { String seed = rs.getString("identity_seed"); String increment = rs.getString("identity_increment"); if (seed == null || increment == null) { 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 d5c1b920f0..bec7a7b5fc 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 @@ -605,14 +605,16 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { "column_comment", "numeric_precision", "numeric_scale", - "character_maximum_length", - "identity_seed", - "identity_increment" + "character_maximum_length" }, new Object[][]{ - {"id", "int", false, null, null, 32, 0, null, "1", "1"}, - {"name", "character varying", true, null, null, null, null, 64, null, null} + {"id", "int", false, null, null, 32, 0, null}, + {"name", "character varying", true, null, null, null, null, 64} } + ), + resultSet( + new String[]{"column_name", "identity_seed", "identity_increment"}, + new Object[][]{{"id", "1", "1"}} ) )); @@ -620,7 +622,26 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertEquals("identity(1,1)", columns.get(0).getExtra()); Assertions.assertNull(columns.get(1).getExtra()); - Assertions.assertTrue(sql.get(1).contains("LEFT JOIN sys.identity_columns"), sql.get(1)); + Assertions.assertFalse(sql.get(1).contains("sys.identity_columns"), sql.get(1)); + Assertions.assertTrue(sql.get(2).contains("FROM sys.identity_columns"), sql.get(2)); + } + + @Test + void sqlServerCompatGetColumnsIgnoresBrokenIdentityCatalog() throws Exception { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + setSqlServerIdentityCatalogMode(agent, true); + TestSupport.setPrivateConnection(agent, sqlServerIdentityFailureConnection(sql)); + + List columns = agent.getColumns("dbo", "orders"); + + Assertions.assertEquals(2, columns.size()); + Assertions.assertEquals("id", columns.get(0).getName()); + Assertions.assertTrue(columns.get(0).getIs_primary_key()); + Assertions.assertNull(columns.get(0).getExtra()); + Assertions.assertFalse(sql.get(1).contains("sys.identity_columns"), sql.get(1)); + Assertions.assertTrue(sql.get(2).contains("FROM sys.identity_columns"), sql.get(2)); + Assertions.assertFalse(isSqlServerIdentityCatalogMode(agent)); } @Test @@ -750,6 +771,45 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { return preparedConnectionWithFailures(sql, List.of(failingSqlFragment), fallback); } + private static Connection sqlServerIdentityFailureConnection(List sql) { + return proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) { + return proxy(Statement.class, (statementMethod, statementArgs) -> { + if ("executeQuery".equals(statementMethod.getName())) { + String query = String.valueOf(statementArgs[0]); + sql.add(query); + if (query.contains("FROM sys.identity_columns")) { + throw new SQLException("ERROR: cannot open file base/14465/t48_3852767: No such file or directory"); + } + if (query.contains("information_schema.table_constraints")) { + return resultSet(new String[]{"column_name"}, new Object[][]{{"id"}}); + } + return resultSet( + new String[]{ + "column_name", + "data_type", + "is_nullable", + "column_default", + "column_comment", + "numeric_precision", + "numeric_scale", + "character_maximum_length" + }, + new Object[][]{ + {"id", "int", false, null, null, 32, 0, null}, + {"name", "character varying", false, null, null, null, null, 64} + } + ); + } + if ("close".equals(statementMethod.getName())) return null; + return defaultValue(statementMethod.getReturnType()); + }); + } + if ("isClosed".equals(method.getName())) return false; + return defaultValue(method.getReturnType()); + }); + } + private static Connection preparedConnectionWithMetadataFailure( List sql, String message,