diff --git a/superset/utils/rls.py b/superset/utils/rls.py index f6643b9fc51..b9fc1c9975e 100644 --- a/superset/utils/rls.py +++ b/superset/utils/rls.py @@ -52,6 +52,7 @@ def apply_rls( method = database.db_engine_spec.get_rls_method() # collect all RLS predicates for all tables in the query + default_catalog = database.get_default_catalog() predicates: dict[Table, list[Any]] = {} for table in parsed_statement.tables: table = table.qualify(catalog=catalog, schema=schema) @@ -60,7 +61,7 @@ def apply_rls( for predicate in get_predicates_for_table( table, database, - database.get_default_catalog(), + default_catalog, exclude_dataset_id=exclude_dataset_id, ) if predicate @@ -98,7 +99,6 @@ def get_predicates_for_table( filters = [ SqlaTable.database_id == database.id, catalog_predicate, - SqlaTable.schema == table.schema, SqlaTable.table_name == table.table, ] # When applying RLS to a virtual dataset's inner SQL, skip a match against @@ -109,7 +109,30 @@ def get_predicates_for_table( if exclude_dataset_id is not None: filters.append(SqlaTable.id != exclude_dataset_id) - dataset = db.session.query(SqlaTable).filter(and_(*filters)).one_or_none() + dataset = ( + db.session.query(SqlaTable) + .filter(and_(*filters, SqlaTable.schema == table.schema)) + .one_or_none() + ) + if not dataset and table.schema: + # A dataset stored without a schema is scoped to the database's default + # schema, so a query resolving to that same schema must still pick up its + # predicates. This mirrors the null-catalog fallback above. + # + # This is a second query rather than an ``OR`` on the first so that an exact + # schema match always wins and neither query can match more than one dataset. + # Resolving the default schema probes the analytic database, so it is deferred + # until a null-schema dataset is known to exist. + null_schema_dataset = ( + db.session.query(SqlaTable) + .filter(and_(*filters, SqlaTable.schema.is_(None))) + .one_or_none() + ) + if null_schema_dataset and table.schema == database.get_default_schema( + table.catalog + ): + dataset = null_schema_dataset + if not dataset: return [] diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index 313c556a945..c554a95f405 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -17,13 +17,16 @@ # pylint: disable=import-outside-toplevel, invalid-name, unused-argument, too-many-locals import json # noqa: TID251 -from unittest.mock import MagicMock +from typing import Any +from unittest.mock import MagicMock, patch from urllib.parse import parse_qs, urlparse from uuid import UUID import pytest from freezegun import freeze_time from pytest_mock import MockerFixture +from sqlalchemy import text +from sqlalchemy.orm import Session from superset.app import SupersetApp from superset.common.db_query_status import QueryStatus @@ -608,6 +611,81 @@ def test_get_predicates_for_table(mocker: MockerFixture) -> None: ) +def test_get_predicates_for_table_null_schema_dataset(session: Session) -> None: + """ + A dataset stored with a NULL schema is scoped to the database's default + schema, mirroring the existing null-catalog fallback. + + A query resolving to that default schema must find the dataset, so its RLS + predicates are applied instead of being silently dropped. A query against a + different schema must not, since the null-schema dataset doesn't describe it. + """ + from superset.connectors.sqla.models import SqlaTable + + SqlaTable.metadata.create_all(session.get_bind()) + + database = Database(database_name="rls_db", sqlalchemy_uri="sqlite://") + # registered without an explicit schema, e.g. via the dataset API + dataset = SqlaTable(table_name="t1", schema=None, catalog=None, database=database) + session.add_all([database, dataset]) + session.flush() + + with ( + patch.object( + SqlaTable, "get_sqla_row_level_filters", return_value=[text("c1 = 1")] + ), + patch.object(Database, "get_default_schema", return_value="public"), + ): + assert get_predicates_for_table( + Table("t1", "public", None), database, None + ) == ["c1 = 1"] + + assert ( + get_predicates_for_table(Table("t1", "sales", None), database, None) == [] + ) + + +def test_get_predicates_for_table_prefers_exact_schema_match(session: Session) -> None: + """ + A dataset stored without a schema and one stored with the default schema can + coexist for the same table. The exact match must win, and the lookup must stay + unambiguous rather than treating both rows as candidates for a single dataset. + """ + from superset.connectors.sqla.models import SqlaTable + + SqlaTable.metadata.create_all(session.get_bind()) + + database = Database(database_name="rls_db_exact", sqlalchemy_uri="sqlite://") + session.add_all( + [ + database, + SqlaTable(table_name="t1", schema=None, catalog=None, database=database), + SqlaTable( + table_name="t1", schema="public", catalog=None, database=database + ), + ] + ) + session.flush() + + def row_level_filters( + self: Any, include_global_guest_rls: bool = True + ) -> list[Any]: + return [text(f"c1 = '{self.schema}'")] + + with ( + patch.object( + SqlaTable, + "get_sqla_row_level_filters", + autospec=True, + side_effect=row_level_filters, + ), + patch.object(Database, "get_default_schema", return_value="public"), + ): + assert get_predicates_for_table( + Table("t1", "public", None), database, None + ) == ["c1 = 'public'"] + + def test_get_predicates_for_table_excludes_self(mocker: MockerFixture) -> None: """ When ``exclude_dataset_id`` is supplied, the lookup query must add an