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>

agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java

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

agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java

@@ -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() {

agents/common/src/main/java/com/dbx/agent/BatchExecutor.java

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

agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java

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

agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java

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

agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java

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

agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java

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

agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java

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