hughhhh commented on code in PR #43870: URL: https://github.com/apache/superset/pull/43870#discussion_r4319669911
########## superset/connectors/sqla/partition_mapping.py: ########## @@ -0,0 +1,1913 @@ +# 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. +""" +Partition filter mapping. + +Datasets on Hadoop-family engines are commonly partitioned on a *technical* +column -- an epoch integer, a lowercased region key -- that no analyst would +filter on. Unless a query carries a predicate on that column the engine scans +every partition. + +A dataset owner names one partition column ``p``, one business column that +filters are mirrored from, and a value transform ``T`` (a SQL expression +containing a ``:value`` placeholder). Superset then appends an equivalent +predicate on ``p`` to every query, so chart authors change nothing and queries +prune. + +The load-bearing assumption +--------------------------- +Everything here reasons about ``T(col) op T(v)``, but what is emitted is +``p op T(v)`` -- a predicate on a *physically different column*. The step from +one to the other is:: + + p = T(mapped_col) for every row in the table + +Superset cannot verify that; it is a property of whatever ETL populates the +partition column. If that job lags, backfills with different logic, or writes +the partition key in a different timezone, mirrored predicates silently drop +real rows. The mapping is only as trustworthy as the pipeline behind it. + +Operator safety +--------------- +A mirrored predicate ``P2`` may only be ``AND``-ed onto a query when the +original predicate ``P1`` *implies* it: + +=========================================== ============================= +Original Safe when +=========================================== ============================= +``col = v``, ``col IN (...)`` ``T`` is a function *and* the + engine compares ``col`` + byte-exactly +``col >=|>|<|<= v``, ``TEMPORAL_RANGE`` only if ``T`` is monotonic +``col != v``, ``NOT IN``, ``LIKE``, ... never +=========================================== ============================= + +Equality under the engine's own rules +------------------------------------- +"``T`` is a function" gives ``col = v`` implies ``T(col) = T(v)`` for *value* +equality. The engine compares with *SQL* equality, and the two part company +wherever that comparison is not byte-exact: under a case-insensitive collation +a stored ``country`` of ``'us'`` satisfies a filter for ``'US'``, while the +mirror ``region_key = hex('US')`` excludes the row and the chart loses it with +nothing to indicate why. Trailing-space-insensitive ``CHAR`` comparison says +the same thing about padding. + +So equality and ``IN`` are withdrawn for a *string* mapped column unless the +engine spec declares ``binary_string_comparison`` -- `equality_mirrors_safely`. +Numeric and temporal comparison is exact everywhere and is untouched, and range +mirroring is untouched too: it already rests on the owner declaring ``T`` +order-preserving *with respect to the column's own order*, which is an +assertion about the very semantics this is checking for. + +What the gate cannot see is a column declaring its own collation +(``country COLLATE NOCASE``) on an otherwise byte-exact engine. Nothing in +SQLAlchemy's reflection or the engine specs exposes that, so there it remains +the owner's assumption, like ``p = T(mapped_col)`` itself. + +Negations are never safe because ``T`` need not be injective: +``lower(:value)`` with ``country != 'US'`` mirrors to ``region_key != 'us'``, +which wrongly excludes rows whose ``country`` is already lowercase ``'us'`` -- +rows the original filter *keeps*. + +Monotonicity is a property of the transform, not of the column's data type: +``hour(:value)``, ``date_format(:value, 'dd')`` and ``dayofweek(:value)`` are +all reasonable transforms on a ``TIMESTAMP`` column and none of them preserve +ordering. It is therefore declared by the owner, not inferred. + +Time grains +----------- +A grained filter -- drill-to-detail, mostly -- compares the *truncated* column, +``trunc(col) op v``, so the raw bounds it carries do not describe the rows it +keeps. A row in the final partial bucket satisfies ``trunc(col) < until`` while +``col < until`` excludes it. + +Which direction a grain rounds is not knowable from the duration alone, and the +obvious guess is wrong: ``WEEK_ENDING_SATURDAY`` rounds *forward* on Hive and +Presto, and Ocient's grains are ``ROUND``, i.e. to nearest. What every grain +does satisfy is a bound on the displacement:: + + |trunc(ts) - ts| < width(grain) + +so widening *both* bounds by one bucket width is no narrower than the real +predicate whichever way the grain rounds -- and it needs no monotonicity of +``trunc`` itself, which is what rescues Hive's oddly-anchored ``P1W``. Width is +a property of the grain, not of the engine, which is what makes this +maintainable; see `grain_bucket_width`. A grain whose SQL an operator supplied +has no known width, and does not mirror at all. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +from collections.abc import Sequence +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, cast, TYPE_CHECKING + +import numpy as np +import pandas as pd +import sqlalchemy as sa +from dateutil.relativedelta import relativedelta +from flask import current_app as app +from flask_babel import lazy_gettext as _ +from sqlalchemy.engine.interfaces import Dialect +from sqlalchemy.sql.elements import ColumnElement + +from superset.constants import LRU_CACHE_MAX_SIZE, TimeGrain +from superset.exceptions import ( + QueryClauseValidationException, + SupersetParseError, + SupersetSecurityException, +) +from superset.extensions import cache_manager, feature_flag_manager +from superset.sql.parse import SQLScript, SQLStatement +from superset.superset_typing import FilterValues +from superset.utils import core as utils, json +from superset.utils.core import FilterOperator + +if TYPE_CHECKING: + from superset.connectors.sqla.models import SqlaTable, TableColumn + from superset.db_engine_specs.base import BaseEngineSpec + from superset.models.core import Database + +logger = logging.getLogger(__name__) + +FEATURE_FLAG = "PARTITION_FILTER_MAPPING" + +#: Longest transform accepted anywhere. A transform is one SQL expression an +#: owner types by hand, so this is generous; the point is that it is bounded. +#: The bound matters because `is_transform_active` parses the stored value on +#: each Explore load, and an unbounded string would make that parse the +#: expensive part of rendering a chart. +#: +#: The typed column field, the import schema and the preview request each +#: enforce it through a marshmallow `Length` validator, which is what produces +#: the per-field message those three callers want. `stored_expression_error` +#: enforces it too, for the paths that load no schema -- the deprecated +#: `POST /datasource/save/` is one, and it writes straight through +#: `update_from_object`. +MAX_TRANSFORM_LENGTH = 1024 + +#: Placeholder the owner writes in the transform, e.g. ``unix_timestamp(:value)``. +#: Matched with word boundaries so ``:values`` is not mistaken for it. +VALUE_PLACEHOLDER_RE = re.compile(r":value\b") + +#: Balanced Jinja blocks. The probe would render these in a different context +#: at a different time from the chart query, so they are rejected at save time. +JINJA_BLOCK_RE = re.compile(r"\{\{.*?\}\}|\{%.*?%\}|\{#.*?#\}", re.DOTALL) + +#: Substituted for ``:value`` before parsing -- sqlglot rejects a bare ``:value`` +#: on most dialects. Mirrors the ``_JINJA_BLOCK_RE`` -> ``NULL`` trick used by +#: ``validate_stored_expression``. +_PARSE_STANDIN = "NULL" + +#: Substituted for ``:value`` to find out *where* the placeholder landed, which +#: `_PARSE_STANDIN` cannot answer: ``NULL`` is a keyword, so after substitution +#: the parse tree no longer records that a placeholder was ever there. An +#: identifier does leave a trace -- a column reference -- and one that is +#: missing is a placeholder the engine will never evaluate. Named so it cannot +#: plausibly collide with a real column; a collision only makes the count +#: disagree, which fails closed. +_PLACEHOLDER_STANDIN = "superset_pfm_value_standin" + +#: Functions whose value depends on wall-clock time or randomness. The probe +#: runs in a different session at a different moment from the chart query and +#: its result is then cached, so any of these freezes a snapshot of probe time +#: into the emitted predicate. +NON_DETERMINISTIC_FUNCTIONS = { + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "NOW", + "RAND", + "RANDOM", + "UUID", +} + +#: Functions that mean "now" only in their zero-argument form. On Hive and +#: Impala ``unix_timestamp()`` is the current time while ``unix_timestamp(x)`` +#: -- the canonical transform for this feature -- is pure. +NON_DETERMINISTIC_WHEN_NILADIC = {"UNIX_TIMESTAMP"} + +#: Safe for any function ``T``. +MIRRORABLE_ALWAYS = {FilterOperator.EQUALS, FilterOperator.IN} + +#: Safe only when ``T`` preserves ordering. +MIRRORABLE_IF_MONOTONIC = { + FilterOperator.GREATER_THAN, + FilterOperator.GREATER_THAN_OR_EQUALS, + FilterOperator.LESS_THAN, + FilterOperator.LESS_THAN_OR_EQUALS, + FilterOperator.TEMPORAL_RANGE, +} + + +#: Every operator that can be mirrored under *some* transform. The preview +#: endpoint accepts these; whether a given one actually mirrors still depends on +#: the monotonicity declaration. +MIRRORABLE_OPERATORS = MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + +#: Operators the preview endpoint can construct. `TEMPORAL_RANGE` is mirrored by +#: the query path, but as *two* bounds that `_collect_partition_mirror_range` +#: decomposes a time range into -- the preview request carries sample values, not +#: a since/until pair, so there is no range for it to build. The editor never +#: asks for one either: `previewOperatorFor` sends `>=`, `=` or `IN`. +PREVIEWABLE_OPERATORS = MIRRORABLE_OPERATORS - {FilterOperator.TEMPORAL_RANGE} + + +def mirrorable_operators( + is_monotonic: bool, *, equality_is_safe: bool = True +) -> set[FilterOperator]: + """ + The operators whose predicates may be mirrored onto the partition column. + + :param is_monotonic: whether the owner declared the transform + order-preserving + :param equality_is_safe: whether the engine's ``=`` on the mapped column + compares the way ``T`` was written against -- see + `equality_mirrors_safely` + + Both call sites -- the query path through `PartitionMapping.mirrors` and the + Explore indicator through `partition_filter_mapping_summary` -- come through + here, which is what keeps the glyph from promising pruning the SQL will not + do. + """ + operators = MIRRORABLE_ALWAYS if equality_is_safe else set() + if is_monotonic: + operators = operators | MIRRORABLE_IF_MONOTONIC + return set(operators) + + +def equality_mirrors_safely( + column: TableColumn | None, + db_engine_spec: type[BaseEngineSpec], +) -> bool: + """ + Whether ``col = v`` on this column implies ``T(col) = T(v)`` in the engine. + + The operator matrix calls equality safe "for any function ``T``" because + ``T`` is a function, so ``col = v`` gives ``T(col) = T(v)``. That reasons + about *value* equality while the engine reasons about *SQL* equality, and + the two part company under any comparison that is not byte-exact: with a + case-insensitive collation a stored ``country`` of ``'us'`` satisfies a + filter for ``'US'``, but the mirror ``region_key = hex('US')`` excludes the + row and the chart loses it with no indication why. Trailing-space-insensitive + ``CHAR`` comparison says the same thing about padding. + + Only string columns are at risk; numeric and temporal comparison is exact + everywhere. And only equality: range mirroring already requires the owner + to declare ``T`` order-preserving *with respect to the column's own order*, + which is an assertion about the same comparison semantics this is checking + for -- so the declaration covers it where equality has nothing to cover it. + + What this cannot see is a column that declares its own collation + (``country COLLATE NOCASE``) on an otherwise byte-exact engine. Neither + SQLAlchemy's reflection nor the engine specs expose it, so on such a column + the assumption stays where ``p = T(mapped_col)`` already is: with the owner. + """ + if column is None: + return True + try: + is_string = column.type_generic == utils.GenericDataType.STRING + except Exception: # pylint: disable=broad-except # noqa: BLE001 + # An unresolvable type is treated as a string: the gate exists to stop + # a silent wrong answer, so it fails closed. + return False + if not is_string: + return True + return bool(db_engine_spec.binary_string_comparison) + + +#: How wide one bucket of each built-in time grain is. +#: +#: Keyed on the ISO duration a filter carries in its ``grain``. Written out +#: rather than parsed: four of the week grains are ISO *intervals* with an +#: anchor (``P1W/1970-01-03T00:00:00Z``) that ``isodate.parse_duration`` +#: rejects outright, and ``PT0.5H`` / ``P0.25Y`` are fractional. A literal +#: table is also the thing a reviewer can check a line at a time. +#: +#: ``relativedelta`` rather than ``timedelta`` so the calendar grains stay +#: calendar arithmetic: a month is not 30 days. +GRAIN_BUCKET_WIDTHS: dict[str, relativedelta] = { + TimeGrain.SECOND: relativedelta(seconds=1), + TimeGrain.FIVE_SECONDS: relativedelta(seconds=5), + TimeGrain.THIRTY_SECONDS: relativedelta(seconds=30), + TimeGrain.MINUTE: relativedelta(minutes=1), + TimeGrain.FIVE_MINUTES: relativedelta(minutes=5), + TimeGrain.TEN_MINUTES: relativedelta(minutes=10), + TimeGrain.FIFTEEN_MINUTES: relativedelta(minutes=15), + TimeGrain.THIRTY_MINUTES: relativedelta(minutes=30), + TimeGrain.HALF_HOUR: relativedelta(minutes=30), + TimeGrain.HOUR: relativedelta(hours=1), + TimeGrain.SIX_HOURS: relativedelta(hours=6), + TimeGrain.DAY: relativedelta(days=1), + TimeGrain.WEEK: relativedelta(days=7), + TimeGrain.WEEK_STARTING_SUNDAY: relativedelta(days=7), + TimeGrain.WEEK_STARTING_MONDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SATURDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SUNDAY: relativedelta(days=7), + TimeGrain.MONTH: relativedelta(months=1), + TimeGrain.QUARTER: relativedelta(months=3), + TimeGrain.QUARTER_YEAR: relativedelta(months=3), + TimeGrain.YEAR: relativedelta(years=1), +} + + +def grain_bucket_width(grain: str | None, engine: str) -> relativedelta | None: + """ + How far a grain's truncation can move a timestamp, or ``None`` if unknown. + + A grained filter compares the *truncated* column, so the raw bounds it + carries do not describe the rows it keeps. Widening both bounds by one + bucket recovers a predicate that is no narrower than the real one -- see + `_collect_partition_mirror_range` for the argument. That only works for a + grain whose bucket width Superset knows, which excludes anything an + operator supplied. + + :param grain: the ISO duration from the filter's ``grain`` + :param engine: the engine spec's ``engine``, to check per-engine overrides + """ + if not grain: + return None + + # An operator-declared grain carries whatever duration string they typed, + # and `TIME_GRAIN_ADDON_EXPRESSIONS` can redefine a *built-in* grain's SQL + # per engine -- `P1D` could be anything at all. Neither has a width we can + # claim to know, so both fall back to not mirroring. + if grain in app.config["TIME_GRAIN_ADDONS"]: + return None + if grain in app.config["TIME_GRAIN_ADDON_EXPRESSIONS"].get(engine, {}): + return None + + return GRAIN_BUCKET_WIDTHS.get(grain) + + +@dataclass(frozen=True) +class PartitionMapping: + """A resolved, usable partition filter mapping.""" + + partition_column: str + mapped_column: str + value_transform: str + is_monotonic: bool + #: Whether the engine compares the mapped column the way ``T`` was written + #: against. Resolved once where the column and the engine are both in hand, + #: rather than re-derived at each `mirrors` call. + equality_is_safe: bool = True + + def mirrors(self, operator: FilterOperator) -> bool: + return operator in mirrorable_operators( + self.is_monotonic, equality_is_safe=self.equality_is_safe + ) + + +def contains_value_placeholder(transform: str | None) -> bool: + """Whether the transform contains the ``:value`` placeholder.""" + return bool(transform) and VALUE_PLACEHOLDER_RE.search(transform or "") is not None + + +def contains_jinja(transform: str | None) -> bool: + """Whether the transform contains a balanced Jinja block.""" + return bool(transform) and JINJA_BLOCK_RE.search(transform or "") is not None + + +def parse_skeleton(transform: str, standin: str = _PARSE_STANDIN) -> str: + """ + The transform with ``:value`` substituted out, ready for a SQL parser. + + ``sanitize_clause`` / sqlglot choke on a bare ``:value`` on most dialects, + so the placeholder is swapped for a benign literal first -- the same trick + ``validate_stored_expression`` uses for Jinja blocks. + + :param standin: what to substitute. `placeholder_is_executable` passes + `_PLACEHOLDER_STANDIN` to find out where the placeholder landed. + """ + return VALUE_PLACEHOLDER_RE.sub(standin, transform) Review Comment: Fixed. The rendered SQL is substituted literally rather than used as a regex replacement string, so escaped backslashes in an array literal survive into the probe. ########## superset/connectors/sqla/partition_mapping.py: ########## @@ -0,0 +1,1314 @@ +# 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. +""" +Partition filter mapping. + +Datasets on Hadoop-family engines are commonly partitioned on a *technical* +column -- an epoch integer, a lowercased region key -- that no analyst would +filter on. Unless a query carries a predicate on that column the engine scans +every partition. + +A dataset owner names one partition column ``p``, one business column that +filters are mirrored from, and a value transform ``T`` (a SQL expression +containing a ``:value`` placeholder). Superset then appends an equivalent +predicate on ``p`` to every query, so chart authors change nothing and queries +prune. + +The load-bearing assumption +--------------------------- +Everything here reasons about ``T(col) op T(v)``, but what is emitted is +``p op T(v)`` -- a predicate on a *physically different column*. The step from +one to the other is:: + + p = T(mapped_col) for every row in the table + +Superset cannot verify that; it is a property of whatever ETL populates the +partition column. If that job lags, backfills with different logic, or writes +the partition key in a different timezone, mirrored predicates silently drop +real rows. The mapping is only as trustworthy as the pipeline behind it. + +Operator safety +--------------- +A mirrored predicate ``P2`` may only be ``AND``-ed onto a query when the +original predicate ``P1`` *implies* it: + +=========================================== ============================= +Original Safe when +=========================================== ============================= +``col = v``, ``col IN (...)`` always -- ``T`` is a function +``col >=|>|<|<= v``, ``TEMPORAL_RANGE`` only if ``T`` is monotonic +``col != v``, ``NOT IN``, ``LIKE``, ... never +=========================================== ============================= + +Negations are never safe because ``T`` need not be injective: +``lower(:value)`` with ``country != 'US'`` mirrors to ``region_key != 'us'``, +which wrongly excludes rows whose ``country`` is already lowercase ``'us'`` -- +rows the original filter *keeps*. + +Monotonicity is a property of the transform, not of the column's data type: +``hour(:value)``, ``date_format(:value, 'dd')`` and ``dayofweek(:value)`` are +all reasonable transforms on a ``TIMESTAMP`` column and none of them preserve +ordering. It is therefore declared by the owner, not inferred. + +Time grains +----------- +A grained filter -- drill-to-detail, mostly -- compares the *truncated* column, +``trunc(col) op v``, so the raw bounds it carries do not describe the rows it +keeps. A row in the final partial bucket satisfies ``trunc(col) < until`` while +``col < until`` excludes it. + +Which direction a grain rounds is not knowable from the duration alone, and the +obvious guess is wrong: ``WEEK_ENDING_SATURDAY`` rounds *forward* on Hive and +Presto, and Ocient's grains are ``ROUND``, i.e. to nearest. What every grain +does satisfy is a bound on the displacement:: + + |trunc(ts) - ts| < width(grain) + +so widening *both* bounds by one bucket width is no narrower than the real +predicate whichever way the grain rounds -- and it needs no monotonicity of +``trunc`` itself, which is what rescues Hive's oddly-anchored ``P1W``. Width is +a property of the grain, not of the engine, which is what makes this +maintainable; see `grain_bucket_width`. A grain whose SQL an operator supplied +has no known width, and does not mirror at all. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, cast, TYPE_CHECKING + +import numpy as np +import pandas as pd +import sqlalchemy as sa +from dateutil.relativedelta import relativedelta +from flask import current_app as app +from flask_babel import lazy_gettext as _ +from sqlalchemy.engine.interfaces import Dialect +from sqlalchemy.sql.elements import ColumnElement + +from superset.constants import LRU_CACHE_MAX_SIZE, TimeGrain +from superset.exceptions import ( + QueryClauseValidationException, + SupersetParseError, + SupersetSecurityException, +) +from superset.extensions import cache_manager, feature_flag_manager +from superset.sql.parse import SQLStatement +from superset.utils import json +from superset.utils.core import FilterOperator + +if TYPE_CHECKING: + from superset.connectors.sqla.models import SqlaTable, TableColumn + from superset.models.core import Database + +logger = logging.getLogger(__name__) + +FEATURE_FLAG = "PARTITION_FILTER_MAPPING" + +#: Longest transform accepted anywhere. A transform is one SQL expression an +#: owner types by hand, so this is generous; the point is that it is bounded. +#: Every entry point enforces it -- the typed column field, the import schema +#: and the preview request -- because `is_transform_active` parses the stored +#: value on each Explore load, and an unbounded string would make that parse +#: the expensive part of rendering a chart. +MAX_TRANSFORM_LENGTH = 1024 + +#: Placeholder the owner writes in the transform, e.g. ``unix_timestamp(:value)``. +#: Matched with word boundaries so ``:values`` is not mistaken for it. +VALUE_PLACEHOLDER_RE = re.compile(r":value\b") + +#: Balanced Jinja blocks. The probe would render these in a different context +#: at a different time from the chart query, so they are rejected at save time. +JINJA_BLOCK_RE = re.compile(r"\{\{.*?\}\}|\{%.*?%\}|\{#.*?#\}", re.DOTALL) + +#: Substituted for ``:value`` before parsing -- sqlglot rejects a bare ``:value`` +#: on most dialects. Mirrors the ``_JINJA_BLOCK_RE`` -> ``NULL`` trick used by +#: ``validate_stored_expression``. +_PARSE_STANDIN = "NULL" + +#: Functions whose value depends on wall-clock time or randomness. The probe +#: runs in a different session at a different moment from the chart query and +#: its result is then cached, so any of these freezes a snapshot of probe time +#: into the emitted predicate. +NON_DETERMINISTIC_FUNCTIONS = { + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "NOW", + "RAND", + "RANDOM", + "UUID", +} + +#: Functions that mean "now" only in their zero-argument form. On Hive and +#: Impala ``unix_timestamp()`` is the current time while ``unix_timestamp(x)`` +#: -- the canonical transform for this feature -- is pure. +NON_DETERMINISTIC_WHEN_NILADIC = {"UNIX_TIMESTAMP"} + +#: Safe for any function ``T``. +MIRRORABLE_ALWAYS = {FilterOperator.EQUALS, FilterOperator.IN} + +#: Safe only when ``T`` preserves ordering. +MIRRORABLE_IF_MONOTONIC = { + FilterOperator.GREATER_THAN, + FilterOperator.GREATER_THAN_OR_EQUALS, + FilterOperator.LESS_THAN, + FilterOperator.LESS_THAN_OR_EQUALS, + FilterOperator.TEMPORAL_RANGE, +} + + +#: Every operator that can be mirrored under *some* transform. The preview +#: endpoint accepts these; whether a given one actually mirrors still depends on +#: the monotonicity declaration. +MIRRORABLE_OPERATORS = MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + +#: Operators the preview endpoint can construct. `TEMPORAL_RANGE` is mirrored by +#: the query path, but as *two* bounds that `_collect_partition_mirror_range` +#: decomposes a time range into -- the preview request carries sample values, not +#: a since/until pair, so there is no range for it to build. The editor never +#: asks for one either: `previewOperatorFor` sends `>=`, `=` or `IN`. +PREVIEWABLE_OPERATORS = MIRRORABLE_OPERATORS - {FilterOperator.TEMPORAL_RANGE} + + +def mirrorable_operators(is_monotonic: bool) -> set[FilterOperator]: + """ + The operators whose predicates may be mirrored onto the partition column. + + :param is_monotonic: whether the owner declared the transform + order-preserving + """ + if is_monotonic: + return MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + return set(MIRRORABLE_ALWAYS) + + +#: How wide one bucket of each built-in time grain is. +#: +#: Keyed on the ISO duration a filter carries in its ``grain``. Written out +#: rather than parsed: four of the week grains are ISO *intervals* with an +#: anchor (``P1W/1970-01-03T00:00:00Z``) that ``isodate.parse_duration`` +#: rejects outright, and ``PT0.5H`` / ``P0.25Y`` are fractional. A literal +#: table is also the thing a reviewer can check a line at a time. +#: +#: ``relativedelta`` rather than ``timedelta`` so the calendar grains stay +#: calendar arithmetic: a month is not 30 days. +GRAIN_BUCKET_WIDTHS: dict[str, relativedelta] = { + TimeGrain.SECOND: relativedelta(seconds=1), + TimeGrain.FIVE_SECONDS: relativedelta(seconds=5), + TimeGrain.THIRTY_SECONDS: relativedelta(seconds=30), + TimeGrain.MINUTE: relativedelta(minutes=1), + TimeGrain.FIVE_MINUTES: relativedelta(minutes=5), + TimeGrain.TEN_MINUTES: relativedelta(minutes=10), + TimeGrain.FIFTEEN_MINUTES: relativedelta(minutes=15), + TimeGrain.THIRTY_MINUTES: relativedelta(minutes=30), + TimeGrain.HALF_HOUR: relativedelta(minutes=30), + TimeGrain.HOUR: relativedelta(hours=1), + TimeGrain.SIX_HOURS: relativedelta(hours=6), + TimeGrain.DAY: relativedelta(days=1), + TimeGrain.WEEK: relativedelta(days=7), + TimeGrain.WEEK_STARTING_SUNDAY: relativedelta(days=7), + TimeGrain.WEEK_STARTING_MONDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SATURDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SUNDAY: relativedelta(days=7), + TimeGrain.MONTH: relativedelta(months=1), + TimeGrain.QUARTER: relativedelta(months=3), + TimeGrain.QUARTER_YEAR: relativedelta(months=3), + TimeGrain.YEAR: relativedelta(years=1), +} + + +def grain_bucket_width(grain: str | None, engine: str) -> relativedelta | None: + """ + How far a grain's truncation can move a timestamp, or ``None`` if unknown. + + A grained filter compares the *truncated* column, so the raw bounds it + carries do not describe the rows it keeps. Widening both bounds by one + bucket recovers a predicate that is no narrower than the real one -- see + `_collect_partition_mirror_range` for the argument. That only works for a + grain whose bucket width Superset knows, which excludes anything an + operator supplied. + + :param grain: the ISO duration from the filter's ``grain`` + :param engine: the engine spec's ``engine``, to check per-engine overrides + """ + if not grain: + return None + + # An operator-declared grain carries whatever duration string they typed, + # and `TIME_GRAIN_ADDON_EXPRESSIONS` can redefine a *built-in* grain's SQL + # per engine -- `P1D` could be anything at all. Neither has a width we can + # claim to know, so both fall back to not mirroring. + if grain in app.config["TIME_GRAIN_ADDONS"]: + return None + if grain in app.config["TIME_GRAIN_ADDON_EXPRESSIONS"].get(engine, {}): + return None + + return GRAIN_BUCKET_WIDTHS.get(grain) + + +@dataclass(frozen=True) +class PartitionMapping: + """A resolved, usable partition filter mapping.""" + + partition_column: str + mapped_column: str + value_transform: str + is_monotonic: bool + + def mirrors(self, operator: FilterOperator) -> bool: + return operator in mirrorable_operators(self.is_monotonic) + + +def contains_value_placeholder(transform: str | None) -> bool: + """Whether the transform contains the ``:value`` placeholder.""" + return bool(transform) and VALUE_PLACEHOLDER_RE.search(transform or "") is not None + + +def contains_jinja(transform: str | None) -> bool: + """Whether the transform contains a balanced Jinja block.""" + return bool(transform) and JINJA_BLOCK_RE.search(transform or "") is not None + + +def parse_skeleton(transform: str) -> str: + """ + The transform with ``:value`` substituted out, ready for a SQL parser. + + ``sanitize_clause`` / sqlglot choke on a bare ``:value`` on most dialects, + so the placeholder is swapped for a benign literal first -- the same trick + ``validate_stored_expression`` uses for Jinja blocks. + """ + return VALUE_PLACEHOLDER_RE.sub(_PARSE_STANDIN, transform) + + +#: Prefix the transform is wrapped in before parsing. Its length is subtracted +#: from any reported column so positions refer to what the owner actually typed. +_SELECT_PREFIX = "SELECT " + + +def _parse_skeleton(transform: str, engine: str) -> SQLStatement | None: + """ + Parse ``SELECT <transform>`` with the placeholder substituted out. + + Returns ``None`` when the transform does not parse. + """ + try: + return SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError: + return None + + +def parse_error_detail(transform: str, engine: str) -> str | None: + """ + Where the parser gave up on the transform. + + Returns ``None`` when it parses, or when the parser offered no position. + Note that sqlglot parses unknown functions happily -- a misspelled function + name is not a parse error, it is an engine error, and surfaces only when the + transform is evaluated. + + The parser's own ``highlight`` is deliberately dropped: it would name the + ``NULL`` we substituted for ``:value``, which is not a token the owner typed. + """ + try: + SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError as ex: + column = (ex.error.extra or {}).get("column") + if not isinstance(column, int): + return None + return str( + _( + "syntax error at position %(position)d.", + position=_position_in_transform(transform, column), + ) + ) + return None + + +def _position_in_transform(transform: str, parsed_column: int) -> int: + """ + Map a column in the parsed skeleton back to the transform as typed. + + Two substitutions stand between them: the ``SELECT`` prefix, and every + ``:value`` that became a shorter ``NULL``. Without unwinding both, a + reported position drifts left by two characters per placeholder ahead of it + -- which is worst exactly where transforms usually break, at the end. + """ + position = max(parsed_column - len(_SELECT_PREFIX), 0) + shift = len(":value") - len(_PARSE_STANDIN) + preceding = sum( + 1 + for index, match in enumerate(VALUE_PLACEHOLDER_RE.finditer(transform)) + if match.start() - index * shift < position + ) + return position + preceding * shift + + +def is_parseable(transform: str | None, engine: str) -> bool: + """ + Whether the transform parses as a single select expression. + + "Single" is the load-bearing word. `SELECT lower(:value), 'x'` parses just + as happily as `SELECT lower(:value)`, but it returns two columns per input + value, and the probe reads one column per input -- so an `IN` filter would + get a predicate built from the wrong halves of the wrong rows. Rejecting the + list here is what lets the probe trust its own column count. + """ + if not transform or not transform.strip(): + return False + statement = _parse_skeleton(transform, engine) + return statement is not None and statement.count_select_expressions() == 1 + + +def find_non_deterministic_functions(transform: str, engine: str) -> set[str]: + """ + Names of non-deterministic functions the transform calls. + + ``UNIX_TIMESTAMP`` is only reported in its zero-argument form, which means + "now" on Hive and Impala; the one-argument form is the canonical temporal + transform and stays allowed. + """ + statement = _parse_skeleton(transform, engine) + if statement is None: + return set() + + found = { + name + for name in NON_DETERMINISTIC_FUNCTIONS + if statement.check_functions_present({name}) + } + return found | _find_niladic_calls(statement) + + +def _find_niladic_calls(statement: SQLStatement) -> set[str]: + """ + Names from ``NON_DETERMINISTIC_WHEN_NILADIC`` called with no arguments. + + Note some dialects resolve the zero-argument form themselves -- Hive parses + ``unix_timestamp()`` straight to ``CURRENT_TIMESTAMP`` -- in which case the + name-based check above has already caught it. This is the backstop for the + dialects that do not. + """ + return NON_DETERMINISTIC_WHEN_NILADIC & statement.get_niladic_functions() + + +def resolve_partition_mapping(datasource: SqlaTable) -> PartitionMapping | None: + """ + Resolve the dataset's mapping, or ``None`` when nothing may be mirrored. + + Every bail-out here is defensive as well as functional: save-time validation + rejects most of these, but rows predating the validation can still violate + the invariants, and a column sync can invalidate a mapping that was fine + when it was written. + + The transform gate is `is_transform_active`, the same function the Explore + indicator reads, which is `validate_transform` with the messages discarded. + Anything narrower here would be a second, weaker statement of the same rule: + a transform calling `now()` that reached storage without passing + `UpdateDatasetCommand` -- through import, or through a bundle written by hand + -- would be reported inactive by the editor and still mirrored by this + function, freezing a snapshot of probe time into the predicate with nothing + on screen to say so. + """ + if not feature_flag_manager.is_feature_enabled(FEATURE_FLAG): + return None + + partition_column = getattr(datasource, "partition_column", None) + if not partition_column: + return None + + columns_by_name = {column.column_name: column for column in datasource.columns} + if partition_column not in columns_by_name: + # The partition column was dropped by a column sync or at the source. + return None + + mapped_column_name = ( + getattr(datasource, "partition_mapped_column", None) or datasource.main_dttm_col + ) + if not mapped_column_name or mapped_column_name not in columns_by_name: + return None + + if mapped_column_name == partition_column: + # Self-mapping: the mirrored predicate would duplicate the original. + return None + + mapped_column = columns_by_name[mapped_column_name] + transform = getattr(mapped_column, "partition_value_transform", None) + if not is_transform_active(transform, datasource.database.backend): + return None + + if has_active_advanced_data_type(mapped_column): + # `translate_filter` builds its own predicate shape from *translated* + # values, so the `(operator, value)` pair the operator matrix reasons + # about does not exist and mirroring would apply the wrong values. + return None + + return PartitionMapping( + partition_column=str(partition_column), + mapped_column=str(mapped_column_name), + value_transform=cast(str, transform), + is_monotonic=bool( + getattr(mapped_column, "partition_transform_is_monotonic", False) + ), + ) + + +def has_active_advanced_data_type(column: TableColumn) -> bool: + """ + Whether the column's advanced data type is configured and switched on. + + Such a column is never mirrored: ``translate_filter`` builds its own + predicate shape from translated values, so the ``(operator, value)`` pair + the operator matrix reasons about does not exist. + """ + advanced_data_type = getattr(column, "advanced_data_type", None) + if not advanced_data_type: + return False + if not feature_flag_manager.is_feature_enabled("ENABLE_ADVANCED_DATA_TYPES"): + return False + return advanced_data_type in app.config.get("ADVANCED_DATA_TYPES", {}) + + +def stored_expression_error( + database: "Database", + catalog: str | None, + schema: str | None, + transform: str, +) -> str | None: + """ + Why this transform may not be stored or run, if there is a reason. + + The same subquery, function-denylist and RLS policy every other stored + expression goes through, applied to a value transform. It matters more + here than the name suggests: `build_probe_sql` binds only `:value` and + splices the rest of the transform in as SQL text, which the engine then + executes, so an ungated transform is arbitrary SQL. A dataset editor + without SQL Lab could store `(SELECT secret FROM protected_table LIMIT 1) + || :value` and read the answer back out of the emitted predicate. + + Returns the engine-agnostic reason as a string rather than raising, + because its four callers disagree about what to do with it: the preview + reports it at the field, a PUT and the legacy datasource save refuse the + write, the importer drops the transform and keeps the dataset, and the + probe simply declines to run. Raising would make three of those four + write a `try` around a question. + + Imported inside the function rather than at module scope: + `connectors.sqla.models` imports this module, so the dependency only runs + one way at import time. + """ + from superset.connectors.sqla.models import ( # pylint: disable=import-outside-toplevel,cyclic-import + validate_stored_expression, + ) + + try: + validate_stored_expression(database, catalog, schema, parse_skeleton(transform)) + except SupersetSecurityException as ex: + return str(ex.error.message) + except QueryClauseValidationException as ex: + return str(ex.message) + return None + + +def build_probe_sql( + transform: str, + values: list[Any], + dialect: Dialect | None = None, + from_suffix: str = "", +) -> str: + """ + Compile a single ``SELECT`` that evaluates the transform at every value. + + Values are attacker-controlled (a Gamma user picks filter values), so they + are bound as parameters and rendered by the dialect's own literal processor + rather than interpolated into the SQL text. + + ``from_suffix`` comes from the engine spec's ``select_without_from_suffix``. + A ``SELECT`` with no ``FROM`` is not universal SQL: Oracle and Db2 need a + one-row table to select from, and without it every probe on those engines + raises -- which `_run_probe` swallows, so the only symptom is a correctly + configured mapping that silently never prunes. + + Note this deliberately does *not* go through ``BaseEngineSpec``'s text + helper, which escapes ``:`` on every engine but Athena and would destroy the + ``:value`` placeholder before it can be bound. + """ + selections = [] + for index, value in enumerate(values): + clause = sa.text(transform).bindparams(sa.bindparam("value", value=value)) + compiled = clause.compile( + dialect=dialect, + compile_kwargs={"literal_binds": True}, Review Comment: Fixed. `_compile_literal` compiles with inline literals and then undoes the paramstyle's percent doubling, guarded on the identifier preparer's `_double_percents` — the same normalization `Database.compile_sqla_query` applies, and for the same reason: the probe executes with no bound parameters, so nothing would undo it later. ########## superset/connectors/sqla/partition_mapping.py: ########## @@ -0,0 +1,2264 @@ +# 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. +""" +Partition filter mapping. + +Datasets on Hadoop-family engines are commonly partitioned on a *technical* +column -- an epoch integer, a lowercased region key -- that no analyst would +filter on. Unless a query carries a predicate on that column the engine scans +every partition. + +A dataset owner names one partition column ``p``, one business column that +filters are mirrored from, and a value transform ``T`` (a SQL expression +containing a ``:value`` placeholder). Superset then appends an equivalent +predicate on ``p`` to every query, so chart authors change nothing and queries +prune. + +The load-bearing assumption +--------------------------- +Everything here reasons about ``T(col) op T(v)``, but what is emitted is +``p op T(v)`` -- a predicate on a *physically different column*. The step from +one to the other is:: + + p = T(mapped_col) for every row in the table + +Superset cannot verify that; it is a property of whatever ETL populates the +partition column. If that job lags, backfills with different logic, or writes +the partition key in a different timezone, mirrored predicates silently drop +real rows. The mapping is only as trustworthy as the pipeline behind it. + +Operator safety +--------------- +A mirrored predicate ``P2`` may only be ``AND``-ed onto a query when the +original predicate ``P1`` *implies* it: + +=========================================== ============================= +Original Safe when +=========================================== ============================= +``col = v``, ``col IN (...)`` ``T`` is a function *and* the + engine compares ``col`` + byte-exactly +``col >=|>|<|<= v``, ``TEMPORAL_RANGE`` only if ``T`` is monotonic +``col != v``, ``NOT IN``, ``LIKE``, ... never +=========================================== ============================= + +Monotonic here means non-decreasing, not strictly increasing: a bucketing +transform such as a day key is monotonic and maps a whole day of timestamps +onto one value. ``col < v`` then only implies ``T(col) <= T(v)``, so a strict +bound is mirrored non-strictly -- see `mirror_operator`. + +Equality under the engine's own rules +------------------------------------- +"``T`` is a function" gives ``col = v`` implies ``T(col) = T(v)`` for *value* +equality. The engine compares with *SQL* equality, and the two part company +wherever that comparison is not byte-exact: under a case-insensitive collation +a stored ``country`` of ``'us'`` satisfies a filter for ``'US'``, while the +mirror ``region_key = hex('US')`` excludes the row and the chart loses it with +nothing to indicate why. Trailing-space-insensitive ``CHAR`` comparison says +the same thing about padding. + +So equality and ``IN`` are withdrawn for a *string* mapped column unless the +engine spec declares ``binary_string_comparison`` -- `equality_mirrors_safely`. +Numeric and temporal comparison is exact everywhere and is untouched, and range +mirroring is untouched too: it already rests on the owner declaring ``T`` +order-preserving *with respect to the column's own order*, which is an +assertion about the very semantics this is checking for. + +What the gate cannot see is a column declaring its own collation +(``country COLLATE NOCASE``) on an otherwise byte-exact engine. Nothing in +SQLAlchemy's reflection or the engine specs exposes that, so there it remains +the owner's assumption, like ``p = T(mapped_col)`` itself. + +Resolution of the value the engine compares +------------------------------------------- +Every row of the table above assumes the engine compares the value the mirror +was built from. It does not always. A ``DATE`` column compared against +``'2026-07-06 10:08:11'`` is compared on the date part alone -- whether the +engine writes the literal itself (``TO_DATE('2026-07-06', 'YYYY-MM-DD')`` on +Postgres, ``DATE '2026-07-06'`` on Presto, Trino, BigQuery and some two dozen +others) or applies an implicit cast to a bound string. The filter therefore +keeps the whole of that day while ``p = T(2026-07-06 10:08:11)`` keeps one +instant of it: narrower than the predicate it stands in for, which is the one +thing a mirror may never be. + +What an engine throws away is detected from the column's type rather than +enumerated per engine -- `ExploreMixin._engine_literal_resolution` -- and what +to do about it is a property of the *operator*, not of the value: + +* a bound rounds outward, backward for a lower one and forward for an upper one, + through `ExploreMixin._round_bounds_outward`. The mirror then reads one extra + bucket at the edge, and the original predicate is still in the query to + exclude the rows that bucket adds. +* an equality, or one member of an ``IN`` list, has nowhere to widen to: these + predicates are ``AND``-ed onto the query, so "or the day either side" is not + something the mirror can express. They decline, which costs the query its + pruning and nothing else. + +So a ``DATE``-typed mapped column mirrors a range at day granularity, and does +not mirror an equality whose value carries a time of day at all. Map a +``TIMESTAMP`` column instead and every operator is available again. + +Negations are never safe because ``T`` need not be injective: +``lower(:value)`` with ``country != 'US'`` mirrors to ``region_key != 'us'``, +which wrongly excludes rows whose ``country`` is already lowercase ``'us'`` -- +rows the original filter *keeps*. + +Monotonicity is a property of the transform, not of the column's data type: +``hour(:value)``, ``date_format(:value, 'dd')`` and ``dayofweek(:value)`` are +all reasonable transforms on a ``TIMESTAMP`` column and none of them preserve +ordering. It is therefore declared by the owner, not inferred. + +Time grains +----------- +A grained filter -- drill-to-detail, mostly -- compares the *truncated* column, +``trunc(col) op v``, so the raw bounds it carries do not describe the rows it +keeps. A row in the final partial bucket satisfies ``trunc(col) < until`` while +``col < until`` excludes it. + +Which direction a grain rounds is not knowable from the duration alone, and the +obvious guess is wrong: ``WEEK_ENDING_SATURDAY`` rounds *forward* on Hive and +Presto, and Ocient's grains are ``ROUND``, i.e. to nearest. What every grain +does satisfy is a bound on the displacement:: + + |trunc(ts) - ts| < width(grain) + +so widening *both* bounds by one bucket width is no narrower than the real +predicate whichever way the grain rounds -- and it needs no monotonicity of +``trunc`` itself, which is what rescues Hive's oddly-anchored ``P1W``. Width is +a property of the grain, not of the engine, which is what makes this +maintainable; see `grain_bucket_width`. A grain whose SQL an operator supplied +has no known width, and does not mirror at all. +""" + +from __future__ import annotations + +import hashlib +import logging +import numbers +import re +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import date, datetime +from functools import lru_cache +from typing import Any, cast, TYPE_CHECKING + +import numpy as np +import pandas as pd +import sqlalchemy as sa +from dateutil.relativedelta import relativedelta +from flask import current_app as app +from flask_babel import lazy_gettext as _ +from sqlalchemy.engine.interfaces import Dialect +from sqlalchemy.sql.elements import ColumnElement + +from superset.constants import LRU_CACHE_MAX_SIZE, TimeGrain +from superset.exceptions import ( + QueryClauseValidationException, + SupersetParseError, + SupersetSecurityException, +) +from superset.extensions import cache_manager, feature_flag_manager +from superset.sql.parse import SQLScript, SQLStatement +from superset.superset_typing import FilterValues +from superset.utils import core as utils, json +from superset.utils.core import FilterOperator + +if TYPE_CHECKING: + from superset.connectors.sqla.models import SqlaTable, TableColumn + from superset.db_engine_specs.base import BaseEngineSpec + from superset.models.core import Database + +logger = logging.getLogger(__name__) + +FEATURE_FLAG = "PARTITION_FILTER_MAPPING" + +#: Longest transform accepted anywhere. A transform is one SQL expression an +#: owner types by hand, so this is generous; the point is that it is bounded. +#: The bound matters because `is_transform_active` parses the stored value on +#: each Explore load, and an unbounded string would make that parse the +#: expensive part of rendering a chart. +#: +#: The typed column field, the import schema and the preview request each +#: enforce it through a marshmallow `Length` validator, which is what produces +#: the per-field message those three callers want. `stored_expression_error` +#: enforces it too, for the paths that load no schema -- the deprecated +#: `POST /datasource/save/` is one, and it writes straight through +#: `update_from_object`. +MAX_TRANSFORM_LENGTH = 1024 + +#: Placeholder the owner writes in the transform, e.g. ``unix_timestamp(:value)``. +#: Matched with word boundaries so ``:values`` is not mistaken for it. +VALUE_PLACEHOLDER_RE = re.compile(r":value\b") + +#: Balanced Jinja blocks. The probe would render these in a different context +#: at a different time from the chart query, so they are rejected at save time. +JINJA_BLOCK_RE = re.compile(r"\{\{.*?\}\}|\{%.*?%\}|\{#.*?#\}", re.DOTALL) + +#: Substituted for ``:value`` before parsing -- sqlglot rejects a bare ``:value`` +#: on most dialects. Mirrors the ``_JINJA_BLOCK_RE`` -> ``NULL`` trick used by +#: ``validate_stored_expression``. +_PARSE_STANDIN = "NULL" + +#: Substituted for ``:value`` to find out *where* the placeholder landed, which +#: `_PARSE_STANDIN` cannot answer: ``NULL`` is a keyword, so after substitution +#: the parse tree no longer records that a placeholder was ever there. An +#: identifier does leave a trace -- a column reference -- and one that is +#: missing is a placeholder the engine will never evaluate. Named so it cannot +#: plausibly collide with a real column; a collision only makes the count +#: disagree, which fails closed. +_PLACEHOLDER_STANDIN = "superset_pfm_value_standin" + +#: Functions whose value depends on wall-clock time or randomness. The probe +#: runs in a different session at a different moment from the chart query and +#: its result is then cached, so any of these freezes a snapshot of probe time +#: into the emitted predicate. +NON_DETERMINISTIC_FUNCTIONS = { + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "NOW", + "RAND", + "RANDOM", + "UUID", +} + +#: Functions that mean "now" only in their zero-argument form. On Hive and +#: Impala ``unix_timestamp()`` is the current time while ``unix_timestamp(x)`` +#: -- the canonical transform for this feature -- is pure. +NON_DETERMINISTIC_WHEN_NILADIC = {"UNIX_TIMESTAMP"} + +#: Safe for any function ``T``. +MIRRORABLE_ALWAYS = {FilterOperator.EQUALS, FilterOperator.IN} + +#: Safe only when ``T`` preserves ordering. +MIRRORABLE_IF_MONOTONIC = { + FilterOperator.GREATER_THAN, + FilterOperator.GREATER_THAN_OR_EQUALS, + FilterOperator.LESS_THAN, + FilterOperator.LESS_THAN_OR_EQUALS, + FilterOperator.TEMPORAL_RANGE, +} + + +#: Every operator that can be mirrored under *some* transform. The preview +#: endpoint accepts these; whether a given one actually mirrors still depends on +#: the monotonicity declaration. +MIRRORABLE_OPERATORS = MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + +#: Operators the preview endpoint can construct. `TEMPORAL_RANGE` is mirrored by +#: the query path, but as *two* bounds that `_collect_partition_mirror_range` +#: decomposes a time range into -- the preview request carries sample values, not +#: a since/until pair, so there is no range for it to build. The editor never +#: asks for one either: `previewOperatorFor` sends `>=`, `=` or `IN`. +PREVIEWABLE_OPERATORS = MIRRORABLE_OPERATORS - {FilterOperator.TEMPORAL_RANGE} + + +#: A transform that preserves ordering need not be *strictly* increasing: a day +#: key such as ``to_char(ts, 'YYYYMMDD')`` maps a whole day of timestamps onto +#: one value. For such a ``T``, ``col < v`` only implies ``T(col) <= T(v)`` and +#: ``col > v`` only implies ``T(col) >= T(v)`` -- the rows in the bucket at the +#: boundary satisfy the non-strict form and not the strict one. A strict mirror +#: therefore drops rows the filter keeps, which is the one thing a mirror may +#: never do, so a strict bound is mirrored non-strictly. +_NON_STRICT_EQUIVALENT = { + FilterOperator.GREATER_THAN: FilterOperator.GREATER_THAN_OR_EQUALS, + FilterOperator.LESS_THAN: FilterOperator.LESS_THAN_OR_EQUALS, +} + + +def mirror_operator(operator: FilterOperator) -> FilterOperator: + """ + The operator a mirrored predicate uses to stand in for ``operator``. + + Strict bounds widen; everything else passes through. Idempotent, so it is + safe wherever a request may already have been widened. + + The widened bound reads one extra partition at the edge, and the original + predicate is still in the query to exclude the rows that partition adds -- + the same trade as the grain widening and the ``NULL``-partition escape. + Declaring the transform strictly increasing would buy that partition back, + which is not worth a second property only the owner can vouch for. + """ + return _NON_STRICT_EQUIVALENT.get(operator, operator) + + +def mirrorable_operators( + is_monotonic: bool, *, equality_is_safe: bool = True +) -> set[FilterOperator]: + """ + The operators whose predicates may be mirrored onto the partition column. + + :param is_monotonic: whether the owner declared the transform + order-preserving + :param equality_is_safe: whether the engine's ``=`` on the mapped column + compares the way ``T`` was written against -- see + `equality_mirrors_safely` + + Both call sites -- the query path through `PartitionMapping.mirrors` and the + Explore indicator through `partition_filter_mapping_summary` -- come through + here, which is what keeps the glyph from promising pruning the SQL will not + do. + """ + operators = MIRRORABLE_ALWAYS if equality_is_safe else set() + if is_monotonic: + operators = operators | MIRRORABLE_IF_MONOTONIC + return set(operators) + + +def equality_mirrors_safely( + column: TableColumn | None, + db_engine_spec: type[BaseEngineSpec], +) -> bool: + """ + Whether ``col = v`` on this column implies ``T(col) = T(v)`` in the engine. + + The operator matrix calls equality safe "for any function ``T``" because + ``T`` is a function, so ``col = v`` gives ``T(col) = T(v)``. That reasons + about *value* equality while the engine reasons about *SQL* equality, and + the two part company under any comparison that is not byte-exact: with a + case-insensitive collation a stored ``country`` of ``'us'`` satisfies a + filter for ``'US'``, but the mirror ``region_key = hex('US')`` excludes the + row and the chart loses it with no indication why. Trailing-space-insensitive + ``CHAR`` comparison says the same thing about padding. + + Only string columns are at risk; numeric and temporal comparison is exact + everywhere. And only equality: range mirroring already requires the owner + to declare ``T`` order-preserving *with respect to the column's own order*, + which is an assertion about the same comparison semantics this is checking + for -- so the declaration covers it where equality has nothing to cover it. + + What this cannot see is a column that declares its own collation + (``country COLLATE NOCASE``) on an otherwise byte-exact engine. Neither + SQLAlchemy's reflection nor the engine specs expose it, so on such a column + the assumption stays where ``p = T(mapped_col)`` already is: with the owner. + """ + if column is None: + return True + try: + is_string = column.type_generic == utils.GenericDataType.STRING + except Exception: # pylint: disable=broad-except # noqa: BLE001 + # An unresolvable type is treated as a string: the gate exists to stop + # a silent wrong answer, so it fails closed. + return False + if not is_string: + return True + return bool(db_engine_spec.binary_string_comparison) + + +#: How wide one bucket of each built-in time grain is. +#: +#: Keyed on the ISO duration a filter carries in its ``grain``. Written out +#: rather than parsed: four of the week grains are ISO *intervals* with an +#: anchor (``P1W/1970-01-03T00:00:00Z``) that ``isodate.parse_duration`` +#: rejects outright, and ``PT0.5H`` / ``P0.25Y`` are fractional. A literal +#: table is also the thing a reviewer can check a line at a time. +#: +#: ``relativedelta`` rather than ``timedelta`` so the calendar grains stay +#: calendar arithmetic: a month is not 30 days. +GRAIN_BUCKET_WIDTHS: dict[str, relativedelta] = { + TimeGrain.SECOND: relativedelta(seconds=1), + TimeGrain.FIVE_SECONDS: relativedelta(seconds=5), + TimeGrain.THIRTY_SECONDS: relativedelta(seconds=30), + TimeGrain.MINUTE: relativedelta(minutes=1), + TimeGrain.FIVE_MINUTES: relativedelta(minutes=5), + TimeGrain.TEN_MINUTES: relativedelta(minutes=10), + TimeGrain.FIFTEEN_MINUTES: relativedelta(minutes=15), + TimeGrain.THIRTY_MINUTES: relativedelta(minutes=30), + TimeGrain.HALF_HOUR: relativedelta(minutes=30), + TimeGrain.HOUR: relativedelta(hours=1), + TimeGrain.SIX_HOURS: relativedelta(hours=6), + TimeGrain.DAY: relativedelta(days=1), + TimeGrain.WEEK: relativedelta(days=7), + TimeGrain.WEEK_STARTING_SUNDAY: relativedelta(days=7), + TimeGrain.WEEK_STARTING_MONDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SATURDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SUNDAY: relativedelta(days=7), + TimeGrain.MONTH: relativedelta(months=1), + TimeGrain.QUARTER: relativedelta(months=3), + TimeGrain.QUARTER_YEAR: relativedelta(months=3), + TimeGrain.YEAR: relativedelta(years=1), +} + + +def grain_bucket_width(grain: str | None, engine: str) -> relativedelta | None: + """ + How far a grain's truncation can move a timestamp, or ``None`` if unknown. + + A grained filter compares the *truncated* column, so the raw bounds it + carries do not describe the rows it keeps. Widening both bounds by one + bucket recovers a predicate that is no narrower than the real one -- see + `_collect_partition_mirror_range` for the argument. That only works for a + grain whose bucket width Superset knows, which excludes anything an + operator supplied. + + :param grain: the ISO duration from the filter's ``grain`` + :param engine: the engine spec's ``engine``, to check per-engine overrides + """ + if not grain: + return None + + # An operator-declared grain carries whatever duration string they typed, + # and `TIME_GRAIN_ADDON_EXPRESSIONS` can redefine a *built-in* grain's SQL + # per engine -- `P1D` could be anything at all. Neither has a width we can + # claim to know, so both fall back to not mirroring. + if grain in app.config["TIME_GRAIN_ADDONS"]: + return None + if grain in app.config["TIME_GRAIN_ADDON_EXPRESSIONS"].get(engine, {}): + return None + + return GRAIN_BUCKET_WIDTHS.get(grain) + + +@dataclass(frozen=True) +class PartitionMapping: + """A resolved, usable partition filter mapping.""" + + partition_column: str + mapped_column: str + value_transform: str + is_monotonic: bool + #: Whether the engine compares the mapped column the way ``T`` was written + #: against. Resolved once where the column and the engine are both in hand, + #: rather than re-derived at each `mirrors` call. + equality_is_safe: bool = True + + def mirrors(self, operator: FilterOperator) -> bool: + return operator in mirrorable_operators( + self.is_monotonic, equality_is_safe=self.equality_is_safe + ) + + +def contains_value_placeholder(transform: str | None) -> bool: + """Whether the transform contains the ``:value`` placeholder.""" + return bool(transform) and VALUE_PLACEHOLDER_RE.search(transform or "") is not None + + +def contains_jinja(transform: str | None) -> bool: + """Whether the transform contains a balanced Jinja block.""" + return bool(transform) and JINJA_BLOCK_RE.search(transform or "") is not None + + +def parse_skeleton(transform: str, standin: str = _PARSE_STANDIN) -> str: + """ + The transform with ``:value`` substituted out, ready for a SQL parser. + + ``sanitize_clause`` / sqlglot choke on a bare ``:value`` on most dialects, + so the placeholder is swapped for a benign literal first -- the same trick + ``validate_stored_expression`` uses for Jinja blocks. + + :param standin: what to substitute. `placeholder_is_executable` passes + `_PLACEHOLDER_STANDIN` to find out where the placeholder landed. + """ + # Substituted through a function so `standin` is taken literally. As a + # replacement *string* `re.sub` would read backslashes in it as escapes, + # and `build_probe_sql` passes rendered SQL here: a filter value containing + # `\n` would probe a newline instead of the two characters the predicate + # compares, and one containing `\1` would raise outright. + return VALUE_PLACEHOLDER_RE.sub(lambda _: standin, transform) + + +#: Prefix the transform is wrapped in before parsing. Its length is subtracted +#: from any reported column so positions refer to what the owner actually typed. +_SELECT_PREFIX = "SELECT " + + +def _parse_skeleton( + transform: str, engine: str, standin: str = _PARSE_STANDIN +) -> SQLStatement | None: + """ + Parse ``SELECT <transform>`` with the placeholder substituted out. + + Returns ``None`` when the transform does not parse. + """ + try: + return SQLStatement( + f"{_SELECT_PREFIX}{parse_skeleton(transform, standin)}", engine + ) + except SupersetParseError: + return None + + +def parse_error_detail(transform: str, engine: str) -> str | None: + """ + Where the parser gave up on the transform. + + Returns ``None`` when it parses, or when the parser offered no position. + Note that sqlglot parses unknown functions happily -- a misspelled function + name is not a parse error, it is an engine error, and surfaces only when the + transform is evaluated. + + The parser's own ``highlight`` is deliberately dropped: it would name the + ``NULL`` we substituted for ``:value``, which is not a token the owner typed. + """ + try: + SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError as ex: + column = (ex.error.extra or {}).get("column") + if not isinstance(column, int): + return None + return str( + _( + "syntax error at position %(position)d.", + position=_position_in_transform(transform, column), + ) + ) + return None + + +def _position_in_transform(transform: str, parsed_column: int) -> int: + """ + Map a column in the parsed skeleton back to the transform as typed. + + Two substitutions stand between them: the ``SELECT`` prefix, and every + ``:value`` that became a shorter ``NULL``. Without unwinding both, a + reported position drifts left by two characters per placeholder ahead of it + -- which is worst exactly where transforms usually break, at the end. + """ + position = max(parsed_column - len(_SELECT_PREFIX), 0) + shift = len(":value") - len(_PARSE_STANDIN) + preceding = sum( + 1 + for index, match in enumerate(VALUE_PLACEHOLDER_RE.finditer(transform)) + if match.start() - index * shift < position + ) + return position + preceding * shift + + +def placeholder_is_executable(transform: str | None, engine: str) -> bool: + """ + Whether every ``:value`` sits where the engine will evaluate it. + + `contains_value_placeholder` is a regex over raw text, and `parse_skeleton` + substitutes before parsing, so neither notices a placeholder that is not in + an executable position. Two shapes reach the probe that way, and both are + worse than the inactive mapping a malformed transform earns: + + ``':value'`` -- the placeholder inside a string literal. `build_probe_sql` + binds it anyway, because SQLAlchemy's own ``text()`` scan is no more + literal-aware than the regex, and the dialect then renders the bind + *including its own quotes* in the middle of the owner's quotes. A value of + ``" || (SELECT secret FROM vault LIMIT 1) || "`` closes the literal and + compiles to a working sub-query -- the very read primitive + `stored_expression_error` refuses when it is written openly. + + ``1 -- :value`` -- the placeholder inside a comment. The probe's selections + are joined on one line, so the comment swallows its own ``AS v0`` alias and + everything after it; a single-value probe still returns one column, the + constant is accepted as the transform's result, and the mirror then emits + ``partition_col = <constant>`` for every filter value. That is wrong rows, + not lost pruning. + + Answered structurally: substitute an identifier rather than a keyword and + count the column references it produced. A literal holds no column + reference and a comment is not parsed at all, so either case comes back + short of the number of placeholders the owner wrote. + """ + if not transform: + return False + expected = len(VALUE_PLACEHOLDER_RE.findall(transform)) + if not expected: + return False + statement = _parse_skeleton(transform, engine, _PLACEHOLDER_STANDIN) + if statement is None: + return False + return statement.count_bare_column_references(_PLACEHOLDER_STANDIN) == expected Review Comment: Fixed. A transform naming the stand-in is refused outright: `if _PLACEHOLDER_STANDIN in transform.lower(): return False`, case-insensitively, ahead of the count check. The compiled probe is also re-read by `probe_sql_is_evaluable`. ########## superset/connectors/sqla/partition_mapping.py: ########## @@ -0,0 +1,1314 @@ +# 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. +""" +Partition filter mapping. + +Datasets on Hadoop-family engines are commonly partitioned on a *technical* +column -- an epoch integer, a lowercased region key -- that no analyst would +filter on. Unless a query carries a predicate on that column the engine scans +every partition. + +A dataset owner names one partition column ``p``, one business column that +filters are mirrored from, and a value transform ``T`` (a SQL expression +containing a ``:value`` placeholder). Superset then appends an equivalent +predicate on ``p`` to every query, so chart authors change nothing and queries +prune. + +The load-bearing assumption +--------------------------- +Everything here reasons about ``T(col) op T(v)``, but what is emitted is +``p op T(v)`` -- a predicate on a *physically different column*. The step from +one to the other is:: + + p = T(mapped_col) for every row in the table + +Superset cannot verify that; it is a property of whatever ETL populates the +partition column. If that job lags, backfills with different logic, or writes +the partition key in a different timezone, mirrored predicates silently drop +real rows. The mapping is only as trustworthy as the pipeline behind it. + +Operator safety +--------------- +A mirrored predicate ``P2`` may only be ``AND``-ed onto a query when the +original predicate ``P1`` *implies* it: + +=========================================== ============================= +Original Safe when +=========================================== ============================= +``col = v``, ``col IN (...)`` always -- ``T`` is a function +``col >=|>|<|<= v``, ``TEMPORAL_RANGE`` only if ``T`` is monotonic +``col != v``, ``NOT IN``, ``LIKE``, ... never +=========================================== ============================= + +Negations are never safe because ``T`` need not be injective: +``lower(:value)`` with ``country != 'US'`` mirrors to ``region_key != 'us'``, +which wrongly excludes rows whose ``country`` is already lowercase ``'us'`` -- +rows the original filter *keeps*. + +Monotonicity is a property of the transform, not of the column's data type: +``hour(:value)``, ``date_format(:value, 'dd')`` and ``dayofweek(:value)`` are +all reasonable transforms on a ``TIMESTAMP`` column and none of them preserve +ordering. It is therefore declared by the owner, not inferred. + +Time grains +----------- +A grained filter -- drill-to-detail, mostly -- compares the *truncated* column, +``trunc(col) op v``, so the raw bounds it carries do not describe the rows it +keeps. A row in the final partial bucket satisfies ``trunc(col) < until`` while +``col < until`` excludes it. + +Which direction a grain rounds is not knowable from the duration alone, and the +obvious guess is wrong: ``WEEK_ENDING_SATURDAY`` rounds *forward* on Hive and +Presto, and Ocient's grains are ``ROUND``, i.e. to nearest. What every grain +does satisfy is a bound on the displacement:: + + |trunc(ts) - ts| < width(grain) + +so widening *both* bounds by one bucket width is no narrower than the real +predicate whichever way the grain rounds -- and it needs no monotonicity of +``trunc`` itself, which is what rescues Hive's oddly-anchored ``P1W``. Width is +a property of the grain, not of the engine, which is what makes this +maintainable; see `grain_bucket_width`. A grain whose SQL an operator supplied +has no known width, and does not mirror at all. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, cast, TYPE_CHECKING + +import numpy as np +import pandas as pd +import sqlalchemy as sa +from dateutil.relativedelta import relativedelta +from flask import current_app as app +from flask_babel import lazy_gettext as _ +from sqlalchemy.engine.interfaces import Dialect +from sqlalchemy.sql.elements import ColumnElement + +from superset.constants import LRU_CACHE_MAX_SIZE, TimeGrain +from superset.exceptions import ( + QueryClauseValidationException, + SupersetParseError, + SupersetSecurityException, +) +from superset.extensions import cache_manager, feature_flag_manager +from superset.sql.parse import SQLStatement +from superset.utils import json +from superset.utils.core import FilterOperator + +if TYPE_CHECKING: + from superset.connectors.sqla.models import SqlaTable, TableColumn + from superset.models.core import Database + +logger = logging.getLogger(__name__) + +FEATURE_FLAG = "PARTITION_FILTER_MAPPING" + +#: Longest transform accepted anywhere. A transform is one SQL expression an +#: owner types by hand, so this is generous; the point is that it is bounded. +#: Every entry point enforces it -- the typed column field, the import schema +#: and the preview request -- because `is_transform_active` parses the stored +#: value on each Explore load, and an unbounded string would make that parse +#: the expensive part of rendering a chart. +MAX_TRANSFORM_LENGTH = 1024 + +#: Placeholder the owner writes in the transform, e.g. ``unix_timestamp(:value)``. +#: Matched with word boundaries so ``:values`` is not mistaken for it. +VALUE_PLACEHOLDER_RE = re.compile(r":value\b") + +#: Balanced Jinja blocks. The probe would render these in a different context +#: at a different time from the chart query, so they are rejected at save time. +JINJA_BLOCK_RE = re.compile(r"\{\{.*?\}\}|\{%.*?%\}|\{#.*?#\}", re.DOTALL) + +#: Substituted for ``:value`` before parsing -- sqlglot rejects a bare ``:value`` +#: on most dialects. Mirrors the ``_JINJA_BLOCK_RE`` -> ``NULL`` trick used by +#: ``validate_stored_expression``. +_PARSE_STANDIN = "NULL" + +#: Functions whose value depends on wall-clock time or randomness. The probe +#: runs in a different session at a different moment from the chart query and +#: its result is then cached, so any of these freezes a snapshot of probe time +#: into the emitted predicate. +NON_DETERMINISTIC_FUNCTIONS = { + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "NOW", + "RAND", + "RANDOM", + "UUID", +} + +#: Functions that mean "now" only in their zero-argument form. On Hive and +#: Impala ``unix_timestamp()`` is the current time while ``unix_timestamp(x)`` +#: -- the canonical transform for this feature -- is pure. +NON_DETERMINISTIC_WHEN_NILADIC = {"UNIX_TIMESTAMP"} + +#: Safe for any function ``T``. +MIRRORABLE_ALWAYS = {FilterOperator.EQUALS, FilterOperator.IN} + +#: Safe only when ``T`` preserves ordering. +MIRRORABLE_IF_MONOTONIC = { + FilterOperator.GREATER_THAN, + FilterOperator.GREATER_THAN_OR_EQUALS, + FilterOperator.LESS_THAN, + FilterOperator.LESS_THAN_OR_EQUALS, + FilterOperator.TEMPORAL_RANGE, +} + + +#: Every operator that can be mirrored under *some* transform. The preview +#: endpoint accepts these; whether a given one actually mirrors still depends on +#: the monotonicity declaration. +MIRRORABLE_OPERATORS = MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + +#: Operators the preview endpoint can construct. `TEMPORAL_RANGE` is mirrored by +#: the query path, but as *two* bounds that `_collect_partition_mirror_range` +#: decomposes a time range into -- the preview request carries sample values, not +#: a since/until pair, so there is no range for it to build. The editor never +#: asks for one either: `previewOperatorFor` sends `>=`, `=` or `IN`. +PREVIEWABLE_OPERATORS = MIRRORABLE_OPERATORS - {FilterOperator.TEMPORAL_RANGE} + + +def mirrorable_operators(is_monotonic: bool) -> set[FilterOperator]: + """ + The operators whose predicates may be mirrored onto the partition column. + + :param is_monotonic: whether the owner declared the transform + order-preserving + """ + if is_monotonic: + return MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + return set(MIRRORABLE_ALWAYS) + + +#: How wide one bucket of each built-in time grain is. +#: +#: Keyed on the ISO duration a filter carries in its ``grain``. Written out +#: rather than parsed: four of the week grains are ISO *intervals* with an +#: anchor (``P1W/1970-01-03T00:00:00Z``) that ``isodate.parse_duration`` +#: rejects outright, and ``PT0.5H`` / ``P0.25Y`` are fractional. A literal +#: table is also the thing a reviewer can check a line at a time. +#: +#: ``relativedelta`` rather than ``timedelta`` so the calendar grains stay +#: calendar arithmetic: a month is not 30 days. +GRAIN_BUCKET_WIDTHS: dict[str, relativedelta] = { + TimeGrain.SECOND: relativedelta(seconds=1), + TimeGrain.FIVE_SECONDS: relativedelta(seconds=5), + TimeGrain.THIRTY_SECONDS: relativedelta(seconds=30), + TimeGrain.MINUTE: relativedelta(minutes=1), + TimeGrain.FIVE_MINUTES: relativedelta(minutes=5), + TimeGrain.TEN_MINUTES: relativedelta(minutes=10), + TimeGrain.FIFTEEN_MINUTES: relativedelta(minutes=15), + TimeGrain.THIRTY_MINUTES: relativedelta(minutes=30), + TimeGrain.HALF_HOUR: relativedelta(minutes=30), + TimeGrain.HOUR: relativedelta(hours=1), + TimeGrain.SIX_HOURS: relativedelta(hours=6), + TimeGrain.DAY: relativedelta(days=1), + TimeGrain.WEEK: relativedelta(days=7), + TimeGrain.WEEK_STARTING_SUNDAY: relativedelta(days=7), + TimeGrain.WEEK_STARTING_MONDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SATURDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SUNDAY: relativedelta(days=7), + TimeGrain.MONTH: relativedelta(months=1), + TimeGrain.QUARTER: relativedelta(months=3), + TimeGrain.QUARTER_YEAR: relativedelta(months=3), + TimeGrain.YEAR: relativedelta(years=1), +} + + +def grain_bucket_width(grain: str | None, engine: str) -> relativedelta | None: + """ + How far a grain's truncation can move a timestamp, or ``None`` if unknown. + + A grained filter compares the *truncated* column, so the raw bounds it + carries do not describe the rows it keeps. Widening both bounds by one + bucket recovers a predicate that is no narrower than the real one -- see + `_collect_partition_mirror_range` for the argument. That only works for a + grain whose bucket width Superset knows, which excludes anything an + operator supplied. + + :param grain: the ISO duration from the filter's ``grain`` + :param engine: the engine spec's ``engine``, to check per-engine overrides + """ + if not grain: + return None + + # An operator-declared grain carries whatever duration string they typed, + # and `TIME_GRAIN_ADDON_EXPRESSIONS` can redefine a *built-in* grain's SQL + # per engine -- `P1D` could be anything at all. Neither has a width we can + # claim to know, so both fall back to not mirroring. + if grain in app.config["TIME_GRAIN_ADDONS"]: + return None + if grain in app.config["TIME_GRAIN_ADDON_EXPRESSIONS"].get(engine, {}): + return None + + return GRAIN_BUCKET_WIDTHS.get(grain) + + +@dataclass(frozen=True) +class PartitionMapping: + """A resolved, usable partition filter mapping.""" + + partition_column: str + mapped_column: str + value_transform: str + is_monotonic: bool + + def mirrors(self, operator: FilterOperator) -> bool: + return operator in mirrorable_operators(self.is_monotonic) + + +def contains_value_placeholder(transform: str | None) -> bool: + """Whether the transform contains the ``:value`` placeholder.""" + return bool(transform) and VALUE_PLACEHOLDER_RE.search(transform or "") is not None + + +def contains_jinja(transform: str | None) -> bool: + """Whether the transform contains a balanced Jinja block.""" + return bool(transform) and JINJA_BLOCK_RE.search(transform or "") is not None + + +def parse_skeleton(transform: str) -> str: + """ + The transform with ``:value`` substituted out, ready for a SQL parser. + + ``sanitize_clause`` / sqlglot choke on a bare ``:value`` on most dialects, + so the placeholder is swapped for a benign literal first -- the same trick + ``validate_stored_expression`` uses for Jinja blocks. + """ + return VALUE_PLACEHOLDER_RE.sub(_PARSE_STANDIN, transform) + + +#: Prefix the transform is wrapped in before parsing. Its length is subtracted +#: from any reported column so positions refer to what the owner actually typed. +_SELECT_PREFIX = "SELECT " + + +def _parse_skeleton(transform: str, engine: str) -> SQLStatement | None: + """ + Parse ``SELECT <transform>`` with the placeholder substituted out. + + Returns ``None`` when the transform does not parse. + """ + try: + return SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError: + return None + + +def parse_error_detail(transform: str, engine: str) -> str | None: + """ + Where the parser gave up on the transform. + + Returns ``None`` when it parses, or when the parser offered no position. + Note that sqlglot parses unknown functions happily -- a misspelled function + name is not a parse error, it is an engine error, and surfaces only when the + transform is evaluated. + + The parser's own ``highlight`` is deliberately dropped: it would name the + ``NULL`` we substituted for ``:value``, which is not a token the owner typed. + """ + try: + SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError as ex: + column = (ex.error.extra or {}).get("column") + if not isinstance(column, int): + return None + return str( + _( + "syntax error at position %(position)d.", + position=_position_in_transform(transform, column), + ) + ) + return None + + +def _position_in_transform(transform: str, parsed_column: int) -> int: + """ + Map a column in the parsed skeleton back to the transform as typed. + + Two substitutions stand between them: the ``SELECT`` prefix, and every + ``:value`` that became a shorter ``NULL``. Without unwinding both, a + reported position drifts left by two characters per placeholder ahead of it + -- which is worst exactly where transforms usually break, at the end. + """ + position = max(parsed_column - len(_SELECT_PREFIX), 0) + shift = len(":value") - len(_PARSE_STANDIN) + preceding = sum( + 1 + for index, match in enumerate(VALUE_PLACEHOLDER_RE.finditer(transform)) + if match.start() - index * shift < position + ) + return position + preceding * shift + + +def is_parseable(transform: str | None, engine: str) -> bool: + """ + Whether the transform parses as a single select expression. + + "Single" is the load-bearing word. `SELECT lower(:value), 'x'` parses just + as happily as `SELECT lower(:value)`, but it returns two columns per input + value, and the probe reads one column per input -- so an `IN` filter would + get a predicate built from the wrong halves of the wrong rows. Rejecting the + list here is what lets the probe trust its own column count. + """ + if not transform or not transform.strip(): + return False + statement = _parse_skeleton(transform, engine) + return statement is not None and statement.count_select_expressions() == 1 + + +def find_non_deterministic_functions(transform: str, engine: str) -> set[str]: + """ + Names of non-deterministic functions the transform calls. + + ``UNIX_TIMESTAMP`` is only reported in its zero-argument form, which means + "now" on Hive and Impala; the one-argument form is the canonical temporal + transform and stays allowed. + """ + statement = _parse_skeleton(transform, engine) + if statement is None: + return set() + + found = { + name + for name in NON_DETERMINISTIC_FUNCTIONS + if statement.check_functions_present({name}) + } + return found | _find_niladic_calls(statement) + + +def _find_niladic_calls(statement: SQLStatement) -> set[str]: + """ + Names from ``NON_DETERMINISTIC_WHEN_NILADIC`` called with no arguments. + + Note some dialects resolve the zero-argument form themselves -- Hive parses + ``unix_timestamp()`` straight to ``CURRENT_TIMESTAMP`` -- in which case the + name-based check above has already caught it. This is the backstop for the + dialects that do not. + """ + return NON_DETERMINISTIC_WHEN_NILADIC & statement.get_niladic_functions() + + +def resolve_partition_mapping(datasource: SqlaTable) -> PartitionMapping | None: + """ + Resolve the dataset's mapping, or ``None`` when nothing may be mirrored. + + Every bail-out here is defensive as well as functional: save-time validation + rejects most of these, but rows predating the validation can still violate + the invariants, and a column sync can invalidate a mapping that was fine + when it was written. + + The transform gate is `is_transform_active`, the same function the Explore + indicator reads, which is `validate_transform` with the messages discarded. + Anything narrower here would be a second, weaker statement of the same rule: + a transform calling `now()` that reached storage without passing + `UpdateDatasetCommand` -- through import, or through a bundle written by hand + -- would be reported inactive by the editor and still mirrored by this + function, freezing a snapshot of probe time into the predicate with nothing + on screen to say so. + """ + if not feature_flag_manager.is_feature_enabled(FEATURE_FLAG): + return None + + partition_column = getattr(datasource, "partition_column", None) + if not partition_column: + return None + + columns_by_name = {column.column_name: column for column in datasource.columns} + if partition_column not in columns_by_name: + # The partition column was dropped by a column sync or at the source. + return None + + mapped_column_name = ( + getattr(datasource, "partition_mapped_column", None) or datasource.main_dttm_col + ) + if not mapped_column_name or mapped_column_name not in columns_by_name: + return None + + if mapped_column_name == partition_column: + # Self-mapping: the mirrored predicate would duplicate the original. + return None + + mapped_column = columns_by_name[mapped_column_name] + transform = getattr(mapped_column, "partition_value_transform", None) + if not is_transform_active(transform, datasource.database.backend): + return None + + if has_active_advanced_data_type(mapped_column): + # `translate_filter` builds its own predicate shape from *translated* + # values, so the `(operator, value)` pair the operator matrix reasons + # about does not exist and mirroring would apply the wrong values. + return None + + return PartitionMapping( + partition_column=str(partition_column), + mapped_column=str(mapped_column_name), + value_transform=cast(str, transform), + is_monotonic=bool( + getattr(mapped_column, "partition_transform_is_monotonic", False) + ), + ) + + +def has_active_advanced_data_type(column: TableColumn) -> bool: + """ + Whether the column's advanced data type is configured and switched on. + + Such a column is never mirrored: ``translate_filter`` builds its own + predicate shape from translated values, so the ``(operator, value)`` pair + the operator matrix reasons about does not exist. + """ + advanced_data_type = getattr(column, "advanced_data_type", None) + if not advanced_data_type: + return False + if not feature_flag_manager.is_feature_enabled("ENABLE_ADVANCED_DATA_TYPES"): + return False + return advanced_data_type in app.config.get("ADVANCED_DATA_TYPES", {}) + + +def stored_expression_error( + database: "Database", + catalog: str | None, + schema: str | None, + transform: str, +) -> str | None: + """ + Why this transform may not be stored or run, if there is a reason. + + The same subquery, function-denylist and RLS policy every other stored + expression goes through, applied to a value transform. It matters more + here than the name suggests: `build_probe_sql` binds only `:value` and + splices the rest of the transform in as SQL text, which the engine then + executes, so an ungated transform is arbitrary SQL. A dataset editor + without SQL Lab could store `(SELECT secret FROM protected_table LIMIT 1) + || :value` and read the answer back out of the emitted predicate. + + Returns the engine-agnostic reason as a string rather than raising, + because its four callers disagree about what to do with it: the preview + reports it at the field, a PUT and the legacy datasource save refuse the + write, the importer drops the transform and keeps the dataset, and the + probe simply declines to run. Raising would make three of those four + write a `try` around a question. + + Imported inside the function rather than at module scope: + `connectors.sqla.models` imports this module, so the dependency only runs + one way at import time. + """ + from superset.connectors.sqla.models import ( # pylint: disable=import-outside-toplevel,cyclic-import + validate_stored_expression, + ) + + try: + validate_stored_expression(database, catalog, schema, parse_skeleton(transform)) + except SupersetSecurityException as ex: + return str(ex.error.message) + except QueryClauseValidationException as ex: + return str(ex.message) + return None + + +def build_probe_sql( + transform: str, + values: list[Any], + dialect: Dialect | None = None, + from_suffix: str = "", +) -> str: + """ + Compile a single ``SELECT`` that evaluates the transform at every value. + + Values are attacker-controlled (a Gamma user picks filter values), so they + are bound as parameters and rendered by the dialect's own literal processor + rather than interpolated into the SQL text. + + ``from_suffix`` comes from the engine spec's ``select_without_from_suffix``. + A ``SELECT`` with no ``FROM`` is not universal SQL: Oracle and Db2 need a + one-row table to select from, and without it every probe on those engines + raises -- which `_run_probe` swallows, so the only symptom is a correctly + configured mapping that silently never prunes. + + Note this deliberately does *not* go through ``BaseEngineSpec``'s text + helper, which escapes ``:`` on every engine but Athena and would destroy the + ``:value`` placeholder before it can be bound. + """ + selections = [] + for index, value in enumerate(values): + clause = sa.text(transform).bindparams(sa.bindparam("value", value=value)) + compiled = clause.compile( + dialect=dialect, + compile_kwargs={"literal_binds": True}, + ) + selections.append(f"{compiled} AS v{index}") + return "SELECT " + ", ".join(selections) + from_suffix + + +def evaluate_transform( + database: Database, + catalog: str | None, + schema: str | None, + transform: str, + values: list[Any], + *, + errors: list[str] | None = None, +) -> list[Any] | None: + """ + Evaluate ``transform`` against the engine once per distinct value. + + Returns one result per input value, positionally aligned with ``values``, or + ``None`` if anything at all goes wrong. Failing open costs pruning, never + correctness: the chart query still runs, it just scans more partitions. + + The probe is pinned to the dataset's catalog and schema so session settings + match the chart query as closely as the connection pool allows. It still + runs in a *different* session, which is why transforms calling + session-dependent functions are rejected at save time. + + :param errors: optional sink for the engine's own account of a failure. The + query path passes nothing and stays silent; the editor's preview passes + a list so it can tell the owner *why* the transform did not evaluate -- + a misspelled function is the common case and sqlglot parses it happily. + """ + if not values: + return None + + # Dedupe so a 200-value `IN` list costs one column, not 200. + distinct: list[Any] = [] + seen: set[Any] = set() + for value in values: + key = _hashable(value) + if key not in seen: + seen.add(key) + distinct.append(value) + + cache_key = _probe_cache_key(database, catalog, schema, transform, distinct) + cached = _cache_get(cache_key) + if cached is None: + cached = _run_probe( + database, catalog, schema, transform, distinct, errors=errors + ) + if cached is None: + # Deliberately not cached: a transient engine blip would otherwise + # keep the dataset pruning-free for the whole cache timeout. + return None + _cache_set(cache_key, cached) + + evaluated = dict( + zip((_hashable(value) for value in distinct), cached, strict=False) + ) + return [evaluated[_hashable(value)] for value in values] + + +def _run_probe( + database: Database, + catalog: str | None, + schema: str | None, + transform: str, + distinct: list[Any], + *, + errors: list[str] | None = None, +) -> list[Any] | None: + # The last gate before the transform becomes SQL the engine runs, and the + # only one that covers a transform already in storage. The write-side + # checks can only speak for rows written after they existed; this speaks + # for every row, including ones a pre-fix release stored and ones a door + # that forgets to ask still lets in. + # + # Declining costs pruning, never correctness -- see `evaluate_transform`. + if reason := stored_expression_error(database, catalog, schema, transform): + logger.warning( + "Refusing to probe a partition transform that is not a storable " + "expression; queries will not prune: %s", + reason, + ) + return None + + try: + sql = build_probe_sql( + transform, + distinct, + _dialect_for(database), + database.db_engine_spec.select_without_from_suffix, + ) + frame = database.get_df(sql=sql, catalog=catalog, schema=schema) + if frame is None or frame.empty: + logger.warning( + "Partition transform probe returned no rows; skipping mirroring" + ) + return None + row = frame.iloc[0] Review Comment: Fixed. The frame is read cell-wise with `frame.iat[0, index]` rather than through `frame.iloc[0]`, so a row mixing an exact int64 with a float no longer unifies to float64 and rounds an integer key above 2^53. ########## superset/connectors/sqla/partition_mapping.py: ########## @@ -0,0 +1,1613 @@ +# 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. +""" +Partition filter mapping. + +Datasets on Hadoop-family engines are commonly partitioned on a *technical* +column -- an epoch integer, a lowercased region key -- that no analyst would +filter on. Unless a query carries a predicate on that column the engine scans +every partition. + +A dataset owner names one partition column ``p``, one business column that +filters are mirrored from, and a value transform ``T`` (a SQL expression +containing a ``:value`` placeholder). Superset then appends an equivalent +predicate on ``p`` to every query, so chart authors change nothing and queries +prune. + +The load-bearing assumption +--------------------------- +Everything here reasons about ``T(col) op T(v)``, but what is emitted is +``p op T(v)`` -- a predicate on a *physically different column*. The step from +one to the other is:: + + p = T(mapped_col) for every row in the table + +Superset cannot verify that; it is a property of whatever ETL populates the +partition column. If that job lags, backfills with different logic, or writes +the partition key in a different timezone, mirrored predicates silently drop +real rows. The mapping is only as trustworthy as the pipeline behind it. + +Operator safety +--------------- +A mirrored predicate ``P2`` may only be ``AND``-ed onto a query when the +original predicate ``P1`` *implies* it: + +=========================================== ============================= +Original Safe when +=========================================== ============================= +``col = v``, ``col IN (...)`` ``T`` is a function *and* the + engine compares ``col`` + byte-exactly +``col >=|>|<|<= v``, ``TEMPORAL_RANGE`` only if ``T`` is monotonic +``col != v``, ``NOT IN``, ``LIKE``, ... never +=========================================== ============================= + +Equality under the engine's own rules +------------------------------------- +"``T`` is a function" gives ``col = v`` implies ``T(col) = T(v)`` for *value* +equality. The engine compares with *SQL* equality, and the two part company +wherever that comparison is not byte-exact: under a case-insensitive collation +a stored ``country`` of ``'us'`` satisfies a filter for ``'US'``, while the +mirror ``region_key = hex('US')`` excludes the row and the chart loses it with +nothing to indicate why. Trailing-space-insensitive ``CHAR`` comparison says +the same thing about padding. + +So equality and ``IN`` are withdrawn for a *string* mapped column unless the +engine spec declares ``binary_string_comparison`` -- `equality_mirrors_safely`. +Numeric and temporal comparison is exact everywhere and is untouched, and range +mirroring is untouched too: it already rests on the owner declaring ``T`` +order-preserving *with respect to the column's own order*, which is an +assertion about the very semantics this is checking for. + +What the gate cannot see is a column declaring its own collation +(``country COLLATE NOCASE``) on an otherwise byte-exact engine. Nothing in +SQLAlchemy's reflection or the engine specs exposes that, so there it remains +the owner's assumption, like ``p = T(mapped_col)`` itself. + +Negations are never safe because ``T`` need not be injective: +``lower(:value)`` with ``country != 'US'`` mirrors to ``region_key != 'us'``, +which wrongly excludes rows whose ``country`` is already lowercase ``'us'`` -- +rows the original filter *keeps*. + +Monotonicity is a property of the transform, not of the column's data type: +``hour(:value)``, ``date_format(:value, 'dd')`` and ``dayofweek(:value)`` are +all reasonable transforms on a ``TIMESTAMP`` column and none of them preserve +ordering. It is therefore declared by the owner, not inferred. + +Time grains +----------- +A grained filter -- drill-to-detail, mostly -- compares the *truncated* column, +``trunc(col) op v``, so the raw bounds it carries do not describe the rows it +keeps. A row in the final partial bucket satisfies ``trunc(col) < until`` while +``col < until`` excludes it. + +Which direction a grain rounds is not knowable from the duration alone, and the +obvious guess is wrong: ``WEEK_ENDING_SATURDAY`` rounds *forward* on Hive and +Presto, and Ocient's grains are ``ROUND``, i.e. to nearest. What every grain +does satisfy is a bound on the displacement:: + + |trunc(ts) - ts| < width(grain) + +so widening *both* bounds by one bucket width is no narrower than the real +predicate whichever way the grain rounds -- and it needs no monotonicity of +``trunc`` itself, which is what rescues Hive's oddly-anchored ``P1W``. Width is +a property of the grain, not of the engine, which is what makes this +maintainable; see `grain_bucket_width`. A grain whose SQL an operator supplied +has no known width, and does not mirror at all. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, cast, TYPE_CHECKING + +import numpy as np +import pandas as pd +import sqlalchemy as sa +from dateutil.relativedelta import relativedelta +from flask import current_app as app +from flask_babel import lazy_gettext as _ +from sqlalchemy.engine.interfaces import Dialect +from sqlalchemy.sql.elements import ColumnElement + +from superset.constants import LRU_CACHE_MAX_SIZE, TimeGrain +from superset.exceptions import ( + QueryClauseValidationException, + SupersetParseError, + SupersetSecurityException, +) +from superset.extensions import cache_manager, feature_flag_manager +from superset.sql.parse import SQLScript, SQLStatement +from superset.utils import core as utils, json +from superset.utils.core import FilterOperator + +if TYPE_CHECKING: + from superset.connectors.sqla.models import SqlaTable, TableColumn + from superset.db_engine_specs.base import BaseEngineSpec + from superset.models.core import Database + +logger = logging.getLogger(__name__) + +FEATURE_FLAG = "PARTITION_FILTER_MAPPING" + +#: Longest transform accepted anywhere. A transform is one SQL expression an +#: owner types by hand, so this is generous; the point is that it is bounded. +#: Every entry point enforces it -- the typed column field, the import schema +#: and the preview request -- because `is_transform_active` parses the stored +#: value on each Explore load, and an unbounded string would make that parse +#: the expensive part of rendering a chart. +MAX_TRANSFORM_LENGTH = 1024 + +#: Placeholder the owner writes in the transform, e.g. ``unix_timestamp(:value)``. +#: Matched with word boundaries so ``:values`` is not mistaken for it. +VALUE_PLACEHOLDER_RE = re.compile(r":value\b") + +#: Balanced Jinja blocks. The probe would render these in a different context +#: at a different time from the chart query, so they are rejected at save time. +JINJA_BLOCK_RE = re.compile(r"\{\{.*?\}\}|\{%.*?%\}|\{#.*?#\}", re.DOTALL) + +#: Substituted for ``:value`` before parsing -- sqlglot rejects a bare ``:value`` +#: on most dialects. Mirrors the ``_JINJA_BLOCK_RE`` -> ``NULL`` trick used by +#: ``validate_stored_expression``. +_PARSE_STANDIN = "NULL" + +#: Functions whose value depends on wall-clock time or randomness. The probe +#: runs in a different session at a different moment from the chart query and +#: its result is then cached, so any of these freezes a snapshot of probe time +#: into the emitted predicate. +NON_DETERMINISTIC_FUNCTIONS = { + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "NOW", + "RAND", + "RANDOM", + "UUID", +} + +#: Functions that mean "now" only in their zero-argument form. On Hive and +#: Impala ``unix_timestamp()`` is the current time while ``unix_timestamp(x)`` +#: -- the canonical transform for this feature -- is pure. +NON_DETERMINISTIC_WHEN_NILADIC = {"UNIX_TIMESTAMP"} + +#: Safe for any function ``T``. +MIRRORABLE_ALWAYS = {FilterOperator.EQUALS, FilterOperator.IN} + +#: Safe only when ``T`` preserves ordering. +MIRRORABLE_IF_MONOTONIC = { + FilterOperator.GREATER_THAN, + FilterOperator.GREATER_THAN_OR_EQUALS, + FilterOperator.LESS_THAN, + FilterOperator.LESS_THAN_OR_EQUALS, + FilterOperator.TEMPORAL_RANGE, +} + + +#: Every operator that can be mirrored under *some* transform. The preview +#: endpoint accepts these; whether a given one actually mirrors still depends on +#: the monotonicity declaration. +MIRRORABLE_OPERATORS = MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC + +#: Operators the preview endpoint can construct. `TEMPORAL_RANGE` is mirrored by +#: the query path, but as *two* bounds that `_collect_partition_mirror_range` +#: decomposes a time range into -- the preview request carries sample values, not +#: a since/until pair, so there is no range for it to build. The editor never +#: asks for one either: `previewOperatorFor` sends `>=`, `=` or `IN`. +PREVIEWABLE_OPERATORS = MIRRORABLE_OPERATORS - {FilterOperator.TEMPORAL_RANGE} + + +def mirrorable_operators( + is_monotonic: bool, *, equality_is_safe: bool = True +) -> set[FilterOperator]: + """ + The operators whose predicates may be mirrored onto the partition column. + + :param is_monotonic: whether the owner declared the transform + order-preserving + :param equality_is_safe: whether the engine's ``=`` on the mapped column + compares the way ``T`` was written against -- see + `equality_mirrors_safely` + + Both call sites -- the query path through `PartitionMapping.mirrors` and the + Explore indicator through `partition_filter_mapping_summary` -- come through + here, which is what keeps the glyph from promising pruning the SQL will not + do. + """ + operators = MIRRORABLE_ALWAYS if equality_is_safe else set() + if is_monotonic: + operators = operators | MIRRORABLE_IF_MONOTONIC + return set(operators) + + +def equality_mirrors_safely( + column: TableColumn | None, + db_engine_spec: type[BaseEngineSpec], +) -> bool: + """ + Whether ``col = v`` on this column implies ``T(col) = T(v)`` in the engine. + + The operator matrix calls equality safe "for any function ``T``" because + ``T`` is a function, so ``col = v`` gives ``T(col) = T(v)``. That reasons + about *value* equality while the engine reasons about *SQL* equality, and + the two part company under any comparison that is not byte-exact: with a + case-insensitive collation a stored ``country`` of ``'us'`` satisfies a + filter for ``'US'``, but the mirror ``region_key = hex('US')`` excludes the + row and the chart loses it with no indication why. Trailing-space-insensitive + ``CHAR`` comparison says the same thing about padding. + + Only string columns are at risk; numeric and temporal comparison is exact + everywhere. And only equality: range mirroring already requires the owner + to declare ``T`` order-preserving *with respect to the column's own order*, + which is an assertion about the same comparison semantics this is checking + for -- so the declaration covers it where equality has nothing to cover it. + + What this cannot see is a column that declares its own collation + (``country COLLATE NOCASE``) on an otherwise byte-exact engine. Neither + SQLAlchemy's reflection nor the engine specs expose it, so on such a column + the assumption stays where ``p = T(mapped_col)`` already is: with the owner. + """ + if column is None: + return True + try: + is_string = column.type_generic == utils.GenericDataType.STRING + except Exception: # pylint: disable=broad-except # noqa: BLE001 + # An unresolvable type is treated as a string: the gate exists to stop + # a silent wrong answer, so it fails closed. + return False + if not is_string: + return True + return bool(db_engine_spec.binary_string_comparison) + + +#: How wide one bucket of each built-in time grain is. +#: +#: Keyed on the ISO duration a filter carries in its ``grain``. Written out +#: rather than parsed: four of the week grains are ISO *intervals* with an +#: anchor (``P1W/1970-01-03T00:00:00Z``) that ``isodate.parse_duration`` +#: rejects outright, and ``PT0.5H`` / ``P0.25Y`` are fractional. A literal +#: table is also the thing a reviewer can check a line at a time. +#: +#: ``relativedelta`` rather than ``timedelta`` so the calendar grains stay +#: calendar arithmetic: a month is not 30 days. +GRAIN_BUCKET_WIDTHS: dict[str, relativedelta] = { + TimeGrain.SECOND: relativedelta(seconds=1), + TimeGrain.FIVE_SECONDS: relativedelta(seconds=5), + TimeGrain.THIRTY_SECONDS: relativedelta(seconds=30), + TimeGrain.MINUTE: relativedelta(minutes=1), + TimeGrain.FIVE_MINUTES: relativedelta(minutes=5), + TimeGrain.TEN_MINUTES: relativedelta(minutes=10), + TimeGrain.FIFTEEN_MINUTES: relativedelta(minutes=15), + TimeGrain.THIRTY_MINUTES: relativedelta(minutes=30), + TimeGrain.HALF_HOUR: relativedelta(minutes=30), + TimeGrain.HOUR: relativedelta(hours=1), + TimeGrain.SIX_HOURS: relativedelta(hours=6), + TimeGrain.DAY: relativedelta(days=1), + TimeGrain.WEEK: relativedelta(days=7), + TimeGrain.WEEK_STARTING_SUNDAY: relativedelta(days=7), + TimeGrain.WEEK_STARTING_MONDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SATURDAY: relativedelta(days=7), + TimeGrain.WEEK_ENDING_SUNDAY: relativedelta(days=7), + TimeGrain.MONTH: relativedelta(months=1), + TimeGrain.QUARTER: relativedelta(months=3), + TimeGrain.QUARTER_YEAR: relativedelta(months=3), + TimeGrain.YEAR: relativedelta(years=1), +} + + +def grain_bucket_width(grain: str | None, engine: str) -> relativedelta | None: + """ + How far a grain's truncation can move a timestamp, or ``None`` if unknown. + + A grained filter compares the *truncated* column, so the raw bounds it + carries do not describe the rows it keeps. Widening both bounds by one + bucket recovers a predicate that is no narrower than the real one -- see + `_collect_partition_mirror_range` for the argument. That only works for a + grain whose bucket width Superset knows, which excludes anything an + operator supplied. + + :param grain: the ISO duration from the filter's ``grain`` + :param engine: the engine spec's ``engine``, to check per-engine overrides + """ + if not grain: + return None + + # An operator-declared grain carries whatever duration string they typed, + # and `TIME_GRAIN_ADDON_EXPRESSIONS` can redefine a *built-in* grain's SQL + # per engine -- `P1D` could be anything at all. Neither has a width we can + # claim to know, so both fall back to not mirroring. + if grain in app.config["TIME_GRAIN_ADDONS"]: + return None + if grain in app.config["TIME_GRAIN_ADDON_EXPRESSIONS"].get(engine, {}): + return None + + return GRAIN_BUCKET_WIDTHS.get(grain) + + +@dataclass(frozen=True) +class PartitionMapping: + """A resolved, usable partition filter mapping.""" + + partition_column: str + mapped_column: str + value_transform: str + is_monotonic: bool + #: Whether the engine compares the mapped column the way ``T`` was written + #: against. Resolved once where the column and the engine are both in hand, + #: rather than re-derived at each `mirrors` call. + equality_is_safe: bool = True + + def mirrors(self, operator: FilterOperator) -> bool: + return operator in mirrorable_operators( + self.is_monotonic, equality_is_safe=self.equality_is_safe + ) + + +def contains_value_placeholder(transform: str | None) -> bool: + """Whether the transform contains the ``:value`` placeholder.""" + return bool(transform) and VALUE_PLACEHOLDER_RE.search(transform or "") is not None + + +def contains_jinja(transform: str | None) -> bool: + """Whether the transform contains a balanced Jinja block.""" + return bool(transform) and JINJA_BLOCK_RE.search(transform or "") is not None + + +def parse_skeleton(transform: str) -> str: + """ + The transform with ``:value`` substituted out, ready for a SQL parser. + + ``sanitize_clause`` / sqlglot choke on a bare ``:value`` on most dialects, + so the placeholder is swapped for a benign literal first -- the same trick + ``validate_stored_expression`` uses for Jinja blocks. + """ + return VALUE_PLACEHOLDER_RE.sub(_PARSE_STANDIN, transform) + + +#: Prefix the transform is wrapped in before parsing. Its length is subtracted +#: from any reported column so positions refer to what the owner actually typed. +_SELECT_PREFIX = "SELECT " + + +def _parse_skeleton(transform: str, engine: str) -> SQLStatement | None: + """ + Parse ``SELECT <transform>`` with the placeholder substituted out. + + Returns ``None`` when the transform does not parse. + """ + try: + return SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError: + return None + + +def parse_error_detail(transform: str, engine: str) -> str | None: + """ + Where the parser gave up on the transform. + + Returns ``None`` when it parses, or when the parser offered no position. + Note that sqlglot parses unknown functions happily -- a misspelled function + name is not a parse error, it is an engine error, and surfaces only when the + transform is evaluated. + + The parser's own ``highlight`` is deliberately dropped: it would name the + ``NULL`` we substituted for ``:value``, which is not a token the owner typed. + """ + try: + SQLStatement(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError as ex: + column = (ex.error.extra or {}).get("column") + if not isinstance(column, int): + return None + return str( + _( + "syntax error at position %(position)d.", + position=_position_in_transform(transform, column), + ) + ) + return None + + +def _position_in_transform(transform: str, parsed_column: int) -> int: + """ + Map a column in the parsed skeleton back to the transform as typed. + + Two substitutions stand between them: the ``SELECT`` prefix, and every + ``:value`` that became a shorter ``NULL``. Without unwinding both, a + reported position drifts left by two characters per placeholder ahead of it + -- which is worst exactly where transforms usually break, at the end. + """ + position = max(parsed_column - len(_SELECT_PREFIX), 0) + shift = len(":value") - len(_PARSE_STANDIN) + preceding = sum( + 1 + for index, match in enumerate(VALUE_PLACEHOLDER_RE.finditer(transform)) + if match.start() - index * shift < position + ) + return position + preceding * shift + + +def is_parseable(transform: str | None, engine: str) -> bool: + """ + Whether the transform parses as a single select expression. + + "Single" is the load-bearing word. `SELECT lower(:value), 'x'` parses just + as happily as `SELECT lower(:value)`, but it returns two columns per input + value, and the probe reads one column per input -- so an `IN` filter would + get a predicate built from the wrong halves of the wrong rows. Rejecting the + list here is what lets the probe trust its own column count. + """ + if not transform or not transform.strip(): + return False + statement = _parse_skeleton(transform, engine) + return statement is not None and statement.count_select_expressions() == 1 + + +def is_bare_expression(transform: str | None, engine: str) -> bool: + """ + Whether the transform is a scalar expression and nothing more. + + `is_parseable` is not enough, because it counts only the projection. + ``secret || :value FROM vault`` holds exactly one select expression, so it + parses as "a single expression" -- and `build_probe_sql` then emits + ``SELECT secret || 'us' AS v0``, where its own alias is read as a table + alias on the FROM clause the transform smuggled in. The probe reads a + column the owner was never granted and hands it back through the predicate + the preview panel renders. + + A transform is a scalar function of ``:value`` by definition -- the whole + feature rests on ``partition_col = T(mapped_col)``, which only type-checks + for scalar ``T`` -- so there is no legitimate transform with a clause of + its own, and none with a sub-query either. + """ + if not transform or not transform.strip(): + return False + statement = _parse_skeleton(transform, engine) + return ( + statement is not None + and statement.is_bare_select_expression() + and not statement.has_subquery() + ) + + +def is_unfinished(transform: str | None, engine: str) -> bool: + """ + Whether the transform is not SQL yet, as opposed to the wrong SQL. + + The distinction matters because the two deserve opposite treatment. A + half-typed ``unix_timestamp(:value`` is what a text input produces on the + way to something valid: a PUT stores it and reports the mapping inactive, + so discarding it would mean an export could not round-trip the dataset it + came from. ``unix_timestamp(:value); DROP TABLE t`` is not on the way to + anything, and has to be refused wherever it is offered. + + Both fail `is_parseable`, and both fail to parse as a single *statement* -- + so neither of those tells them apart. Parsing as a *script* does: the + multi-statement form is two valid statements, while unfinished text is no + statement at all. + """ + if not transform or not transform.strip(): + return False + try: + SQLScript(f"{_SELECT_PREFIX}{parse_skeleton(transform)}", engine) + except SupersetParseError: + return True + return False + + +def find_non_deterministic_functions(transform: str, engine: str) -> set[str]: + """ + Names of non-deterministic functions the transform calls. + + ``UNIX_TIMESTAMP`` is only reported in its zero-argument form, which means + "now" on Hive and Impala; the one-argument form is the canonical temporal + transform and stays allowed. + """ + statement = _parse_skeleton(transform, engine) + if statement is None: + return set() + + found = { + name + for name in NON_DETERMINISTIC_FUNCTIONS + if statement.check_functions_present({name}) + } + return found | _find_niladic_calls(statement) + + +def _find_niladic_calls(statement: SQLStatement) -> set[str]: + """ + Names from ``NON_DETERMINISTIC_WHEN_NILADIC`` called with no arguments. + + Note some dialects resolve the zero-argument form themselves -- Hive parses + ``unix_timestamp()`` straight to ``CURRENT_TIMESTAMP`` -- in which case the + name-based check above has already caught it. This is the backstop for the + dialects that do not. + """ + return NON_DETERMINISTIC_WHEN_NILADIC & statement.get_niladic_functions() + + +def resolve_partition_mapping(datasource: SqlaTable) -> PartitionMapping | None: + """ + Resolve the dataset's mapping, or ``None`` when nothing may be mirrored. + + Every bail-out here is defensive as well as functional: save-time validation + rejects most of these, but rows predating the validation can still violate + the invariants, and a column sync can invalidate a mapping that was fine + when it was written. + + The transform gate is `is_transform_active`, the same function the Explore + indicator reads, which is `validate_transform` with the messages discarded. + Anything narrower here would be a second, weaker statement of the same rule: + a transform calling `now()` that reached storage without passing + `UpdateDatasetCommand` -- through import, or through a bundle written by hand + -- would be reported inactive by the editor and still mirrored by this + function, freezing a snapshot of probe time into the predicate with nothing + on screen to say so. + """ + if not feature_flag_manager.is_feature_enabled(FEATURE_FLAG): + return None + + partition_column = getattr(datasource, "partition_column", None) + if not partition_column: + return None + + columns_by_name = {column.column_name: column for column in datasource.columns} + if partition_column not in columns_by_name: + # The partition column was dropped by a column sync or at the source. + return None + + mapped_column_name = ( + getattr(datasource, "partition_mapped_column", None) or datasource.main_dttm_col + ) + if not mapped_column_name or mapped_column_name not in columns_by_name: + return None + + if mapped_column_name == partition_column: + # Self-mapping: the mirrored predicate would duplicate the original. + return None + + mapped_column = columns_by_name[mapped_column_name] + transform = getattr(mapped_column, "partition_value_transform", None) + if not is_transform_active(transform, datasource.database.backend): + return None + + if has_active_advanced_data_type(mapped_column): + # `translate_filter` builds its own predicate shape from *translated* + # values, so the `(operator, value)` pair the operator matrix reasons + # about does not exist and mirroring would apply the wrong values. + return None + + return PartitionMapping( + partition_column=str(partition_column), + mapped_column=str(mapped_column_name), + value_transform=cast(str, transform), + is_monotonic=bool( + getattr(mapped_column, "partition_transform_is_monotonic", False) + ), + equality_is_safe=equality_mirrors_safely( + mapped_column, datasource.database.db_engine_spec + ), + ) + + +def has_active_advanced_data_type(column: TableColumn) -> bool: + """ + Whether the column's advanced data type is configured and switched on. + + Such a column is never mirrored: ``translate_filter`` builds its own + predicate shape from translated values, so the ``(operator, value)`` pair + the operator matrix reasons about does not exist. + """ + advanced_data_type = getattr(column, "advanced_data_type", None) + if not advanced_data_type: + return False + if not feature_flag_manager.is_feature_enabled("ENABLE_ADVANCED_DATA_TYPES"): + return False + return advanced_data_type in app.config.get("ADVANCED_DATA_TYPES", {}) + + +def _denylist_engine_key(database: "Database") -> str: + """ + The name ``DISALLOWED_SQL_*`` is keyed by for this database. + + The engine spec's own ``engine`` first, because that is what every other + denylist gate uses (`_raise_for_disallowed_sql`, `sql_lab`) and what an + operator writing the config reads in the documentation; a spec covering + several SQLAlchemy backends through ``engine_aliases`` reports the one + canonical name under which its entry is written. + + Falling back to the URL's backend, because resolving the spec loads the + SQLAlchemy dialect entrypoint, which imports the driver package. On a + deployment missing an optional driver that raises rather than merely being + absent, asking for the spec here would turn a dataset import into a hard + failure -- and for every engine without aliases the two names are the same + string anyway. + """ + try: + return database.db_engine_spec.engine + except Exception: # pylint: disable=broad-except # noqa: BLE001 + return database.backend + + +def stored_expression_error( + database: "Database", + catalog: str | None, + schema: str | None, + transform: str, +) -> str | None: + """ + Why this transform may not be stored or run, if there is a reason. + + Stricter than the policy a general stored expression goes through, because + a value transform is a narrower thing: `build_probe_sql` binds only + `:value` and splices the rest of the transform in as SQL text, which the + engine then executes, so an ungated transform is arbitrary SQL. A dataset + editor without SQL Lab could store `(SELECT secret FROM protected_table + LIMIT 1) || :value`, or `secret || :value FROM protected_table`, and read + the answer back out of the emitted predicate the preview panel renders. + + So the transform must be a bare scalar expression: no clause of its own and + no sub-query, whatever `ALLOW_ADHOC_SUBQUERY` says. That is not a + restriction the feature pays for -- it rests on + ``partition_col = T(mapped_col)``, which only type-checks for scalar ``T`` + -- and it is what makes the table denylist and the RLS rewrite moot here, + since neither has a table reference left to govern. + + Returns the engine-agnostic reason as a string rather than raising, + because its four callers disagree about what to do with it: the preview + reports it at the field, a PUT and the legacy datasource save refuse the + write, the importer drops the transform and keeps the dataset, and the + probe simply declines to run. Raising would make three of those four + write a `try` around a question. + + Imported inside the function rather than at module scope: + `connectors.sqla.models` imports this module, so the dependency only runs + one way at import time. + """ + from superset.connectors.sqla.models import ( # pylint: disable=import-outside-toplevel,cyclic-import + validate_stored_expression, + ) + + statement = _parse_skeleton(transform, database.backend) + if statement is None: + return str( + _("A partition value transform must parse as a single SQL expression.") + ) + + # The shape gate again, even though `validate_transform` blocks the same + # thing at save time. This is the last door before the transform becomes SQL + # an engine runs, and the only door a row written by an earlier release -- + # or by an importer that ran with the feature flag off -- still passes + # through. + if not statement.is_bare_select_expression(): + return str( + _( + "A partition value transform must be a single SQL expression, " + "with no FROM, WHERE or other clause." + ) + ) + + # Unconditionally, not under `ALLOW_ADHOC_SUBQUERY`. A transform is a scalar Review Comment: Fixed. The `schema_to_xml` / `query_to_xml` / `table_to_xml` family (and their `_xmlschema` and `_xml_and_xmlschema` spellings) are in the default `DISALLOWED_SQL_FUNCTIONS` for PostgreSQL, and the denylist is enforced on save rather than only at probe time. -- 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]
