From 808ba668ffd851cdb862334f399720f9feb6156c Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Wed, 17 Dec 2025 11:02:25 -0500 Subject: [PATCH] Integrate apply_data_access_rules() into SQL Lab query execution --- superset/sql_lab.py | 8 +++ tests/unit_tests/sql_lab_test.py | 118 +++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 8ce20da5bb6..68dbcb86f41 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -67,6 +67,7 @@ from superset.utils.core import ( from superset.utils.dates import now_as_float from superset.utils.decorators import stats_timing from superset.utils.rls import apply_rls +from superset.data_access_rules.utils import apply_data_access_rules if TYPE_CHECKING: from superset.models.core import Database @@ -419,6 +420,13 @@ def execute_sql_statements( # noqa: C901 for statement in parsed_script.statements: apply_rls(query.database, query.catalog, default_schema, statement) + if is_feature_enabled("DATA_ACCESS_RULES"): + default_schema = query.database.get_default_schema_for_query(query) + for statement in parsed_script.statements: + apply_data_access_rules( + query.database, query.catalog, default_schema, statement + ) + if query.select_as_cta: # CTAS is valid when the last statement is a SELECT, while CVAS is valid when # there is only a single statement which must be a SELECT. diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index 2f165df009f..2fc16118174 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -36,7 +36,9 @@ from superset.sql_lab import ( get_sql_results, ) from superset.utils.rls import apply_rls, get_predicates_for_table +from superset.data_access_rules.utils import apply_data_access_rules from tests.conftest import with_config +from tests.unit_tests.conftest import with_feature_flags from tests.unit_tests.models.core_test import oauth2_client_info @@ -301,3 +303,119 @@ def test_get_predicates_for_table(mocker: MockerFixture) -> None: table = Table("t1", "public", "examples") assert get_predicates_for_table(table, database, "examples") == ["c1 = 1"] + + +def test_apply_data_access_rules(mocker: MockerFixture) -> None: + """ + Test the ``apply_data_access_rules`` helper function. + """ + from superset.data_access_rules.utils import ( + AccessCheckResult, + RLSPredicate, + TableAccessInfo, + ) + + database = mocker.MagicMock() + database.database_name = "test_db" + database.get_default_schema_for_query.return_value = "public" + database.get_default_catalog.return_value = "examples" + database.db_engine_spec = PostgresEngineSpec + + # Mock get_user_rules to return a rule with RLS predicates + get_user_rules = mocker.patch( + "superset.data_access_rules.utils.get_user_rules", + return_value=[], + ) + + # Mock check_table_access to return allowed access with RLS predicates + mocker.patch( + "superset.data_access_rules.utils.check_table_access", + return_value=TableAccessInfo( + access=AccessCheckResult.ALLOWED, + rls_predicates=[ + RLSPredicate(predicate="org_id = 1", group_key=None), + ], + cls_rules={}, + ), + ) + + # Mock is_feature_enabled + mocker.patch( + "superset.data_access_rules.utils.is_feature_enabled", + return_value=True, + ) + + parsed_statement = SQLStatement("SELECT * FROM t1", "postgresql") + + apply_data_access_rules(database, "examples", "public", parsed_statement) + + # Since we mocked the feature flag check, the function should have processed + get_user_rules.assert_called_once() + + +@with_feature_flags(DATA_ACCESS_RULES=True) +@with_config( + { + "SQLLAB_PAYLOAD_MAX_MB": 50, + "DISALLOWED_SQL_FUNCTIONS": {}, + "SQLLAB_CTAS_NO_LIMIT": False, + "SQL_MAX_ROW": 100000, + "QUERY_LOGGER": None, + "TROUBLESHOOTING_LINK": None, + "STATS_LOGGER": MagicMock(), + } +) +def test_execute_sql_statements_with_data_access_rules( + mocker: MockerFixture, app +) -> None: + """ + Test that `execute_sql_statements` calls `apply_data_access_rules` + when the DATA_ACCESS_RULES feature flag is enabled. + """ + # Mock apply_data_access_rules to track calls + mock_apply_dar = mocker.patch("superset.sql_lab.apply_data_access_rules") + + # Mock the query object and database + query = mocker.MagicMock() + query.limit = 1 + query.database = mocker.MagicMock() + query.database.cache_timeout = 100 + query.status = "RUNNING" + query.select_as_cta = False + query.database.allow_run_async = True + query.database.allow_dml = False + query.catalog = "examples" + query.database.get_default_schema_for_query.return_value = "public" + + # Mock get_query to return our mocked query object + mocker.patch("superset.sql_lab.get_query", return_value=query) + + # Mock db.session.refresh + mocker.patch("superset.sql_lab.db.session.refresh", return_value=None) + + # Mock the results backend + mocker.patch("superset.sql_lab.results_backend", return_value=True) + + # Mock sys.getsizeof to simulate a small payload + mocker.patch("sys.getsizeof", return_value=1000) + + # Mock _serialize_payload + mocker.patch( + "superset.sql_lab._serialize_payload", return_value="serialized_payload" + ) + + try: + execute_sql_statements( + query_id=1, + rendered_query="SELECT 42 AS answer", + return_results=True, + store_results=True, + start_time=None, + expand_data=False, + log_params={}, + ) + except Exception: + pass # We're just testing that apply_data_access_rules is called + + # Verify apply_data_access_rules was called + assert mock_apply_dar.called