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

fix(kingbase): set pg_catalog last in pg-compat search_path 

* fix(kingbase): prioritize selected schema in pg catalog mode

* Fix：安全解析 Kingbase 所选模式对象

---------

Co-authored-by: zipg <4047349+zipg@users.noreply.github.com>
This commit is contained in:
zipg 2026-07-16 00:54:33 +08:00 committed by GitHub
parent d9076f4689
commit 12ab08f543
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 376 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java
View File

@ -142,6 +142,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options.getMaxRows(),
options.getFetchSize(),
options.getTimeoutSecs(),
@ -158,6 +159,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
resultValueReader()
);
@ -182,6 +184,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
resultValueReader()
);
@ -199,7 +202,13 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
@Override
public QueryResult executeTransaction(List<String> statements, String schema) {
return TransactionExecutor.executeUpdateStatements(requireConnected(), statements, schema, this::setSchemaSQL);
return TransactionExecutor.executeUpdateStatements(
requireConnected(),
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
}
@Override

18
agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java
View File

@ -59,6 +59,7 @@ public abstract class BaseDatabaseAgent implements DatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
JdbcExecutor.current()::defaultResultValue
);
@ -81,6 +82,7 @@ public abstract class BaseDatabaseAgent implements DatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
JdbcExecutor.current()::defaultResultValue
);
@ -98,12 +100,24 @@ public abstract class BaseDatabaseAgent implements DatabaseAgent {
@Override
public QueryResult executeTransaction(List<String> statements, String schema) {
return TransactionExecutor.executeUpdateStatements(requireConnected(), statements, schema, this::setSchemaSQL);
return TransactionExecutor.executeUpdateStatements(
requireConnected(),
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
}
@Override
public QueryResult executeBatch(List<String> statements, String schema) {
return BatchExecutor.executeBatchStatements(requireConnected(), statements, schema, this::setSchemaSQL);
return BatchExecutor.executeBatchStatements(
requireConnected(),
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
}
protected Connection requireConnected() {

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

@ -7,6 +7,7 @@ import java.sql.Statement;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.function.Supplier;
public final class BatchExecutor {
private BatchExecutor() {
@ -17,10 +18,20 @@ public final class BatchExecutor {
List<String> statements,
String schema,
Function<String, String> setSchemaSql
) {
return executeBatchStatements(conn, statements, schema, setSchemaSql, () -> "");
}
public static QueryResult executeBatchStatements(
Connection conn,
List<String> statements,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql
) {
return unchecked(() -> {
long start = System.currentTimeMillis();
applySchema(conn, schema, setSchemaSql);
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
long totalAffected = 0;
int statementCount = 0;
try (Statement stmt = conn.createStatement()) {
@ -78,8 +89,13 @@ public final class BatchExecutor {
}
}
private static void applySchema(Connection conn, String schema, Function<String, String> setSchemaSql) throws Exception {
JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql);
private static void applySchema(
Connection conn,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql
) throws Exception {
JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql, resetSchemaSql);
}
private static <T> T unchecked(ThrowingSupplier<T> supplier) {

16
agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java
View File

@ -167,6 +167,7 @@ public interface DatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
AgentExecutionContext.jdbcExecutor()::defaultResultValue
);
@ -190,6 +191,7 @@ public interface DatabaseAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
AgentExecutionContext.jdbcExecutor()::defaultResultValue
);
@ -222,7 +224,13 @@ public interface DatabaseAgent {
if (conn == null) {
throw new IllegalStateException("Not connected");
}
return TransactionExecutor.executeUpdateStatements(conn, statements, schema, this::setSchemaSQL);
return TransactionExecutor.executeUpdateStatements(
conn,
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
}
default QueryResult executeBatch(List<String> statements, String schema) {
@ -230,13 +238,17 @@ public interface DatabaseAgent {
if (conn == null) {
throw new IllegalStateException("Not connected");
}
return BatchExecutor.executeBatchStatements(conn, statements, schema, this::setSchemaSQL);
return BatchExecutor.executeBatchStatements(conn, statements, schema, this::setSchemaSQL, this::resetSchemaSQL);
}
default String setSchemaSQL(String schema) {
return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema);
}
default String resetSchemaSQL() {
return "";
}
static String buildTableDdl(
String schema,
String table,

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

@ -16,6 +16,7 @@ import java.util.Locale;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.function.Supplier;
public final class JdbcExecutor {
public static final JdbcExecutor INSTANCE = new JdbcExecutor();
@ -69,12 +70,26 @@ public final class JdbcExecutor {
Integer fetchSize,
int timeoutSecs,
ResultValueReader valueReader
) {
return execute(conn, sql, schema, setSchemaSql, () -> "", maxRows, fetchSize, timeoutSecs, valueReader);
}
public QueryResult execute(
Connection conn,
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
int maxRows,
Integer fetchSize,
int timeoutSecs,
ResultValueReader valueReader
) {
return unchecked(() -> {
String trimmedSql = trimSql(sql);
long start = System.currentTimeMillis();
applySchema(conn, schema, setSchemaSql);
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
try (Statement stmt = conn.createStatement()) {
activeStatements.add(stmt);
@ -178,7 +193,19 @@ public final class JdbcExecutor {
QueryPageOptions options,
ResultValueReader valueReader
) {
return executePage(conn, sql, schema, setSchemaSql, options, valueReader, sessions);
return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, sessions);
}
public QueryPageResult executePage(
Connection conn,
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
QueryPageOptions options,
ResultValueReader valueReader
) {
return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, sessions);
}
public QueryPageResult startTableRead(
@ -189,7 +216,19 @@ public final class JdbcExecutor {
QueryPageOptions options,
ResultValueReader valueReader
) {
return executePage(conn, sql, schema, setSchemaSql, options, valueReader, tableReadSessions);
return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, tableReadSessions);
}
public QueryPageResult startTableRead(
Connection conn,
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
QueryPageOptions options,
ResultValueReader valueReader
) {
return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, tableReadSessions);
}
private QueryPageResult executePage(
@ -197,6 +236,7 @@ public final class JdbcExecutor {
String sql,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
QueryPageOptions options,
ResultValueReader valueReader,
ConcurrentHashMap<String, QuerySession> targetSessions
@ -206,7 +246,7 @@ public final class JdbcExecutor {
String trimmedSql = trimSql(sql);
long start = System.currentTimeMillis();
applySchema(conn, schema, setSchemaSql);
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
Statement stmt = conn.createStatement();
activeStatements.add(stmt);
@ -592,9 +632,14 @@ public final class JdbcExecutor {
return Math.min(requestedRows, 1024);
}
private void applySchema(Connection conn, String schema, Function<String, String> setSchemaSql) throws SQLException {
private void applySchema(
Connection conn,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql
) throws SQLException {
try {
JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql);
JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql, resetSchemaSql);
} catch (SQLException e) {
throw e;
} catch (Exception e) {

45
agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java
View File

@ -3,31 +3,59 @@ package com.dbx.agent;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.Map;
import java.util.WeakHashMap;
import java.util.function.Function;
import java.util.function.Supplier;
final class JdbcSchemaSwitcher {
private static final Map<Connection, String> RESET_SQL_BY_CONNECTION =
Collections.synchronizedMap(new WeakHashMap<>());
private JdbcSchemaSwitcher() {
}
static void apply(Connection conn, String schema, Function<String, String> setSchemaSql) throws Exception {
apply(conn, schema, setSchemaSql, () -> "");
}
static void apply(
Connection conn,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql
) throws Exception {
if (schema == null || schema.trim().isEmpty()) {
String resetSql = RESET_SQL_BY_CONNECTION.remove(conn);
if (resetSql == null) {
return;
}
SchemaSqlResult resetResult = applySchemaSql(conn, () -> resetSql);
if (resetResult.attempted && resetResult.error != null) {
RESET_SQL_BY_CONNECTION.put(conn, resetSql);
throw resetResult.error;
}
return;
}
SchemaSqlResult schemaSqlResult = applySchemaSql(conn, schema, setSchemaSql);
SchemaSqlResult schemaSqlResult = applySchemaSql(conn, () -> setSchemaSql.apply(schema));
Exception schemaSqlError = schemaSqlResult.error;
if (schemaSqlError == null) {
rememberResetSql(conn, resetSchemaSql);
return;
}
try {
conn.setSchema(schema);
rememberResetSql(conn, resetSchemaSql);
return;
} catch (SQLException | AbstractMethodError ignored) {
// Some JDBC drivers only expose schema switching through SQL.
}
try {
conn.setCatalog(schema);
rememberResetSql(conn, resetSchemaSql);
return;
} catch (SQLException | AbstractMethodError ignored) {
// Last fallback failed as well; surface the SQL-switch error below.
@ -38,10 +66,21 @@ final class JdbcSchemaSwitcher {
}
}
private static SchemaSqlResult applySchemaSql(Connection conn, String schema, Function<String, String> setSchemaSql) {
private static void rememberResetSql(Connection conn, Supplier<String> resetSchemaSql) {
try {
String sql = resetSchemaSql.get();
if (sql != null && !sql.trim().isEmpty()) {
RESET_SQL_BY_CONNECTION.put(conn, sql);
}
} catch (RuntimeException ignored) {
// Drivers without a reset command keep the legacy no-op behavior.
}
}
private static SchemaSqlResult applySchemaSql(Connection conn, Supplier<String> schemaSqlSupplier) {
String schemaSql;
try {
schemaSql = setSchemaSql.apply(schema);
schemaSql = schemaSqlSupplier.get();
} catch (RuntimeException e) {
return new SchemaSqlResult(true, e);
}

32
agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java
View File

@ -40,6 +40,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options.getMaxRows(),
options.getFetchSize(),
options.getTimeoutSecs(),
@ -54,6 +55,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
geometryAwareResolver()
);
@ -66,6 +68,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
sql,
schema,
this::setSchemaSQL,
this::resetSchemaSQL,
options,
geometryAwareResolver()
);
@ -133,7 +136,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
public List<DatabaseInfo> listDatabases() {
return unchecked(() -> {
List<DatabaseInfo> result = new ArrayList<>();
try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname");
try (java.sql.PreparedStatement stmt = requireConnection().prepareStatement("SELECT datname FROM pg_catalog.pg_database WHERE datistemplate = false ORDER BY datname");
ResultSet rs = stmt.executeQuery()) {
while (rs.next()) {
result.add(new DatabaseInfo(rs.getString("datname")));
@ -333,15 +336,15 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
String upperType = objectType.toUpperCase();
String sql;
if ("VIEW".equals(upperType) || "MATERIALIZED VIEW".equals(upperType)) {
sql = "SELECT pg_get_viewdef(to_regclass(?), true)";
sql = "SELECT pg_catalog.pg_get_viewdef(pg_catalog.to_regclass(?), true)";
} else if ("FUNCTION".equals(upperType)) {
sql = "SELECT pg_get_functiondef(p.oid)\n" +
"FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace\n" +
sql = "SELECT pg_catalog.pg_get_functiondef(p.oid)\n" +
"FROM pg_catalog.pg_proc p JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace\n" +
"WHERE n.nspname = ? AND p.proname = ? AND p.prokind = 'f'\n" +
"ORDER BY p.oid LIMIT 1";
} else if ("PROCEDURE".equals(upperType)) {
sql = "SELECT pg_get_functiondef(p.oid)\n" +
"FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace\n" +
sql = "SELECT pg_catalog.pg_get_functiondef(p.oid)\n" +
"FROM pg_catalog.pg_proc p JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace\n" +
"WHERE n.nspname = ? AND p.proname = ? AND p.prokind = 'p'\n" +
"ORDER BY p.oid LIMIT 1";
} else {
@ -451,13 +454,13 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
String sql = "SELECT i.relname AS index_name, am.amname AS index_type, " +
"ix.indisunique AS is_unique, ix.indisprimary AS is_primary, " +
"array_agg(a.attname ORDER BY k.n) AS columns " +
"FROM pg_index ix " +
"JOIN pg_class t ON t.oid = ix.indrelid " +
"JOIN pg_class i ON i.oid = ix.indexrelid " +
"JOIN pg_namespace n ON n.oid = t.relnamespace " +
"JOIN pg_am am ON am.oid = i.relam " +
"FROM pg_catalog.pg_index ix " +
"JOIN pg_catalog.pg_class t ON t.oid = ix.indrelid " +
"JOIN pg_catalog.pg_class i ON i.oid = ix.indexrelid " +
"JOIN pg_catalog.pg_namespace n ON n.oid = t.relnamespace " +
"JOIN pg_catalog.pg_am am ON am.oid = i.relam " +
"JOIN LATERAL (SELECT unnest(ix.indkey) AS attnum, generate_series(1, array_length(ix.indkey, 1)) AS n) AS k ON true " +
"JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum " +
"JOIN pg_catalog.pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum " +
"WHERE n.nspname = ? AND t.relname = ? " +
"GROUP BY i.relname, am.amname, ix.indisunique, ix.indisprimary " +
"ORDER BY i.relname";
@ -565,6 +568,11 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
return "SET search_path TO " + quoteIdentifier(schema);
}
@Override
public String resetSchemaSQL() {
return "RESET search_path";
}
private java.sql.Connection requireConnection() {
return requireConnected();
}

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

@ -5,6 +5,7 @@ import java.sql.Statement;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.function.Supplier;
public final class TransactionExecutor {
private TransactionExecutor() {
@ -16,7 +17,17 @@ public final class TransactionExecutor {
String schema,
Function<String, String> setSchemaSql
) {
return executeStatements(conn, statements, schema, setSchemaSql, new StatementRunner() {
return executeUpdateStatements(conn, statements, schema, setSchemaSql, () -> "");
}
public static QueryResult executeUpdateStatements(
Connection conn,
List<String> statements,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql
) {
return executeStatements(conn, statements, schema, setSchemaSql, resetSchemaSql, new StatementRunner() {
@Override
public long run(Statement stmt, String sql) throws Exception {
return stmt.executeUpdate(sql);
@ -30,18 +41,29 @@ public final class TransactionExecutor {
String schema,
Function<String, String> setSchemaSql,
StatementRunner runner
) {
return executeStatements(conn, statements, schema, setSchemaSql, () -> "", runner);
}
public static QueryResult executeStatements(
Connection conn,
List<String> statements,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
StatementRunner runner
) {
return unchecked(() -> {
long start = System.currentTimeMillis();
if (!supportsTransactions(conn)) {
long totalAffected = executeAll(conn, statements, schema, setSchemaSql, runner);
long totalAffected = executeAll(conn, statements, schema, setSchemaSql, resetSchemaSql, runner);
return result(totalAffected, start);
}
boolean savedAutoCommit = conn.getAutoCommit();
conn.setAutoCommit(false);
try {
long totalAffected = executeAll(conn, statements, schema, setSchemaSql, runner);
long totalAffected = executeAll(conn, statements, schema, setSchemaSql, resetSchemaSql, runner);
conn.commit();
return result(totalAffected, start);
} catch (Exception e) {
@ -58,9 +80,10 @@ public final class TransactionExecutor {
List<String> statements,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
StatementRunner runner
) throws Exception {
applySchema(conn, schema, setSchemaSql);
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
long totalAffected = 0;
for (String statement : statements) {
try (Statement stmt = conn.createStatement()) {
@ -70,8 +93,13 @@ public final class TransactionExecutor {
return totalAffected;
}
private static void applySchema(Connection conn, String schema, Function<String, String> setSchemaSql) throws Exception {
JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql);
private static void applySchema(
Connection conn,
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql
) throws Exception {
JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql, resetSchemaSql);
}
private static boolean supportsTransactions(Connection conn) {

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

@ -85,6 +85,17 @@ class JdbcExecutorTest {
assertEquals(Arrays.asList("execute:USE APP", "setSchema:APP"), calls);
}
@Test
void schemaSwitcherRestoresOriginalContextWhenNextQueryHasNoSchema() throws Exception {
List<String> calls = new ArrayList<>();
Connection connection = schemaConnection(calls, false);
JdbcSchemaSwitcher.apply(connection, "APP", schema -> "SET search_path TO " + schema, () -> "RESET search_path");
JdbcSchemaSwitcher.apply(connection, null, schema -> "SET search_path TO " + schema, () -> "RESET search_path");
assertEquals(Arrays.asList("execute:SET search_path TO APP", "execute:RESET search_path"), calls);
}
private static ResultSet resultSet(byte[] bytes, StringSupplier stringSupplier) {
return resultSet(bytes, stringSupplier, false);
}

8
agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java
View File

@ -23,11 +23,14 @@ class PostgresLikeAgentTest {
TestPostgresLikeAgent agent = new TestPostgresLikeAgent();
agent.connect(new ConnectParams());
agent.listDatabases();
agent.listSchemas();
agent.listTables("app");
agent.listObjects("app");
agent.getObjectSource("app", "refresh_orders", "FUNCTION");
agent.getColumns("app", "orders");
agent.listCheckConstraintsForTest("app", "orders");
agent.listIndexes("app", "orders");
agent.listForeignKeys("app", "orders");
agent.listTriggers("app", "orders");
@ -35,13 +38,18 @@ class PostgresLikeAgentTest {
assertFalse(sql.contains("FROM information_schema"), sql);
assertFalse(sql.contains("JOIN information_schema"), sql);
assertTrue(sql.contains("pg_catalog.pg_database"), sql);
assertTrue(sql.contains("pg_catalog.pg_namespace"), sql);
assertTrue(sql.contains("pg_catalog.pg_class"), sql);
assertTrue(sql.contains("pg_catalog.pg_proc"), sql);
assertTrue(sql.contains("pg_catalog.pg_index"), sql);
assertTrue(sql.contains("pg_catalog.pg_attribute"), sql);
assertTrue(sql.contains("pg_catalog.pg_constraint"), sql);
assertTrue(sql.contains("pg_catalog.pg_get_constraintdef"), sql);
assertTrue(sql.contains("pg_catalog.pg_trigger"), sql);
assertFalse(sql.contains("FROM pg_database"), sql);
assertFalse(sql.contains("FROM pg_proc"), sql);
assertFalse(sql.contains("FROM pg_index"), sql);
assertFalse(sql.contains(" AS key "), sql);
assertFalse(sql.contains(" key."), sql);
}

2
agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java
View File

@ -37,7 +37,7 @@ class HighgoAgentTest extends JdbcFakeExecutionBehaviorTest {
Assertions.assertEquals(
Arrays.asList(
"SELECT pg_get_viewdef(to_regclass(?), true)",
"SELECT pg_catalog.pg_get_viewdef(pg_catalog.to_regclass(?), true)",
"param:1=\"bad\"\"schema\".\"view's name\""
),
JdbcMetadataSqlFake.statements

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

@ -79,7 +79,7 @@ public final class KingbaseAgent extends PostgresLikeAgent {
private static boolean mysqlSqlModeExists(Connection connection) {
try (Statement stmt = connection.createStatement();
ResultSet rs = stmt.executeQuery("SELECT 1 FROM sys_settings WHERE LOWER(name) = 'sql_mode'")) {
ResultSet rs = stmt.executeQuery("SELECT 1 FROM sys_catalog.sys_settings WHERE LOWER(name) = 'sql_mode'")) {
return rs.next();
} catch (Exception ignored) {
return false;
@ -102,7 +102,7 @@ public final class KingbaseAgent extends PostgresLikeAgent {
return unchecked(() -> {
for (String sql : List.of(
"SELECT datname AS database_name FROM sys_catalog.sys_database WHERE datistemplate = false AND datallowconn = true ORDER BY datname",
"SELECT datname AS database_name FROM pg_database WHERE datistemplate = false AND datallowconn = true ORDER BY datname"
"SELECT datname AS database_name FROM pg_catalog.pg_database WHERE datistemplate = false AND datallowconn = true ORDER BY datname"
)) {
try {
List<DatabaseInfo> result = queryDatabases(sql);
@ -147,7 +147,7 @@ public final class KingbaseAgent extends PostgresLikeAgent {
"AND UPPER(schema_name) NOT LIKE 'XLOG%' " +
"ORDER BY schema_name"
: "SELECT nspname AS schema_name " +
"FROM sys_namespace " +
"FROM sys_catalog.sys_namespace " +
"WHERE nspname NOT LIKE 'sys_temp_%' " +
"AND nspname NOT LIKE 'sys_toast_temp_%' " +
"ORDER BY nspname";
@ -637,10 +637,9 @@ public final class KingbaseAgent extends PostgresLikeAgent {
@Override
public String setSchemaSQL(String schema) {
if (postgresCatalogMode) return super.setSchemaSQL(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";
// Keep sys_catalog's implicit priority for functions, types, and
// operators. User table references are schema-qualified before execution.
return "SET search_path TO " + JdbcIdentifiers.INSTANCE.doubleQuote(effectiveSchema(schema));
}
@Override

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

@ -51,11 +51,11 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
}
@Test
void schemaSwitchPlacesSysCatalogAfterSelectedSchema() {
void schemaSwitchKeepsSystemCatalogImplicitlyFirst() {
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"));
Assertions.assertEquals("SET search_path TO \"app\"", agent.setSchemaSQL("app"));
Assertions.assertEquals("SET search_path TO \"app\"\"prod\"", agent.setSchemaSQL("app\"prod"));
}
@Test
@ -84,14 +84,14 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
agent.setMysqlCompatMode(true);
TestSupport.setPrivateConnection(agent, preparedConnectionWithFailures(
sql,
List.of("sys_catalog.sys_database", "FROM pg_database"),
List.of("sys_catalog.sys_database", "FROM pg_catalog.pg_database"),
resultSet(new String[]{"database_name"}, new Object[][]{{"TEST"}})
));
Assertions.assertEquals("TEST", agent.listDatabases().get(0).getName());
Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_database"), sql.get(0));
Assertions.assertTrue(sql.get(0).contains("datallowconn = true"), sql.get(0));
Assertions.assertTrue(sql.get(1).contains("FROM pg_database"), sql.get(1));
Assertions.assertTrue(sql.get(1).contains("FROM pg_catalog.pg_database"), sql.get(1));
Assertions.assertTrue(sql.get(1).contains("datallowconn = true"), sql.get(1));
Assertions.assertEquals("SELECT current_database() AS database_name", sql.get(2));
}
@ -127,7 +127,7 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
Assertions.assertEquals(1, databases.size());
Assertions.assertEquals("test", databases.get(0).getName());
Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_database"), sql.get(0));
Assertions.assertTrue(sql.get(1).contains("FROM pg_database"), sql.get(1));
Assertions.assertTrue(sql.get(1).contains("FROM pg_catalog.pg_database"), sql.get(1));
Assertions.assertTrue(sql.get(1).contains("datallowconn = true"), sql.get(1));
}
@ -141,7 +141,7 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
)));
Assertions.assertEquals(Arrays.asList("public", "sys_catalog"), agent.listSchemas());
Assertions.assertTrue(sql.get(0).contains("FROM sys_namespace"), sql.get(0));
Assertions.assertTrue(sql.get(0).contains("FROM sys_catalog.sys_namespace"), sql.get(0));
Assertions.assertFalse(sql.get(0).contains("SYS%"), sql.get(0));
}
@ -164,6 +164,7 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
Assertions.assertEquals("SELECT 1 FROM pg_catalog.pg_namespace WHERE 1 = 0", sql.get(1));
Assertions.assertTrue(sql.get(2).contains("FROM pg_catalog.pg_namespace"), sql.get(2));
Assertions.assertEquals("SET search_path TO \"app\"", agent.setSchemaSQL("app"));
Assertions.assertEquals("RESET search_path", agent.resetSchemaSQL());
}
@Test
@ -179,7 +180,7 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
Assertions.assertTrue(agent.isMysqlCompatMode());
Assertions.assertEquals("`", agent.getIdentifierQuote());
Assertions.assertEquals("SELECT 1 FROM sys_catalog.sys_namespace WHERE 1 = 0", sql.get(0));
Assertions.assertEquals("SELECT 1 FROM sys_settings WHERE LOWER(name) = 'sql_mode'", sql.get(1));
Assertions.assertEquals("SELECT 1 FROM sys_catalog.sys_settings WHERE LOWER(name) = 'sql_mode'", sql.get(1));
}
@Test
@ -194,7 +195,7 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest {
Assertions.assertTrue(isSqlServerIdentityCatalogMode(agent));
Assertions.assertEquals("SELECT 1 FROM sys_catalog.sys_namespace WHERE 1 = 0", sql.get(0));
Assertions.assertEquals("SELECT 1 FROM sys_settings WHERE LOWER(name) = 'sql_mode'", sql.get(1));
Assertions.assertEquals("SELECT 1 FROM sys_catalog.sys_settings WHERE LOWER(name) = 'sql_mode'", sql.get(1));
Assertions.assertEquals("SELECT 1 FROM sys.identity_columns WHERE 1 = 0", sql.get(2));
}

153
crates/dbx-core/src/query.rs
View File

@ -5,7 +5,9 @@ use duckdb::types::{TimeUnit, Value, ValueRef};
use futures::StreamExt;
use mysql_async::prelude::Queryable;
use serde::Serialize;
use sqlparser::ast::{visit_relations_mut, Ident, ObjectName, ObjectNamePart, ObjectType, Statement};
use sqlparser::ast::{
visit_relations_mut, Ident, ObjectName, ObjectNamePart, ObjectType, Statement, TableFactor, VisitMut, VisitorMut,
};
use sqlparser::dialect::{GenericDialect, PostgreSqlDialect};
use sqlparser::parser::Parser;
use std::collections::HashSet;
@ -218,12 +220,16 @@ fn schema_for_execution_context(db_type: Option<DatabaseType>, schema: Option<&s
}
fn sql_for_execution_context(db_type: Option<DatabaseType>, sql: &str, schema: Option<&str>) -> String {
if matches!(db_type, Some(DatabaseType::Iris)) {
if let Some(schema) = schema.map(str::trim).filter(|schema| !schema.is_empty()) {
return qualify_iris_unqualified_dml(sql, schema).unwrap_or_else(|| sql.to_string());
let Some(schema) = schema.map(str::trim).filter(|schema| !schema.is_empty()) else {
return sql.to_string();
};
match db_type {
Some(DatabaseType::Iris) => qualify_iris_unqualified_dml(sql, schema).unwrap_or_else(|| sql.to_string()),
Some(DatabaseType::Kingbase) => {
qualify_kingbase_unqualified_relations(sql, schema).unwrap_or_else(|| sql.to_string())
}
_ => sql.to_string(),
}
sql.to_string()
}
fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option<String> {
@ -235,12 +241,12 @@ fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option<String> {
let mut changed = false;
for statement in &mut statements {
if !iris_statement_uses_schema_search_path(statement) {
if !statement_uses_schema_context(statement) {
continue;
}
let cte_names = iris_statement_cte_names(statement);
let cte_names = statement_cte_names(statement);
let _ = visit_relations_mut(statement, |name| {
if qualify_iris_relation_name(name, schema, &cte_names) {
if qualify_unqualified_relation_name(name, schema, &cte_names) {
changed = true;
}
ControlFlow::<()>::Continue(())
@ -250,7 +256,63 @@ fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option<String> {
changed.then(|| statements.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))
}
fn iris_statement_uses_schema_search_path(statement: &Statement) -> bool {
fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str) -> Option<String> {
let dialect = PostgreSqlDialect {};
let mut statements = Parser::parse_sql(&dialect, sql).ok()?;
if statements.is_empty() {
return None;
}
let mut changed = false;
for statement in &mut statements {
if !statement_uses_schema_context(statement) {
continue;
}
let cte_names = statement_cte_names(statement);
let mut qualifier =
KingbaseRelationQualifier { schema, cte_names: &cte_names, parameterized_table_depth: 0, changed: false };
let _ = statement.visit(&mut qualifier);
changed |= qualifier.changed;
}
changed.then(|| statements.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))
}
struct KingbaseRelationQualifier<'a> {
schema: &'a str,
cte_names: &'a HashSet<String>,
parameterized_table_depth: usize,
changed: bool,
}
impl VisitorMut for KingbaseRelationQualifier<'_> {
type Break = ();
fn pre_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow<Self::Break> {
if matches!(table_factor, TableFactor::Table { args: Some(_), .. }) {
self.parameterized_table_depth += 1;
}
ControlFlow::Continue(())
}
fn post_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow<Self::Break> {
if matches!(table_factor, TableFactor::Table { args: Some(_), .. }) {
self.parameterized_table_depth = self.parameterized_table_depth.saturating_sub(1);
}
ControlFlow::Continue(())
}
fn post_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow<Self::Break> {
if self.parameterized_table_depth == 0
&& qualify_unqualified_relation_name(relation, self.schema, self.cte_names)
{
self.changed = true;
}
ControlFlow::Continue(())
}
}
fn statement_uses_schema_context(statement: &Statement) -> bool {
matches!(
statement,
Statement::Query(_)
@ -261,7 +323,7 @@ fn iris_statement_uses_schema_search_path(statement: &Statement) -> bool {
)
}
fn qualify_iris_relation_name(name: &mut ObjectName, schema: &str, cte_names: &HashSet<String>) -> bool {
fn qualify_unqualified_relation_name(name: &mut ObjectName, schema: &str, cte_names: &HashSet<String>) -> bool {
let [ObjectNamePart::Identifier(table)] = name.0.as_slice() else {
return false;
};
@ -274,33 +336,37 @@ fn qualify_iris_relation_name(name: &mut ObjectName, schema: &str, cte_names: &H
true
}
fn iris_statement_cte_names(statement: &Statement) -> HashSet<String> {
fn statement_cte_names(statement: &Statement) -> HashSet<String> {
let mut names = HashSet::new();
collect_iris_statement_cte_names(statement, &mut names);
collect_statement_cte_names(statement, &mut names);
names
}
fn collect_iris_statement_cte_names(statement: &Statement, names: &mut HashSet<String>) {
fn collect_statement_cte_names(statement: &Statement, names: &mut HashSet<String>) {
match statement {
Statement::Query(query) => collect_iris_query_cte_names(query, names),
Statement::Query(query) => collect_query_cte_names(query, names),
Statement::Insert(insert) => {
if let Some(source) = &insert.source {
collect_iris_query_cte_names(source, names);
collect_query_cte_names(source, names);
}
}
_ => {}
}
}
fn collect_iris_query_cte_names(query: &sqlparser::ast::Query, names: &mut HashSet<String>) {
fn collect_query_cte_names(query: &sqlparser::ast::Query, names: &mut HashSet<String>) {
if let Some(with) = &query.with {
for cte in &with.cte_tables {
names.insert(cte.alias.name.value.to_ascii_uppercase());
collect_iris_query_cte_names(&cte.query, names);
collect_query_cte_names(&cte.query, names);
}
}
}
fn qualifies_unqualified_agent_relations(db_type: Option<DatabaseType>) -> bool {
matches!(db_type, Some(DatabaseType::Iris | DatabaseType::Kingbase))
}
#[derive(Clone, Debug, Default)]
pub struct QueryExecutionOptions {
pub max_rows: Option<usize>,
@ -2240,7 +2306,7 @@ pub async fn execute_statements(
let db_type = connection_database_type_for_pool_key(state, &pool_key).await;
let execution_schema = schema_for_execution_context(db_type, schema);
let rewritten_statements;
let statements = if matches!(db_type, Some(DatabaseType::Iris)) {
let statements = if qualifies_unqualified_agent_relations(db_type) {
rewritten_statements =
statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::<Vec<_>>();
rewritten_statements.as_slice()
@ -2651,7 +2717,7 @@ async fn exec_tx_explicit_inner(
let db_type = connection_database_type_for_pool_key(state, pool_key).await;
let execution_schema = schema_for_execution_context(db_type, schema);
let rewritten_statements;
let statements = if matches!(db_type, Some(DatabaseType::Iris)) {
let statements = if qualifies_unqualified_agent_relations(db_type) {
rewritten_statements =
statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::<Vec<_>>();
rewritten_statements.as_slice()
@ -4363,6 +4429,55 @@ mod tests {
);
}
#[test]
fn kingbase_execution_context_qualifies_only_unqualified_relations() {
assert_eq!(
sql_for_execution_context(Some(DatabaseType::Kingbase), "SELECT * FROM sys_user", Some("app")),
"SELECT * FROM \"app\".sys_user"
);
assert_eq!(
sql_for_execution_context(Some(DatabaseType::Kingbase), "SELECT * FROM other_schema.sys_user", Some("app")),
"SELECT * FROM other_schema.sys_user"
);
let mixed = sql_for_execution_context(
Some(DatabaseType::Kingbase),
"SELECT pg_typeof(u.id) FROM generate_series(1, 2) AS n JOIN sys_user u ON true",
Some("app"),
);
assert!(mixed.contains("FROM generate_series(1, 2) AS n"), "{mixed}");
assert!(mixed.contains("JOIN \"app\".sys_user u"), "{mixed}");
assert!(mixed.contains("pg_typeof(u.id)"), "{mixed}");
}
#[test]
fn kingbase_execution_context_preserves_ctes_functions_types_and_unsupported_sql() {
assert_eq!(
sql_for_execution_context(
Some(DatabaseType::Kingbase),
"WITH current_user AS (SELECT * FROM sys_user) SELECT * FROM current_user",
Some("APP")
),
"WITH current_user AS (SELECT * FROM \"APP\".sys_user) SELECT * FROM current_user"
);
assert_eq!(
sql_for_execution_context(
Some(DatabaseType::Kingbase),
"SELECT pg_typeof(1::int), current_user",
Some("APP")
),
"SELECT pg_typeof(1::int), current_user"
);
assert_eq!(
sql_for_execution_context(Some(DatabaseType::Kingbase), "CREATE TABLE sys_user (id INT)", Some("APP")),
"CREATE TABLE sys_user (id INT)"
);
assert_eq!(
sql_for_execution_context(Some(DatabaseType::Kingbase), "SELECT * FROM", Some("APP")),
"SELECT * FROM"
);
}
#[test]
fn parses_postgres_drop_database_target() {
assert_eq!(parse_drop_database_target("DROP DATABASE vaultwarden;"), Some("vaultwarden".to_string()));