diff --git a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java index d48c0fc98..7e6dedce2 100644 --- a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java +++ b/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 statements, String schema) { - return TransactionExecutor.executeUpdateStatements(requireConnected(), statements, schema, this::setSchemaSQL); + return TransactionExecutor.executeUpdateStatements( + requireConnected(), + statements, + schema, + this::setSchemaSQL, + this::resetSchemaSQL + ); } @Override diff --git a/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java index d8c107622..7731bcaec 100644 --- a/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java +++ b/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 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 statements, String schema) { - return BatchExecutor.executeBatchStatements(requireConnected(), statements, schema, this::setSchemaSQL); + return BatchExecutor.executeBatchStatements( + requireConnected(), + statements, + schema, + this::setSchemaSQL, + this::resetSchemaSQL + ); } protected Connection requireConnected() { 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 134ada07c..8075fc9b8 100644 --- a/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java +++ b/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 statements, String schema, Function setSchemaSql + ) { + return executeBatchStatements(conn, statements, schema, setSchemaSql, () -> ""); + } + + public static QueryResult executeBatchStatements( + Connection conn, + List statements, + String schema, + Function setSchemaSql, + Supplier 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 setSchemaSql) throws Exception { - JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql); + private static void applySchema( + Connection conn, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql + ) throws Exception { + JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql, resetSchemaSql); } private static T unchecked(ThrowingSupplier supplier) { diff --git a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java index d33011cba..7f7e79095 100644 --- a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java +++ b/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 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, 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 039980de5..e21ea8680 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/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 setSchemaSql, + Supplier 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 setSchemaSql, + Supplier 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 setSchemaSql, + Supplier 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 setSchemaSql, + Supplier resetSchemaSql, QueryPageOptions options, ResultValueReader valueReader, ConcurrentHashMap 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 setSchemaSql) throws SQLException { + private void applySchema( + Connection conn, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql + ) throws SQLException { try { - JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql); + JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql, resetSchemaSql); } catch (SQLException e) { throw e; } catch (Exception 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 index 76ffa5ab8..32d428259 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java +++ b/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 RESET_SQL_BY_CONNECTION = + Collections.synchronizedMap(new WeakHashMap<>()); + private JdbcSchemaSwitcher() { } static void apply(Connection conn, String schema, Function setSchemaSql) throws Exception { + apply(conn, schema, setSchemaSql, () -> ""); + } + + static void apply( + Connection conn, + String schema, + Function setSchemaSql, + Supplier 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 setSchemaSql) { + private static void rememberResetSql(Connection conn, Supplier 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 schemaSqlSupplier) { String schemaSql; try { - schemaSql = setSchemaSql.apply(schema); + schemaSql = schemaSqlSupplier.get(); } catch (RuntimeException e) { return new SchemaSqlResult(true, e); } diff --git a/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java index 36f70a761..220a25449 100644 --- a/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java +++ b/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 listDatabases() { return unchecked(() -> { List 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(); } 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 2bd12ed46..3f1450701 100644 --- a/agents/common/src/main/java/com/dbx/agent/TransactionExecutor.java +++ b/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 setSchemaSql ) { - return executeStatements(conn, statements, schema, setSchemaSql, new StatementRunner() { + return executeUpdateStatements(conn, statements, schema, setSchemaSql, () -> ""); + } + + public static QueryResult executeUpdateStatements( + Connection conn, + List statements, + String schema, + Function setSchemaSql, + Supplier 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 setSchemaSql, StatementRunner runner + ) { + return executeStatements(conn, statements, schema, setSchemaSql, () -> "", runner); + } + + public static QueryResult executeStatements( + Connection conn, + List statements, + String schema, + Function setSchemaSql, + Supplier 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 statements, String schema, Function setSchemaSql, + Supplier 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 setSchemaSql) throws Exception { - JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql); + private static void applySchema( + Connection conn, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql + ) throws Exception { + JdbcSchemaSwitcher.apply(conn, schema, setSchemaSql, resetSchemaSql); } private static boolean supportsTransactions(Connection conn) { 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 6d54dfc0f..3210f0b9d 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -85,6 +85,17 @@ class JdbcExecutorTest { assertEquals(Arrays.asList("execute:USE APP", "setSchema:APP"), calls); } + @Test + void schemaSwitcherRestoresOriginalContextWhenNextQueryHasNoSchema() throws Exception { + List 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); } diff --git a/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java b/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java index f3d427917..0960ed6bd 100644 --- a/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java +++ b/agents/common/src/test/java/com/dbx/agent/PostgresLikeAgentTest.java @@ -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); } diff --git a/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java b/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java index 218631226..a7928ecf9 100644 --- a/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java +++ b/agents/drivers/highgo/src/test/java/com/dbx/agent/highgo/HighgoAgentTest.java @@ -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 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 5e6c22d0b..0623abd48 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 @@ -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 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 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 c390d4222..cbd7a65d8 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 @@ -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)); } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 2c06d41e9..7b0353455 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -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, schema: Option<&s } fn sql_for_execution_context(db_type: Option, 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 { @@ -235,12 +241,12 @@ fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option { 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 { changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) } -fn iris_statement_uses_schema_search_path(statement: &Statement) -> bool { +fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str) -> Option { + 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::>().join("; ")) +} + +struct KingbaseRelationQualifier<'a> { + schema: &'a str, + cte_names: &'a HashSet, + parameterized_table_depth: usize, + changed: bool, +} + +impl VisitorMut for KingbaseRelationQualifier<'_> { + type Break = (); + + fn pre_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow { + 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 { + 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 { + 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) -> bool { +fn qualify_unqualified_relation_name(name: &mut ObjectName, schema: &str, cte_names: &HashSet) -> 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 { +fn statement_cte_names(statement: &Statement) -> HashSet { 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) { +fn collect_statement_cte_names(statement: &Statement, names: &mut HashSet) { 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) { +fn collect_query_cte_names(query: &sqlparser::ast::Query, names: &mut HashSet) { 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) -> bool { + matches!(db_type, Some(DatabaseType::Iris | DatabaseType::Kingbase)) +} + #[derive(Clone, Debug, Default)] pub struct QueryExecutionOptions { pub max_rows: Option, @@ -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::>(); 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::>(); 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()));