xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(kingbase): prioritize selected schema search path

This commit is contained in:
t8y2 2026-07-06 19:42:49 +08:00
parent 2b5aaa2a54
commit ece47c85d6
9 changed files with 146 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
agents/common/src/main/java/com/dbx/agent/BatchExecutor.java
View File

@ -79,26 +79,7 @@ public final class BatchExecutor {
}
private static void applySchema(Connection conn, String schema, Function<String, String> 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> T unchecked(ThrowingSupplier<T> supplier) {

29
agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java
View File

@ -565,31 +565,12 @@ public final class JdbcExecutor {
}
private void applySchema(Connection conn, String schema, Function<String, String> 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);
}
}

68
agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java Normal file
View File

@ -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<String, String> 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<String, String> 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;
}
}
}

25
agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java
View File

@ -71,30 +71,7 @@ public final class TransactionExecutor {
}
private static void applySchema(Connection conn, String schema, Function<String, String> 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) {

4
agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java
View File

@ -67,7 +67,7 @@ class AbstractJdbcAgentTest {
assertEquals(Collections.singletonList("VALUE"), result.getColumns());
assertEquals(Collections.singletonList(Collections.<Object>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",

2
agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java
View File

@ -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
);
}

56
agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java
View File

@ -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<String> 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<String> 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<String> 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<String> 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};

5
agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java
View File

@ -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

8
agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java
View File

@ -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<String> sql = new ArrayList<>();