From 2da2db6c7c0464fffd00663e56adaea7ffc6e4b6 Mon Sep 17 00:00:00 2001 From: Shaitan <105581038+sha174n@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:57:45 +0100 Subject: [PATCH] feat(sql): schema-qualified table denylist + information_schema/lo_* defaults (#41120) Co-authored-by: Claude Opus 4.8 (1M context) --- UPDATING.md | 11 + superset/commands/sql_lab/estimate.py | 79 ++- superset/config.py | 42 ++ superset/models/core.py | 40 ++ superset/models/helpers.py | 76 ++- superset/sql/execution/executor.py | 62 ++- superset/sql/parse.py | 227 ++++++++- superset/sql_lab.py | 35 +- .../commands/sql_lab/test_estimate.py | 78 +-- tests/unit_tests/models/core_test.py | 39 ++ tests/unit_tests/models/helpers_test.py | 55 +++ .../unit_tests/sql/execution/test_executor.py | 86 +++- tests/unit_tests/sql/parse_tests.py | 452 ++++++++++++++++++ 13 files changed, 1143 insertions(+), 139 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 5a41479ac9e..8eb28cfbec0 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -24,6 +24,17 @@ assists people when migrating to a new version. ## Next +### SQL Lab denies large-object and information_schema access by default + +`DISALLOWED_SQL_FUNCTIONS` and `DISALLOWED_SQL_TABLES` now ship with additional default entries, so SQL Lab and chart-data queries that reference them are rejected where they were previously allowed: + +- PostgreSQL large-object routines (`lo_from_bytea`, `lo_export`, `lo_import`, `lo_put`, `lo_create`, `lo_creat`, `lowrite`, `lo_get`, `loread`, `lo_unlink`), which read and write bytes on the database server's filesystem. +- The SQL-standard `information_schema` views (`tables`, `columns`, `routines`, `views`, the privilege/grant views, etc.), which expose table, column, privilege, and view-definition metadata across the whole database. + +Deployments that legitimately query these (for example tooling that introspects `information_schema`) can restore the previous behavior by overriding `DISALLOWED_SQL_FUNCTIONS` / `DISALLOWED_SQL_TABLES` in `superset_config.py` to drop the entries they need. + +Because the denylist now resolves the effective schema through the query-aware path, PostgreSQL queries that change the `search_path` (e.g. `SET search_path = ...`) are rejected on the SQL Lab execution and cost-estimate paths whenever any `DISALLOWED_SQL_TABLES` entry is configured (the default for PostgreSQL), matching the behavior previously applied only when `RLS_IN_SQLLAB` was enabled. + ### SQL parser input length cap (SQL_MAX_PARSE_LENGTH) The SQL parser now rejects scripts whose UTF-8 byte length exceeds the new diff --git a/superset/commands/sql_lab/estimate.py b/superset/commands/sql_lab/estimate.py index 36bae9de412..7d7ef228c2e 100644 --- a/superset/commands/sql_lab/estimate.py +++ b/superset/commands/sql_lab/estimate.py @@ -34,7 +34,6 @@ from superset.exceptions import ( ) from superset.jinja_context import get_template_processor from superset.models.core import Database -from superset.models.sql_lab import Query from superset.sql.parse import SQLScript from superset.utils import core as utils from superset.utils.rls import apply_rls @@ -102,57 +101,43 @@ class QueryEstimationCommand(BaseCommand): db_engine_spec.engine, set(), ) - if disallowed_tables and parsed_script.check_tables_present(disallowed_tables): - found_tables = set() - for statement in parsed_script.statements: - present = {table.table.lower() for table in statement.tables} - for table in disallowed_tables: - if table.lower() in present: - found_tables.add(table) - raise SupersetDisallowedSQLTableException(found_tables or disallowed_tables) + rls_enabled = is_feature_enabled("RLS_IN_SQLLAB") + + # Resolve the effective per-query schema once, the same way the execution + # path does (``sql_lab.execute_sql_statements``), but only when a control + # below actually needs it. Going through ``get_default_schema_for_query`` + # rather than the static ``get_default_schema`` runs engine-specific + # per-query security gates too — e.g. ``PostgresEngineSpec`` rejects a + # query that sets ``search_path`` — and resolves unqualified references to + # the schema the engine uses at runtime, so both the denylist check and + # RLS injection match the execution path exactly. + catalog: str | None = None + effective_schema = "" + if disallowed_tables or rls_enabled: + catalog = self._catalog or self._database.get_default_catalog() + resolved_schema = self._database.resolve_query_default_schema( + self._sql, self._schema, catalog, self._template_params + ) + # An explicit schema still wins for matching/RLS targeting; otherwise + # fall back to the runtime-resolved default. + effective_schema = self._schema or resolved_schema or "" + + if disallowed_tables: + # Honors schema-qualified denylist entries (e.g. + # ``information_schema.tables``) and reports only the tables + # actually referenced by the query. + found_tables = parsed_script.get_disallowed_tables( + disallowed_tables, effective_schema + ) + if found_tables: + raise SupersetDisallowedSQLTableException(found_tables) if parsed_script.has_mutation() and not self._database.allow_dml: raise SupersetDMLNotAllowedException() - if is_feature_enabled("RLS_IN_SQLLAB"): - # Resolve the default catalog/schema the same way the execution path - # does (``sql_lab.execute_sql_statements``) before injecting RLS. - # Crucially this goes through ``get_default_schema_for_query`` rather - # than the plain ``get_default_schema``, so engine-specific per-query - # security gates run too — e.g. ``PostgresEngineSpec`` rejects a query - # that sets ``search_path``. Resolving against the static default - # schema instead would both skip that gate and let unqualified tables - # dodge the RLS predicates the real query enforces, defeating the - # security parity this command exists to provide. - catalog = self._catalog or self._database.get_default_catalog() - # Build a transient (unsaved) Query so the engine spec can resolve the - # effective per-query schema exactly as the executor does. Mirror the - # probe built in ``SupersetSecurityManager.raise_for_access``: set a - # ``client_id`` (the column is ``nullable=False``) and expunge it, so - # the ``database`` backref's ``cascade="all, delete-orphan"`` cannot - # autoflush this incomplete row into the session when ``apply_rls`` - # issues its own ``db.session`` query below. - probe_query = Query( - database=self._database, - sql=self._sql, - schema=self._schema or None, - catalog=catalog, - client_id=utils.shortid()[:10], - user_id=utils.get_user_id(), - ) - db.session.expunge(probe_query) - # Always resolve through ``get_default_schema_for_query`` — even when - # the caller pinned a schema — so the engine's per-query security gate - # runs (e.g. ``PostgresEngineSpec`` rejects a query that sets - # ``search_path``), exactly as the executor does unconditionally. Only - # the resulting value falls back to the resolved default; an explicit - # schema still wins for the RLS predicate target. - resolved_schema = self._database.get_default_schema_for_query( - probe_query, self._template_params - ) - schema = self._schema or resolved_schema or "" + if rls_enabled: for statement in parsed_script.statements: - apply_rls(self._database, catalog, schema, statement) + apply_rls(self._database, catalog, effective_schema, statement) return parsed_script.format() return sql diff --git a/superset/config.py b/superset/config.py index 76f9485e2ef..0f918224fcc 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1958,6 +1958,24 @@ DISALLOWED_SQL_FUNCTIONS: dict[str, set[str]] = { "pg_read_file", "pg_ls_dir", "pg_read_binary_file", + # PostgreSQL large-object functions: writers can plant arbitrary + # bytes on the server filesystem (lo_export, lo_from_bytea, lowrite, + # lo_put, lo_create, lo_creat, lo_import), readers can pull bytes back + # out (lo_get, loread), lo_truncate/lo_truncate64 shrink existing + # large objects, and lo_unlink deletes them outright. Defense-in-depth + # on top of is_mutating()'s function-name check. + "lo_from_bytea", + "lo_export", + "lo_import", + "lo_put", + "lo_create", + "lo_creat", + "lowrite", + "lo_get", + "loread", + "lo_truncate", + "lo_truncate64", + "lo_unlink", # XML functions that can execute SQL "database_to_xml", "database_to_xmlschema", @@ -2048,6 +2066,30 @@ DISALLOWED_SQL_TABLES: dict[str, set[str]] = { "pg_stat_replication", "pg_stat_wal_receiver", "pg_user", + # The SQL-standard `information_schema` views expose table / + # column / privilege / view-definition metadata across the entire + # database role the connection user can see. Entries are + # schema-qualified so `check_tables_present` only matches when the + # reference resolves to `information_schema.` -- either written + # explicitly or as an unqualified name under an `information_schema` + # search_path -- not any user table that happens to share a name. + "information_schema.tables", + "information_schema.columns", + "information_schema.schemata", + "information_schema.views", + "information_schema.routines", + "information_schema.role_table_grants", + "information_schema.role_column_grants", + "information_schema.role_routine_grants", + "information_schema.table_privileges", + "information_schema.column_privileges", + "information_schema.usage_privileges", + "information_schema.key_column_usage", + "information_schema.table_constraints", + "information_schema.referential_constraints", + "information_schema.view_table_usage", + "information_schema.applicable_roles", + "information_schema.enabled_roles", }, "mysql": { "mysql.user", diff --git a/superset/models/core.py b/superset/models/core.py index 38884c75f73..5a8c6e6e173 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -698,6 +698,46 @@ class Database(CoreDatabase, AuditMixinNullable, ImportExportMixin): # pylint: self, query, template_params ) + def resolve_query_default_schema( + self, + sql: str, + schema: str | None, + catalog: str | None, + template_params: Optional[dict[str, Any]] = None, + ) -> str | None: + """ + Resolve the effective per-query default schema for a SQL string through + the query-aware :meth:`get_default_schema_for_query`. + + Builds a transient (unsaved) ``Query`` probe so the engine spec resolves + the schema exactly as execution does -- running engine-specific per-query + security gates too -- then expunges it so the ``database`` backref's + ``cascade="all, delete-orphan"`` cannot autoflush this incomplete row + into the session. Centralizes the probe construction shared by the SQL + Lab executor, the cost-estimate command, and datasource denylist checks + so the schema-resolution behavior stays in sync across these + security-sensitive paths. + + :param sql: Original (pre-render) SQL the query will execute + :param schema: Explicit per-query schema, if any + :param catalog: Resolved catalog + :param template_params: Jinja template parameters, if any + :returns: The runtime-resolved default schema, or None + """ + from superset.models.sql_lab import Query + + probe_query = Query( + database=self, + sql=sql, + schema=schema or None, + catalog=catalog, + client_id=utils.shortid()[:10], + user_id=utils.get_user_id(), + ) + if probe_query in db.session: + db.session.expunge(probe_query) + return self.get_default_schema_for_query(probe_query, template_params) + @staticmethod def post_process_df(df: pd.DataFrame) -> pd.DataFrame: def column_needs_conversion(df_series: pd.Series) -> bool: diff --git a/superset/models/helpers.py b/superset/models/helpers.py index cf44392bbb8..77963bdd288 100644 --- a/superset/models/helpers.py +++ b/superset/models/helpers.py @@ -1187,6 +1187,31 @@ class ExploreMixin: # pylint: disable=too-many-public-methods # for datasources of type query return [] + def _resolve_denylist_schema(self, sql: str) -> Optional[str]: + """ + Resolve the effective default schema for ``DISALLOWED_SQL_TABLES`` + checks, memoized per datasource instance. + + An explicit datasource schema always wins. Otherwise the schema is + resolved through the query-aware ``get_default_schema_for_query`` (not + the static ``get_default_schema``) so the value matches what the engine + uses at runtime -- honoring dynamic-schema engines and URI/connect-arg + schema rules -- exactly like the SQL Lab / executor denylist gate. + + The result is cached on the instance so adhoc-expression validation, + which runs once per selected column/metric/order-by, resolves the schema + at most once instead of issuing a fallback inspector round-trip per + expression. + """ + if self.schema: + return self.schema + if not hasattr(self, "_denylist_default_schema"): + catalog = self.catalog or self.database.get_default_catalog() + self._denylist_default_schema = self.database.resolve_query_default_schema( + sql, None, catalog + ) + return self._denylist_default_schema + def _process_sql_expression( # pylint: disable=too-many-arguments self, expression: Optional[str], @@ -1236,24 +1261,23 @@ class ExploreMixin: # pylint: disable=too-many-public-methods disallowed_functions ): raise SupersetDisallowedSQLFunctionException(disallowed_functions) - if disallowed_tables and parsed.check_tables_present(disallowed_tables): + if disallowed_tables: # Report only the tables actually found in the expression, # mirroring the canonical execution-time gate in # `superset.sql_lab._validate_query` so the user-facing - # error doesn't echo the operator's full denylist. - present_tables = { - table.table.lower() - for statement in parsed.statements - for table in statement.tables - } - found_tables = { - table - for table in disallowed_tables - if table.lower() in present_tables - } - raise SupersetDisallowedSQLTableException( - found_tables or disallowed_tables + # error doesn't echo the operator's full denylist. Honors + # schema-qualified denylist entries (e.g. + # ``information_schema.tables``) and resolves unqualified + # references against the same runtime schema the SQL Lab / + # executor gate resolves them to (via the query-aware + # resolver, not the static inspector default), so both + # surfaces enforce the denylist consistently. + effective_schema = self._resolve_denylist_schema(sql_to_check) + found_tables = parsed.get_disallowed_tables( + disallowed_tables, effective_schema ) + if found_tables: + raise SupersetDisallowedSQLTableException(found_tables) return expression def _process_select_expression( @@ -1480,19 +1504,21 @@ class ExploreMixin: # pylint: disable=too-many-public-methods disallowed_functions ): raise SupersetDisallowedSQLFunctionException(disallowed_functions) - if disallowed_tables and parsed_script.check_tables_present(disallowed_tables): + if disallowed_tables: # Report only the tables actually found in the query, mirroring the # canonical execution-time gate so the user-facing error doesn't - # echo the operator's full denylist. - present_tables = { - table.table.lower() - for statement in parsed_script.statements - for table in statement.tables - } - found_tables = { - table for table in disallowed_tables if table.lower() in present_tables - } - raise SupersetDisallowedSQLTableException(found_tables or disallowed_tables) + # echo the operator's full denylist. Honors schema-qualified + # denylist entries (e.g. ``information_schema.tables``) and resolves + # unqualified references against the same runtime schema the SQL Lab + # / executor gate resolves them to (via the query-aware resolver, + # not the static inspector default), so both surfaces enforce the + # denylist consistently. + effective_schema = self._resolve_denylist_schema(sql) + found_tables = parsed_script.get_disallowed_tables( + disallowed_tables, effective_schema + ) + if found_tables: + raise SupersetDisallowedSQLTableException(found_tables) def query(self, query_obj: QueryObjectDict) -> QueryResult: """ diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py index 13e12fc1a47..f13db9ae901 100644 --- a/superset/sql/execution/executor.py +++ b/superset/sql/execution/executor.py @@ -231,7 +231,7 @@ class SQLExecutor: ) # 2. Security checks on transformed script - self._check_security(transformed_script) + self._check_security(transformed_script, schema) # 3. Get mutation status and format SQL has_mutation = transformed_script.has_mutation() @@ -355,7 +355,7 @@ class SQLExecutor: ) # 2. Security checks on transformed script - self._check_security(transformed_script) + self._check_security(transformed_script, schema) # 3. Get mutation status and format SQL has_mutation = transformed_script.has_mutation() @@ -436,9 +436,21 @@ class SQLExecutor: rendered_sql, self.database.db_engine_spec.engine ) - # 4. Get catalog and schema + # 4. Get catalog and the effective per-query schema. Resolve the schema + # through the query-aware ``get_default_schema_for_query`` rather than the + # static ``get_default_schema``, the same way ``execute_sql_statements`` + # and the estimate path do: it resolves an unqualified reference to the + # schema the engine actually uses at runtime (engines without + # dynamic-schema support ignore the request's selected schema) and runs + # engine-specific per-query security gates (e.g. ``PostgresEngineSpec`` + # rejects a query that sets ``search_path``), so the denylist check and + # RLS injection match execution instead of a schema that may never apply. catalog = opts.catalog or self.database.get_default_catalog() - schema = opts.schema or self.database.get_default_schema(catalog) + # Resolve unconditionally, even when an explicit schema is supplied, so + # the engine's per-query security gate always runs (parity with the + # estimate path); the explicit schema still wins as the effective target. + resolved_schema = self._resolve_query_schema(sql, opts, catalog) + schema = opts.schema or resolved_schema # 5. Apply RLS to transformed script only self._apply_rls_to_script(transformed_script, catalog, schema) @@ -449,11 +461,12 @@ class SQLExecutor: return original_script, transformed_script, catalog, schema - def _check_security(self, script: SQLScript) -> None: + def _check_security(self, script: SQLScript, schema: str | None = None) -> None: """ Perform security checks on prepared SQL script. :param script: Prepared SQLScript + :param schema: Effective schema unqualified references resolve to :raises SupersetSecurityException: If security checks fail """ # Check disallowed functions @@ -469,7 +482,7 @@ class SQLExecutor: ) # Check disallowed tables - if disallowed_tables := self._check_disallowed_tables(script): + if disallowed_tables := self._check_disallowed_tables(script, schema): raise SupersetSecurityException( SupersetError( message=f"Disallowed SQL tables: {', '.join(disallowed_tables)}", @@ -702,11 +715,32 @@ class SQLExecutor: return found if found else None - def _check_disallowed_tables(self, script: SQLScript) -> set[str] | None: + def _resolve_query_schema( + self, sql: str, opts: QueryOptions, catalog: str | None + ) -> str | None: + """ + Resolve the effective per-query default schema through the query-aware + ``get_default_schema_for_query`` so the denylist check and RLS injection + match the schema the engine uses at runtime, and engine-specific + per-query security gates run on this path too. + + :param sql: Original (pre-render) SQL the query will execute + :param opts: Query options (supplies schema, template params) + :param catalog: Resolved catalog + :returns: The runtime-resolved default schema, or None + """ + return self.database.resolve_query_default_schema( + sql, opts.schema, catalog, opts.template_params + ) + + def _check_disallowed_tables( + self, script: SQLScript, schema: str | None = None + ) -> set[str] | None: """ Check for disallowed SQL tables/views. :param script: Parsed SQL script + :param schema: Effective schema unqualified references resolve to :returns: Set of disallowed tables found, or None if none found """ disallowed_config = app.config.get("DISALLOWED_SQL_TABLES", {}) @@ -717,15 +751,11 @@ class SQLExecutor: if not engine_disallowed: return None - # Single-pass AST-based table detection - found: set[str] = set() - for statement in script.statements: - present = {table.table.lower() for table in statement.tables} - for table in engine_disallowed: - if table.lower() in present: - found.add(table) - - return found or None + # Honors schema-qualified denylist entries (e.g. + # ``information_schema.tables``) as well as bare names. The effective + # schema lets an unqualified reference that resolves to it at runtime + # (via the connection ``search_path``) match too. + return script.get_disallowed_tables(engine_disallowed, schema) or None def _apply_rls_to_script( self, script: SQLScript, catalog: str | None, schema: str | None diff --git a/superset/sql/parse.py b/superset/sql/parse.py index 8d1330403ea..c470e43e2da 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -559,15 +559,49 @@ class BaseSQLStatement(Generic[InternalRepresentation]): """ raise NotImplementedError() - def check_tables_present(self, tables: set[str]) -> bool: + def check_tables_present( + self, tables: set[str], default_schema: str | None = None + ) -> bool: """ Check if any of the given tables are present in the statement. :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Schema unqualified references resolve to at + runtime (e.g. the session ``search_path`` / selected schema) :return: True if any of the tables are present """ raise NotImplementedError() + def changes_search_path(self) -> bool: + """ + Check if the statement changes the session ``search_path``. + + Defaults to ``False``; engines whose statements can rebind unqualified + schema resolution override this. + + :return: True if the statement changes the session ``search_path`` + """ + return False + + def get_disallowed_tables( + self, + tables: set[str], + default_schema: str | None = None, + schema_indeterminate: bool = False, + ) -> set[str]: + """ + Return the subset of ``tables`` referenced by this statement. + + :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Schema unqualified references resolve to at + runtime (e.g. the session ``search_path`` / selected schema) + :param schema_indeterminate: When True, unqualified references are + matched against schema-qualified entries too (see + :meth:`SQLStatement.get_disallowed_tables`) + :return: The matched entries, in their original denylist form + """ + raise NotImplementedError() + def get_limit_value(self) -> int | None: """ Get the limit value of the statement. @@ -670,7 +704,10 @@ class SQLStatement(BaseSQLStatement[exp.Expression]): "LO_IMPORT", "LO_PUT", "LO_CREATE", + "LO_CREAT", "LOWRITE", + "LO_TRUNCATE", + "LO_TRUNCATE64", "LO_UNLINK", # PostgreSQL sequence mutators. `SELECT setval('seq', N)` and # `SELECT nextval('seq')` look like reads but change sequence state @@ -1060,15 +1097,135 @@ class SQLStatement(BaseSQLStatement[exp.Expression]): return any(function.upper() in present for function in functions) - def check_tables_present(self, tables: set[str]) -> bool: + def check_tables_present( + self, tables: set[str], default_schema: str | None = None + ) -> bool: """ Check if any of the given tables are present in the statement. + Denylist entries may be bare (``pg_stat_activity``) or + schema-qualified (``information_schema.tables``). Bare entries + match by table name regardless of schema; qualified entries + require the schema to match too. This lets us block all access + to ``information_schema`` without also blocking any + user-authored table that happens to be named ``tables``. + :param tables: Set of table names to check for (case-insensitive) - :return: True if any of the tables are present + :param default_schema: Schema unqualified references resolve to at + runtime (e.g. the session ``search_path`` / selected schema) + :return: True if any of the given tables is referenced """ - present = {table.table.lower() for table in self.tables} - return any(table.lower() in present for table in tables) + return bool(self.get_disallowed_tables(tables, default_schema)) + + def changes_search_path(self) -> bool: + """ + Return True if the statement changes the session ``search_path``. + + A ``SET search_path = ...`` makes unqualified references in later + statements resolve to a schema other than the caller's + ``default_schema``, so denylist matching against ``default_schema`` + alone becomes unreliable once such a statement is present. + """ + # `SET search_path = schema` (and the `TO`/`SESSION`/`LOCAL` variants) + # parse as a structured exp.Set, surfaced by get_settings(). Strip any + # identifier quoting so `SET "search_path" = ...` (equivalent to the + # unquoted form in Postgres) is still recognized. + if any(key.strip('"').lower() == "search_path" for key in self.get_settings()): + return True + # `set_config('search_path', ...)` rebinds the search path through a + # function call rather than a SET statement, so it never reaches + # get_settings() and must be detected on the parsed tree. + for func in self._parsed.find_all(exp.Anonymous): + if ( + func.name.lower() == "set_config" + and func.expressions + and isinstance(func.expressions[0], exp.Literal) + and func.expressions[0].name.lower() == "search_path" + ): + return True + # Exotic forms (e.g. `SET search_path TO "$user", public`) fall back to + # an opaque exp.Command. Match the leading setting name rather than + # scanning the whole expression, so `SET ROLE my_search_path_role` + # (whose value merely contains the substring) is not misclassified. + parsed = self._parsed + if isinstance(parsed, exp.Command) and parsed.name.upper() == "SET": + tokens = str(parsed.expression).replace("=", " ").split() + while tokens and tokens[0].upper() in {"SESSION", "LOCAL"}: + tokens.pop(0) + return bool(tokens) and tokens[0].strip('"').lower() == "search_path" + return False + + def get_disallowed_tables( + self, + tables: set[str], + default_schema: str | None = None, + schema_indeterminate: bool = False, + ) -> set[str]: + """ + Return the subset of ``tables`` referenced by this statement. + + Matching mirrors :meth:`check_tables_present`: bare entries match by + table name regardless of schema, while schema-qualified entries + require the schema to match too. Entries are returned in their + original denylist form so callers can report exactly which + denylisted tables were hit. + + A reference without an explicit schema is resolved against + ``default_schema`` when one is supplied, so an unqualified ``tables`` + run under ``search_path = information_schema`` still matches the + ``information_schema.tables`` entry, while the same name under a + user schema does not. + + :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Schema unqualified references resolve to at + runtime (e.g. the session ``search_path`` / selected schema) + :param schema_indeterminate: When True, the effective ``search_path`` + cannot be pinned to ``default_schema`` (e.g. the script contains a + ``SET search_path``), so an unqualified reference is matched + against the bare-name portion of schema-qualified entries too, to + avoid bypassing a qualified denylist entry + :return: The matched entries, in their original denylist form + """ + fallback = default_schema.lower() if default_schema else None + present_bare: set[str] = set() + present_qualified: set[str] = set() + present_unqualified: set[str] = set() + for t in self.tables: + bare = t.table.lower() + present_bare.add(bare) + if t.schema: + present_qualified.add(f"{t.schema.lower()}.{bare}") + # Also index the fully-qualified (catalog.schema.table) form so a + # three-part denylist entry can match; without this, qualified + # entries deeper than schema.table would silently never match. + if t.catalog: + present_qualified.add( + f"{t.catalog.lower()}.{t.schema.lower()}.{bare}" + ) + else: + present_unqualified.add(bare) + # An unqualified reference can only be resolved against a known + # default schema. When ``default_schema`` is None (the runtime + # search_path is unknown to us), a qualified denylist entry is + # matched only via the ``schema_indeterminate`` bare-name + # fallback below, never here, this is an inherent limit of static + # analysis without the live search_path. + if fallback: + present_qualified.add(f"{fallback}.{bare}") + found: set[str] = set() + for entry in tables: + needle = entry.lower() + if "." in needle: + if needle in present_qualified: + found.add(entry) + elif ( + schema_indeterminate + and needle.rsplit(".", 1)[1] in present_unqualified + ): + found.add(entry) + elif needle in present_bare: + found.add(entry) + return found def get_limit_value(self) -> int | None: """ @@ -1500,16 +1657,36 @@ class KustoKQLStatement(BaseSQLStatement[str]): logger.warning("Kusto KQL doesn't support checking for functions present.") return False - def check_tables_present(self, tables: set[str]) -> bool: + def check_tables_present( + self, tables: set[str], default_schema: str | None = None + ) -> bool: """ Check if any of the given tables are present in the statement. :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Unused; accepted for interface parity :return: True if any of the tables are present """ logger.warning("Kusto KQL doesn't support checking for tables present.") return False + def get_disallowed_tables( + self, + tables: set[str], + default_schema: str | None = None, + schema_indeterminate: bool = False, + ) -> set[str]: + """ + Return the subset of ``tables`` referenced by this statement. + + :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Unused; accepted for interface parity + :param schema_indeterminate: Unused; accepted for interface parity + :return: The matched entries, in their original denylist form + """ + logger.warning("Kusto KQL doesn't support checking for tables present.") + return set() + def get_limit_value(self) -> int | None: """ Get the limit value of the statement. @@ -1681,16 +1858,46 @@ class SQLScript: for statement in self.statements ) - def check_tables_present(self, tables: set[str]) -> bool: + def check_tables_present( + self, tables: set[str], default_schema: str | None = None + ) -> bool: """ Check if any of the given tables are present in the script. :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Schema unqualified references resolve to at + runtime (e.g. the session ``search_path`` / selected schema) :return: True if any of the tables are present """ - return any( - statement.check_tables_present(tables) for statement in self.statements - ) + return bool(self.get_disallowed_tables(tables, default_schema)) + + def get_disallowed_tables( + self, tables: set[str], default_schema: str | None = None + ) -> set[str]: + """ + Return the subset of ``tables`` referenced anywhere in the script. + + :param tables: Set of table names to check for (case-insensitive) + :param default_schema: Schema unqualified references resolve to at + runtime (e.g. the session ``search_path`` / selected schema) + :return: The matched entries, in their original denylist form + """ + # A `SET search_path` only affects statements that run *after* it, so + # track the indeterminate state in statement order: an unqualified + # reference is matched conservatively (against the bare-name portion of + # qualified denylist entries) only once a preceding statement has + # rebound the search path. This keeps a qualified denylist entry from + # being bypassed (e.g. `SET search_path = information_schema; SELECT * + # FROM tables`) without penalizing statements that ran beforehand. + found: set[str] = set() + schema_indeterminate = False + for statement in self.statements: + found |= statement.get_disallowed_tables( + tables, default_schema, schema_indeterminate + ) + if statement.changes_search_path(): + schema_indeterminate = True + return found def is_valid_ctas(self) -> bool: """ diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 34611d38a33..518a9b48e6b 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -433,23 +433,34 @@ def execute_sql_statements( # noqa: C901 db_engine_spec.engine, set(), ) - if disallowed_tables and parsed_script.check_tables_present(disallowed_tables): - # Report only the tables actually found in the query - found_tables = set() - for statement in parsed_script.statements: - present = {table.table.lower() for table in statement.tables} - for table in disallowed_tables: - if table.lower() in present: - found_tables.add(table) - raise SupersetDisallowedSQLTableException(found_tables or disallowed_tables) + rls_enabled = is_feature_enabled("RLS_IN_SQLLAB") + + # Resolve the effective per-query schema once and share it between the + # denylist check and RLS injection, but only when a control below needs it. + # Going through the query-aware ``get_default_schema_for_query`` (rather than + # the static ``get_default_schema``) resolves an unqualified reference to the + # schema the engine actually uses at runtime -- engines without dynamic-schema + # support ignore the request's selected schema -- so both controls match the + # execution path instead of a schema that may never apply. + effective_schema = "" + if disallowed_tables or rls_enabled: + effective_schema = database.get_default_schema_for_query(query) + + if disallowed_tables: + # Report only the denylisted tables actually referenced in the query, + # honoring schema-qualified entries (e.g. ``information_schema.tables``). + found_tables = parsed_script.get_disallowed_tables( + disallowed_tables, effective_schema + ) + if found_tables: + raise SupersetDisallowedSQLTableException(found_tables) if parsed_script.has_mutation() and not database.allow_dml: raise SupersetDMLNotAllowedException() - if is_feature_enabled("RLS_IN_SQLLAB"): - default_schema = query.database.get_default_schema_for_query(query) + if rls_enabled: for statement in parsed_script.statements: - apply_rls(query.database, query.catalog, default_schema, statement) + apply_rls(query.database, query.catalog, effective_schema, statement) if query.select_as_cta: # CTAS is valid when the last statement is a SELECT, while CVAS is valid when diff --git a/tests/unit_tests/commands/sql_lab/test_estimate.py b/tests/unit_tests/commands/sql_lab/test_estimate.py index bb28e299222..61afd2bd97b 100644 --- a/tests/unit_tests/commands/sql_lab/test_estimate.py +++ b/tests/unit_tests/commands/sql_lab/test_estimate.py @@ -181,8 +181,15 @@ def test_apply_sql_security_allows_dml_when_enabled(mock_app: MagicMock) -> None assert command._apply_sql_security("INSERT INTO t VALUES (1)") +@patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=False) @patch("superset.commands.sql_lab.estimate.app") -def test_apply_sql_security_blocks_disallowed_table(mock_app: MagicMock) -> None: +def test_apply_sql_security_blocks_disallowed_table( + mock_app: MagicMock, + mock_is_feature_enabled: MagicMock, +) -> None: + """A query referencing a table on ``DISALLOWED_SQL_TABLES`` for the engine + is rejected on the estimate path with ``SupersetDisallowedSQLTableException``, + mirroring the execution-time denylist gate.""" mock_app.config = { "DISALLOWED_SQL_FUNCTIONS": {}, "DISALLOWED_SQL_TABLES": {"postgresql": {"secrets"}}, @@ -190,10 +197,40 @@ def test_apply_sql_security_blocks_disallowed_table(mock_app: MagicMock) -> None from superset.exceptions import SupersetDisallowedSQLTableException command = _make_command_with_db("SELECT * FROM secrets", allow_dml=True) + cast( + MagicMock, command._database + ).resolve_query_default_schema.return_value = "public" with pytest.raises(SupersetDisallowedSQLTableException): command._apply_sql_security("SELECT * FROM secrets") +@patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=False) +@patch("superset.commands.sql_lab.estimate.app") +def test_apply_sql_security_denylist_runs_schema_gate( + mock_app: MagicMock, + mock_is_feature_enabled: MagicMock, +) -> None: + """The denylist check resolves the effective schema through the shared + query-aware ``resolve_query_default_schema`` (not the static + ``get_default_schema``), so the engine's per-query security gate — e.g. the + Postgres ``search_path`` check — runs on the estimate path even with RLS + disabled, matching ``sql_lab.execute_sql_statements``.""" + mock_app.config = { + "DISALLOWED_SQL_FUNCTIONS": {}, + "DISALLOWED_SQL_TABLES": {"postgresql": {"secrets"}}, + } + command = _make_command_with_db( + "SET search_path = secret; SELECT * FROM t", allow_dml=True + ) + database = cast(MagicMock, command._database) + database.resolve_query_default_schema.side_effect = _security_exception() + + with pytest.raises(SupersetSecurityException): + command._apply_sql_security("SET search_path = secret; SELECT * FROM t") + + database.resolve_query_default_schema.assert_called_once() + + @patch("superset.commands.sql_lab.estimate.app") def test_apply_sql_security_blocks_disallowed_function(mock_app: MagicMock) -> None: """A disallowed function cannot be probed via cost estimation either.""" @@ -218,15 +255,11 @@ def test_apply_sql_security_allows_benign_select(mock_app: MagicMock) -> None: @patch("superset.commands.sql_lab.estimate.apply_rls") -@patch("superset.commands.sql_lab.estimate.Query") -@patch("superset.commands.sql_lab.estimate.db") @patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=True) @patch("superset.commands.sql_lab.estimate.app") def test_apply_sql_security_injects_rls_when_enabled( mock_app: MagicMock, mock_is_feature_enabled: MagicMock, - mock_db: MagicMock, - mock_query: MagicMock, mock_apply_rls: MagicMock, ) -> None: """With RLS_IN_SQLLAB enabled, RLS predicates are applied per statement so @@ -238,14 +271,13 @@ def test_apply_sql_security_injects_rls_when_enabled( mock_is_feature_enabled.assert_called_with("RLS_IN_SQLLAB") mock_apply_rls.assert_called_once() - # The transient probe Query is expunged so its (deliberately incomplete) - # row can't autoflush into the session when apply_rls queries below. - mock_db.session.expunge.assert_called_once_with(mock_query.return_value) + # Effective schema is resolved through the shared query-aware + # ``Database.resolve_query_default_schema`` (which builds and expunges the + # transient probe), keeping parity with the execution path. + cast(MagicMock, command._database).resolve_query_default_schema.assert_called_once() assert isinstance(result, str) -@patch("superset.commands.sql_lab.estimate.Query") -@patch("superset.commands.sql_lab.estimate.db") @patch("superset.commands.sql_lab.estimate.apply_rls") @patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=True) @patch("superset.commands.sql_lab.estimate.app") @@ -253,8 +285,6 @@ def test_apply_sql_security_resolves_default_schema_for_rls( mock_app: MagicMock, mock_is_feature_enabled: MagicMock, mock_apply_rls: MagicMock, - mock_db: MagicMock, - mock_query: MagicMock, ) -> None: """When no catalog/schema is supplied, RLS must be applied against the database's *resolved* default catalog/schema — mirroring the execution path @@ -269,16 +299,16 @@ def test_apply_sql_security_resolves_default_schema_for_rls( command._schema = "" command._catalog = None database.get_default_catalog.return_value = "default_catalog" - database.get_default_schema_for_query.return_value = "public" + database.resolve_query_default_schema.return_value = "public" command._apply_sql_security("SELECT * FROM t") # Default catalog/schema are resolved before injection, in the same order # as the executor (catalog first, then schema derived per-query). The schema - # goes through ``get_default_schema_for_query`` so engine-specific per-query - # security gates (e.g. the Postgres ``search_path`` check) run as well. + # goes through the shared ``resolve_query_default_schema`` so engine-specific + # per-query security gates (e.g. the Postgres ``search_path`` check) run too. database.get_default_catalog.assert_called_once_with() - database.get_default_schema_for_query.assert_called_once() + database.resolve_query_default_schema.assert_called_once() # RLS is applied with the *resolved* values, never the raw ""/None. # apply_rls(database, catalog, schema, statement) @@ -287,8 +317,6 @@ def test_apply_sql_security_resolves_default_schema_for_rls( assert call_args[2] == "public" -@patch("superset.commands.sql_lab.estimate.Query") -@patch("superset.commands.sql_lab.estimate.db") @patch("superset.commands.sql_lab.estimate.apply_rls") @patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=True) @patch("superset.commands.sql_lab.estimate.app") @@ -296,12 +324,10 @@ def test_apply_sql_security_respects_explicit_catalog_schema( mock_app: MagicMock, mock_is_feature_enabled: MagicMock, mock_apply_rls: MagicMock, - mock_db: MagicMock, - mock_query: MagicMock, ) -> None: """An explicitly supplied catalog short-circuits default-catalog resolution, and the explicit schema wins as the RLS target — but the schema resolver - ``get_default_schema_for_query`` is still invoked so the engine's per-query + ``resolve_query_default_schema`` is still invoked so the engine's per-query security gate runs even when a schema is pinned (parity with the executor, which calls it unconditionally).""" mock_app.config = {"DISALLOWED_SQL_FUNCTIONS": {}, "DISALLOWED_SQL_TABLES": {}} @@ -317,14 +343,12 @@ def test_apply_sql_security_respects_explicit_catalog_schema( # ...but the schema gate must run even when a schema is pinned, otherwise an # explicit-schema estimate could smuggle a ``SET search_path`` past the gate # the executor enforces. - database.get_default_schema_for_query.assert_called_once() + database.resolve_query_default_schema.assert_called_once() call_args = mock_apply_rls.call_args.args assert call_args[1] == "my_catalog" assert call_args[2] == "my_schema" -@patch("superset.commands.sql_lab.estimate.Query") -@patch("superset.commands.sql_lab.estimate.db") @patch("superset.commands.sql_lab.estimate.apply_rls") @patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=True) @patch("superset.commands.sql_lab.estimate.app") @@ -332,10 +356,8 @@ def test_apply_sql_security_propagates_engine_schema_gate( mock_app: MagicMock, mock_is_feature_enabled: MagicMock, mock_apply_rls: MagicMock, - mock_db: MagicMock, - mock_query: MagicMock, ) -> None: - """Default-schema resolution goes through ``get_default_schema_for_query``, + """Default-schema resolution goes through ``resolve_query_default_schema``, so an engine-specific per-query security gate (e.g. the Postgres ``search_path`` check that rejects ``SET search_path = ...``) is enforced on the estimate path too, rather than being silently bypassed. @@ -348,7 +370,7 @@ def test_apply_sql_security_propagates_engine_schema_gate( command._schema = "" command._catalog = None database.get_default_catalog.return_value = "default_catalog" - database.get_default_schema_for_query.side_effect = _security_exception() + database.resolve_query_default_schema.side_effect = _security_exception() with pytest.raises(SupersetSecurityException): command._apply_sql_security("SET search_path = secret; SELECT * FROM t") diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index 9a369f9ece6..c4b5ae6c18d 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -1560,6 +1560,45 @@ def test_database_execute_async_without_options(mocker: MockerFixture) -> None: assert result == mock_handle +def test_resolve_query_default_schema_builds_probe(mocker: MockerFixture) -> None: + """``resolve_query_default_schema`` builds a transient probe Query carrying + the request's SQL/schema/catalog and resolves through the query-aware + ``get_default_schema_for_query`` (which runs per-query engine security + gates), forwarding template params.""" + database = Database(database_name="test_db", sqlalchemy_uri="sqlite://") + resolve_mock = mocker.patch.object( + database, "get_default_schema_for_query", return_value="resolved_schema" + ) + + result = database.resolve_query_default_schema( + "SELECT 1", "explicit", "cat", {"p": 1} + ) + + assert result == "resolved_schema" + probe, template_params = resolve_mock.call_args.args + assert probe.sql == "SELECT 1" + assert probe.schema == "explicit" + assert probe.catalog == "cat" + assert template_params == {"p": 1} + + +def test_resolve_query_default_schema_normalizes_blank_schema( + mocker: MockerFixture, +) -> None: + """A blank request schema reaches the probe as ``None`` so the engine spec + resolves it, rather than matching on an empty string.""" + database = Database(database_name="test_db", sqlalchemy_uri="sqlite://") + resolve_mock = mocker.patch.object( + database, "get_default_schema_for_query", return_value="public" + ) + + database.resolve_query_default_schema("SELECT 1", "", None) + + probe = resolve_mock.call_args.args[0] + assert probe.schema is None + assert probe.catalog is None + + def test_clear_bootstrap_cache_logs_warning_on_failure( mocker: MockerFixture, ) -> None: diff --git a/tests/unit_tests/models/helpers_test.py b/tests/unit_tests/models/helpers_test.py index fab19f5f22e..58411fcfae7 100644 --- a/tests/unit_tests/models/helpers_test.py +++ b/tests/unit_tests/models/helpers_test.py @@ -3195,6 +3195,61 @@ def test_process_sql_expression_no_gate_when_denylists_empty( assert result is not None +def test_resolve_denylist_schema_uses_query_aware_resolution( + mocker: MockerFixture, database: Database +) -> None: + """A datasource without an explicit schema resolves the denylist schema + through the query-aware ``get_default_schema_for_query`` (matching the SQL + Lab / executor gate), not the static inspector-based ``get_default_schema``.""" + from superset.connectors.sqla.models import SqlaTable + + query_aware = mocker.patch.object( + database, "get_default_schema_for_query", return_value="resolved" + ) + static = mocker.patch.object(database, "get_default_schema") + table = SqlaTable(database=database, schema=None, table_name="t") + + assert table._resolve_denylist_schema("SELECT 1") == "resolved" + query_aware.assert_called_once() + static.assert_not_called() + + +def test_resolve_denylist_schema_memoizes_across_expressions( + mocker: MockerFixture, database: Database +) -> None: + """The resolved schema is cached per datasource so adhoc-expression + validation, which runs once per column/metric/order-by, does not re-resolve + (and re-probe) the schema for every expression.""" + from superset.connectors.sqla.models import SqlaTable + + query_aware = mocker.patch.object( + database, "get_default_schema_for_query", return_value="resolved" + ) + table = SqlaTable(database=database, schema=None, table_name="t") + + # Distinct expression shapes (column / metric / order-by) mirror a real + # request, where validation runs once per selected field with different + # SQL each time. Caching must be keyed on the datasource instance, not the + # SQL string, so a single resolution still covers all of them. + for sql in ("price", "SUM(price)", "created_at DESC"): + assert table._resolve_denylist_schema(sql) == "resolved" + query_aware.assert_called_once() + + +def test_resolve_denylist_schema_explicit_schema_skips_probe( + mocker: MockerFixture, database: Database +) -> None: + """An explicit datasource schema is returned directly, with no probe Query + and no inspector round-trip.""" + from superset.connectors.sqla.models import SqlaTable + + query_aware = mocker.patch.object(database, "get_default_schema_for_query") + table = SqlaTable(database=database, schema="analytics", table_name="t") + + assert table._resolve_denylist_schema("SELECT 1") == "analytics" + query_aware.assert_not_called() + + def test_get_sqla_query_dotted_struct_column_bigquery( mocker: MockerFixture, session: Session, diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py index fdaf52f8a79..5ee25b8cf90 100644 --- a/tests/unit_tests/sql/execution/test_executor.py +++ b/tests/unit_tests/sql/execution/test_executor.py @@ -1374,8 +1374,10 @@ def test_execute_uses_default_catalog_and_schema( get_default_catalog_mock = mocker.patch.object( database, "get_default_catalog", return_value="main" ) + # Schema is resolved through the query-aware ``get_default_schema_for_query`` + # (so per-query engine gates run), not the static ``get_default_schema``. get_default_schema_mock = mocker.patch.object( - database, "get_default_schema", return_value="public" + database, "get_default_schema_for_query", return_value="public" ) mocker.patch.dict( current_app.config, @@ -1395,6 +1397,88 @@ def test_execute_uses_default_catalog_and_schema( get_default_schema_mock.assert_called() +def test_resolve_query_schema_uses_query_aware_resolution( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """``_resolve_query_schema`` resolves through the query-aware + ``get_default_schema_for_query`` (which runs per-query engine security gates), + handing it a transient probe Query that carries the request's SQL, schema, + catalog and template params.""" + from superset.models.sql_lab import Query + from superset.sql.execution.executor import SQLExecutor + + resolve_mock = mocker.patch.object( + database, "get_default_schema_for_query", return_value="resolved_schema" + ) + + executor = SQLExecutor(database) + options = QueryOptions(schema="explicit", template_params={"p": 1}) + + result = executor._resolve_query_schema("SELECT 1", options, "cat") + + assert result == "resolved_schema" + probe, template_params = resolve_mock.call_args.args + assert isinstance(probe, Query) + assert probe.sql == "SELECT 1" + assert probe.schema == "explicit" + assert probe.catalog == "cat" + assert template_params == {"p": 1} + + +def test_resolve_query_schema_omits_blank_schema( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """An unset request schema reaches the probe as ``None`` so the engine spec + resolves the runtime default instead of matching on an empty string.""" + resolve_mock = mocker.patch.object( + database, "get_default_schema_for_query", return_value="public" + ) + + from superset.sql.execution.executor import SQLExecutor + + executor = SQLExecutor(database) + + result = executor._resolve_query_schema("SELECT 1", QueryOptions(), None) + + assert result == "public" + probe = resolve_mock.call_args.args[0] + assert probe.schema is None + assert probe.catalog is None + + +def test_prepare_sql_runs_schema_gate_with_explicit_schema( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """The per-query schema gate must run even when an explicit schema is + supplied, so an explicit-schema request cannot smuggle a ``SET search_path`` + past the gate the resolver enforces (parity with the estimate path, which + resolves unconditionally).""" + from superset.errors import ErrorLevel, SupersetError, SupersetErrorType + from superset.exceptions import SupersetSecurityException + from superset.sql.execution.executor import SQLExecutor + + gate = mocker.patch.object( + database, + "get_default_schema_for_query", + side_effect=SupersetSecurityException( + SupersetError( + message="blocked", + error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR, + level=ErrorLevel.ERROR, + ) + ), + ) + executor = SQLExecutor(database) + + with pytest.raises(SupersetSecurityException): + executor._prepare_sql( + "SET search_path = secret; SELECT 1", + QueryOptions(schema="explicit"), + ) + + gate.assert_called_once() + + # ============================================================================= # Async Query Status and Result Tests # ============================================================================= diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 55a530266be..5bb8a9e8ab5 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -28,6 +28,7 @@ from superset.exceptions import QueryClauseValidationException, SupersetParseErr from superset.jinja_context import JinjaTemplateProcessor from superset.sql.parse import ( _check_script_length, + BaseSQLStatement, CTASMethod, extract_tables_from_statement, JinjaSQLResult, @@ -1458,7 +1459,12 @@ def test_is_mutating_anonymous_block(sql: str, expected: bool) -> None: ("SELECT lo_import('/etc/passwd')", True), ("SELECT lo_put(12345, 0, decode('00', 'hex'))", True), ("SELECT lo_create(0)", True), + # lo_creat is the legacy large-object creator (distinct from lo_create). + ("SELECT lo_creat(-1)", True), ("SELECT lowrite(12345, decode('00', 'hex'))", True), + # lo_truncate/lo_truncate64 shrink an existing large object: a write. + ("SELECT lo_truncate(12345, 0)", True), + ("SELECT lo_truncate64(12345, 0)", True), # lo_unlink deletes a large object outright. ("SELECT lo_unlink(12345)", True), # PostgreSQL sequence mutators. setval()/nextval() look like reads but @@ -3510,6 +3516,441 @@ def test_check_tables_present(sql: str, engine: str, expected: bool) -> None: assert SQLScript(sql, engine).check_tables_present(tables) == expected +@pytest.mark.parametrize( + "engine, sql, denylist, expected", + [ + # Postgres: schema-qualified denylist entry matches schema-qualified + # reference. + ( + "postgresql", + "SELECT * FROM information_schema.tables", + {"information_schema.tables"}, + True, + ), + # ... and is case-insensitive. + ( + "postgresql", + "SELECT * FROM INFORMATION_SCHEMA.TABLES", + {"information_schema.tables"}, + True, + ), + # Schema-qualified denylist entry does NOT match a bare-name table + # of the same name in another schema. A user table named `tables` + # remains queryable. + ( + "postgresql", + "SELECT * FROM public.tables", + {"information_schema.tables"}, + False, + ), + ( + "postgresql", + "SELECT * FROM tables", + {"information_schema.tables"}, + False, + ), + # Bare-name denylist entry still matches by table name only + # (existing behavior, schema-agnostic). + ( + "postgresql", + "SELECT * FROM pg_stat_activity", + {"pg_stat_activity"}, + True, + ), + ( + "postgresql", + "SELECT * FROM pg_catalog.pg_stat_activity", + {"pg_stat_activity"}, + True, + ), + # Mixed entries: one schema-qualified, one bare. Match either. + ( + "postgresql", + "SELECT * FROM information_schema.columns", + {"information_schema.tables", "information_schema.columns"}, + True, + ), + ( + "postgresql", + "SELECT * FROM pg_roles", + {"information_schema.tables", "pg_roles"}, + True, + ), + # Negative control. + ( + "postgresql", + "SELECT * FROM my_table", + {"information_schema.tables", "pg_roles"}, + False, + ), + # MySQL: the shipped DISALLOWED_SQL_TABLES['mysql'] entries are all + # schema-qualified (`mysql.user`, `performance_schema.threads`, + # `performance_schema.processlist`). Without schema-aware matching + # the entries are dead config. These cases pin the fix. + ( + "mysql", + "SELECT user, host, authentication_string FROM mysql.user", + {"mysql.user"}, + True, + ), + ( + "mysql", + "SELECT * FROM performance_schema.threads", + {"performance_schema.threads"}, + True, + ), + ( + "mysql", + "SELECT * FROM performance_schema.processlist", + {"performance_schema.processlist"}, + True, + ), + # MySQL must NOT block a user-authored table that shares the leaf + # name with the system view. + ( + "mysql", + "SELECT * FROM mydb.user", + {"mysql.user"}, + False, + ), + # MSSQL: same shape, `sys.*` entries are schema-qualified. + ( + "mssql", + "SELECT name, password_hash FROM sys.sql_logins", + {"sys.sql_logins"}, + True, + ), + ( + "mssql", + "SELECT name, sid FROM sys.server_principals", + {"sys.server_principals"}, + True, + ), + ( + "mssql", + "SELECT * FROM sys.configurations", + {"sys.configurations"}, + True, + ), + # MSSQL must NOT block a user-authored table sharing the leaf name. + ( + "mssql", + "SELECT * FROM mydb.sql_logins", + {"sys.sql_logins"}, + False, + ), + # Three-part (catalog.schema.table) denylist entries match the + # fully-qualified reference, the multi-dot form is indexed rather than + # silently dead. + ( + "trino", + "SELECT * FROM cat.sys.sql_logins", + {"cat.sys.sql_logins"}, + True, + ), + # ... and a different catalog must NOT match. + ( + "trino", + "SELECT * FROM other.sys.sql_logins", + {"cat.sys.sql_logins"}, + False, + ), + ], +) +def test_check_tables_present_schema_qualified( + engine: str, sql: str, denylist: set[str], expected: bool +) -> None: + """ + `check_tables_present` must distinguish schema-qualified denylist + entries (e.g. `information_schema.tables`, `mysql.user`, + `sys.sql_logins`) from bare-name entries (e.g. `pg_stat_activity`). + Schema-qualified entries only match schema-qualified references in + the SQL; bare entries match the table name regardless of schema. + + Covers Postgres, MySQL, and MSSQL dialects so the shipped + DISALLOWED_SQL_TABLES entries for each remain effective. + """ + assert SQLScript(sql, engine).check_tables_present(denylist) == expected + + +@pytest.mark.parametrize( + "engine, sql, denylist, expected", + [ + # A schema-qualified match is reported in its original denylist form, + # not collapsed to the bare leaf name and not the whole denylist. + ( + "postgresql", + "SELECT * FROM information_schema.tables", + {"information_schema.tables", "information_schema.columns", "pg_roles"}, + {"information_schema.tables"}, + ), + # Bare-name match is reported as-is. + ( + "postgresql", + "SELECT * FROM pg_catalog.pg_stat_activity", + {"pg_stat_activity", "pg_roles"}, + {"pg_stat_activity"}, + ), + # Multiple references across statements union their matches. + ( + "postgresql", + "SELECT * FROM information_schema.tables; SELECT * FROM pg_roles", + {"information_schema.tables", "pg_roles", "pg_settings"}, + {"information_schema.tables", "pg_roles"}, + ), + # No match returns an empty set. + ( + "postgresql", + "SELECT * FROM my_table", + {"information_schema.tables", "pg_roles"}, + set(), + ), + # A three-part (catalog.schema.table) denylist entry matches a + # fully-qualified reference, reported in its original form. + ( + "trino", + "SELECT * FROM cat.sys.sql_logins", + {"cat.sys.sql_logins"}, + {"cat.sys.sql_logins"}, + ), + # ... but only when the catalog lines up: a different catalog does not + # match the three-part entry. + ( + "trino", + "SELECT * FROM other.sys.sql_logins", + {"cat.sys.sql_logins"}, + set(), + ), + ], +) +def test_get_disallowed_tables( + engine: str, sql: str, denylist: set[str], expected: set[str] +) -> None: + """ + `get_disallowed_tables` returns exactly the denylist entries referenced, + in their original (possibly schema-qualified) form, so callers can report + precisely which tables were hit instead of echoing the whole denylist. + """ + assert SQLScript(sql, engine).get_disallowed_tables(denylist) == expected + + +@pytest.mark.parametrize( + "sql, default_schema, denylist, expected", + [ + # Unqualified reference resolves to the default schema, so it matches + # a schema-qualified denylist entry when the schemas line up (e.g. a + # connection whose search_path is `information_schema`). + ( + "SELECT * FROM tables", + "information_schema", + {"information_schema.tables"}, + {"information_schema.tables"}, + ), + # ... case-insensitively. + ( + "SELECT * FROM tables", + "INFORMATION_SCHEMA", + {"information_schema.tables"}, + {"information_schema.tables"}, + ), + # The same unqualified name under a user schema must NOT match: a user + # table named `tables` stays queryable. + ( + "SELECT * FROM tables", + "public", + {"information_schema.tables"}, + set(), + ), + # An explicit schema on the reference wins over the default schema. + ( + "SELECT * FROM public.tables", + "information_schema", + {"information_schema.tables"}, + set(), + ), + # Without a default schema, behavior is unchanged: unqualified + # references never match schema-qualified entries. + ( + "SELECT * FROM tables", + None, + {"information_schema.tables"}, + set(), + ), + # Bare-name denylist entries are schema-agnostic and unaffected by the + # default schema. + ( + "SELECT * FROM pg_stat_activity", + "information_schema", + {"pg_stat_activity"}, + {"pg_stat_activity"}, + ), + # The default schema is forwarded to every statement in a script, so an + # unqualified reference in a later statement is resolved too. + ( + "SELECT * FROM my_table; SELECT * FROM tables", + "information_schema", + {"information_schema.tables"}, + {"information_schema.tables"}, + ), + ], +) +def test_get_disallowed_tables_default_schema( + sql: str, + default_schema: str | None, + denylist: set[str], + expected: set[str], +) -> None: + """ + `get_disallowed_tables` resolves an unqualified reference against the + supplied default schema, so a denylisted system view (e.g. + `information_schema.tables`) is still caught when reached without an + explicit schema under that search_path, without blocking a same-named + user table under a different schema. + """ + assert ( + SQLScript(sql, "postgresql").get_disallowed_tables(denylist, default_schema) + == expected + ) + + +@pytest.mark.parametrize( + "sql, default_schema, denylist, expected", + [ + # `SET search_path` rebinds where an unqualified reference resolves, so + # the static default schema can no longer be trusted. A qualified + # denylist entry must still match the later unqualified reference, + # otherwise the block is trivially bypassable. + ( + "SET search_path = information_schema; SELECT * FROM tables", + "public", + {"information_schema.tables"}, + {"information_schema.tables"}, + ), + # `SET search_path TO "$user", ...` falls back to an exp.Command (it is + # not a structured exp.Set), exercising the same conservative matching + # via the command-name detection branch. + ( + 'SET search_path TO "$user", information_schema; SELECT * FROM tables', + "public", + {"information_schema.tables"}, + {"information_schema.tables"}, + ), + # `set_config('search_path', ...)` rebinds the search path through a + # function call and must trigger the same conservative matching. + ( + "SELECT set_config('search_path', 'information_schema', true);" + " SELECT * FROM tables", + "public", + {"information_schema.tables"}, + {"information_schema.tables"}, + ), + # The search-path change only affects later statements: a statement that + # runs before it keeps resolving against the original default schema, so + # its unqualified reference must NOT be widened. + ( + "SELECT * FROM tables; SET search_path = information_schema", + "public", + {"information_schema.tables"}, + set(), + ), + # An explicitly qualified reference is unambiguous and must NOT be + # widened to match a different schema's denylist entry. + ( + "SET search_path = information_schema; SELECT * FROM public.tables", + "public", + {"information_schema.tables"}, + set(), + ), + # Without a search_path change, matching is unchanged: an unqualified + # reference under a user schema does not match the qualified entry. + ( + "SELECT * FROM tables", + "public", + {"information_schema.tables"}, + set(), + ), + ], +) +def test_get_disallowed_tables_search_path_change( + sql: str, + default_schema: str | None, + denylist: set[str], + expected: set[str], +) -> None: + """ + A `SET search_path` in the script makes unqualified references resolve to a + schema other than the caller's default, so `get_disallowed_tables` matches + them against schema-qualified entries too, closing a denylist bypass. + """ + assert ( + SQLScript(sql, "postgresql").get_disallowed_tables(denylist, default_schema) + == expected + ) + + +@pytest.mark.parametrize( + "sql, expected", + [ + # Structured `SET search_path` (exp.Set), surfaced via get_settings(). + ("SET search_path = information_schema", True), + # Exotic form that falls back to exp.Command; the leading setting name + # is `search_path`. + ('SET search_path TO "$user", public', True), + # `SET SESSION ...` can also fall back to exp.Command; the optional + # SESSION/LOCAL qualifier is skipped before matching the setting name. + ("SET SESSION search_path FROM CURRENT", True), + # A quoted identifier is equivalent to the unquoted form in Postgres, + # so it must be recognized too (both the exp.Set and exp.Command forms). + ('SET "search_path" = information_schema', True), + ('SET "search_path" TO "$user", public', True), + # A `SET` whose value merely contains the substring `search_path` must + # not be misclassified (the setting being changed is `ROLE`). + ("SET ROLE app_search_path_user", False), + # `set_config('search_path', ...)` rebinds the path via a function call. + ("SELECT set_config('search_path', 'information_schema', true)", True), + # A different setting changed through `set_config` is not a search-path + # change. + ("SELECT set_config('statement_timeout', '0', true)", False), + # An unrelated (non-`set_config`) function call is not a change either. + ("SELECT my_custom_func(1)", False), + ("SELECT 1", False), + ], +) +def test_changes_search_path(sql: str, expected: bool) -> None: + """ + `changes_search_path` detects search-path rebinds (via `SET` or + `set_config`) without misclassifying unrelated `SET` statements. + """ + assert SQLStatement(sql, "postgresql").changes_search_path() == expected + + +@pytest.mark.parametrize( + "sql, denylist, expected", + [ + ("SELECT * FROM pg_stat_activity", {"pg_stat_activity"}, True), + ("SELECT * FROM my_table", {"pg_stat_activity"}, False), + ], +) +def test_statement_check_tables_present( + sql: str, denylist: set[str], expected: bool +) -> None: + """ + `SQLStatement.check_tables_present` is the per-statement entry point that + `SQLScript` no longer routes through (it calls `get_disallowed_tables` + directly), so exercise it on its own to keep the override covered. + """ + assert SQLStatement(sql, "postgresql").check_tables_present(denylist) == expected + + +def test_kustokql_statement_check_tables_present() -> None: + """ + `KustoKQLStatement.check_tables_present` is unsupported and always reports + False; exercise it directly so the override stays covered. + """ + statement = KustoKQLStatement("foo | take 100", "kustokql") + assert statement.check_tables_present({"foo"}) is False + + @pytest.mark.parametrize( "kql, expected", [ @@ -3701,6 +4142,17 @@ def test_backtick_invalid_sql_still_fails() -> None: SQLScript(sql, "base") +def test_base_sql_statement_is_destructive_raises_not_implemented() -> None: + """ + BaseSQLStatement.is_destructive is abstract; both concrete subclasses + (SQLStatement and KustoKQLStatement) override it, so calling the base + implementation directly must raise. This exercises the abstract stub + so it stays exercised under coverage. + """ + with pytest.raises(NotImplementedError): + BaseSQLStatement.is_destructive(object()) # type: ignore[arg-type] + + # --------------------------------------------------------------------------- # SQL_MAX_PARSE_LENGTH gate # ---------------------------------------------------------------------------