From ece47c85d6ddf781df2a1be57e6d373b0a895731 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 6 Jul 2026 19:42:49 +0800 Subject: [PATCH] fix(kingbase): prioritize selected schema search path --- .../java/com/dbx/agent/BatchExecutor.java | 21 +----- .../main/java/com/dbx/agent/JdbcExecutor.java | 29 ++------ .../com/dbx/agent/JdbcSchemaSwitcher.java | 68 +++++++++++++++++++ .../com/dbx/agent/TransactionExecutor.java | 25 +------ .../com/dbx/agent/AbstractJdbcAgentTest.java | 4 +- .../agent/CommonJavaCompatibilityTest.java | 2 +- .../java/com/dbx/agent/JdbcExecutorTest.java | 56 +++++++++++++++ .../com/dbx/agent/kingbase/KingbaseAgent.java | 5 +- .../dbx/agent/kingbase/KingbaseAgentTest.java | 8 +++ 9 files changed, 146 insertions(+), 72 deletions(-) create mode 100644 agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java diff --git a/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java b/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java index b8c973f17..134ada07c 100644 --- a/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java @@ -79,26 +79,7 @@ public final class BatchExecutor { } private static void applySchema(Connection conn, String schema, Function setSchemaSql) throws Exception { - if (schema == null || schema.trim().isEmpty()) { - return; - } - try { - conn.setSchema(schema); - return; - } catch (Exception | AbstractMethodError ignored) { - } - try { - conn.setCatalog(schema); - return; - } catch (Exception | AbstractMethodError ignored) { - } - String schemaSql = setSchemaSql.apply(schema); - if (schemaSql == null || schemaSql.trim().isEmpty()) { - return; - } - try (Statement stmt = conn.createStatement()) { - stmt.execute(schemaSql); - } + JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql); } private static T unchecked(ThrowingSupplier supplier) { diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java index 0b11f87d8..cbd77dcf2 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java @@ -565,31 +565,12 @@ public final class JdbcExecutor { } private void applySchema(Connection conn, String schema, Function setSchemaSql) throws SQLException { - if (schema == null || schema.trim().isEmpty()) { - return; - } - // Prefer JDBC standard APIs over database-specific SQL. - // setSchema (JDBC 4.1) and setCatalog (JDBC 1.0) work universally - // across all JDBC drivers without needing to know the database dialect. try { - conn.setSchema(schema); - return; - } catch (SQLException | AbstractMethodError ignored) { - // setSchema not supported by this driver - } - try { - conn.setCatalog(schema); - return; - } catch (SQLException | AbstractMethodError ignored) { - // setCatalog not supported either - } - // Fallback: execute database-specific SQL (e.g. USE, SET SCHEMA, etc.) - String schemaSql = setSchemaSql.apply(schema); - if (schemaSql == null || schemaSql.trim().isEmpty()) { - return; - } - try (Statement stmt = conn.createStatement()) { - stmt.execute(schemaSql); + JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql); + } catch (SQLException e) { + throw e; + } catch (Exception e) { + throw new SQLException(e); } } diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java b/agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java new file mode 100644 index 000000000..76ffa5ab8 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java @@ -0,0 +1,68 @@ +package com.dbx.agent; + +import java.sql.Connection; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.function.Function; + +final class JdbcSchemaSwitcher { + private JdbcSchemaSwitcher() { + } + + static void apply(Connection conn, String schema, Function setSchemaSql) throws Exception { + if (schema == null || schema.trim().isEmpty()) { + return; + } + + SchemaSqlResult schemaSqlResult = applySchemaSql(conn, schema, setSchemaSql); + Exception schemaSqlError = schemaSqlResult.error; + if (schemaSqlError == null) { + return; + } + + try { + conn.setSchema(schema); + return; + } catch (SQLException | AbstractMethodError ignored) { + // Some JDBC drivers only expose schema switching through SQL. + } + try { + conn.setCatalog(schema); + return; + } catch (SQLException | AbstractMethodError ignored) { + // Last fallback failed as well; surface the SQL-switch error below. + } + + if (schemaSqlResult.attempted) { + throw schemaSqlError; + } + } + + private static SchemaSqlResult applySchemaSql(Connection conn, String schema, Function setSchemaSql) { + String schemaSql; + try { + schemaSql = setSchemaSql.apply(schema); + } catch (RuntimeException e) { + return new SchemaSqlResult(true, e); + } + if (schemaSql == null || schemaSql.trim().isEmpty()) { + return new SchemaSqlResult(false, new SQLException("No schema switch SQL provided")); + } + try (Statement stmt = conn.createStatement()) { + stmt.execute(schemaSql); + return new SchemaSqlResult(true, null); + } catch (SQLException | AbstractMethodError e) { + return new SchemaSqlResult(true, e instanceof SQLException ? (SQLException) e : new SQLException(e)); + } + } + + private static final class SchemaSqlResult { + private final boolean attempted; + private final Exception error; + + private SchemaSqlResult(boolean attempted, Exception error) { + this.attempted = attempted; + this.error = error; + } + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java b/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java index 9ba7b877d..2bd12ed46 100644 --- a/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java @@ -71,30 +71,7 @@ public final class TransactionExecutor { } private static void applySchema(Connection conn, String schema, Function setSchemaSql) throws Exception { - if (schema == null || schema.trim().isEmpty()) { - return; - } - // Prefer JDBC standard APIs over database-specific SQL. - try { - conn.setSchema(schema); - return; - } catch (Exception | AbstractMethodError ignored) { - // setSchema not supported by this driver - } - try { - conn.setCatalog(schema); - return; - } catch (Exception | AbstractMethodError ignored) { - // setCatalog not supported either - } - // Fallback: execute database-specific SQL (e.g. USE, SET SCHEMA, etc.) - String schemaSql = setSchemaSql.apply(schema); - if (schemaSql == null || schemaSql.trim().isEmpty()) { - return; - } - try (Statement stmt = conn.createStatement()) { - stmt.execute(schemaSql); - } + JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql); } private static boolean supportsTransactions(Connection conn) { diff --git a/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java b/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java index 5c03e7464..2043c6388 100644 --- a/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java +++ b/agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java @@ -67,7 +67,7 @@ class AbstractJdbcAgentTest { assertEquals(Collections.singletonList("VALUE"), result.getColumns()); assertEquals(Collections.singletonList(Collections.singletonList("row-value")), result.getRows()); assertFalse(result.getTruncated()); - assertEquals(Arrays.asList("setSchema:APP", "setMaxRows:26", "setFetchSize:7", "execute:SELECT VALUE"), tracking.calls); + assertEquals(Arrays.asList("execute:USE APP", "setMaxRows:26", "setFetchSize:7", "execute:SELECT VALUE"), tracking.calls); } @Test @@ -153,7 +153,7 @@ class AbstractJdbcAgentTest { assertEquals( Arrays.asList( "setAutoCommit:false", - "setSchema:APP", + "execute:USE APP", "executeUpdate:UPDATE A", "executeUpdate:UPDATE B", "commit", diff --git a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java index e415de7cc..92ff463ab 100644 --- a/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java +++ b/agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java @@ -334,7 +334,7 @@ class CommonJavaCompatibilityTest { assertEquals(2L, result.getAffected_rows()); assertEquals( - Arrays.asList("supportsTransactions", "setSchema:APP", "executeUpdate:UPDATE A SET ID = 1", "executeUpdate:UPDATE B SET ID = 2"), + Arrays.asList("supportsTransactions", "execute:SET SCHEMA \"APP\"", "executeUpdate:UPDATE A SET ID = 1", "executeUpdate:UPDATE B SET ID = 2"), calls ); } diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java index f6c125523..6d54dfc0f 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -5,10 +5,15 @@ import org.junit.jupiter.api.Test; import java.lang.reflect.InvocationHandler; import java.lang.reflect.Method; import java.lang.reflect.Proxy; +import java.sql.Connection; import java.sql.ResultSet; import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; import java.sql.Types; +import java.util.ArrayList; import java.util.Arrays; +import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import javax.sql.rowset.serial.SerialBlob; @@ -62,10 +67,61 @@ class JdbcExecutorTest { assertEquals(2, fixture.getColumnTypeCalls()); } + @Test + void schemaSwitcherPrefersDriverSpecificSql() throws Exception { + List calls = new ArrayList<>(); + + JdbcSchemaSwitcher.apply(schemaConnection(calls, false), "APP", schema -> "USE " + schema); + + assertEquals(Arrays.asList("execute:USE APP"), calls); + } + + @Test + void schemaSwitcherFallsBackToSetSchemaWhenSqlFails() throws Exception { + List calls = new ArrayList<>(); + + JdbcSchemaSwitcher.apply(schemaConnection(calls, true), "APP", schema -> "USE " + schema); + + assertEquals(Arrays.asList("execute:USE APP", "setSchema:APP"), calls); + } + private static ResultSet resultSet(byte[] bytes, StringSupplier stringSupplier) { return resultSet(bytes, stringSupplier, false); } + private static Connection schemaConnection(List calls, boolean failSchemaSql) { + InvocationHandler handler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "createStatement": + return schemaStatement(calls, failSchemaSql); + case "setSchema": + calls.add("setSchema:" + args[0]); + return null; + default: + return defaultValue(method.getReturnType()); + } + }; + return (Connection) Proxy.newProxyInstance(Connection.class.getClassLoader(), new Class[]{Connection.class}, handler); + } + + private static Statement schemaStatement(List calls, boolean failSchemaSql) { + InvocationHandler handler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "execute": + calls.add("execute:" + args[0]); + if (failSchemaSql) { + throw new SQLException("unsupported schema SQL"); + } + return false; + case "close": + return null; + default: + return defaultValue(method.getReturnType()); + } + }; + return (Statement) Proxy.newProxyInstance(Statement.class.getClassLoader(), new Class[]{Statement.class}, handler); + } + private static CountingResultSetFixture countingResultSet(Object[][] rows) { String[] labels = {"id", "name"}; int[] sqlTypes = {Types.INTEGER, Types.VARCHAR}; 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 3fd839a10..87b12f90e 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 @@ -542,7 +542,10 @@ public final class KingbaseAgent extends PostgresLikeAgent { @Override public String setSchemaSQL(String schema) { - return "SET search_path TO " + JdbcIdentifiers.INSTANCE.doubleQuote(effectiveSchema(schema)); + // Kingbase searches sys_catalog implicitly before user schemas unless it + // is listed explicitly. Put it after the selected schema so business + // tables named like system tables (for example sys_config) win. + return "SET search_path TO " + JdbcIdentifiers.INSTANCE.doubleQuote(effectiveSchema(schema)) + ", sys_catalog"; } @Override 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 eef795655..a29c7e36e 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 @@ -45,6 +45,14 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertEquals("jdbc:kingbase8://{host}:{port}/{database}", agent.getProfile().getUrlTemplate()); } + @Test + void schemaSwitchPlacesSysCatalogAfterSelectedSchema() { + KingbaseAgent agent = new KingbaseAgent(); + + Assertions.assertEquals("SET search_path TO \"app\", sys_catalog", agent.setSchemaSQL("app")); + Assertions.assertEquals("SET search_path TO \"app\"\"prod\", sys_catalog", agent.setSchemaSQL("app\"prod")); + } + @Test void mysqlCompatListDatabasesUsesCurrentDatabase() { List sql = new ArrayList<>();