This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch
fix-virtual-dataset-excess-db-connections
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to
refs/heads/fix-virtual-dataset-excess-db-connections by this push:
new d143dac1160 fix(sqla): collect expression RLS keys without building
chart queries
d143dac1160 is described below
commit d143dac1160081327b63028da3fbd448ca9f4f9e
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Mon Oct 26 17:45:05 2026 +0000
fix(sqla): collect expression RLS keys without building chart queries
---
docs/admin_docs/security/security.mdx | 2 +
superset/connectors/sqla/models.py | 9 +-
superset/models/helpers.py | 240 +++++++++----
superset/models/sql_lab.py | 2 +-
.../sqla/virtual_dataset_connections_test.py | 376 +++++++++++++++++++++
5 files changed, 560 insertions(+), 69 deletions(-)
diff --git a/docs/admin_docs/security/security.mdx
b/docs/admin_docs/security/security.mdx
index 78369dfa053..7d4a41350ac 100644
--- a/docs/admin_docs/security/security.mdx
+++ b/docs/admin_docs/security/security.mdx
@@ -666,6 +666,8 @@ When `ALLOW_ADHOC_SUBQUERY` is enabled, subqueries in chart
columns, metrics and
also receive the referenced datasets' RLS filters, and those filters
contribute to the
chart's cache key. Unqualified table names use the connection's execution
schema on
engines that cannot change schema per query, not an ignored dataset schema
setting.
+Expression RLS cache keys are collected without building the chart query or
running
+series-limit prequeries; schema discovery is shared with the dataset SQL's RLS
lookup.
For example, if you have:
diff --git a/superset/connectors/sqla/models.py
b/superset/connectors/sqla/models.py
index 54b38100dcd..f36be7def67 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -2886,9 +2886,12 @@ class SqlaTable(
filtered_query_obj = {
k: v for k, v in query_obj.items() if k in SQLA_QUERY_KEYS
}
- sqla_query = self.get_sqla_query(**cast(Any, filtered_query_obj))
+ sqla_query = self.get_sqla_query(
+ **cast(Any, filtered_query_obj), execute_prequeries=False
+ )
extra_cache_keys += sqla_query.extra_cache_keys
+ default_schema = None
# For virtual datasets, include RLS predicates in the cache key
if self.is_virtual and self.sql:
default_schema = self._get_rls_schema()
@@ -2902,7 +2905,9 @@ class SqlaTable(
# Add each predicate as a separate cache key component
extra_cache_keys.extend(rls_predicates)
- extra_cache_keys.extend(self._get_query_rls_cache_keys(query_obj))
+ extra_cache_keys.extend(
+ self._get_query_rls_cache_keys(query_obj, default_schema)
+ )
return list(set(extra_cache_keys))
@property
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
index c1b8af12646..095cbd28716 100644
--- a/superset/models/helpers.py
+++ b/superset/models/helpers.py
@@ -4028,22 +4028,120 @@ class ExploreMixin: # pylint:
disable=too-many-public-methods
self.template_params_dict,
)
- def _get_query_rls_cache_keys(self, query_obj: QueryObjectDict) ->
list[str]:
- """Include expression subqueries' RLS, not just the datasource SQL's
rules.
+ def _get_query_sql_expressions(
+ self, query_obj: QueryObjectDict, template_processor:
BaseTemplateProcessor
+ ) -> list[str]:
+ """Find SQL-bearing chart fields without building a chart or its
prequeries."""
+ # Real column names take precedence over verbose names in filter
lookup.
+ columns = {col.verbose_name: col.expression for col in self.columns}
+ columns.update({col.column_name: col.expression for col in
self.columns})
+ metrics = {metric.metric_name: metric.expression for metric in
self.metrics}
+ expressions: list[str] = []
+
+ def add(value: Any, metric: bool = False) -> None:
+ """Resolve stored references in the namespace used by the query
sink."""
+ if isinstance(value, str):
+ references = metrics if metric else columns
+ expressions.append(references.get(value) or value)
+ elif isinstance(value, dict):
+ if sql := value.get("sqlExpression"):
+ sql = template_processor.process_template(sql)
+ # Adhoc columns can name a saved calculated column after
+ # rendering. SQL metrics use their SQL literally instead.
+ expressions.append(
+ sql if "expressionType" in value else columns.get(sql)
or sql
+ )
+ add(value.get("column"))
+ add(value.get("column_name"))
- Build (but do not execute) the query to cover rendered expressions in
columns,
- metrics, filters, ordering and stored expressions through their
query-time
- validation paths. These subqueries are not constrained by the outer
dataset's
- RLS, so neither its id nor global guest rules may be excluded here.
+ for value in [
+ *(query_obj.get("columns") or []),
+ *(query_obj.get("groupby") or []),
+ *(query_obj.get("series_columns") or []),
+ ]:
+ add(value)
+ for value in query_obj.get("metrics") or []:
+ add(value, metric=True)
+ for value, _ascending in query_obj.get("orderby") or []:
+ add(value, metric=isinstance(value, str) and value in metrics)
+ for flt in query_obj.get("filter") or []:
+ value = flt.get("col")
+ add(value, metric=isinstance(value, str) and value not in columns)
+ add(query_obj.get("series_limit_metric"), metric=True)
+ add(query_obj.get("timeseries_limit_metric"), metric=True)
+ add(query_obj.get("granularity"))
+ add(self.main_dttm_col if query_obj.get("granularity") else None)
+ extras = query_obj.get("extras") or {}
+ expressions.extend(
+ extras[key] for key in ("where", "having") if extras.get(key)
+ )
+ if (
+ query_obj.get("apply_fetch_values_predicate")
+ and self.fetch_values_predicate
+ ):
+ expressions.append(self.fetch_values_predicate)
+ return expressions
+
+ def _get_query_rls_cache_keys(
+ self, query_obj: QueryObjectDict, default_schema: str | None = None
+ ) -> list[str]:
+ """Collect expression-only RLS rules without building or executing SQL.
+
+ Only subquery expressions need table discovery. Reuse the
+ datasource SQL's resolved schema, or resolve it once for physical
datasets.
+ Do not exclude the dataset's own id: an independent expression
subquery is
+ not constrained by the outer query's RLS.
"""
from superset.utils.rls import collect_rls_predicates_for_sql
if not query_obj or not is_feature_enabled("ALLOW_ADHOC_SUBQUERY"):
return []
- sql = self.get_query_str_extended(query_obj, mutate=False).full_sql
- return collect_rls_predicates_for_sql(
- sql, self.database, self.catalog, self._get_rls_schema()
+ extras = query_obj.get("extras") or {}
+ granularity = query_obj.get("granularity")
+ if granularity is not None and granularity not in self.dttm_cols:
+ granularity = self.main_dttm_col
+ template_kwargs = {
+ key: query_obj.get(key)
+ for key in (
+ "columns",
+ "groupby",
+ "metrics",
+ "row_limit",
+ "row_offset",
+ "filter",
+ )
+ }
+ template_kwargs.update(
+ time_column=granularity,
+ time_grain=extras.get("time_grain_sqla"),
+ table_columns=[col.column_name for col in self.columns],
)
+ template_kwargs.update(self.template_params_dict)
+ template_processor = self.get_template_processor(**template_kwargs)
+ predicates: set[str] = set()
+ for expression in self._get_query_sql_expressions(
+ query_obj, template_processor
+ ):
+ expression = template_processor.process_template(expression)
+ sql = f"SELECT {expression}"
+ try:
+ parsed = SQLStatement(sql, self.db_engine_spec.engine)
+ except SupersetParseError:
+ # The collector supplies a per-user failure key for
unparseable SQL.
+ pass
+ else:
+ if not parsed.has_subquery():
+ continue
+ if default_schema is None:
+ default_schema = self._get_rls_schema()
+ # A query root lets scope traversal find tables in scalar
expressions
+ # and select-list fragments such as DISTINCT (SELECT ...).
+ predicates.update(
+ collect_rls_predicates_for_sql(
+ sql, self.database, self.catalog, default_schema
+ )
+ )
+ return sorted(predicates)
def get_from_clause(
self, template_processor: Optional[BaseTemplateProcessor] = None
@@ -5034,8 +5132,13 @@ class ExploreMixin: # pylint:
disable=too-many-public-methods
timeseries_limit: Optional[int] = None,
timeseries_limit_metric: Optional[Metric] = None,
time_shift: Optional[str] = None,
+ execute_prequeries: bool = True,
) -> SqlaQuery:
- """Querying any sqla table from this common interface"""
+ """Build a chart query, optionally omitting execution of series
prequeries.
+
+ Cache-key extraction disables prequeries: only template side effects
are
+ needed, not the top-series rows used to constrain an executed chart.
+ """
if granularity not in self.dttm_cols and granularity is not None:
granularity = self.main_dttm_col
@@ -6199,67 +6302,72 @@ class ExploreMixin: # pylint:
disable=too-many-public-methods
)
]
- # run prequery to get top groups
- prequery_obj: QueryObjectDict = {
- "is_timeseries": False,
- "row_limit": series_limit,
- "metrics": metrics,
- "granularity": granularity,
- "groupby": groupby,
- "from_dttm": inner_from_dttm or from_dttm,
- "to_dttm": inner_to_dttm or to_dttm,
- "filter": filter or [],
- "orderby": orderby,
- "extras": extras,
- "columns": get_non_base_axis_columns(columns),
- "order_desc": True,
- }
-
- result = self.query(prequery_obj)
- prequeries.append(result.query)
- dimensions = [
- c
- for c in result.df.columns
- if c not in metrics and c in groupby_series_columns
- ]
- top_groups = self._get_top_groups(
- result.df, dimensions, groupby_series_columns,
columns_by_name
- )
-
- if group_others_when_limit_reached:
- # Apply Others grouping using the refactored method
- def _create_top_groups_condition(col_name: str, expr: Any)
-> Any:
- return top_groups
+ if execute_prequeries:
+ # run prequery to get top groups
+ prequery_obj: QueryObjectDict = {
+ "is_timeseries": False,
+ "row_limit": series_limit,
+ "metrics": metrics,
+ "granularity": granularity,
+ "groupby": groupby,
+ "from_dttm": inner_from_dttm or from_dttm,
+ "to_dttm": inner_to_dttm or to_dttm,
+ "filter": filter or [],
+ "orderby": orderby,
+ "extras": extras,
+ "columns": get_non_base_axis_columns(columns),
+ "order_desc": True,
+ }
- select_exprs, groupby_all_columns = (
- self._apply_series_others_grouping(
- select_exprs,
- groupby_all_columns,
- groupby_series_columns,
- _create_top_groups_condition,
- )
+ result = self.query(prequery_obj)
+ prequeries.append(result.query)
+ dimensions = [
+ c
+ for c in result.df.columns
+ if c not in metrics and c in groupby_series_columns
+ ]
+ top_groups = self._get_top_groups(
+ result.df, dimensions, groupby_series_columns,
columns_by_name
)
- # Reconstruct query with modified expressions
- qry = sa.select(*select_exprs)
- if groupby_all_columns:
- qry = qry.group_by(
-
*self._groupby_clause_exprs(groupby_all_columns.values())
+ if group_others_when_limit_reached:
+ # Apply Others grouping using the refactored method
+ def _create_top_groups_condition(
+ col_name: str, expr: Any
+ ) -> Any:
+ return top_groups
+
+ select_exprs, groupby_all_columns = (
+ self._apply_series_others_grouping(
+ select_exprs,
+ groupby_all_columns,
+ groupby_series_columns,
+ _create_top_groups_condition,
+ )
)
- # Re-apply WHERE and HAVING clauses lost during query
reconstruction
- qry = self._reapply_query_filters(
- qry,
- apply_fetch_values_predicate,
- template_processor,
- granularity,
- time_filters,
- where_clause_and,
- having_clause_and,
- )
- else:
- # Original behavior: filter to only top groups
- qry = qry.where(top_groups)
+ # Reconstruct query with modified expressions
+ qry = sa.select(*select_exprs)
+ if groupby_all_columns:
+ qry = qry.group_by(
+ *self._groupby_clause_exprs(
+ groupby_all_columns.values()
+ )
+ )
+
+ # Restore filters lost during query reconstruction.
+ qry = self._reapply_query_filters(
+ qry,
+ apply_fetch_values_predicate,
+ template_processor,
+ granularity,
+ time_filters,
+ where_clause_and,
+ having_clause_and,
+ )
+ else:
+ # Original behavior: filter to only top groups
+ qry = qry.where(top_groups)
qry = qry.select_from(tbl)
self.rename_shadowing_aliases(qry)
diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py
index b806d471f9e..467199dec62 100644
--- a/superset/models/sql_lab.py
+++ b/superset/models/sql_lab.py
@@ -417,7 +417,7 @@ class Query(
self.catalog,
default_schema,
)
- predicates.extend(self._get_query_rls_cache_keys(query_obj))
+ predicates.extend(self._get_query_rls_cache_keys(query_obj,
default_schema))
return list(dict.fromkeys(predicates))
def get_time_grains(self) -> list[TimeGrainDict]:
diff --git
a/tests/unit_tests/connectors/sqla/virtual_dataset_connections_test.py
b/tests/unit_tests/connectors/sqla/virtual_dataset_connections_test.py
index bc80ad44f94..579a193de51 100644
--- a/tests/unit_tests/connectors/sqla/virtual_dataset_connections_test.py
+++ b/tests/unit_tests/connectors/sqla/virtual_dataset_connections_test.py
@@ -680,3 +680,379 @@ def test_virtual_dataset_without_uri_database(
extractor.assert_called_once()
assert extractor.call_args.args[0].database is None
+
+
[email protected]("datasource_class", [SqlaTable, Query])
[email protected]("dynamic_schema", [False, True])
[email protected]("adhoc", [False, True])
+@with_feature_flags(ALLOW_ADHOC_SUBQUERY=True)
+def test_cache_hit_schema_discovery_connection_budget(
+ mocker: MockerFixture,
+ tmp_path: Path,
+ app_context: None,
+ datasource_class: type[SqlaTable] | type[Query],
+ dynamic_schema: bool,
+ adhoc: bool,
+) -> None:
+ """Expression cache keys must reuse the datasource's single inspector
lookup."""
+ uri = f"sqlite:///{tmp_path / 'analytic.db'}"
+ engine = create_engine(uri, poolclass=NullPool)
+ connections = MagicMock()
+ event.listen(engine, "connect", connections)
+ database = Database(id=1, database_name="analytic", sqlalchemy_uri=uri)
+ mocker.patch.object(database, "_get_sqla_engine", return_value=engine)
+ mocker.patch.object(
+ database.db_engine_spec, "supports_dynamic_schema", dynamic_schema
+ )
+ # SQLite needs its inspector to discover "main". The static-engine case
+ # deliberately disagrees with the dataset's configured schema.
+ virtual = datasource_class(
+ id=2,
+ sql="SELECT col1 FROM physical",
+ schema=None if dynamic_schema else "not_the_default",
+ database=database,
+ )
+ if isinstance(virtual, SqlaTable):
+ virtual.table_name = "virtual"
+ virtual.columns = [TableColumn(column_name="col1", type="INTEGER")]
+ else:
+ virtual.extra = {
+ "columns": [{"column_name": "col1", "type": "INTEGER", "is_dttm":
False}]
+ }
+ physical = SqlaTable(id=1, table_name="physical", schema="main",
database=database)
+ rule = MagicMock(clause="col1 = 1", group_key=None)
+ mocker.patch("superset.utils.rls.find_datasets", return_value=[physical])
+ mocker.patch(
+ "superset.connectors.sqla.models.security_manager.get_rls_filters",
+ side_effect=lambda dataset: [rule] if dataset.id == 1 else [],
+ )
+ mocker.patch(
+
"superset.connectors.sqla.models.security_manager.get_guest_rls_filters",
+ return_value=[],
+ )
+ mocker.patch(
+
"superset.common.query_context_processor.security_manager.get_rls_cache_key",
+ return_value=[],
+ )
+ cache = MagicMock(is_loaded=True, status="success",
bq_memory_limited=False)
+ cache.df = pd.DataFrame({"restricted": [1]})
+ mocker.patch(
+ "superset.common.query_context_processor.QueryCacheManager.get",
+ return_value=cache,
+ )
+ processor = QueryContextProcessor(
+ MagicMock(
+ datasource=virtual,
+ force=False,
+ is_async_execution=False,
+ custom_cache_timeout=300,
+ result_type="full",
+ )
+ )
+ query = QueryObject(
+ datasource=virtual,
+ columns=[
+ {"sqlExpression": "(SELECT MAX(col1) FROM physical)", "label":
"restricted"}
+ ]
+ if adhoc
+ else ["col1"],
+ metrics=[],
+ )
+ cache_key = mocker.spy(query, "cache_key")
+ result = processor.get_df_payload_result(query)
+
+ assert result.payload["status"] == "success"
+ assert cache_key.call_args.kwargs["extra_cache_keys"] == ["(col1 = 1)"]
+ assert connections.call_count == 1
+ cache.set_query_result.assert_not_called()
+ engine.dispose()
+
+
[email protected](
+ "static_schema_datasource,extra_cache_macro",
+ [(SqlaTable, False), (SqlaTable, True), (Query, False)],
+ indirect=["static_schema_datasource"],
+)
+@with_feature_flags(ALLOW_ADHOC_SUBQUERY=True, ENABLE_TEMPLATE_PROCESSING=True)
+def test_cache_hit_series_limit_does_not_execute_prequery(
+ static_schema_datasource: tuple[SqlaTable | Query, MagicMock],
+ mocker: MockerFixture,
+ extra_cache_macro: bool,
+) -> None:
+ """No-join engines must not execute top-series queries for cache keys."""
+ virtual, _ = static_schema_datasource
+ # Use Druid's no-join planner with the fixture's SQLite transport.
+ mocker.patch.object(
+ virtual.database, "get_db_engine_spec", return_value=DruidEngineSpec
+ )
+ mocker.patch.object(DruidEngineSpec, "allows_joins", False)
+ if extra_cache_macro:
+ virtual.sql += " WHERE 1 = {{ cache_key_wrapper(1) }}"
+ execute = mocker.spy(virtual, "query")
+ cache = MagicMock(is_loaded=True, status="success",
bq_memory_limited=False)
+ cache.df = pd.DataFrame({"col1": [1], "total": [1]})
+ mocker.patch(
+ "superset.common.query_context_processor.QueryCacheManager.get",
+ return_value=cache,
+ )
+ mocker.patch(
+
"superset.common.query_context_processor.security_manager.get_rls_cache_key",
+ return_value=[],
+ )
+ processor = QueryContextProcessor(
+ MagicMock(
+ datasource=virtual,
+ force=False,
+ is_async_execution=False,
+ custom_cache_timeout=300,
+ result_type="full",
+ )
+ )
+ query = QueryObject(
+ datasource=virtual,
+ columns=["col1"],
+ metrics=[
+ {"expressionType": "SQL", "sqlExpression": "SUM(col1)", "label":
"total"}
+ ],
+ series_columns=["col1"],
+ series_limit=1,
+ )
+
+ result = processor.get_df_payload_result(query)
+
+ assert result.payload["status"] == "success"
+ execute.assert_not_called()
+ cache.set_query_result.assert_not_called()
+
+
[email protected]("static_schema_datasource", [SqlaTable],
indirect=True)
[email protected](
+ "fields",
+ [
+ {"columns": ["calculated"]},
+ {
+ "columns": [
+ {
+ "sqlExpression": "calculated",
+ "label": "axis",
+ "columnType": "BASE_AXIS",
+ "timeGrain": "P1D",
+ }
+ ]
+ },
+ {"groupby": ["calculated"]},
+ {"columns": ["(SELECT MAX(col1) FROM physical)"], "metrics":
["restricted"]},
+ {"metrics": ["restricted"]},
+ {
+ "metrics": [
+ {
+ "expressionType": "SIMPLE",
+ "aggregate": "MAX",
+ "column": {"column_name": "calculated"},
+ "label": "total",
+ }
+ ]
+ },
+ {"orderby": [("restricted", True)]},
+ {"orderby": [("calculated", True)]},
+ {"filter": [{"col": "calculated", "op": "==", "val": 1}]},
+ {"filter": [{"col": "caption", "op": "==", "val": 1}]},
+ {"filter": [{"col": "restricted", "op": "==", "val": 1}]},
+ {
+ "series_limit": 1,
+ "series_columns": ["col1"],
+ "series_limit_metric": "restricted",
+ },
+ {"timeseries_limit": 1, "timeseries_limit_metric": "restricted"},
+ {
+ "series_limit": 1,
+ "series_columns": ["col1"],
+ "series_limit_metric": {
+ "expressionType": "SQL",
+ "sqlExpression": "MAX({{ subquery }})",
+ "label": "total",
+ },
+ },
+ {"granularity": "calculated", "is_timeseries": True},
+ {"extras": {"where": "1 = {{ subquery }}"}},
+ {"extras": {"having": "1 = {{ subquery }}"}},
+ {"apply_fetch_values_predicate": True},
+ ],
+)
+@with_feature_flags(ALLOW_ADHOC_SUBQUERY=True, ENABLE_TEMPLATE_PROCESSING=True)
+def test_expression_cache_key_sibling_fields(
+ static_schema_datasource: tuple[SqlaTable | Query, MagicMock],
+ mocker: MockerFixture,
+ fields: QueryObjectDict,
+) -> None:
+ """Every expression entry point contributes RLS without a full query
build."""
+ virtual, rule = static_schema_datasource
+ assert isinstance(virtual, SqlaTable)
+ virtual.sql = "SELECT 1 AS col1"
+ virtual.template_params = '{"subquery": "(SELECT MAX(col1) FROM
physical)"}'
+ virtual.columns.append(
+ TableColumn(
+ column_name="calculated",
+ verbose_name="caption",
+ is_dttm=True,
+ expression="{{ subquery }}",
+ )
+ )
+ virtual.metrics.append(
+ SqlMetric(
+ metric_name="restricted",
+ expression="MAX({{ subquery }})",
+ )
+ )
+ virtual.fetch_values_predicate = "1 = (SELECT MAX(col1) FROM physical)"
+ build = mocker.spy(virtual, "get_sqla_query")
+ query: QueryObjectDict = {
+ "columns": ["col1"],
+ "metrics": [],
+ "is_timeseries": False,
+ }
+ query.update(fields)
+
+ assert virtual.get_extra_cache_keys(query) == ["(col1 = 1)"]
+ rule.clause = "col1 = 2"
+ assert virtual.get_extra_cache_keys(query) == ["(col1 = 2)"]
+ build.assert_not_called()
+
+
+def test_series_limit_execution_still_runs_prequery(
+ static_schema_datasource: tuple[SqlaTable | Query, MagicMock],
+ mocker: MockerFixture,
+) -> None:
+ """Skipping prequeries for keys must not disable top-series execution."""
+ virtual, _ = static_schema_datasource
+ mocker.patch.object(
+ virtual.database, "get_db_engine_spec", return_value=DruidEngineSpec
+ )
+ mocker.patch.object(DruidEngineSpec, "allows_joins", False)
+ execute = mocker.spy(virtual, "query")
+ query = QueryObject(
+ datasource=virtual,
+ columns=["col1"],
+ metrics=[
+ {"expressionType": "SQL", "sqlExpression": "SUM(col1)", "label":
"total"}
+ ],
+ series_columns=["col1"],
+ series_limit=1,
+ )
+
+ result = virtual.get_query_result(query)
+
+ assert result.status == "success"
+ assert result.df.to_dict(orient="list") == {"col1": [1], "total": [1]}
+ assert execute.call_count == 2
+
+
[email protected]("static_schema_datasource", [SqlaTable],
indirect=True)
[email protected]("sink", ["metric", "column", "verbose",
"rendered_reference"])
+@with_feature_flags(ALLOW_ADHOC_SUBQUERY=True, ENABLE_TEMPLATE_PROCESSING=True)
+def test_expression_cache_key_reference_precedence(
+ static_schema_datasource: tuple[SqlaTable | Query, MagicMock],
+ mocker: MockerFixture,
+ sink: str,
+) -> None:
+ """Name collisions and rendered column references retain the executed
rules."""
+ virtual, rule = static_schema_datasource
+ assert isinstance(virtual, SqlaTable)
+ virtual.sql = "SELECT 1 AS col1"
+ protected = "(SELECT MAX(col1) FROM physical)"
+ unrelated = "(SELECT MAX(col1) FROM unrelated)"
+ virtual.columns.append(
+ TableColumn(
+ column_name="restricted",
+ expression=unrelated if sink == "metric" else protected,
+ )
+ )
+ virtual.columns.append(
+ TableColumn(
+ column_name="another",
+ verbose_name="restricted",
+ expression=unrelated,
+ )
+ )
+ virtual.metrics.append(
+ SqlMetric(
+ metric_name="restricted",
+ expression=protected if sink == "metric" else unrelated,
+ )
+ )
+ query: QueryObjectDict = {
+ "columns": ["restricted"],
+ "metrics": [],
+ "is_timeseries": False,
+ }
+ if sink == "metric":
+ query.update(columns=[], metrics=["restricted"])
+ elif sink == "verbose":
+ query.update(
+ columns=["col1"], filter=[{"col": "restricted", "op": "==", "val":
1}]
+ )
+ elif sink == "rendered_reference":
+ query["columns"] = [{"sqlExpression": "{{ 'restricted' }}", "label":
"result"}]
+ build = mocker.spy(virtual, "get_sqla_query")
+
+ assert virtual.get_extra_cache_keys(query) == ["(col1 = 1)"]
+ rule.clause = "col1 = 2"
+ assert virtual.get_extra_cache_keys(query) == ["(col1 = 2)"]
+ build.assert_not_called()
+
+
[email protected]("static_schema_datasource", [SqlaTable],
indirect=True)
[email protected]("subquery", [False, True])
+@with_feature_flags(ALLOW_ADHOC_SUBQUERY=True)
+def test_physical_expression_cache_key_schema_budget(
+ static_schema_datasource: tuple[SqlaTable | Query, MagicMock],
+ mocker: MockerFixture,
+ subquery: bool,
+) -> None:
+ """Physical datasets discover a schema only once, and only for
subqueries."""
+ physical, _ = static_schema_datasource
+ physical.sql = None
+ expression = "(SELECT MAX(col1) FROM physical)" if subquery else "col1 + 1"
+ query: QueryObjectDict = {
+ "columns": [
+ {"sqlExpression": expression, "label": str(index)} for index in
range(3)
+ ],
+ "metrics": [],
+ "is_timeseries": False,
+ }
+ discover = mocker.spy(physical.database, "resolve_query_default_schema")
+ build = mocker.spy(physical, "get_sqla_query")
+
+ assert physical.get_extra_cache_keys(query) == (["(col1 = 1)"] if subquery
else [])
+ assert discover.call_count == (1 if subquery else 0)
+ build.assert_not_called()
+
+
+@with_feature_flags(ALLOW_ADHOC_SUBQUERY=True)
+def test_expression_cache_key_parse_failure_is_user_scoped(
+ static_schema_datasource: tuple[SqlaTable | Query, MagicMock],
+ mocker: MockerFixture,
+) -> None:
+ """Unparseable expressions must retain the collector's per-user
sentinel."""
+ virtual, _ = static_schema_datasource
+ virtual.sql = "SELECT 1 AS col1"
+ user_id = mocker.patch("superset.utils.rls.get_user_id", return_value=41)
+ mocker.patch(
+ "superset.utils.rls.security_manager.get_current_guest_user_if_guest",
+ return_value=None,
+ )
+ query: QueryObjectDict = {
+ "columns": [{"sqlExpression": "(SELECT FROM !)", "label": "broken"}],
+ "metrics": [],
+ "is_timeseries": False,
+ }
+ build = mocker.spy(virtual, "get_sqla_query")
+
+ assert virtual.get_extra_cache_keys(query) == [
+ "rls-predicate-parse-failed-for-user-41"
+ ]
+ user_id.return_value = 42
+ assert virtual.get_extra_cache_keys(query) == [
+ "rls-predicate-parse-failed-for-user-42"
+ ]
+ build.assert_not_called()
