hughhhh commented on code in PR #43870:
URL: https://github.com/apache/superset/pull/43870#discussion_r4319677594
##########
tests/unit_tests/commands/dataset/update_test.py:
##########
@@ -1366,3 +1375,254 @@ def
test_update_dataset_rejects_malicious_fetch_values_predicate(
and "fetch_values_predicate" in (exc.field_name or "")
for exc in excinfo.value._exceptions
)
+
+
+def _mapping_command(
+ mocker: MockerFixture,
+ transform: str | None,
+ *,
+ partition_mapped_column: str | None = None,
+) -> UpdateDatasetCommand:
+ """A command whose stored dataset maps `event_time` onto `dt_epoch`."""
+ mapped_column = mocker.MagicMock()
+ mapped_column.column_name = "event_time"
+ mapped_column.partition_value_transform = transform
+ partition_column = mocker.MagicMock()
+ partition_column.column_name = "dt_epoch"
+ # The partition column is the target of a mapping, never the source of one.
+ partition_column.partition_value_transform = None
+
+ mock_dataset = mocker.MagicMock(is_managed_externally=False)
+ mock_dataset.database.backend = "sqlite"
+ mock_dataset.catalog = None
+ mock_dataset.schema = "main"
+ mock_dataset.columns = [mapped_column, partition_column]
+ mock_dataset.main_dttm_col = "event_time"
+ mock_dataset.partition_column = "dt_epoch"
+ mock_dataset.partition_mapped_column = partition_mapped_column
+
+ command = UpdateDatasetCommand(1, {})
+ command._model = mock_dataset
+ return command
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_an_unparseable_transform_does_not_block_the_save(
+ mocker: MockerFixture,
+) -> None:
+ """
+ An unparseable transform is a Tier-2 issue: the mapping saves and stays
+ inactive. `validate_stored_expression` rejects anything it cannot parse, so
+ running it here would turn that into a blocking error and cost the owner
the
+ rest of their edits.
+ """
+ gate =
mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value")
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert exceptions == []
+ gate.assert_not_called()
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_a_parseable_transform_still_goes_through_the_stored_expression_gate(
+ mocker: MockerFixture,
+) -> None:
+ """The parser gate that governs every other stored expression still
runs."""
+ gate =
mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value)")
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert exceptions == []
+ gate.assert_called_once()
+ assert ":value" not in gate.call_args.args[-1]
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
[email protected](
+ "transform",
+ [
+ "unix_timestamp(:value) UNION ALL SELECT password FROM ab_user",
+ "unix_timestamp(:value); DROP TABLE ab_user",
+ "password || :value FROM ab_user",
+ ],
+ ids=["set-operation", "multi-statement", "bare-from"],
+)
+def test_a_transform_that_is_the_wrong_sql_still_reaches_the_gate(
+ mocker: MockerFixture, transform: str
+) -> None:
+ """
+ The companion to the test above, and the reason the skip condition is "not
+ SQL yet" rather than "does not parse". All three of these fail to parse as
+ a single statement, exactly as a half-typed transform does -- so a skip
+ keyed on parseability would have waved them through. None of them is on the
+ way to a valid transform.
+ """
+ gate =
mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, transform)
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ # Either the shape issue reports it or the gate does; what matters is that
+ # it is not silently accepted the way an unfinished transform is.
+ assert exceptions or gate.called
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_a_jinja_transform_is_still_rejected(mocker: MockerFixture) -> None:
+ """
+ Skipping the gate for unparseable transforms is not a hole for templating:
+ Jinja is a blocking issue of its own, reported before the gate is reached.
+ """
+ mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp({{ current_user_id()
}})")
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert [exc.field_name for exc in exceptions] ==
["partition_value_transform"]
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_a_description_only_put_survives_an_implicit_self_mapping(
+ mocker: MockerFixture,
+) -> None:
+ """
+ `fetch_metadata` could choose the partition column as the default datetime
+ column without anyone asking, and a blocking self-mapping error then made
+ the dataset unsaveable forever -- including by the very write that would
+ have set the override that fixes it, and including a write that changes
+ nothing but the description.
+
+ The mapping is inert in that state either way, since
+ `resolve_partition_mapping` bails out on it, so the issue is reported
+ rather than blocking.
+ """
+ mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value)")
+ command._model.main_dttm_col = "dt_epoch"
Review Comment:
Fixed — `pre-commit run mypy` passes on this file at the current head.
##########
tests/unit_tests/connectors/sqla/partition_mapping_test.py:
##########
@@ -0,0 +1,2512 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for ``superset.connectors.sqla.partition_mapping``."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from decimal import Decimal
+from importlib import import_module
+from typing import Any, cast
+from unittest.mock import MagicMock, patch, PropertyMock
+
+import pandas as pd
+import pytest
+from dateutil.relativedelta import relativedelta
+from flask import Flask
+from sqlalchemy.dialects import mysql, postgresql
+
+from superset.config import DISALLOWED_SQL_FUNCTIONS
+from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.connectors.sqla.partition_mapping import (
+ _mirror_verdict_cache_key,
+ _placeholder_is_bindable,
+ _probe_cache_key,
+ _reads_as_temporal,
+ _render_literal,
+ build_mirrored_predicates,
+ build_probe_sql,
+ contains_executable_comment,
+ contains_jinja,
+ contains_value_placeholder,
+ drop_unmapped_value_transforms,
+ equality_mirrors_safely,
+ evaluate_transform,
+ find_non_deterministic_functions,
+ grain_bucket_width,
+ GRAIN_BUCKET_WIDTHS,
+ is_bare_expression,
+ is_parseable,
+ is_transform_active,
+ is_unfinished,
+ known_mirror_verdict,
+ MappingValidationIssue,
+ mirror_operator,
+ MIRRORABLE_ALWAYS,
+ MIRRORABLE_IF_MONOTONIC,
+ mirrorable_operators,
+ parse_error_detail,
+ placeholder_is_executable,
+ probe_sql_is_evaluable,
+ probed_value_type_error,
+ RawProbeValue,
+ resolve_partition_mapping,
+ stored_expression_error,
+ validate_partition_mapping,
+ validate_transform,
+)
+from superset.constants import TimeGrain
+from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.oracle import OracleEngineSpec
+from superset.models.core import Database
+from superset.sql.parse import SQLStatement
+from superset.utils.core import FilterOperator, override_user
+
+
[email protected](autouse=True)
+def real_probe_cache(app: Flask) -> Any:
+ """
+ The test app runs a null cache, which would make every cache assertion here
+ vacuously pass. Swap in a real in-memory one for the duration.
+ """
+ from flask_caching import Cache
+
+ from superset.extensions import cache_manager
+
+ cache = Cache(config={"CACHE_TYPE": "SimpleCache",
"CACHE_DEFAULT_TIMEOUT": 300})
+ cache.init_app(app)
+ original = cache_manager._cache # noqa: SLF001
+ cache_manager._cache = cache # noqa: SLF001
+ yield
+ cache_manager._cache = original # noqa: SLF001
+
+
[email protected](autouse=True)
+def enable_partition_filter_mapping(app: Flask) -> Any:
+ """The feature ships off; every test here exercises it on."""
+ original =
app.config["DEFAULT_FEATURE_FLAGS"].get("PARTITION_FILTER_MAPPING")
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = True
+ yield
+ if original is None:
+ del app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"]
+ else:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] =
original
+
+
+def _table(**kwargs: Any) -> SqlaTable:
+ database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
+ defaults: dict[str, Any] = {
+ "table_name": "web_events",
+ "database": database,
+ "main_dttm_col": "event_time",
+ "columns": [
+ TableColumn(column_name="event_time", is_dttm=True,
type="TIMESTAMP"),
+ TableColumn(
+ column_name="dt_epoch",
+ type="BIGINT",
+ partition_value_transform=None,
+ ),
+ TableColumn(column_name="country", type="VARCHAR"),
+ TableColumn(column_name="region_key", type="VARCHAR"),
+ ],
+ }
+ defaults.update(kwargs)
+ return SqlaTable(**defaults)
+
+
+def _mapped_table(
+ *,
+ partition_column: str = "dt_epoch",
+ mapped_column: str = "event_time",
+ transform: str | None = "unix_timestamp(:value)",
+ monotonic: bool = True,
+ partition_mapped_column: str | None = None,
+ main_dttm_col: str | None = "event_time",
+) -> SqlaTable:
+ table = _table(main_dttm_col=main_dttm_col)
+ table.partition_column = partition_column
+ table.partition_mapped_column = partition_mapped_column
+ for column in table.columns:
+ if column.column_name == mapped_column:
+ column.partition_value_transform = transform
+ column.partition_transform_is_monotonic = monotonic
+ return table
+
+
+# ---------------------------------------------------------------------------
+# §2 — operator safety matrix
+# ---------------------------------------------------------------------------
+
+
+def test_equality_and_in_are_always_mirrorable() -> None:
+ """``=`` and ``IN`` are safe for any function ``T``."""
+ assert MIRRORABLE_ALWAYS == {FilterOperator.EQUALS, FilterOperator.IN}
+
+
+def test_range_operators_require_a_monotonic_transform() -> None:
+ assert MIRRORABLE_IF_MONOTONIC == {
+ FilterOperator.GREATER_THAN,
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ FilterOperator.LESS_THAN,
+ FilterOperator.LESS_THAN_OR_EQUALS,
+ FilterOperator.TEMPORAL_RANGE,
+ }
+
+
+def test_mirrorable_operators_excludes_ranges_when_not_monotonic() -> None:
+ assert mirrorable_operators(is_monotonic=False) == MIRRORABLE_ALWAYS
+
+
+def test_mirrorable_operators_includes_ranges_when_monotonic() -> None:
+ assert mirrorable_operators(is_monotonic=True) == (
+ MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC
+ )
+
+
[email protected](
+ "operator",
+ [
+ FilterOperator.NOT_EQUALS,
+ FilterOperator.NOT_IN,
+ FilterOperator.LIKE,
+ FilterOperator.ILIKE,
+ FilterOperator.NOT_LIKE,
+ FilterOperator.NOT_ILIKE,
+ FilterOperator.IS_NULL,
+ FilterOperator.IS_NOT_NULL,
+ FilterOperator.IS_TRUE,
+ FilterOperator.IS_FALSE,
+ ],
+)
+def test_negations_and_pattern_matches_are_never_mirrorable(
+ operator: FilterOperator,
+) -> None:
+ """
+ ``T`` is not injective, so ``col != v`` does **not** imply ``T(col) !=
T(v)``:
+ mirroring it would drop rows the original filter keeps.
+ """
+ assert operator not in mirrorable_operators(is_monotonic=True)
+
+
+def test_equality_is_dropped_when_the_engine_does_not_compare_byte_exactly()
-> None:
+ """
+ "Safe for any function ``T``" reasons about value equality: ``T`` is a
+ function, so ``col = v`` gives ``T(col) = T(v)``. The engine reasons about
+ *SQL* equality, and the two part company under a case-insensitive
+ collation -- a stored ``'us'`` satisfies a filter for ``'US'`` while the
+ mirror ``hex('US')`` excludes the row, and the chart loses it silently.
+ """
+ assert mirrorable_operators(is_monotonic=False, equality_is_safe=False) ==
set()
+
+
+def test_ranges_survive_an_unsafe_equality() -> None:
+ """
+ Range mirroring already rests on the owner declaring ``T``
+ order-preserving with respect to the column's own order -- an assertion
+ about the very comparison semantics this gate checks for. Equality has no
+ such declaration behind it, which is why only equality is withdrawn.
+ """
+ assert (
+ mirrorable_operators(is_monotonic=True, equality_is_safe=False)
+ == MIRRORABLE_IF_MONOTONIC
+ )
+
+
[email protected](
+ "operator,expected",
+ [
+ (FilterOperator.GREATER_THAN, FilterOperator.GREATER_THAN_OR_EQUALS),
+ (FilterOperator.LESS_THAN, FilterOperator.LESS_THAN_OR_EQUALS),
+ (
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ ),
+ (FilterOperator.LESS_THAN_OR_EQUALS,
FilterOperator.LESS_THAN_OR_EQUALS),
+ (FilterOperator.EQUALS, FilterOperator.EQUALS),
+ (FilterOperator.IN, FilterOperator.IN),
+ ],
+)
+def test_a_strict_bound_mirrors_non_strictly(
+ operator: FilterOperator, expected: FilterOperator
+) -> None:
+ """
+ Order-preserving means non-decreasing, not injective, so ``col < v``
+ implies only ``T(col) <= T(v)``. A day-key transform maps a whole day onto
+ one value and a strict mirror drops the boundary bucket. The non-strict
+ cases are already safe and pass through, which is what makes applying this
+ twice harmless.
+ """
+ assert mirror_operator(operator) == expected
+
+
+def test_relaxing_a_bound_stays_inside_the_mirrorable_set() -> None:
+ """
+ The substitution is about strictness, not about which filters mirror. The
+ Explore indicator reads `mirrorable_operators`, so leaving that set alone
is
+ what keeps the glyph from promising pruning the SQL does not do -- and it
is
+ also what guarantees `handle_comparison_filter` is never handed an operator
+ it has no case for.
+ """
+ assert {
+ mirror_operator(operator) for operator in MIRRORABLE_IF_MONOTONIC
+ } <= MIRRORABLE_IF_MONOTONIC
+
+
+def _column(type_: str) -> TableColumn:
+ database = Database(database_name="t", sqlalchemy_uri="sqlite://")
+ table = SqlaTable(table_name="t", database=database)
+ return TableColumn(column_name="c", type=type_, table=table)
+
+
[email protected](
+ "type_, binary, expected",
+ [
+ ("VARCHAR", True, True),
+ ("VARCHAR", False, False),
+ # Numeric and temporal comparison is exact everywhere.
+ ("BIGINT", False, True),
+ ("TIMESTAMP", False, True),
+ ("DOUBLE", False, True),
+ ],
+)
+def test_equality_mirrors_safely_only_guards_string_columns(
+ app: Flask, type_: str, binary: bool, expected: bool
+) -> None:
+ class _Spec(BaseEngineSpec):
+ binary_string_comparison = binary
+
+ with app.app_context():
+ assert equality_mirrors_safely(_column(type_), _Spec) is expected
+
+
+def test_equality_mirrors_safely_fails_closed_on_an_unresolvable_type(
+ app: Flask,
+) -> None:
+ """
+ The gate exists to stop a silent wrong answer, so a column whose type it
+ cannot read is treated as the risky case rather than waved through.
+ """
+ column = MagicMock()
+ type(column).type_generic = PropertyMock(side_effect=ValueError("no type"))
+
+ with app.app_context():
+ assert equality_mirrors_safely(column, BaseEngineSpec) is False
+
+
+def test_no_mapped_column_leaves_the_matrix_alone(app: Flask) -> None:
+ """Nothing to reason about, and the resolver bails out on it anyway."""
+ with app.app_context():
+ assert equality_mirrors_safely(None, BaseEngineSpec) is True
+
+
[email protected](
+ "module, spec_name",
+ [
+ ("presto", "PrestoEngineSpec"),
+ ("hive", "HiveEngineSpec"),
+ ("trino", "TrinoEngineSpec"),
+ ("impala", "ImpalaEngineSpec"),
+ ("postgres", "PostgresEngineSpec"),
+ ("bigquery", "BigQueryEngineSpec"),
+ ("sqlite", "SqliteEngineSpec"),
+ ],
+)
+def test_the_byte_exact_engines_opt_in(module: str, spec_name: str) -> None:
+ """
+ The engines this feature targets compare text byte-exactly, so the gate
+ must not cost them the second canonical mapping (`lower(:value)` onto a
+ lowercased key).
+ """
+ spec = getattr(import_module(f"superset.db_engine_specs.{module}"),
spec_name)
+ assert spec.binary_string_comparison is True
+
+
[email protected]("module, spec_name", [("mysql", "MySQLEngineSpec")])
+def test_a_case_insensitive_engine_does_not_opt_in(module: str, spec_name:
str) -> None:
+ """
+ MySQL's default collation is `utf8mb4_0900_ai_ci`, which is exactly the
+ case the reviewer raised.
+ """
+ spec = getattr(import_module(f"superset.db_engine_specs.{module}"),
spec_name)
+ assert spec.binary_string_comparison is False
+
+
+def test_the_base_spec_does_not_opt_in() -> None:
+ """A spec that has not said so does not mirror string equality."""
+ assert BaseEngineSpec.binary_string_comparison is False
+
+
+# ---------------------------------------------------------------------------
+# Time grain bucket widths
+# ---------------------------------------------------------------------------
+
+
[email protected]("grain", list(TimeGrain))
+def test_every_builtin_grain_has_a_known_bucket_width(
+ app: Flask, grain: TimeGrain
+) -> None:
+ """
+ The widening is only sound for a grain whose bucket width we know, so a
+ grain added upstream without a width here would silently stop mirroring.
+ This is the test that notices.
+ """
+ with app.app_context():
+ assert grain_bucket_width(grain.value, "sqlite") is not None
+
+
+def test_a_calendar_grain_widens_by_a_calendar_month(app: Flask) -> None:
+ """A month is not thirty days, and the month it lands in decides its
width."""
+ with app.app_context():
+ width = grain_bucket_width(TimeGrain.MONTH.value, "sqlite")
+
+ assert width is not None
+ assert datetime(2026, 1, 31) + width == datetime(2026, 2, 28)
+ assert datetime(2026, 2, 1) + width == datetime(2026, 3, 1)
+
+
+def test_a_fractional_grain_has_a_width(app: Flask) -> None:
+ """
+ ``PT0.5H`` and ``P0.25Y`` are why the widths are written out rather than
+ parsed -- and the anchored week grains are intervals ``parse_duration``
+ rejects outright.
+ """
+ with app.app_context():
+ assert grain_bucket_width(TimeGrain.HALF_HOUR.value, "sqlite") ==
relativedelta(
+ minutes=30
+ )
+ assert grain_bucket_width(
+ TimeGrain.WEEK_ENDING_SATURDAY.value, "sqlite"
+ ) == relativedelta(days=7)
+
+
+def test_an_unknown_grain_has_no_known_width(app: Flask) -> None:
+ with app.app_context():
+ assert grain_bucket_width("P13X", "sqlite") is None
+ assert grain_bucket_width(None, "sqlite") is None
+
+
+def test_an_addon_grain_has_no_known_width(app: Flask) -> None:
+ """An operator-declared grain carries whatever duration string they
typed."""
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDONS"] = {"PT2S": "2 second"}
+ try:
+ assert grain_bucket_width("PT2S", "sqlite") is None
+ finally:
+ app.config["TIME_GRAIN_ADDONS"] = {}
+
+
+def test_a_builtin_grain_an_operator_redefined_has_no_known_width(
+ app: Flask,
+) -> None:
+ """
+ ``TIME_GRAIN_ADDON_EXPRESSIONS`` replaces a built-in grain's SQL for one
+ engine, so its duration no longer bounds the displacement *there* -- but it
+ still does on every other engine.
+ """
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {
+ "sqlite": {"P1D": "DATETIME({col}, 'start of year')"}
+ }
+ try:
+ assert grain_bucket_width(TimeGrain.DAY.value, "sqlite") is None
+ assert grain_bucket_width(TimeGrain.DAY.value, "hive") is not None
+ finally:
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {}
+
+
+def test_the_width_table_is_keyed_on_grain_durations(app: Flask) -> None:
+ """Keys are the ISO duration a filter carries, not the enum member name."""
+ assert GRAIN_BUCKET_WIDTHS[TimeGrain.DAY] == relativedelta(days=1)
+ assert set(GRAIN_BUCKET_WIDTHS) == {grain.value for grain in TimeGrain}
+
+
+# ---------------------------------------------------------------------------
+# §4.1 — resolving the effective mapping, and the defensive bail-outs
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_returns_the_mapping_when_everything_lines_up(app: Flask) ->
None:
+ with app.app_context():
+ mapping = resolve_partition_mapping(_mapped_table())
+
+ assert mapping is not None
+ assert mapping.partition_column == "dt_epoch"
+ assert mapping.mapped_column == "event_time"
+ assert mapping.value_transform == "unix_timestamp(:value)"
+ assert mapping.is_monotonic is True
+
+
+def test_effective_mapped_column_follows_main_dttm_col(app: Flask) -> None:
+ """No explicit override: the mapping follows the default datetime
column."""
+ with app.app_context():
+ mapping = resolve_partition_mapping(_mapped_table())
+
+ assert mapping is not None
+ assert mapping.mapped_column == "event_time"
+
+
+def test_explicit_override_wins_over_main_dttm_col(app: Flask) -> None:
+ table = _mapped_table(
+ mapped_column="country",
+ transform="lower(:value)",
+ monotonic=False,
+ partition_mapped_column="country",
+ )
+ table.partition_column = "region_key"
+
+ with app.app_context():
+ mapping = resolve_partition_mapping(table)
+
+ assert mapping is not None
+ assert mapping.mapped_column == "country"
+ assert mapping.partition_column == "region_key"
+ assert mapping.is_monotonic is False
+
+
+def test_resolve_returns_none_when_the_feature_flag_is_off(app: Flask) -> None:
+ table = _mapped_table()
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ return_value=False,
+ ):
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_without_a_partition_column(app: Flask) -> None:
+ table = _mapped_table()
+ table.partition_column = None
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_partition_column_no_longer_exists(
+ app: Flask,
+) -> None:
+ """A column sync can drop the physical partition column out from under
us."""
+ table = _mapped_table()
+ table.partition_column = "dt_epoch_gone"
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_mapped_column_no_longer_exists(
+ app: Flask,
+) -> None:
+ table = _mapped_table(main_dttm_col="vanished")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_on_self_mapping(app: Flask) -> None:
+ """
+ ``partition_column == effective mapped column`` mirrors a column onto
itself.
+ Save-time validation rejects it, but rows predating that validation exist.
+ """
+ table = _mapped_table(
+ partition_column="event_time",
+ mapped_column="event_time",
+ )
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
[email protected]("transform", [None, "", " "])
+def test_resolve_returns_none_without_a_transform(
+ app: Flask, transform: str | None
+) -> None:
+ table = _mapped_table(transform=transform)
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_transform_lacks_the_placeholder(
+ app: Flask,
+) -> None:
+ table = _mapped_table(transform="unix_timestamp(event_time)")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_a_non_deterministic_transform(app: Flask) ->
None:
+ """
+ The query-time gate is `is_transform_active`, the same one the Explore
+ indicator reads. Anything narrower would let a transform that reached
+ storage without passing `UpdateDatasetCommand` -- through import, or a
+ hand-written bundle -- mirror a filter with a value frozen at probe time,
+ while the editor reported the mapping as inactive and nothing on screen
+ said otherwise.
+ """
+ table = _mapped_table(transform="unix_timestamp(:value) + rand()")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_a_multi_expression_transform(app: Flask) ->
None:
+ """A select list is not an expression; the probe cannot align its
columns."""
+ table = _mapped_table(transform="lower(:value), 'x'")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_an_unparseable_transform(app: Flask) -> None:
+ table = _mapped_table(transform="unix_timestamp(:value")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_mapped_column_has_an_advanced_data_type(
+ app: Flask,
+) -> None:
+ """
+ ``translate_filter`` builds its own predicate shape from *translated*
values,
+ so the ``(operator, value)`` pair the operator matrix reasons about does
not
+ exist. Mirroring anyway would silently apply the wrong values (§4.1).
+ """
+ table = _mapped_table()
+ for column in table.columns:
+ if column.column_name == "event_time":
+ column.advanced_data_type = "port"
+
+ with app.app_context():
+ with patch.dict(app.config["ADVANCED_DATA_TYPES"], {"port":
MagicMock()}):
+ with patch(
+
"superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ side_effect=lambda flag: (
+ flag in {"PARTITION_FILTER_MAPPING",
"ENABLE_ADVANCED_DATA_TYPES"}
+ ),
+ ):
+ assert resolve_partition_mapping(table) is None
+
+
+def test_advanced_data_type_does_not_block_when_the_flag_is_off(app: Flask) ->
None:
+ """An inert ``advanced_data_type`` is not a reason to skip mirroring."""
+ table = _mapped_table()
+ for column in table.columns:
+ if column.column_name == "event_time":
+ column.advanced_data_type = "port"
+
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ side_effect=lambda flag: flag == "PARTITION_FILTER_MAPPING",
+ ):
+ assert resolve_partition_mapping(table) is not None
+
+
+# ---------------------------------------------------------------------------
+# §5 — transform inspection helpers
+# ---------------------------------------------------------------------------
+
+
[email protected](
+ "transform,expected",
+ [
+ ("unix_timestamp(:value)", True),
+ ("lower(:value)", True),
+ ("CAST(:value AS BIGINT)", True),
+ ("unix_timestamp(event_time)", False),
+ (":values", False),
+ ("", False),
+ (None, False),
+ ],
+)
+def test_contains_value_placeholder(transform: str | None, expected: bool) ->
None:
+ assert contains_value_placeholder(transform) is expected
+
+
[email protected](
+ "transform,expected",
+ [
+ ("unix_timestamp(:value)", False),
+ ("{{ current_username() }}", True),
+ ("lower({% if x %}:value{% endif %})", True),
+ ("lower(:value) -- {# comment #}", True),
+ ],
+)
+def test_contains_jinja(transform: str, expected: bool) -> None:
+ assert contains_jinja(transform) is expected
+
+
[email protected](
+ ("transform", "expected"),
+ [
+ ("lower(:value)", False),
+ ("lower(:value) /* an ordinary comment */", False),
+ ("lower(:value) -- a line comment", False),
+ (":value /*!50000 + 1 */", True),
+ (":value /*M!50000 + 1 */", True),
+ (":value /*! + 1 */", True),
+ ("", False),
+ (None, False),
+ ],
+)
+def test_contains_executable_comment(transform: str | None, expected: bool) ->
None:
+ assert contains_executable_comment(transform) is expected
+
+
[email protected](
+ "transform",
+ [
+ "unix_timestamp(:value)",
+ "lower(:value)",
+ "CAST(:value AS BIGINT)",
+ "date_format(:value, 'yyyyMMdd')",
+ ],
+)
+def test_pure_transforms_report_no_non_deterministic_functions(
+ transform: str,
+) -> None:
+ assert find_non_deterministic_functions(transform, "hive") == set()
+
+
[email protected](
+ "transform,expected_name",
+ [
+ ("date_diff(:value, now())", "NOW"),
+ ("CAST(:value AS BIGINT) + rand()", "RAND"),
+ ("CAST(:value AS DATE) - current_date", "CURRENT_DATE"),
+ ],
+)
+def test_non_deterministic_functions_are_reported(
+ transform: str, expected_name: str
+) -> None:
+ """
+ The probe runs at a different moment and in a different session from the
+ chart query, and its result is cached, so anything time- or
+ randomness-dependent freezes a snapshot of probe time into the predicate.
+ """
+ assert expected_name in find_non_deterministic_functions(transform, "hive")
+
+
+def test_niladic_unix_timestamp_is_rejected_but_the_unary_form_is_not() ->
None:
+ """
+ On Hive/Impala ``unix_timestamp()`` means "now" while ``unix_timestamp(x)``
+ -- the canonical temporal transform -- is pure. The distinction is the
whole
+ reason this check inspects arity rather than just the name.
+ """
+ assert find_non_deterministic_functions("unix_timestamp(:value)", "hive")
== set()
+ assert find_non_deterministic_functions(
+ "unix_timestamp(:value) - unix_timestamp()", "hive"
+ )
+
+
+# ---------------------------------------------------------------------------
+# §4.2 — probe evaluation
+# ---------------------------------------------------------------------------
+
+
+def _user(username: str) -> Any:
+ """A stand-in for `override_user`, which only needs `.username` here."""
+ user = MagicMock()
+ user.username = username
+ user.is_anonymous = False
+ return user
+
+
+def _database_returning(values: list[Any]) -> Database:
+ """A real ``Database`` (so the dialect is real) with only ``get_df``
stubbed."""
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=pd.DataFrame(
+ [values], columns=[f"v{i}" for i in range(len(values))]
+ )
+ )
+ return database
+
+
+def _database_returning_frame(frame: pd.DataFrame) -> Database:
+ """
+ ``_database_returning`` for a probe result whose column dtypes matter.
+
+ The flat-list sibling builds its frame from one Python list, so pandas
+ infers a single dtype and the row-versus-cell distinction cannot be seen.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=frame
+ )
+ return database
+
+
+def _probe(database: Database) -> MagicMock:
+ """
+ The stubbed ``get_df``, typed for call assertions.
+
+ ``Database.get_df`` is a real method, so mypy reads its declared signature
+ rather than the mock that replaced it and rejects ``.call_args`` and
+ friends.
+ """
+ return cast(MagicMock, database.get_df)
+
+
+def test_evaluate_transform_returns_one_value_per_input(app: Flask) -> None:
+ database = _database_returning([1767225600, 1769904000])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database,
+ None,
+ "default",
+ "unix_timestamp(:value)",
+ ["2026-01-01 00:00:00", "2026-02-01 00:00:00"],
+ )
+
+ assert result == [1767225600, 1769904000]
+ _probe(database).assert_called_once()
+
+
+def test_evaluate_transform_binds_values_rather_than_interpolating(
+ app: Flask,
+) -> None:
+ """
+ Filter values are attacker-controlled (a Gamma user picks them), so they
must
+ be bound and escaped, never pasted into the probe SQL.
+ """
+ database = _database_returning(["o''brien"])
+
+ with app.app_context():
+ evaluate_transform(database, None, None, "lower(:value)", ["O'Brien"])
+
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert "O'Brien" not in sql
+ assert "O''Brien" in sql
+
+
[email protected](
+ ("transform", "expected"),
+ [
+ ("lower(:value)", True),
+ ("CAST(:value AS BIGINT)", True),
+ ("date_format(:value, '%Y%m%d')", True),
+ # The cast shorthand: a trailing colon stops SQLAlchemy's bind scan.
+ (":value::bigint", False),
+ ("(:value::date)::text", False),
+ ],
+)
+def test_a_postgres_cast_shorthand_is_not_bindable(
+ transform: str, expected: bool
+) -> None:
+ """
+ `VALUE_PLACEHOLDER_RE` is `:value\b`, which matches inside
`:value::bigint`.
+ SQLAlchemy's own scan ends with `(?![:\\w$])`, so the following colon stops
+ it: `text(":value::bigint")` reports a parameter named `valu`, and asking
it
+ to bind `value` raises.
+
+ Pinned because the two regexes disagreeing is the whole bug: every
+ write-side gate accepted the transform and the probe then died on an
+ `ArgumentError` that `_probe` swallows.
+ """
+ assert _placeholder_is_bindable(transform) is expected
+
+
+def test_a_cast_shorthand_transform_still_probes(app: Flask) -> None:
+ """
+ The transform is valid SQL, passes every gate, and used to prune nothing:
+ `build_probe_sql` raised, `_probe` swallowed it, and the only symptom was a
+ mapping that silently never mirrored -- with the preview reporting an
engine
+ failure for SQL that was never sent.
+ """
+ database = _database_returning([20260115])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database, None, None, ":value::bigint", ["20260115"]
+ )
+
+ assert result == [20260115]
+ assert _probe(database).call_count == 1
+ assert "::bigint" in _probe(database).call_args.kwargs["sql"]
+
+
+def test_a_cast_shorthand_value_is_still_escaped_by_the_dialect(
+ app: Flask,
+) -> None:
+ """
+ The fallback substitutes where the bind path binds, so it has to earn the
+ same guarantee its sibling above pins: the value is rendered by the
+ dialect's literal processor, not pasted in.
+
+ Binding was only ever the vehicle for reaching that processor --
+ `_compile_literal` inlines the result either way -- so this renders first
+ and substitutes second. Same string, different order.
+ """
+ database = _database_returning(["x"])
+
+ with app.app_context():
+ evaluate_transform(database, None, None, ":value::text", ["O'Brien"])
+
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert "O'Brien" not in sql
+ assert "O''Brien" in sql
+
+
+def test_evaluate_transform_dedupes_repeated_values(app: Flask) -> None:
+ """Three inputs, two distinct: the probe evaluates the transform twice."""
+ database = _database_returning(["us", "us"])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database, None, None, "lower(:value)", ["US", "US", "us"]
+ )
+
+ assert result == ["us", "us", "us"]
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert sql.count("lower") == 2
+
+
+def test_evaluate_transform_fails_open_on_a_short_result_row(app: Flask) ->
None:
+ """A row narrower than the probe asked for means the results cannot be
+ aligned back to their inputs; pruning is skipped rather than guessed at."""
+ database = _database_returning(["us"])
+
+ with app.app_context():
+ assert (
+ evaluate_transform(database, None, None, "lower(:value)", ["US",
"CA"])
+ is None
+ )
+
+
+def test_evaluate_transform_fails_open_on_a_wide_result_row(app: Flask) ->
None:
+ """
+ Too many columns is the dangerous direction. A transform whose select list
+ holds two expressions returns 2N columns for N inputs, and reading the
first
+ N interleaves the expressions instead of taking one per value -- a
predicate
+ built from the wrong values rather than one that is merely short. A frame
of
+ `('us', 'x', 'ca', 'x')` read as the first two columns would have yielded
+ `('us', 'x')` and dropped every CA row.
+
+ The transform is a perfectly ordinary `lower(:value)` and the width comes
+ from the stubbed frame, which is the only way to reach the guard: a
+ transform that really does carry two expressions -- `lower(:value), 'x'` --
+ is refused by `stored_expression_error` for not being a single expression,
+ so `get_df` is never called and `None` comes back for an unrelated reason.
+ Asserting the probe *ran* is what separates the two, and it is what makes
+ relaxing the guard from `!=` to `<` fail here rather than pass: four
columns
+ for two inputs is not fewer than two.
+ """
+ database = _database_returning(["us", "x", "ca", "x"])
+
+ with app.app_context():
+ assert (
+ evaluate_transform(database, None, None, "lower(:value)", ["US",
"CA"])
+ is None
+ )
+
+ assert _probe(database).call_count == 1
+
+
+def test_the_probe_select_carries_the_engine_s_from_clause(app: Flask) -> None:
+ """
+ A `SELECT` with no `FROM` is not universal SQL: Oracle and Db2 need a
+ one-row table to select from. Without this every probe on those engines
+ raises, which `evaluate_transform` swallows -- so the only symptom is a
+ correctly configured mapping that silently never prunes.
+ """
+ assert build_probe_sql("lower(:value)", ["US"]) == "SELECT lower('US') AS
v0"
+ assert build_probe_sql("lower(:value)", ["US"], None, " FROM
DUAL").endswith(
+ " FROM DUAL"
+ )
+ assert OracleEngineSpec.select_without_from_suffix == " FROM DUAL"
+ assert BaseEngineSpec.select_without_from_suffix == ""
+
+
[email protected]("dialect", [mysql.dialect(), postgresql.dialect()])
+def test_the_probe_select_does_not_double_percent_signs(dialect: Any) -> None:
+ """
+ `TextClause` compilation runs the dialect's `post_process_text`, which
+ doubles every `%` on any paramstyle whose DBAPI later interpolates
+ parameters -- MySQL and Postgres, and every spec inheriting from them. The
+ probe has no such interpolation pass: it executes with no bound parameters
+ at all. Left doubled, `date_format(:value, '%Y%m%d')` asks the warehouse
+ for the literal string `%%Y%%m%%d`, so the probe returns that instead of a
+ date and the mirrored predicate drops every row the filter keeps.
+ """
+ sql = build_probe_sql("date_format(:value, '%Y%m%d')", ["2026-01-15"],
dialect)
+
+ assert "'%Y%m%d'" in sql
+ assert "%%" not in sql
+
+
+def test_the_probe_select_is_unchanged_on_a_dialect_that_does_not_double() ->
None:
+ """
+ Pins the no-dialect path, which is also the blind spot: the default dialect
+ does not double percent signs, so no assertion made through it can see the
+ bug the two cases above cover.
+ """
+ assert (
+ build_probe_sql("date_format(:value, '%Y%m%d')", ["2026-01-15"])
+ == "SELECT date_format('2026-01-15', '%Y%m%d') AS v0"
+ )
+
+
+def test_a_raw_probe_value_is_substituted_literally() -> None:
+ r"""
+ `RawProbeValue.sql` is rendered from a filter value, so it reaches
+ `parse_skeleton` as the *replacement* for `:value`. Read as a replacement
+ template, a backslash in it is an escape: `array('a\nb')` -- which is how a
+ dialect renders the two characters `\` and `n` -- would become a real
+ newline, so the engine is asked about a different array than the predicate
+ compares and the mirror can exclude rows the filter keeps.
+ """
+ sql = build_probe_sql("toYYYYMM(:value)",
[RawProbeValue(r"array('a\nb')")])
+
+ assert sql == r"SELECT toYYYYMM(array('a\nb')) AS v0"
+
+
+def test_a_raw_probe_value_does_not_re_inject_the_placeholder() -> None:
+ r"""
+ `\g<0>` in a replacement template expands to the whole match, which would
+ put `:value` back into the probe -- an unbound parameter the engine
refuses.
+ """
+ sql = build_probe_sql("toYYYYMM(:value)", [RawProbeValue(r"'a\g<0>b'")])
+
+ assert sql == r"SELECT toYYYYMM('a\g<0>b') AS v0"
+ assert ":value" not in sql
+
+
+def test_a_raw_probe_value_with_a_group_reference_does_not_raise() -> None:
+ r"""
+ `\1` is an invalid group reference for this pattern, so as a replacement
+ template it raised `re.PatternError` rather than probing anything.
+ """
+ sql = build_probe_sql("toYYYYMM(:value)", [RawProbeValue(r"'a\1b'")])
+
+ assert sql == r"SELECT toYYYYMM('a\1b') AS v0"
+
+
+def test_a_probed_literal_keeps_its_percent_sign(app: Flask) -> None:
+ """
+ The preview panel renders probed values through the same dialect, so a
+ value that merely contains a `%` would be shown doubled.
+ """
+ database = Database(database_name="pct_db",
sqlalchemy_uri="postgresql://h/d")
+
+ with app.app_context():
+ assert _render_literal(database, "100%") == "'100%'"
+
+
+def test_the_probe_reads_each_cell_without_row_wide_dtype_coercion(
+ app: Flask,
+) -> None:
+ """
+ A row is a Series, so it carries one dtype across every probe column: a row
+ mixing an exact int64 with a float unifies to float64 and an integer
+ partition key above 2^53 comes back rounded. The mirror then asks for a key
+ the warehouse does not hold, and the row the filter matched is dropped.
+ """
+ database = _database_returning_frame(
+ pd.DataFrame({"v0": [9007199254740993], "v1": [1.5]})
+ )
+
+ with app.app_context():
+ probed = evaluate_transform(
+ database, None, None, "CAST(:value AS NUMERIC)",
["9007199254740993", "1.5"]
+ )
+
+ assert probed == [9007199254740993, 1.5]
+
+
+def test_evaluate_transform_pins_catalog_and_schema(app: Flask) -> None:
+ """
+ The probe runs in a different session from the chart query; pinning the
+ catalog and schema keeps session settings as close as the pool allows.
+ """
+ database = _database_returning([1])
+
+ with app.app_context():
+ evaluate_transform(database, "prod", "analytics", "lower(:value)",
["x"])
+
+ kwargs = _probe(database).call_args.kwargs
+ assert kwargs["catalog"] == "prod"
+ assert kwargs["schema"] == "analytics"
+
+
+def test_evaluate_transform_fails_open_when_the_probe_raises(app: Flask) ->
None:
+ """A wedged engine must not break the chart — it just stops pruning."""
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ side_effect=RuntimeError("connection reset")
+ )
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)",
["x"]) is None
+
+
+def test_evaluate_transform_fails_open_on_an_empty_result(app: Flask) -> None:
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=pd.DataFrame()
+ )
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)",
["x"]) is None
+
+
+def test_evaluate_transform_returns_none_for_no_values(app: Flask) -> None:
+ database = _database_returning([])
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)", [])
is None
+
+ _probe(database).assert_not_called()
+
+
[email protected](
+ "transform",
+ [
+ "(SELECT password FROM ab_user LIMIT 1) || :value",
+ "lower(:value) UNION ALL SELECT password FROM ab_user",
+ "lower(:value); DROP TABLE ab_user",
+ "password || :value FROM ab_user",
+ "lower(:value) FROM ab_user WHERE 1 = 1",
+ ],
+ ids=[
+ "subquery",
+ "set-operation",
+ "multi-statement",
+ "bare-from",
+ "from-and-where",
+ ],
+)
+def test_a_transform_that_is_not_a_storable_expression_never_runs(
+ app: Flask, transform: str
+) -> None:
+ """
+ The probe is where a transform stops being text and becomes SQL.
+
+ `build_probe_sql` binds `:value` and splices the rest in verbatim, so a
+ transform that was never a single harmless expression is arbitrary SQL
+ against the dataset's database -- read back out of the emitted predicate
+ by anyone who can open "View query". The write-side gates refuse these,
+ but they can only speak for rows written after they existed, so the last
+ word belongs here.
+ """
+ database = _database_returning(["irrelevant"])
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, transform, ["us"]) is
None
+
+ _probe(database).assert_not_called()
+
+
[email protected](
+ "transform, unfinished",
+ [
+ ("unix_timestamp(:value", True),
+ ("unix_timestamp(:value))", True),
+ ("unix_timestamp(:value) +", True),
+ ("unix_timestamp(:value)", False),
+ ("unix_timestamp(:value); DROP TABLE ab_user", False),
+ ("unix_timestamp(:value) UNION ALL SELECT password FROM ab_user",
False),
+ ("password || :value FROM ab_user", False),
+ ],
+)
+def test_is_unfinished_separates_not_sql_yet_from_the_wrong_sql(
+ transform: str, unfinished: bool
+) -> None:
+ """
+ The two deserve opposite treatment -- a half-typed transform is stored and
+ parked inactive, the wrong SQL is refused -- and neither `is_parseable` nor
+ parsing as a single statement tells them apart, because the multi-statement
+ form fails both exactly as unfinished text does. Parsing as a script does.
+ """
+ assert is_unfinished(transform, "hive") is unfinished
+
+
+def test_a_transform_smuggling_its_own_from_clause_is_refused(app: Flask) ->
None:
+ """
+ The hole this gate was added for.
+
+ `is_parseable` counts only the projection, so `password || :value FROM
+ ab_user` reads as a single select expression and passed every write-side
+ gate -- including with `ALLOW_ADHOC_SUBQUERY` off, since a top-level FROM
+ is not a sub-query. `build_probe_sql` then emitted `SELECT password || 'us'
+ AS v0 FROM ab_user`, where its own alias is read as a table alias on the
+ FROM clause the transform brought with it: valid SQL, one row, and the
+ secret handed back through the predicate the preview panel renders.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ transform = "password || :value FROM ab_user"
+
+ assert is_parseable(transform, "sqlite") is True
+ assert is_bare_expression(transform, "sqlite") is False
+
+ with app.app_context():
+ reason = stored_expression_error(database, None, None, transform)
+ assert reason is not None
+
+ blocking = _blocking(validate_transform(transform, "sqlite"))
+ assert len(blocking) == 1
+ assert "single SQL expression" in str(blocking[0].message)
+ assert is_transform_active(transform, "sqlite") is False
+
+
+def test_a_subquery_transform_is_refused_even_with_adhoc_subquery_enabled(
+ app: Flask,
+) -> None:
+ """
+ `validate_adhoc_subquery` admits a sub-query under `ALLOW_ADHOC_SUBQUERY`
+ and returns RLS-rewritten SQL -- which this path has no way to propagate,
+ because `build_probe_sql` splices the original transform text. So the probe
+ ran a sub-query with no row-level-security predicates at all and reported
+ the hidden row in `emitted_predicate`. A transform is a scalar function of
+ `:value`, so refusing the sub-query outright is both correct and what makes
+ the missing RLS rewrite moot.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ transform = "(SELECT secret FROM vault LIMIT 1) || :value"
+
+ with app.app_context():
+ app.config["DEFAULT_FEATURE_FLAGS"]["ALLOW_ADHOC_SUBQUERY"] = True
+ try:
+ reason = stored_expression_error(database, None, None, transform)
+ finally:
+ del app.config["DEFAULT_FEATURE_FLAGS"]["ALLOW_ADHOC_SUBQUERY"]
+
+ assert reason is not None
+ assert "sub-query" in reason
+
+
[email protected](
+ "comment",
+ ["/*!50000 + (SELECT secret FROM vault LIMIT 1) */", "/*M!50000 + 1 */"],
+)
+def test_an_executable_comment_is_refused_by_every_gate(
+ app: Flask, comment: str
+) -> None:
+ """
+ The one construct that made the shape gates describe a statement the engine
+ would not run.
+
+ sqlglot reads a MySQL executable comment as comment data: the pinned parse
+ below reports no sub-query and no clause, so `stored_expression_error` and
+ `probe_sql_is_evaluable` both passed it, while MySQL and MariaDB execute
+ what is inside. An owner with dataset write and no SQL Lab could put one in
+ a transform, call the preview and read a table back out of
+ `emitted_predicate`, with no RLS on the hidden query.
+
+ Refused on every engine, not only the two that honour it -- hence SQLite
+ here. There is no legitimate transform carrying one, and a per-engine gate
+ would store on Postgres what detonates after the dataset is re-pointed.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ transform = f"(:value {comment})"
+
+ # The premise: no parse-tree gate can see it.
+ statement = SQLStatement(f"SELECT (1 {comment}) AS v0", "mysql")
+ assert statement.has_subquery() is False
+ assert statement.get_clause_names() == set()
+
+ with app.app_context():
+ assert stored_expression_error(database, None, None, transform) is not
None
+
+ blocking = [
+ issue for issue in validate_transform(transform, "mysql") if
issue.blocking
+ ]
+ assert len(blocking) == 1
+ assert "executable comment" in str(blocking[0].message)
+
+
+def test_an_executable_comment_stops_the_probe_before_the_engine(app: Flask)
-> None:
+ """
+ The gate above is the save path. This is the row an earlier release stored:
+ `_probe` consults `stored_expression_error` on every call, so a transform
+ already in the metadata database never reaches `get_df`.
+ """
+ database = _database_returning(["x"])
+ transform = "(:value /*!50000 + (SELECT secret FROM vault LIMIT 1) */)"
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, transform, ["US"]) is
None
+
+ assert _probe(database).call_count == 0
+
+
+def test_a_plain_comment_is_still_allowed_through_the_shape_gates(
+ app: Flask,
+) -> None:
+ """
+ Only the executable form is refused. An ordinary comment is inert on every
+ engine, so rejecting it would cost an owner a transform that works -- and
+ the placeholder gate already covers the case that matters, a `:value` the
+ comment swallows.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+
+ with app.app_context():
+ assert (
+ stored_expression_error(database, None, None, "lower(:value) /* ok
*/")
+ is None
+ )
+
+
+def test_a_transform_calling_a_disallowed_function_is_refused(app: Flask) ->
None:
+ """
+ The probe executes the transform, so the operator's function denylist has
+ to apply to it for the same reason it applies in SQL Lab.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ transform = "version() || :value"
+
+ with app.app_context():
+ original = app.config["DISALLOWED_SQL_FUNCTIONS"]
+ try:
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = {"sqlite": {"version"}}
+ reason = stored_expression_error(database, None, None, transform)
+ assert reason is not None
+ assert "version" in reason
+
+ # Named, rather than echoing the operator's whole denylist back.
+ assert "current_user" not in reason
+
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = {}
+ assert stored_expression_error(database, None, None, transform) is
None
+ finally:
+ # Restored: `app.config` is shared, and leaving it emptied silently
+ # disabled this gate for every test that ran afterwards.
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = original
+
+
[email protected](
+ "function",
+ [
+ "database_to_xml",
+ "database_to_xmlschema",
+ "database_to_xml_and_xmlschema",
+ "schema_to_xml",
+ "schema_to_xmlschema",
+ "schema_to_xml_and_xmlschema",
+ "table_to_xml",
+ "table_to_xmlschema",
+ "table_to_xml_and_xmlschema",
+ "query_to_xml",
+ "query_to_xmlschema",
+ "query_to_xml_and_xmlschema",
+ ],
+)
+def test_the_whole_postgres_xml_family_is_denied_by_default(function: str) ->
None:
+ """
+ These read tables with no FROM clause and no sub-query, so every gate that
+ reasons about table references is blind to them -- the denylist is the only
+ thing that can refuse them, and it has to name each shape. `schema_to_xml`
+ and three siblings were missing, so a bare scalar call cleared
+ `stored_expression_error` entirely and returned connection-user-readable
+ schema contents through preview.
+
+ Asserted against the shipped default rather than through the gate, because
+ the gap was an incomplete enumeration rather than broken wiring -- and
+ because `app.config` is ambient here, which the behavioural test below
+ handles by setting the denylist itself.
+ """
+ assert function in DISALLOWED_SQL_FUNCTIONS["postgresql"]
+
+
+def test_a_denied_read_capable_function_is_refused_despite_having_no_from(
+ app: Flask,
+) -> None:
+ """
+ The wiring half: a bare scalar call passes the length, single-statement,
+ bare-expression, sub-query and placeholder gates, so the denylist is what
+ has to stop it.
+ """
+ database = Database(database_name="probe_db",
sqlalchemy_uri="postgresql://u@h/d")
+ transform = "schema_to_xml('public', true, false, '') || :value"
+
+ with app.app_context():
+ original = app.config["DISALLOWED_SQL_FUNCTIONS"]
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = {"postgresql":
{"schema_to_xml"}}
+ try:
+ reason = stored_expression_error(database, None, None, transform)
+ finally:
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = original
+
+ assert reason is not None
+ assert "schema_to_xml" in reason
+
+
[email protected]("function", ["ts_rewrite", "ts_stat"])
+def test_the_postgres_full_text_query_functions_are_denied_by_default(
+ function: str,
+) -> None:
+ """
+ `ts_stat('select ...')` and `ts_rewrite(q, 'select ...')` run their *text*
+ argument as a query. The query is a string, so these have the same blind
+ spot the XML family above does: no FROM clause and no sub-query for a
+ table-reference gate to reason about, which leaves the denylist as the only
+ thing that can refuse them.
+ """
+ assert function in DISALLOWED_SQL_FUNCTIONS["postgresql"]
+
+
+def test_a_denied_query_running_text_function_is_refused(app: Flask) -> None:
+ """
+ The wiring half, and the only thing that proves the *name* matching works
+ for these two: neither is a function sqlglot models, so each reaches
+ `get_disallowed_functions` through its ANONYMOUS branch rather than as a
+ typed node. A denylist entry that only matched modelled functions would
+ pass the enumeration test above and still let these through.
+ """
+ database = Database(database_name="probe_db",
sqlalchemy_uri="postgresql://u@h/d")
+ transform = "ts_stat('select 1') || :value"
+
+ with app.app_context():
+ original = app.config["DISALLOWED_SQL_FUNCTIONS"]
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = {"postgresql": {"ts_stat"}}
+ try:
+ reason = stored_expression_error(database, None, None, transform)
+ finally:
+ app.config["DISALLOWED_SQL_FUNCTIONS"] = original
+
+ assert reason is not None
+ assert "ts_stat" in reason
+
+
+def test_the_function_denylist_is_keyed_on_the_engine_spec_name(app: Flask) ->
None:
+ """
+ Every other denylist gate in the codebase keys on the engine spec's own
+ name, which is also the name the config documentation tells an operator to
+ write, so this one does too -- a spec covering several SQLAlchemy backends
+ through `engine_aliases` reports the one canonical name under which its
+ entry exists.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
Review Comment:
Fixed. The test uses an aliased engine — backend `postgres`, spec
`postgresql`, asserted explicitly — so it now distinguishes the two and fails
if the gate keys on `database.backend`.
##########
tests/unit_tests/connectors/sqla/partition_mapping_test.py:
##########
@@ -0,0 +1,1973 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for ``superset.connectors.sqla.partition_mapping``."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from decimal import Decimal
+from importlib import import_module
+from typing import Any, cast
+from unittest.mock import MagicMock, patch, PropertyMock
+
+import pandas as pd
+import pytest
+from dateutil.relativedelta import relativedelta
+from flask import Flask
+from sqlalchemy.dialects import mysql, postgresql
+
+from superset.config import DISALLOWED_SQL_FUNCTIONS
+from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.connectors.sqla.partition_mapping import (
+ _probe_cache_key,
+ _render_literal,
+ build_probe_sql,
+ contains_jinja,
+ contains_value_placeholder,
+ equality_mirrors_safely,
+ evaluate_transform,
+ find_non_deterministic_functions,
+ grain_bucket_width,
+ GRAIN_BUCKET_WIDTHS,
+ is_bare_expression,
+ is_parseable,
+ is_transform_active,
+ is_unfinished,
+ known_mirror_verdict,
+ MappingValidationIssue,
+ mirror_operator,
+ MIRRORABLE_ALWAYS,
+ MIRRORABLE_IF_MONOTONIC,
+ mirrorable_operators,
+ parse_error_detail,
+ placeholder_is_executable,
+ probe_sql_is_evaluable,
+ probed_value_type_error,
+ RawProbeValue,
+ resolve_partition_mapping,
+ stored_expression_error,
+ validate_partition_mapping,
+ validate_transform,
+)
+from superset.constants import TimeGrain
+from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.oracle import OracleEngineSpec
+from superset.models.core import Database
+from superset.utils.core import FilterOperator
+
+
[email protected](autouse=True)
+def real_probe_cache(app: Flask) -> Any:
+ """
+ The test app runs a null cache, which would make every cache assertion here
+ vacuously pass. Swap in a real in-memory one for the duration.
+ """
+ from flask_caching import Cache
+
+ from superset.extensions import cache_manager
+
+ cache = Cache(config={"CACHE_TYPE": "SimpleCache",
"CACHE_DEFAULT_TIMEOUT": 300})
+ cache.init_app(app)
+ original = cache_manager._cache # noqa: SLF001
+ cache_manager._cache = cache # noqa: SLF001
+ yield
+ cache_manager._cache = original # noqa: SLF001
+
+
[email protected](autouse=True)
+def enable_partition_filter_mapping(app: Flask) -> Any:
+ """The feature ships off; every test here exercises it on."""
+ original =
app.config["DEFAULT_FEATURE_FLAGS"].get("PARTITION_FILTER_MAPPING")
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = True
+ yield
+ if original is None:
+ del app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"]
+ else:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] =
original
+
+
+def _table(**kwargs: Any) -> SqlaTable:
+ database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
+ defaults: dict[str, Any] = {
+ "table_name": "web_events",
+ "database": database,
+ "main_dttm_col": "event_time",
+ "columns": [
+ TableColumn(column_name="event_time", is_dttm=True,
type="TIMESTAMP"),
+ TableColumn(
+ column_name="dt_epoch",
+ type="BIGINT",
+ partition_value_transform=None,
+ ),
+ TableColumn(column_name="country", type="VARCHAR"),
+ TableColumn(column_name="region_key", type="VARCHAR"),
+ ],
+ }
+ defaults.update(kwargs)
+ return SqlaTable(**defaults)
+
+
+def _mapped_table(
+ *,
+ partition_column: str = "dt_epoch",
+ mapped_column: str = "event_time",
+ transform: str | None = "unix_timestamp(:value)",
+ monotonic: bool = True,
+ partition_mapped_column: str | None = None,
+ main_dttm_col: str | None = "event_time",
+) -> SqlaTable:
+ table = _table(main_dttm_col=main_dttm_col)
+ table.partition_column = partition_column
+ table.partition_mapped_column = partition_mapped_column
+ for column in table.columns:
+ if column.column_name == mapped_column:
+ column.partition_value_transform = transform
+ column.partition_transform_is_monotonic = monotonic
+ return table
+
+
+# ---------------------------------------------------------------------------
+# §2 — operator safety matrix
+# ---------------------------------------------------------------------------
+
+
+def test_equality_and_in_are_always_mirrorable() -> None:
+ """``=`` and ``IN`` are safe for any function ``T``."""
+ assert MIRRORABLE_ALWAYS == {FilterOperator.EQUALS, FilterOperator.IN}
+
+
+def test_range_operators_require_a_monotonic_transform() -> None:
+ assert MIRRORABLE_IF_MONOTONIC == {
+ FilterOperator.GREATER_THAN,
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ FilterOperator.LESS_THAN,
+ FilterOperator.LESS_THAN_OR_EQUALS,
+ FilterOperator.TEMPORAL_RANGE,
+ }
+
+
+def test_mirrorable_operators_excludes_ranges_when_not_monotonic() -> None:
+ assert mirrorable_operators(is_monotonic=False) == MIRRORABLE_ALWAYS
+
+
+def test_mirrorable_operators_includes_ranges_when_monotonic() -> None:
+ assert mirrorable_operators(is_monotonic=True) == (
+ MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC
+ )
+
+
[email protected](
+ "operator",
+ [
+ FilterOperator.NOT_EQUALS,
+ FilterOperator.NOT_IN,
+ FilterOperator.LIKE,
+ FilterOperator.ILIKE,
+ FilterOperator.NOT_LIKE,
+ FilterOperator.NOT_ILIKE,
+ FilterOperator.IS_NULL,
+ FilterOperator.IS_NOT_NULL,
+ FilterOperator.IS_TRUE,
+ FilterOperator.IS_FALSE,
+ ],
+)
+def test_negations_and_pattern_matches_are_never_mirrorable(
+ operator: FilterOperator,
+) -> None:
+ """
+ ``T`` is not injective, so ``col != v`` does **not** imply ``T(col) !=
T(v)``:
+ mirroring it would drop rows the original filter keeps.
+ """
+ assert operator not in mirrorable_operators(is_monotonic=True)
+
+
+def test_equality_is_dropped_when_the_engine_does_not_compare_byte_exactly()
-> None:
+ """
+ "Safe for any function ``T``" reasons about value equality: ``T`` is a
+ function, so ``col = v`` gives ``T(col) = T(v)``. The engine reasons about
+ *SQL* equality, and the two part company under a case-insensitive
+ collation -- a stored ``'us'`` satisfies a filter for ``'US'`` while the
+ mirror ``hex('US')`` excludes the row, and the chart loses it silently.
+ """
+ assert mirrorable_operators(is_monotonic=False, equality_is_safe=False) ==
set()
+
+
+def test_ranges_survive_an_unsafe_equality() -> None:
+ """
+ Range mirroring already rests on the owner declaring ``T``
+ order-preserving with respect to the column's own order -- an assertion
+ about the very comparison semantics this gate checks for. Equality has no
+ such declaration behind it, which is why only equality is withdrawn.
+ """
+ assert (
+ mirrorable_operators(is_monotonic=True, equality_is_safe=False)
+ == MIRRORABLE_IF_MONOTONIC
+ )
+
+
[email protected](
+ "operator,expected",
+ [
+ (FilterOperator.GREATER_THAN, FilterOperator.GREATER_THAN_OR_EQUALS),
+ (FilterOperator.LESS_THAN, FilterOperator.LESS_THAN_OR_EQUALS),
+ (
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ ),
+ (FilterOperator.LESS_THAN_OR_EQUALS,
FilterOperator.LESS_THAN_OR_EQUALS),
+ (FilterOperator.EQUALS, FilterOperator.EQUALS),
+ (FilterOperator.IN, FilterOperator.IN),
+ ],
+)
+def test_a_strict_bound_mirrors_non_strictly(
+ operator: FilterOperator, expected: FilterOperator
+) -> None:
+ """
+ Order-preserving means non-decreasing, not injective, so ``col < v``
+ implies only ``T(col) <= T(v)``. A day-key transform maps a whole day onto
+ one value and a strict mirror drops the boundary bucket. The non-strict
+ cases are already safe and pass through, which is what makes applying this
+ twice harmless.
+ """
+ assert mirror_operator(operator) == expected
+
+
+def test_relaxing_a_bound_stays_inside_the_mirrorable_set() -> None:
+ """
+ The substitution is about strictness, not about which filters mirror. The
+ Explore indicator reads `mirrorable_operators`, so leaving that set alone
is
+ what keeps the glyph from promising pruning the SQL does not do -- and it
is
+ also what guarantees `handle_comparison_filter` is never handed an operator
+ it has no case for.
+ """
+ assert {
+ mirror_operator(operator) for operator in MIRRORABLE_IF_MONOTONIC
+ } <= MIRRORABLE_IF_MONOTONIC
+
+
+def _column(type_: str) -> TableColumn:
+ database = Database(database_name="t", sqlalchemy_uri="sqlite://")
+ table = SqlaTable(table_name="t", database=database)
+ return TableColumn(column_name="c", type=type_, table=table)
+
+
[email protected](
+ "type_, binary, expected",
+ [
+ ("VARCHAR", True, True),
+ ("VARCHAR", False, False),
+ # Numeric and temporal comparison is exact everywhere.
+ ("BIGINT", False, True),
+ ("TIMESTAMP", False, True),
+ ("DOUBLE", False, True),
+ ],
+)
+def test_equality_mirrors_safely_only_guards_string_columns(
+ app: Flask, type_: str, binary: bool, expected: bool
+) -> None:
+ class _Spec(BaseEngineSpec):
+ binary_string_comparison = binary
+
+ with app.app_context():
+ assert equality_mirrors_safely(_column(type_), _Spec) is expected
+
+
+def test_equality_mirrors_safely_fails_closed_on_an_unresolvable_type(
+ app: Flask,
+) -> None:
+ """
+ The gate exists to stop a silent wrong answer, so a column whose type it
+ cannot read is treated as the risky case rather than waved through.
+ """
+ column = MagicMock()
+ type(column).type_generic = PropertyMock(side_effect=ValueError("no type"))
+
+ with app.app_context():
+ assert equality_mirrors_safely(column, BaseEngineSpec) is False
+
+
+def test_no_mapped_column_leaves_the_matrix_alone(app: Flask) -> None:
+ """Nothing to reason about, and the resolver bails out on it anyway."""
+ with app.app_context():
+ assert equality_mirrors_safely(None, BaseEngineSpec) is True
+
+
[email protected](
+ "module, spec_name",
+ [
+ ("presto", "PrestoEngineSpec"),
+ ("hive", "HiveEngineSpec"),
+ ("trino", "TrinoEngineSpec"),
+ ("impala", "ImpalaEngineSpec"),
+ ("postgres", "PostgresEngineSpec"),
+ ("bigquery", "BigQueryEngineSpec"),
+ ("sqlite", "SqliteEngineSpec"),
+ ],
+)
+def test_the_byte_exact_engines_opt_in(module: str, spec_name: str) -> None:
+ """
+ The engines this feature targets compare text byte-exactly, so the gate
+ must not cost them the second canonical mapping (`lower(:value)` onto a
+ lowercased key).
+ """
+ spec = getattr(import_module(f"superset.db_engine_specs.{module}"),
spec_name)
+ assert spec.binary_string_comparison is True
+
+
[email protected]("module, spec_name", [("mysql", "MySQLEngineSpec")])
+def test_a_case_insensitive_engine_does_not_opt_in(module: str, spec_name:
str) -> None:
+ """
+ MySQL's default collation is `utf8mb4_0900_ai_ci`, which is exactly the
+ case the reviewer raised.
+ """
+ spec = getattr(import_module(f"superset.db_engine_specs.{module}"),
spec_name)
+ assert spec.binary_string_comparison is False
+
+
+def test_the_base_spec_does_not_opt_in() -> None:
+ """A spec that has not said so does not mirror string equality."""
+ assert BaseEngineSpec.binary_string_comparison is False
+
+
+# ---------------------------------------------------------------------------
+# Time grain bucket widths
+# ---------------------------------------------------------------------------
+
+
[email protected]("grain", list(TimeGrain))
+def test_every_builtin_grain_has_a_known_bucket_width(
+ app: Flask, grain: TimeGrain
+) -> None:
+ """
+ The widening is only sound for a grain whose bucket width we know, so a
+ grain added upstream without a width here would silently stop mirroring.
+ This is the test that notices.
+ """
+ with app.app_context():
+ assert grain_bucket_width(grain.value, "sqlite") is not None
+
+
+def test_a_calendar_grain_widens_by_a_calendar_month(app: Flask) -> None:
+ """A month is not thirty days, and the month it lands in decides its
width."""
+ with app.app_context():
+ width = grain_bucket_width(TimeGrain.MONTH.value, "sqlite")
+
+ assert width is not None
+ assert datetime(2026, 1, 31) + width == datetime(2026, 2, 28)
+ assert datetime(2026, 2, 1) + width == datetime(2026, 3, 1)
+
+
+def test_a_fractional_grain_has_a_width(app: Flask) -> None:
+ """
+ ``PT0.5H`` and ``P0.25Y`` are why the widths are written out rather than
+ parsed -- and the anchored week grains are intervals ``parse_duration``
+ rejects outright.
+ """
+ with app.app_context():
+ assert grain_bucket_width(TimeGrain.HALF_HOUR.value, "sqlite") ==
relativedelta(
+ minutes=30
+ )
+ assert grain_bucket_width(
+ TimeGrain.WEEK_ENDING_SATURDAY.value, "sqlite"
+ ) == relativedelta(days=7)
+
+
+def test_an_unknown_grain_has_no_known_width(app: Flask) -> None:
+ with app.app_context():
+ assert grain_bucket_width("P13X", "sqlite") is None
+ assert grain_bucket_width(None, "sqlite") is None
+
+
+def test_an_addon_grain_has_no_known_width(app: Flask) -> None:
+ """An operator-declared grain carries whatever duration string they
typed."""
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDONS"] = {"PT2S": "2 second"}
+ try:
+ assert grain_bucket_width("PT2S", "sqlite") is None
+ finally:
+ app.config["TIME_GRAIN_ADDONS"] = {}
+
+
+def test_a_builtin_grain_an_operator_redefined_has_no_known_width(
+ app: Flask,
+) -> None:
+ """
+ ``TIME_GRAIN_ADDON_EXPRESSIONS`` replaces a built-in grain's SQL for one
+ engine, so its duration no longer bounds the displacement *there* -- but it
+ still does on every other engine.
+ """
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {
+ "sqlite": {"P1D": "DATETIME({col}, 'start of year')"}
+ }
+ try:
+ assert grain_bucket_width(TimeGrain.DAY.value, "sqlite") is None
+ assert grain_bucket_width(TimeGrain.DAY.value, "hive") is not None
+ finally:
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {}
+
+
+def test_the_width_table_is_keyed_on_grain_durations(app: Flask) -> None:
+ """Keys are the ISO duration a filter carries, not the enum member name."""
+ assert GRAIN_BUCKET_WIDTHS[TimeGrain.DAY] == relativedelta(days=1)
+ assert set(GRAIN_BUCKET_WIDTHS) == {grain.value for grain in TimeGrain}
+
+
+# ---------------------------------------------------------------------------
+# §4.1 — resolving the effective mapping, and the defensive bail-outs
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_returns_the_mapping_when_everything_lines_up(app: Flask) ->
None:
+ with app.app_context():
+ mapping = resolve_partition_mapping(_mapped_table())
+
+ assert mapping is not None
+ assert mapping.partition_column == "dt_epoch"
+ assert mapping.mapped_column == "event_time"
+ assert mapping.value_transform == "unix_timestamp(:value)"
+ assert mapping.is_monotonic is True
+
+
+def test_effective_mapped_column_follows_main_dttm_col(app: Flask) -> None:
+ """No explicit override: the mapping follows the default datetime
column."""
+ with app.app_context():
+ mapping = resolve_partition_mapping(_mapped_table())
+
+ assert mapping is not None
+ assert mapping.mapped_column == "event_time"
+
+
+def test_explicit_override_wins_over_main_dttm_col(app: Flask) -> None:
+ table = _mapped_table(
+ mapped_column="country",
+ transform="lower(:value)",
+ monotonic=False,
+ partition_mapped_column="country",
+ )
+ table.partition_column = "region_key"
+
+ with app.app_context():
+ mapping = resolve_partition_mapping(table)
+
+ assert mapping is not None
+ assert mapping.mapped_column == "country"
+ assert mapping.partition_column == "region_key"
+ assert mapping.is_monotonic is False
+
+
+def test_resolve_returns_none_when_the_feature_flag_is_off(app: Flask) -> None:
+ table = _mapped_table()
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ return_value=False,
+ ):
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_without_a_partition_column(app: Flask) -> None:
+ table = _mapped_table()
+ table.partition_column = None
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_partition_column_no_longer_exists(
+ app: Flask,
+) -> None:
+ """A column sync can drop the physical partition column out from under
us."""
+ table = _mapped_table()
+ table.partition_column = "dt_epoch_gone"
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_mapped_column_no_longer_exists(
+ app: Flask,
+) -> None:
+ table = _mapped_table(main_dttm_col="vanished")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_on_self_mapping(app: Flask) -> None:
+ """
+ ``partition_column == effective mapped column`` mirrors a column onto
itself.
+ Save-time validation rejects it, but rows predating that validation exist.
+ """
+ table = _mapped_table(
+ partition_column="event_time",
+ mapped_column="event_time",
+ )
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
[email protected]("transform", [None, "", " "])
+def test_resolve_returns_none_without_a_transform(
+ app: Flask, transform: str | None
+) -> None:
+ table = _mapped_table(transform=transform)
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_transform_lacks_the_placeholder(
+ app: Flask,
+) -> None:
+ table = _mapped_table(transform="unix_timestamp(event_time)")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_a_non_deterministic_transform(app: Flask) ->
None:
+ """
+ The query-time gate is `is_transform_active`, the same one the Explore
+ indicator reads. Anything narrower would let a transform that reached
+ storage without passing `UpdateDatasetCommand` -- through import, or a
+ hand-written bundle -- mirror a filter with a value frozen at probe time,
+ while the editor reported the mapping as inactive and nothing on screen
+ said otherwise.
+ """
+ table = _mapped_table(transform="unix_timestamp(:value) + rand()")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_a_multi_expression_transform(app: Flask) ->
None:
+ """A select list is not an expression; the probe cannot align its
columns."""
+ table = _mapped_table(transform="lower(:value), 'x'")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_an_unparseable_transform(app: Flask) -> None:
+ table = _mapped_table(transform="unix_timestamp(:value")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_mapped_column_has_an_advanced_data_type(
+ app: Flask,
+) -> None:
+ """
+ ``translate_filter`` builds its own predicate shape from *translated*
values,
+ so the ``(operator, value)`` pair the operator matrix reasons about does
not
+ exist. Mirroring anyway would silently apply the wrong values (§4.1).
+ """
+ table = _mapped_table()
+ for column in table.columns:
+ if column.column_name == "event_time":
+ column.advanced_data_type = "port"
+
+ with app.app_context():
+ with patch.dict(app.config["ADVANCED_DATA_TYPES"], {"port":
MagicMock()}):
+ with patch(
+
"superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ side_effect=lambda flag: (
+ flag in {"PARTITION_FILTER_MAPPING",
"ENABLE_ADVANCED_DATA_TYPES"}
+ ),
+ ):
+ assert resolve_partition_mapping(table) is None
+
+
+def test_advanced_data_type_does_not_block_when_the_flag_is_off(app: Flask) ->
None:
+ """An inert ``advanced_data_type`` is not a reason to skip mirroring."""
+ table = _mapped_table()
+ for column in table.columns:
+ if column.column_name == "event_time":
+ column.advanced_data_type = "port"
+
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ side_effect=lambda flag: flag == "PARTITION_FILTER_MAPPING",
+ ):
+ assert resolve_partition_mapping(table) is not None
+
+
+# ---------------------------------------------------------------------------
+# §5 — transform inspection helpers
+# ---------------------------------------------------------------------------
+
+
[email protected](
+ "transform,expected",
+ [
+ ("unix_timestamp(:value)", True),
+ ("lower(:value)", True),
+ ("CAST(:value AS BIGINT)", True),
+ ("unix_timestamp(event_time)", False),
+ (":values", False),
+ ("", False),
+ (None, False),
+ ],
+)
+def test_contains_value_placeholder(transform: str | None, expected: bool) ->
None:
+ assert contains_value_placeholder(transform) is expected
+
+
[email protected](
+ "transform,expected",
+ [
+ ("unix_timestamp(:value)", False),
+ ("{{ current_username() }}", True),
+ ("lower({% if x %}:value{% endif %})", True),
+ ("lower(:value) -- {# comment #}", True),
+ ],
+)
+def test_contains_jinja(transform: str, expected: bool) -> None:
+ assert contains_jinja(transform) is expected
+
+
[email protected](
+ "transform",
+ [
+ "unix_timestamp(:value)",
+ "lower(:value)",
+ "CAST(:value AS BIGINT)",
+ "date_format(:value, 'yyyyMMdd')",
+ ],
+)
+def test_pure_transforms_report_no_non_deterministic_functions(
+ transform: str,
+) -> None:
+ assert find_non_deterministic_functions(transform, "hive") == set()
+
+
[email protected](
+ "transform,expected_name",
+ [
+ ("date_diff(:value, now())", "NOW"),
+ ("CAST(:value AS BIGINT) + rand()", "RAND"),
+ ("CAST(:value AS DATE) - current_date", "CURRENT_DATE"),
+ ],
+)
+def test_non_deterministic_functions_are_reported(
+ transform: str, expected_name: str
+) -> None:
+ """
+ The probe runs at a different moment and in a different session from the
+ chart query, and its result is cached, so anything time- or
+ randomness-dependent freezes a snapshot of probe time into the predicate.
+ """
+ assert expected_name in find_non_deterministic_functions(transform, "hive")
+
+
+def test_niladic_unix_timestamp_is_rejected_but_the_unary_form_is_not() ->
None:
+ """
+ On Hive/Impala ``unix_timestamp()`` means "now" while ``unix_timestamp(x)``
+ -- the canonical temporal transform -- is pure. The distinction is the
whole
+ reason this check inspects arity rather than just the name.
+ """
+ assert find_non_deterministic_functions("unix_timestamp(:value)", "hive")
== set()
+ assert find_non_deterministic_functions(
+ "unix_timestamp(:value) - unix_timestamp()", "hive"
+ )
+
+
+# ---------------------------------------------------------------------------
+# §4.2 — probe evaluation
+# ---------------------------------------------------------------------------
+
+
+def _database_returning(values: list[Any]) -> Database:
+ """A real ``Database`` (so the dialect is real) with only ``get_df``
stubbed."""
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=pd.DataFrame(
+ [values], columns=[f"v{i}" for i in range(len(values))]
+ )
+ )
+ return database
+
+
+def _database_returning_frame(frame: pd.DataFrame) -> Database:
+ """
+ ``_database_returning`` for a probe result whose column dtypes matter.
+
+ The flat-list sibling builds its frame from one Python list, so pandas
+ infers a single dtype and the row-versus-cell distinction cannot be seen.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=frame
+ )
+ return database
+
+
+def _probe(database: Database) -> MagicMock:
+ """
+ The stubbed ``get_df``, typed for call assertions.
+
+ ``Database.get_df`` is a real method, so mypy reads its declared signature
+ rather than the mock that replaced it and rejects ``.call_args`` and
+ friends.
+ """
+ return cast(MagicMock, database.get_df)
+
+
+def test_evaluate_transform_returns_one_value_per_input(app: Flask) -> None:
+ database = _database_returning([1767225600, 1769904000])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database,
+ None,
+ "default",
+ "unix_timestamp(:value)",
+ ["2026-01-01 00:00:00", "2026-02-01 00:00:00"],
+ )
+
+ assert result == [1767225600, 1769904000]
+ _probe(database).assert_called_once()
+
+
+def test_evaluate_transform_binds_values_rather_than_interpolating(
+ app: Flask,
+) -> None:
+ """
+ Filter values are attacker-controlled (a Gamma user picks them), so they
must
+ be bound and escaped, never pasted into the probe SQL.
+ """
+ database = _database_returning(["o''brien"])
+
+ with app.app_context():
+ evaluate_transform(database, None, None, "lower(:value)", ["O'Brien"])
+
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert "O'Brien" not in sql
+ assert "O''Brien" in sql
+
+
+def test_evaluate_transform_dedupes_repeated_values(app: Flask) -> None:
+ """Three inputs, two distinct: the probe evaluates the transform twice."""
+ database = _database_returning(["us", "us"])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database, None, None, "lower(:value)", ["US", "US", "us"]
+ )
+
+ assert result == ["us", "us", "us"]
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert sql.count("lower") == 2
+
+
+def test_evaluate_transform_fails_open_on_a_short_result_row(app: Flask) ->
None:
+ """A row narrower than the probe asked for means the results cannot be
+ aligned back to their inputs; pruning is skipped rather than guessed at."""
+ database = _database_returning(["us"])
+
+ with app.app_context():
+ assert (
+ evaluate_transform(database, None, None, "lower(:value)", ["US",
"CA"])
+ is None
+ )
+
+
+def test_evaluate_transform_fails_open_on_a_wide_result_row(app: Flask) ->
None:
+ """
+ Too many columns is the dangerous direction. A transform whose select list
+ holds two expressions returns 2N columns for N inputs, and reading the
first
+ N interleaves the expressions instead of taking one per value -- a
predicate
+ built from the wrong values rather than one that is merely short. Here
+ `lower('US'), 'x', lower('CA'), 'x'` would have yielded `('us', 'x')` and
+ dropped every CA row.
+ """
+ database = _database_returning(["us", "x", "ca", "x"])
+
+ with app.app_context():
+ assert (
+ evaluate_transform(database, None, None, "lower(:value), 'x'",
["US", "CA"])
Review Comment:
Fixed. The test uses a valid transform so `get_df` is actually reached and
the interleaving case the docstring describes is the one being asserted;
changing the width guard from `!=` to `<` now fails it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
