You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by hu...@apache.org on 2026/09/30 03:42:38 UTC
(superset) 01/01: feat(datasets): partition filter mapping, stored in tables.extra
This is an automated email from the ASF dual-hosted git repository.
hughhhh pushed a commit to branch hughhhh/tables-extra-no-migration
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 79bbbd76b33365e71fc057184e6666168dba1b3e
Author: Hugh A Miles II <hu...@gmail.com>
AuthorDate: Tue Sep 29 23:39:30 2026 -0400
feat(datasets): partition filter mapping, stored in tables.extra
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, and the only workaround today is hand-writing the predicate as
custom SQL in a virtual dataset, which pushes a performance concern onto every
chart author and takes the dataset out of the physical/syncable path.
A dataset owner names one partition column, one business column whose filters
are mirrored onto it, a value transform (a SQL expression containing `:value`),
and whether that transform preserves ordering. Superset then appends an
equivalent predicate on the partition column to every query: chart authors
change nothing and queries prune. Behind `PARTITION_FILTER_MAPPING`, off by
default.
This consolidates what was a four-PR stack (#43757, #43758, #43759, #43760) and
replaces its four new database columns with a storage interface, so the feature
ships with no Alembic migration and the columns land later as an additive
follow-up.
Storage
-------
The four values live in `tables.extra` under `partition_filter_mapping`, behind
`PartitionMappingStore`: `load`/`save`, one adapter today
(`ExtraJsonPartitionMappingStore`) and one that activates with the migration,
selected by `PARTITION_MAPPING_STORE`. The model, the DAO, validation, the query
rewriter and the preview endpoint are written against the port and never learn
which store answered.
`tables.extra` is a reasonable home rather than only an expedient one: it is
already in `SqlaTable.export_fields` and already one of `ExportDatasetsCommand`'s
`JSON_KEYS`, so a mapping travels through import/export with no bundle-format
change; Continuum already versions the column, so dataset version history and
restore work untouched; and `certification` and `warning_markdown` already prove
the pattern of surfacing typed settings out of that blob through model
properties. It has to be the dataset-level blob, though, never
`table_columns.extra` -- the editor's `buildExtraJsonObject` rebuilds every
column's `extra` from two keys on each save.
Two details worth a look:
- **Transforms are keyed by column name, not flattened onto the mapping.**
Flattening is unsound: an owner who declares a transform on `A` and later
re-points `main_dttm_col` to `B` would have `B`'s values mirrored through
`A`'s expression. Keyed, `B` has no transform and the mapping goes inactive --
which is what four real columns would do.
- **The model's four attributes are read-only properties, and `DatasetDAO`
extracts the mapping from a payload before anything is applied and writes it
after everything is.** The editor's PUT carries both the typed
`partition_column` field and a verbatim copy of `extra`, and under this store
those are the same bytes -- so a `setattr` loop that happens to apply `extra`
last silently discards the typed field, and which order that is depends on
marshmallow's field ordering. Read-only properties make that loss impossible
rather than merely unlikely; `set_partition_mapping` is the one supported
write.
Two bugs in the original stack, fixed here
------------------------------------------
- `_transform_is_usable` omitted the non-deterministic denylist, so a transform
calling `now()` that reached storage without passing `UpdateDatasetCommand` --
via create, via import, or via the free-text Extra box, which every dataset
owner can reach -- was mirrored at query time and its result cached. Replaced
by `is_transform_active`, which is `validate_transform` with the messages
discarded and so cannot drift from the save path.
- `columns.partition_value_transform` was missing from the dataset GET payload
while `DatasourceModal` PUT it back for every column, so every dataset-editor
save cleared the transform.
Also new: `partition_column` is a physical column name, so it joins
`DASHBOARD_DATASET_INACCESSIBLE_FIELDS`; `DatasetPutSchema.extra` gains a
structural validator for the mapping key, because the Extra box is a second door
into the same storage and a mistyped value there would otherwise be silently
ignored; and `SqlaTable.update_from_object` preserves the mapping when a legacy
payload omits `extra`.
No authoring UI is added, as in the original stack: a mapping is configured
through the REST API or the Extra box, and the preview endpoint ships available
but uncalled.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
.../configuration/partition-filter-mapping.mdx | 224 +++++
docs/static/feature-flags.json | 6 +
docs/static/resources/openapi.json | 165 ++++
.../superset-ui-chart-controls/src/types.ts | 12 +
.../Datasource/DatasourceModal/index.tsx | 6 +
.../DatasourceEditor/DatasourceEditor.tsx | 39 +-
.../src/components/Datasource/utils/index.ts | 43 +
.../Datasource/utils/partitionMapping.test.ts | 118 +++
superset-frontend/src/features/datasets/types.ts | 6 +
superset/commands/dataset/update.py | 120 +++
superset/config.py | 26 +
superset/connectors/sqla/models.py | 160 ++++
superset/connectors/sqla/partition_mapping.py | 977 ++++++++++++++++++++
.../connectors/sqla/partition_mapping_storage.py | 606 +++++++++++++
superset/daos/dataset.py | 117 ++-
superset/dashboards/api.py | 6 +
superset/dashboards/schemas.py | 3 +
superset/datasets/api.py | 160 +++-
superset/datasets/schemas.py | 150 +++-
superset/mcp_service/common/schema_discovery.py | 19 +
superset/models/helpers.py | 288 +++++-
superset/sql/parse.py | 38 +
superset/superset_typing.py | 9 +
tests/integration_tests/datasets/api_tests.py | 1 +
tests/unit_tests/commands/dataset/update_test.py | 141 +++
.../sqla/partition_mapping_storage_test.py | 544 ++++++++++++
.../connectors/sqla/partition_mapping_test.py | 984 +++++++++++++++++++++
tests/unit_tests/dao/dataset_test.py | 197 +++++
tests/unit_tests/dashboards/datasets_test.py | 3 +
.../datasets/partition_mapping_preview_test.py | 406 +++++++++
.../partition_mapping_serialization_test.py | 504 +++++++++++
.../unit_tests/models/partition_mirroring_test.py | 908 +++++++++++++++++++
tests/unit_tests/sql/parse_tests.py | 35 +
33 files changed, 6980 insertions(+), 41 deletions(-)
diff --git a/docs/admin_docs/configuration/partition-filter-mapping.mdx b/docs/admin_docs/configuration/partition-filter-mapping.mdx
new file mode 100644
index 00000000000..9025001a3fb
--- /dev/null
+++ b/docs/admin_docs/configuration/partition-filter-mapping.mdx
@@ -0,0 +1,224 @@
+---
+title: Partition Filter Mapping
+hide_title: true
+sidebar_position: 15
+version: 1
+---
+
+<!--
+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
+
+Tables on Hadoop-family engines are often partitioned on a *technical* column — an epoch
+integer, a lowercased region key — that no analyst would ever filter on. Unless a query
+carries a predicate on that column, the engine scans every partition.
+
+Partition filter mapping makes this a dataset setting instead of a per-chart chore. A
+dataset owner names the partition column, the business column whose filters should be
+mirrored onto it, and a value transform. Superset then appends an equivalent predicate on
+the partition column to every query. Chart authors change nothing; queries prune.
+
+:::caution Experimental
+This feature is behind the `PARTITION_FILTER_MAPPING` feature flag and is off by default.
+:::
+
+## Enabling it
+
+```python
+FEATURE_FLAGS = {
+ "PARTITION_FILTER_MAPPING": True,
+}
+```
+
+Configure it as a **static boolean**. `FEATURE_FLAGS` also accepts per-request callables,
+but a flag that resolves differently per user or tenant would let a user with the feature
+off read a cached chart result that was produced from pruned SQL by a user with it on.
+
+## Configuring a mapping
+
+:::note No editor controls yet
+There are no form fields for this in the dataset editor. A mapping is configured through the
+REST API or by hand in the editor's **Extra** JSON box, both described below.
+:::
+
+A mapping names four things:
+
+- the **partition column**, the physical column the engine partitions on;
+- the **mapped column**, whose filters are mirrored onto it. Leave it unset and the mapping
+ follows the dataset's default datetime column, so re-pointing that column moves the mapping
+ with it; set it explicitly to pin the mapping to a different column;
+- the **value transform**, a SQL expression containing a `:value` placeholder, which stands
+ for the filter value being mirrored;
+- whether that transform **preserves ordering**, which gates range mirroring.
+
+Through the API:
+
+```bash
+curl -X PUT "$SUPERSET/api/v1/dataset/$DATASET_ID" \
+ -H 'Content-Type: application/json' \
+ -d '{
+ "partition_column": "dt_epoch",
+ "columns": [
+ {
+ "id": 42,
+ "column_name": "event_time",
+ "partition_value_transform": "unix_timestamp(:value)",
+ "partition_transform_is_monotonic": true
+ }
+ ]
+ }'
+```
+
+Note the `columns` list replaces the dataset's columns, so send every column you want to
+keep — the same rule that already applies to any dataset PUT.
+
+| Mapped column | Partition column | Transform |
+|---|---|---|
+| `event_time` (`TIMESTAMP`) | `dt_epoch` (`BIGINT`) | `unix_timestamp(:value)` |
+| `country` (`VARCHAR`) | `region_key` (`VARCHAR`) | `lower(:value)` |
+
+A filter of `event_time >= '2026-01-01'` then adds `dt_epoch >= 1767225600` to the query.
+The added predicate is an ordinary `WHERE` clause and shows up in **View query**.
+
+### Transform preserves ordering
+
+Range filters — including the Explore time range, the most important case — are only
+mirrored when the transform is declared order-preserving
+(`partition_transform_is_monotonic`).
+
+Monotonicity is a property of the *transform*, not of the column's data type.
+`unix_timestamp(:value)` preserves ordering. `hour(:value)`,
+`date_format(:value, 'dd')` and `dayofweek(:value)` are all perfectly reasonable
+partition transforms on a `TIMESTAMP` column and none of them do: `hour('2026-01-01 23:00')`
+is greater than `hour('2026-01-02 01:00')` even though the first instant is earlier. Mirroring
+a range through one of those would silently return wrong numbers, so Superset asks you to
+declare it rather than guessing.
+
+Left undeclared, `=` and `IN` filters still mirror; ranges do not.
+
+## What is and isn't mirrored
+
+| Filter | Mirrored |
+|---|---|
+| `=`, `IN` | Always |
+| `>`, `>=`, `<`, `<=`, time ranges | Only when the transform preserves ordering |
+| `!=`, `NOT IN`, `LIKE`, `ILIKE`, `IS NULL`, `IS TRUE` | Never |
+
+Negations are never safe. A transform need not be injective: `lower(:value)` with
+`country != 'US'` would mirror to `region_key != 'us'`, which excludes rows whose `country`
+is already lowercase `'us'` — rows the original filter *keeps*.
+
+Known gaps, all of which are out of scope rather than bugs:
+
+- **Filter-value dropdowns do not prune.** Populating a filter's value list runs its own
+ `SELECT DISTINCT`, which never goes through the chart query path. There is no filter to
+ mirror from.
+- **Row-level security predicates do not mirror.** They are stored as raw SQL and appended
+ downstream of the structured filters.
+- **Custom SQL `WHERE` clauses do not mirror**, for the same reason.
+- **Columns with an active advanced data type do not mirror.** Those build their own
+ predicate shape from translated values, so there is no operator/value pair to mirror.
+- Dashboard native filters and cross-filters *do* mirror — they arrive as ordinary filters —
+ they just carry no visual indicator in the filter bar.
+
+## The assumption this rests on
+
+Superset emits a predicate on the partition column that stands in for one on the mapped
+column. That substitution is only valid if, for every row in the table:
+
+```
+partition_column = <transform>(mapped_column)
+```
+
+**Superset cannot verify this.** 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 than the transform resolves, mirrored predicates silently drop real rows
+and charts show quietly wrong numbers. Confirm the invariant with whoever owns the pipeline
+before enabling a mapping on a production dataset.
+
+Related: a predicate like `dt_epoch >= X` also drops rows where `dt_epoch` is `NULL`.
+Partition keys in Hive and Impala are non-null by construction, so this is accepted rather
+than defended against.
+
+## How the transform is evaluated
+
+The transform is evaluated against the engine — pinned to the dataset's database, catalog
+and schema — and the result is emitted as a literal constant. Results are cached
+(`PARTITION_TRANSFORM_PROBE_CACHE_TIMEOUT`, 24 hours by default), which matters because this
+adds a round trip to the chart query path. Day-aligned ranges like "Last month" hit the cache
+constantly; second-granularity relative ranges like "Last 24 hours" essentially never do.
+
+If the evaluation fails for any reason, no predicate is added: the query still runs and is
+still correct, it just scans more partitions.
+
+Because the evaluation happens in a **different session** from the chart query, transforms
+that call non-deterministic functions are rejected when you save. That includes `now()`,
+`current_date`, `current_timestamp`, `rand()` and the zero-argument `unix_timestamp()`, which
+means "now" on Hive and Impala. The one-argument `unix_timestamp(:value)` is fine.
+
+Session-dependent behaviour that Superset cannot detect is still your responsibility:
+`unix_timestamp()` is timezone-dependent on Hive and Impala, so if the evaluating session and
+the query session resolve different timezones the emitted bounds will disagree with the
+timestamp bounds they mirror. Prefer explicitly-anchored transforms.
+
+Jinja templating is not supported in a transform. The template would render in a different
+context and at a different time from the chart query.
+
+## Related configuration
+
+| Setting | Default | Purpose |
+|---|---|---|
+| `PARTITION_TRANSFORM_PROBE_CACHE_TIMEOUT` | 24 hours | How long an evaluated transform stays cached |
+| `PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT` | 30 | Per-user, per-dataset preview requests per minute; the preview endpoint runs a real query |
+| `PARTITION_MAPPING_STORE` | `"extra"` | Where the mapping is persisted. See below. |
+
+## Where a mapping is stored
+
+A mapping lives in the dataset's `extra` JSON, under `partition_filter_mapping`, alongside
+`certification`, `warning_markdown` and `timezone`:
+
+```json
+{
+ "partition_filter_mapping": {
+ "partition_column": "dt_epoch",
+ "mapped_column": null,
+ "column_transforms": {
+ "event_time": {
+ "value_transform": "unix_timestamp(:value)",
+ "is_monotonic": true
+ }
+ }
+ }
+}
+```
+
+That is the same box the dataset editor's **Extra** field edits, so a mapping can be written
+by hand there. Editing it by hand skips the typed API fields but not validation: the same
+checks run on whatever the two doors between them produce. A mapping that is structurally
+wrong — a column name that isn't a string, an `is_monotonic` that isn't a boolean — is
+rejected on save rather than silently ignored.
+
+Storing it here means it needs no schema change, and it also means a mapping travels with the
+dataset through **import/export** and appears in **dataset version history** for free, since
+`extra` already does both.
+
+`PARTITION_MAPPING_STORE` exists for the migration that moves these four values into real
+database columns. Leave it at `"extra"`; setting `"columns"` before that migration has been
+applied raises rather than silently reading columns that are not there.
diff --git a/docs/static/feature-flags.json b/docs/static/feature-flags.json
index 97c49320e79..5afac50e88b 100644
--- a/docs/static/feature-flags.json
+++ b/docs/static/feature-flags.json
@@ -81,6 +81,12 @@
"lifecycle": "development",
"description": "Try to optimize SQL queries \u2014 for now only predicate pushdown is supported"
},
+ {
+ "name": "PARTITION_FILTER_MAPPING",
+ "default": false,
+ "lifecycle": "development",
+ "description": "Configure and validate a mapping from a dataset's business column onto its physical partition column, so engines that require an explicit partition predicate can prune. For now this only stores and validates the mapping; no query consumer emits the mirrored predicate yet."
+ },
{
"name": "PRESTO_EXPAND_DATA",
"default": false,
diff --git a/docs/static/resources/openapi.json b/docs/static/resources/openapi.json
index 5ea742a07eb..e4c1bf1c733 100644
--- a/docs/static/resources/openapi.json
+++ b/docs/static/resources/openapi.json
@@ -3941,6 +3941,16 @@
"parent": {
"type": "object"
},
+ "partition_column": {
+ "type": "string"
+ },
+ "partition_filter_mapping": {
+ "nullable": true,
+ "type": "object"
+ },
+ "partition_mapped_column": {
+ "type": "string"
+ },
"perm": {
"type": "string"
},
@@ -5877,6 +5887,15 @@
"nullable": true,
"type": "boolean"
},
+ "partition_transform_is_monotonic": {
+ "nullable": true,
+ "type": "boolean"
+ },
+ "partition_value_transform": {
+ "description": "SQL expression containing a :value placeholder. Filters on this column are mirrored onto the dataset's partition column with the value passed through this transform.",
+ "nullable": true,
+ "type": "string"
+ },
"python_date_format": {
"maxLength": 255,
"minLength": 1,
@@ -6350,6 +6369,12 @@
"order_by_choices": {
"readOnly": true
},
+ "partition_column": {
+ "readOnly": true
+ },
+ "partition_mapped_column": {
+ "readOnly": true
+ },
"schema": {
"maxLength": 255,
"nullable": true,
@@ -6578,6 +6603,12 @@
"nullable": true,
"type": "boolean"
},
+ "partition_transform_is_monotonic": {
+ "readOnly": true
+ },
+ "partition_value_transform": {
+ "readOnly": true
+ },
"python_date_format": {
"maxLength": 255,
"nullable": true,
@@ -6818,6 +6849,18 @@
"default": false,
"type": "boolean"
},
+ "partition_column": {
+ "maxLength": 250,
+ "minLength": 0,
+ "nullable": true,
+ "type": "string"
+ },
+ "partition_mapped_column": {
+ "maxLength": 250,
+ "minLength": 0,
+ "nullable": true,
+ "type": "string"
+ },
"schema": {
"maxLength": 250,
"minLength": 0,
@@ -6944,6 +6987,18 @@
"nullable": true,
"type": "integer"
},
+ "partition_column": {
+ "maxLength": 250,
+ "minLength": 0,
+ "nullable": true,
+ "type": "string"
+ },
+ "partition_mapped_column": {
+ "maxLength": 250,
+ "minLength": 0,
+ "nullable": true,
+ "type": "string"
+ },
"schema": {
"maxLength": 255,
"minLength": 0,
@@ -7864,6 +7919,35 @@
},
"type": "object"
},
+ "PartitionMappingPreviewSchema": {
+ "properties": {
+ "mapped_column": {
+ "description": "Column whose filters would be mirrored",
+ "maxLength": 250,
+ "minLength": 1,
+ "type": "string"
+ },
+ "sample_value": {
+ "description": "Value to evaluate the transform at",
+ "maxLength": 250,
+ "minLength": 1,
+ "type": "string"
+ },
+ "value_transform": {
+ "description": "SQL expression containing a :value placeholder",
+ "maxLength": 1024,
+ "minLength": 1,
+ "nullable": true,
+ "type": "string"
+ }
+ },
+ "required": [
+ "mapped_column",
+ "sample_value",
+ "value_transform"
+ ],
+ "type": "object"
+ },
"PermissionApi.get": {
"properties": {
"id": {
@@ -24566,6 +24650,87 @@
]
}
},
+ "/api/v1/dataset/{pk}/partition_mapping/preview/": {
+ "post": {
+ "description": "Evaluate a partition value transform at a sample value and return the predicate that would be appended to queries. Parse and denylist checks run before anything reaches the engine, so an unparseable transform costs no query.",
+ "parameters": [
+ {
+ "description": "The dataset ID",
+ "in": "path",
+ "name": "pk",
+ "required": true,
+ "schema": {
+ "type": "integer"
+ }
+ }
+ ],
+ "requestBody": {
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/PartitionMappingPreviewSchema"
+ }
+ }
+ },
+ "required": true
+ },
+ "responses": {
+ "200": {
+ "content": {
+ "application/json": {
+ "schema": {
+ "properties": {
+ "result": {
+ "properties": {
+ "emitted_predicate": {
+ "type": "string"
+ },
+ "error": {
+ "type": "string"
+ },
+ "valid": {
+ "type": "boolean"
+ }
+ },
+ "type": "object"
+ }
+ },
+ "type": "object"
+ }
+ }
+ },
+ "description": "Preview result"
+ },
+ "400": {
+ "$ref": "#/components/responses/400"
+ },
+ "401": {
+ "$ref": "#/components/responses/401"
+ },
+ "404": {
+ "$ref": "#/components/responses/404"
+ },
+ "429": {
+ "$ref": "#/components/responses/400"
+ },
+ "500": {
+ "$ref": "#/components/responses/500"
+ }
+ },
+ "security": [
+ {
+ "jwt": []
+ },
+ {
+ "api_key": []
+ }
+ ],
+ "summary": "Preview a partition filter mapping",
+ "tags": [
+ "Datasets"
+ ]
+ }
+ },
"/api/v1/dataset/{pk}/refresh": {
"put": {
"parameters": [
diff --git a/superset-frontend/packages/superset-ui-chart-controls/src/types.ts b/superset-frontend/packages/superset-ui-chart-controls/src/types.ts
index 2e874cc8324..324b318bdde 100644
--- a/superset-frontend/packages/superset-ui-chart-controls/src/types.ts
+++ b/superset-frontend/packages/superset-ui-chart-controls/src/types.ts
@@ -97,6 +97,12 @@ export interface Dataset {
database?: Record<string, unknown>;
normalize_columns?: boolean;
always_filter_main_dttm?: boolean;
+ partition_column?: string | null;
+ partition_mapped_column?: string | null;
+ // Self-contained summary for the partition pruning indicator. Kept separate
+ // from `columns` because the dashboard payload prunes columns no chart
+ // references, and the partition column is typically referenced by none.
+ partition_filter_mapping?: PartitionFilterMapping | null;
extra?: object | string;
/**
* Stable string values of the features a semantic view's provider declares.
@@ -105,6 +111,12 @@ export interface Dataset {
semantic_view_features?: string[];
}
+export interface PartitionFilterMapping {
+ partition_column: string;
+ mapped_column: string | null;
+ active: boolean;
+}
+
export interface ControlPanelState {
slice: {
slice_id: number;
diff --git a/superset-frontend/src/components/Datasource/DatasourceModal/index.tsx b/superset-frontend/src/components/Datasource/DatasourceModal/index.tsx
index d70a1ad2c34..e296d3c1359 100644
--- a/superset-frontend/src/components/Datasource/DatasourceModal/index.tsx
+++ b/superset-frontend/src/components/Datasource/DatasourceModal/index.tsx
@@ -183,6 +183,8 @@ const DatasourceModal: FunctionComponent<DatasourceModalProps> = ({
currency_code_column: datasource.currency_code_column ?? null,
normalize_columns: datasource.normalize_columns,
always_filter_main_dttm: datasource.always_filter_main_dttm,
+ partition_column: datasource.partition_column ?? null,
+ partition_mapped_column: datasource.partition_mapped_column ?? null,
offset: datasource.offset,
default_endpoint: datasource.default_endpoint,
cache_timeout:
@@ -231,6 +233,10 @@ const DatasourceModal: FunctionComponent<DatasourceModalProps> = ({
is_active: column.is_active,
is_dttm: column.is_dttm,
python_date_format: column.python_date_format || null,
+ partition_value_transform: column.partition_value_transform || null,
+ partition_transform_is_monotonic: Boolean(
+ column.partition_transform_is_monotonic,
+ ),
uuid: column.uuid,
extra: buildExtraJsonObject(column),
}),
diff --git a/superset-frontend/src/components/Datasource/components/DatasourceEditor/DatasourceEditor.tsx b/superset-frontend/src/components/Datasource/components/DatasourceEditor/DatasourceEditor.tsx
index ba148714d03..c77c51a3de0 100644
--- a/superset-frontend/src/components/Datasource/components/DatasourceEditor/DatasourceEditor.tsx
+++ b/superset-frontend/src/components/Datasource/components/DatasourceEditor/DatasourceEditor.tsx
@@ -94,7 +94,11 @@ import SpatialControl from 'src/explore/components/controls/SpatialControl';
import CollectionTable from '../CollectionTable';
import Fieldset from '../Fieldset';
import Field from '../Field';
-import { fetchSyncedColumns, updateColumns } from '../../utils';
+import {
+ clearDanglingPartitionMapping,
+ fetchSyncedColumns,
+ updateColumns,
+} from '../../utils';
import DatasetUsageTab from './components/DatasetUsageTab';
import {
DEFAULT_COLUMNS_FOLDER_UUID,
@@ -156,6 +160,8 @@ interface Column {
certified_by?: string;
certification_details?: string;
is_certified?: boolean;
+ partition_value_transform?: string | null;
+ partition_transform_is_monotonic?: boolean;
}
interface Database {
@@ -198,6 +204,8 @@ interface DatasourceObject {
cache_timeout?: number;
normalize_columns?: boolean;
always_filter_main_dttm?: boolean;
+ partition_column?: string | null;
+ partition_mapped_column?: string | null;
template_params?: string;
spatials?: SpatialConfig[];
all_cols?: string[];
@@ -1327,6 +1335,27 @@ function DatasourceEditor({
) as Column[],
});
+ // A sync can remove the partition column at the source. Leave the editor
+ // showing a mapping that points at a column the table no longer has and
+ // the owner has no way to tell why pruning stopped.
+ const clearedMapping = clearDanglingPartitionMapping(
+ datasource,
+ columnChanges.finalColumns,
+ );
+ if (clearedMapping) {
+ onDatasourcePropChange(
+ 'partition_column',
+ clearedMapping.partition_column,
+ );
+ onDatasourcePropChange(
+ 'partition_mapped_column',
+ clearedMapping.partition_mapped_column,
+ );
+ addSuccessToast(
+ t('The partition filter mapping was cleared: its column is gone'),
+ );
+ }
+
if (datasource.id !== undefined) {
clearDatasetCache(datasource.id);
}
@@ -1350,7 +1379,13 @@ function DatasourceEditor({
} finally {
abortControllers.current.syncMetadata = null;
}
- }, [datasource, addSuccessToast, addDangerToast, setColumns]);
+ }, [
+ datasource,
+ addSuccessToast,
+ addDangerToast,
+ setColumns,
+ onDatasourcePropChange,
+ ]);
// After a physical table change, refresh columns from the new table, matching
// the legacy class component's tableChangeAndSyncMetadata path. Declared after
diff --git a/superset-frontend/src/components/Datasource/utils/index.ts b/superset-frontend/src/components/Datasource/utils/index.ts
index 2740ed92111..0f645241bac 100644
--- a/superset-frontend/src/components/Datasource/utils/index.ts
+++ b/superset-frontend/src/components/Datasource/utils/index.ts
@@ -39,6 +39,13 @@ interface ColumnMetadata {
groupby?: boolean;
filterable?: boolean;
expression?: string;
+ partition_value_transform?: string | null;
+ partition_transform_is_monotonic?: boolean;
+}
+
+interface PartitionMappingFields {
+ partition_column?: string | null;
+ partition_mapped_column?: string | null;
}
interface ColumnChanges {
@@ -223,6 +230,42 @@ export function updateColumns(
* @param signal - Optional AbortSignal to cancel the request
* @returns Promise Array of column metadata objects
*/
+/**
+ * The parts of a partition filter mapping whose columns a sync just removed.
+ *
+ * Returns the fields to patch onto the datasource, or `null` when nothing needs
+ * to change. The backend clears a dangling mapping too, on the authoritative
+ * `override_columns` path -- this keeps the editor's own state honest between
+ * a sync and the next save, so the owner isn't looking at a mapping that points
+ * at a column the table no longer has.
+ */
+export function clearDanglingPartitionMapping(
+ datasource: PartitionMappingFields,
+ finalColumns: Pick<ColumnMetadata, 'column_name'>[],
+): PartitionMappingFields | null {
+ if (!datasource.partition_column) {
+ return null;
+ }
+
+ const surviving = new Set(finalColumns.map(col => col.column_name));
+
+ if (!surviving.has(datasource.partition_column)) {
+ return { partition_column: null, partition_mapped_column: null };
+ }
+
+ if (
+ datasource.partition_mapped_column &&
+ !surviving.has(datasource.partition_mapped_column)
+ ) {
+ return {
+ partition_column: datasource.partition_column,
+ partition_mapped_column: null,
+ };
+ }
+
+ return null;
+}
+
export async function fetchSyncedColumns(
datasource: DatasourceForSync,
signal?: AbortSignal,
diff --git a/superset-frontend/src/components/Datasource/utils/partitionMapping.test.ts b/superset-frontend/src/components/Datasource/utils/partitionMapping.test.ts
new file mode 100644
index 00000000000..b049a9fbb1c
--- /dev/null
+++ b/superset-frontend/src/components/Datasource/utils/partitionMapping.test.ts
@@ -0,0 +1,118 @@
+/**
+ * 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.
+ */
+
+import { clearDanglingPartitionMapping, updateColumns } from '.';
+
+const addSuccessToast = jest.fn();
+
+beforeEach(() => {
+ addSuccessToast.mockClear();
+});
+
+test('a column sync preserves the partition value transform', () => {
+ // The transform is set by hand in the row-expand section, so a sync that
+ // reports the same column must not quietly discard it.
+ const prevCols = [
+ {
+ column_name: 'event_time',
+ type: 'TIMESTAMP',
+ is_dttm: true,
+ partition_value_transform: 'unix_timestamp(:value)',
+ partition_transform_is_monotonic: true,
+ },
+ ];
+ const newCols = [
+ { column_name: 'event_time', type: 'TIMESTAMP', is_dttm: true },
+ ];
+
+ const result = updateColumns(prevCols, newCols, addSuccessToast);
+
+ expect(result.finalColumns[0]).toMatchObject({
+ partition_value_transform: 'unix_timestamp(:value)',
+ partition_transform_is_monotonic: true,
+ });
+});
+
+test('a column sync preserves the transform when the type changes', () => {
+ const prevCols = [
+ {
+ column_name: 'event_time',
+ type: 'TIMESTAMP',
+ is_dttm: true,
+ partition_value_transform: 'unix_timestamp(:value)',
+ },
+ ];
+ const newCols = [
+ { column_name: 'event_time', type: 'DATETIME', is_dttm: true },
+ ];
+
+ const result = updateColumns(prevCols, newCols, addSuccessToast);
+
+ expect(result.modified).toEqual(['event_time']);
+ expect(result.finalColumns[0]).toMatchObject({
+ partition_value_transform: 'unix_timestamp(:value)',
+ });
+});
+
+test('a sync that removes the partition column clears the mapping', () => {
+ const datasource = {
+ partition_column: 'dt_epoch',
+ partition_mapped_column: null,
+ };
+
+ expect(
+ clearDanglingPartitionMapping(datasource, [
+ { column_name: 'event_time' },
+ { column_name: 'country' },
+ ]),
+ ).toEqual({ partition_column: null, partition_mapped_column: null });
+});
+
+test('a sync that removes the mapped column clears only the override', () => {
+ // The partition column is still real, so the designation survives; the
+ // mapping just goes inactive until a column is chosen again.
+ const datasource = {
+ partition_column: 'dt_epoch',
+ partition_mapped_column: 'event_time',
+ };
+
+ expect(
+ clearDanglingPartitionMapping(datasource, [{ column_name: 'dt_epoch' }]),
+ ).toEqual({ partition_column: 'dt_epoch', partition_mapped_column: null });
+});
+
+test('a sync that keeps both columns leaves the mapping alone', () => {
+ const datasource = {
+ partition_column: 'dt_epoch',
+ partition_mapped_column: 'event_time',
+ };
+
+ expect(
+ clearDanglingPartitionMapping(datasource, [
+ { column_name: 'dt_epoch' },
+ { column_name: 'event_time' },
+ ]),
+ ).toBeNull();
+});
+
+test('a dataset with no partition column needs no change', () => {
+ expect(
+ clearDanglingPartitionMapping({}, [{ column_name: 'event_time' }]),
+ ).toBeNull();
+});
diff --git a/superset-frontend/src/features/datasets/types.ts b/superset-frontend/src/features/datasets/types.ts
index b182c3f78dd..fca5d4481c8 100644
--- a/superset-frontend/src/features/datasets/types.ts
+++ b/superset-frontend/src/features/datasets/types.ts
@@ -37,6 +37,10 @@ export type ColumnObject = {
certification_details?: string;
warning_markdown?: string;
advanced_data_type?: string;
+ // Partition filter mapping: filters on this column are mirrored onto the
+ // dataset's `partition_column` with the value passed through this transform.
+ partition_value_transform?: string | null;
+ partition_transform_is_monotonic?: boolean;
};
type MetricObject = {
@@ -81,6 +85,8 @@ export type DatasetObject = {
is_managed_externally: boolean;
normalize_columns: boolean;
always_filter_main_dttm: boolean;
+ partition_column?: string | null;
+ partition_mapped_column?: string | null;
type: DatasourceType;
column_formats: Record<string, string>;
datasource_name: string | null;
diff --git a/superset/commands/dataset/update.py b/superset/commands/dataset/update.py
index a23a10ca4ca..b507e642405 100644
--- a/superset/commands/dataset/update.py
+++ b/superset/commands/dataset/update.py
@@ -47,6 +47,17 @@ from superset.commands.dataset.exceptions import (
)
from superset.commands.utils import compute_subjects, raise_if_managed_externally
from superset.connectors.sqla.models import SqlaTable, validate_stored_expression
+from superset.connectors.sqla.partition_mapping import (
+ is_parseable,
+ parse_skeleton,
+ validate_partition_mapping,
+)
+from superset.connectors.sqla.partition_mapping_storage import (
+ load_partition_mapping,
+ parse_mapping,
+ read_partition_mapping_patch,
+ StoredPartitionMapping,
+)
from superset.daos.dataset import DatasetDAO
from superset.datasets.schemas import FolderSchema
from superset.exceptions import (
@@ -308,6 +319,8 @@ class UpdateDatasetCommand(UpdateMixin, BaseCommand):
if predicate := self._properties.get("fetch_values_predicate"):
self._validate_fetch_values_predicate(predicate, exceptions)
+ self._validate_partition_mapping(exceptions)
+
if folders := self._properties.get("folders"):
valid_uuids: set[UUID] = set()
if metrics:
@@ -416,6 +429,113 @@ class UpdateDatasetCommand(UpdateMixin, BaseCommand):
)
)
+ def _validate_partition_mapping(self, exceptions: list[ValidationError]) -> None:
+ """
+ Validate the dataset's partition filter mapping.
+
+ Only the blocking (Tier 1) issues become validation errors. Tier 2
+ issues -- an unparseable transform, a transform missing `:value` --
+ deliberately let the save through and leave the mapping inactive, per
+ the PRD, so a half-written transform doesn't cost the owner the rest of
+ their edits. They are surfaced by the editor, not by rejecting the PUT.
+
+ The transform is authored by a dataset owner, the same principal and
+ trust level as a calculated-column expression, so it also goes through
+ `validate_stored_expression` -- the parser gate that already governs
+ stored expressions. That gate only applies to a transform that parses:
+ it rejects an unparseable expression outright, which would turn a Tier-2
+ issue into a blocking one and undo the paragraph above. Skipping it
+ there costs nothing -- an unparseable transform is never emitted into a
+ query -- and it is not a hole for templating, because Jinja in a
+ transform is already a Tier-1 blocking issue of its own.
+
+ Every path that *reads* a mapping is gated on the feature flag, so this
+ one is too: with the flag off nothing mirrors, and rejecting a save over
+ a mapping that can never be consumed would be a validation error the
+ owner has no way to act on.
+ """
+ if not is_feature_enabled("PARTITION_FILTER_MAPPING"):
+ return
+
+ self._model = cast(SqlaTable, self._model)
+
+ columns = self._properties.get("columns")
+ column_names = (
+ {column["column_name"] for column in columns}
+ if columns is not None
+ else {column.column_name for column in self._model.columns}
+ )
+
+ main_dttm_col = self._properties.get("main_dttm_col", self._model.main_dttm_col)
+ mapping = self._merged_partition_mapping()
+ if not mapping.partition_column:
+ return
+
+ database = self._properties.get("database") or self._model.database
+ catalog = self._properties.get("catalog", self._model.catalog)
+ schema = self._properties.get("schema", self._model.schema)
+
+ transform = mapping.transform_for(
+ mapping.effective_mapped_column(main_dttm_col)
+ ).value_transform
+
+ for issue in validate_partition_mapping(
+ mapping,
+ column_names=column_names,
+ main_dttm_col=main_dttm_col,
+ engine=database.backend,
+ ):
+ if issue.blocking:
+ exceptions.append(
+ ValidationError(str(issue.message), field_name=issue.field)
+ )
+
+ if transform and is_parseable(transform, database.backend):
+ try:
+ validate_stored_expression(
+ database, catalog, schema, parse_skeleton(transform)
+ )
+ except (SupersetSecurityException, QueryClauseValidationException) as ex:
+ message = (
+ ex.error.message
+ if isinstance(ex, SupersetSecurityException)
+ else ex.message
+ )
+ exceptions.append(
+ ValidationError(
+ message,
+ field_name="partition_value_transform",
+ )
+ )
+
+ def _merged_partition_mapping(self) -> StoredPartitionMapping:
+ """
+ The mapping this request would leave behind, as the store would see it.
+
+ Exactly what `DatasetDAO.save_partition_mapping` will persist, derived
+ the same way, so validation cannot pass a mapping that differs from the
+ one written. That matters because there are two doors into this setting
+ and they have to agree:
+
+ - the typed API fields (`partition_column`, and
+ `partition_value_transform` on a column), which arrive as a patch;
+ - the dataset editor's free-text Extra box, which under the extra-JSON
+ store writes the same bytes directly.
+
+ A request carrying `extra` replaces the stored blob wholesale, so that is
+ the base when present; the typed fields patch on top either way. A PUT
+ that changes only `partition_column` is therefore still validated against
+ the transform already stored.
+ """
+ model = cast(SqlaTable, self._model)
+ stored = (
+ parse_mapping(self._properties["extra"])
+ if "extra" in self._properties
+ else load_partition_mapping(model)
+ )
+ patch = read_partition_mapping_patch(model, self._properties)
+ return patch.apply_to(stored) if patch else stored
+
def _validate_fetch_values_predicate(
self,
predicate: str,
diff --git a/superset/config.py b/superset/config.py
index 482bdaad047..6cfb8166dea 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -422,6 +422,26 @@ SCHEDULED_QUERIES: dict[str, Any] = {}
# parameters here: https://flask-limiter.readthedocs.io/en/stable/configuration.html
RATELIMIT_ENABLED = os.environ.get("SUPERSET_ENV") == "production"
RATELIMIT_APPLICATION = "50 per second"
+# Where a dataset's partition filter mapping is persisted: "extra" keeps it as
+# one key in ``tables.extra``, "columns" reads and writes four real columns.
+#
+# The feature ships on "extra" so it needs no schema change, and the migration
+# that adds the columns flips this, backfills, and swaps the model's properties
+# for the columns in one revision. Setting "columns" before that migration has
+# run raises, rather than silently reading four columns that are not there.
+PARTITION_MAPPING_STORE = "extra"
+
+# How long an evaluated partition value transform stays cached. The probe is a
+# synchronous round trip on the chart-query path, so the hit rate is what keeps
+# it off the hot path: day-aligned ranges ("Last month") hit constantly,
+# second-granularity relative ranges ("Last 24 hours") essentially never.
+PARTITION_TRANSFORM_PROBE_CACHE_TIMEOUT = int(timedelta(days=1).total_seconds())
+
+# Per-user, per-dataset budget for the partition mapping preview endpoint, which
+# fires a real warehouse query from a text input in the dataset editor. Counted
+# in a fixed 60-second window; set to 0 to disable.
+PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT = 30
+
AUTH_RATE_LIMITED = True
AUTH_RATE_LIMIT = "5 per second"
@@ -741,6 +761,12 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = {
# Try to optimize SQL queries — for now only predicate pushdown is supported
# @lifecycle: development
"OPTIMIZE_SQL": False,
+ # Configure and validate a mapping from a dataset's business column onto its
+ # physical partition column, so engines that require an explicit partition
+ # predicate can prune. For now this only stores and validates the mapping;
+ # no query consumer emits the mirrored predicate yet.
+ # @lifecycle: development
+ "PARTITION_FILTER_MAPPING": False,
# Expand nested types in Presto into extra columns/arrays. Experimental,
# doesn't work with all nested types.
# @lifecycle: development
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 8f4e82f688b..4fc22fd5589 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -60,6 +60,7 @@ from sqlalchemy.orm import (
relationship,
RelationshipProperty,
)
+from sqlalchemy.orm.exc import DetachedInstanceError
from sqlalchemy.orm.mapper import Mapper
from sqlalchemy.schema import UniqueConstraint
from sqlalchemy.sql import column, ColumnElement, literal_column, quoted_name, table
@@ -71,6 +72,15 @@ from superset_core.common.models import Dataset as CoreDataset
from superset import db, is_feature_enabled, security_manager
from superset.common.db_query_status import QueryStatus
+from superset.connectors.sqla.partition_mapping import (
+ is_transform_active,
+ resolve_partition_mapping,
+)
+from superset.connectors.sqla.partition_mapping_storage import (
+ ColumnTransform,
+ load_partition_mapping,
+ partition_mapping_store,
+)
from superset.connectors.sqla.utils import (
get_columns_description,
get_physical_table_metadata,
@@ -1146,6 +1156,57 @@ class TableColumn(AuditMixinNullable, ImportExportMixin, CertificationMixin, Mod
def __repr__(self) -> str:
return str(self.column_name)
+ # Partition filter mapping (§ PARTITION_FILTER_MAPPING). The transform is a
+ # SQL expression containing a `:value` placeholder; filters on this column
+ # are mirrored onto the dataset's `partition_column` as
+ # `partition_column <op> <transform evaluated at :value>`.
+ #
+ # Read-only properties rather than columns, because the mapping is persisted
+ # by `PARTITION_MAPPING_STORE` -- today one key in the dataset's `extra`,
+ # four real columns after the migration that adds them. They carry the
+ # post-migration column names on purpose: that migration deletes these
+ # properties, declares the columns, and nothing else in the tree moves.
+ #
+ # No setters. Under the extra-JSON store the mapping and `extra` are the same
+ # bytes, so a setter's write is lost without trace whenever `extra` is
+ # assigned afterwards -- and which of the two runs last depends on
+ # marshmallow's field ordering. Writes go through
+ # `partition_mapping_storage.set_partition_mapping`, or the DAO's funnel on
+ # the API path. This is also why they can't be in `export_fields`: `extra`
+ # already is, and already carries them.
+ @property
+ def partition_value_transform(self) -> str | None:
+ return self._partition_transform().value_transform
+
+ @property
+ def partition_transform_is_monotonic(self) -> bool:
+ return self._partition_transform().is_monotonic
+
+ def _partition_transform(self) -> ColumnTransform:
+ """
+ The transform declared on this column, or an empty one.
+
+ Scoped to this column rather than to the dataset so the pre- and
+ post-migration shapes are identical: each column reports what its own
+ row would say, and a column with no transform reports none even while
+ another column on the same dataset has one.
+
+ Degrades to "not declared" when the parent dataset is unreachable -- a
+ column built but not yet attached (`DatasetDAO._override_columns`
+ constructs `TableColumn(table_id=...)` without the relationship), or one
+ detached from its session. That is the safe direction: mirroring stops
+ rather than starting wrongly. It costs no SQL on the normal path, since
+ the many-to-one loader resolves `self.table` out of the identity map and
+ every caller reached this column through its parent.
+ """
+ try:
+ table = self.table
+ except DetachedInstanceError:
+ return ColumnTransform()
+ if table is None:
+ return ColumnTransform()
+ return load_partition_mapping(table).transform_for(self.column_name)
+
@property
def is_boolean(self) -> bool:
"""
@@ -1416,6 +1477,8 @@ class TableColumn(AuditMixinNullable, ImportExportMixin, CertificationMixin, Mod
"type_generic",
"verbose_name",
"warning_markdown",
+ "partition_value_transform",
+ "partition_transform_is_monotonic",
)
return {s: getattr(self, s) for s in attrs if hasattr(self, s)}
@@ -1924,8 +1987,84 @@ class SqlaTable(
data_["extra"] = self.extra
data_["always_filter_main_dttm"] = self.always_filter_main_dttm
data_["normalize_columns"] = self.normalize_columns
+ data_["partition_column"] = self.partition_column
+ data_["partition_mapped_column"] = self.partition_mapped_column
+ data_["partition_filter_mapping"] = self.partition_filter_mapping_summary
return data_
+ # Partition filter mapping (§ PARTITION_FILTER_MAPPING). Read-only
+ # properties over `PARTITION_MAPPING_STORE` rather than columns, for the
+ # reasons spelled out on `TableColumn.partition_value_transform` and in
+ # `superset.connectors.sqla.partition_mapping_storage`. Writes go through
+ # `set_partition_mapping`, or `DatasetDAO`'s funnel on the API path.
+ @property
+ def partition_column(self) -> str | None:
+ """Physical column the engine partitions on."""
+ return load_partition_mapping(self).partition_column
+
+ @property
+ def partition_mapped_column(self) -> str | None:
+ """
+ Explicit override for the column whose filters are mirrored.
+
+ ``None`` means "follow ``main_dttm_col``", so re-pointing the default
+ datetime column moves the mapping with it.
+ """
+ return load_partition_mapping(self).mapped_column
+
+ def update_from_object(self, obj: dict[str, Any]) -> None:
+ """
+ Sync from the legacy datasource editor's payload, mapping intact.
+
+ `super()` writes `obj.get(attr)` for every name in
+ `update_from_object_fields`, and `extra` is one of them -- so a payload
+ that omits it writes NULL and takes the partition mapping (and the
+ dataset's certification) down with it. Re-applying afterwards preserves
+ the mapping when the payload is silent about `extra`, and leaves an
+ `extra` that *does* carry one exactly as sent.
+ """
+ mapping = load_partition_mapping(self)
+ super().update_from_object(obj)
+ if "extra" not in obj and not mapping.is_empty:
+ partition_mapping_store().save(self, mapping)
+
+ @property
+ def partition_filter_mapping_summary(self) -> dict[str, Any] | None:
+ """
+ Self-contained summary of the mapping for the Explore indicator.
+
+ Deliberately not a lookup into `columns`: `data_for_slices` prunes
+ columns no chart references, and the partition column is typically
+ referenced by none of them, so anything reading it out of
+ `datasource.columns` would work in Explore and break on dashboards.
+
+ `active` is the save path's own verdict rather than an approximation of
+ it. A transform that fails validation -- one missing `:value`, one that
+ does not parse -- is saved inactive on purpose, so a cheaper signal here
+ would advertise a mapping that never mirrors a filter. The parse this
+ costs is memoized on `(transform, engine)` in `is_transform_active`, and
+ datasets without a partition column never reach it.
+ """
+ if not self.partition_column:
+ return None
+
+ columns_by_name = {column.column_name: column for column in self.columns}
+ mapped_column_name = self.partition_mapped_column or self.main_dttm_col
+ mapped_column = columns_by_name.get(mapped_column_name or "")
+ active = bool(
+ self.partition_column in columns_by_name
+ and mapped_column is not None
+ and mapped_column_name != self.partition_column
+ and is_transform_active(
+ mapped_column.partition_value_transform, self.database.backend
+ )
+ )
+ return {
+ "partition_column": self.partition_column,
+ "mapped_column": mapped_column_name,
+ "active": active,
+ }
+
@property
def extra_dict(self) -> dict[str, Any]:
try:
@@ -2638,6 +2777,27 @@ class SqlaTable(
# Add each predicate as a separate cache key component
extra_cache_keys.extend(rls_predicates)
+ # An active partition filter mapping changes the SQL a cached result came
+ # from, so it has to participate in the key or a mapping fix leaves stale
+ # pruned results behind. Only appended when the mapping is actually
+ # active, so keys don't churn for the entire installed base over a
+ # feature nobody has enabled.
+ #
+ # Note `PARTITION_FILTER_MAPPING` must be configured as a static boolean.
+ # `FEATURE_FLAGS` also accepts per-request callables, and a flag that
+ # resolves per user or per tenant would let a flag-off user read a cache
+ # entry written from pruned SQL by a flag-on user.
+ if mapping := resolve_partition_mapping(self):
+ extra_cache_keys.append(
+ (
+ "partition_filter_mapping",
+ mapping.partition_column,
+ mapping.mapped_column,
+ mapping.value_transform,
+ mapping.is_monotonic,
+ )
+ )
+
return list(set(extra_cache_keys))
@property
diff --git a/superset/connectors/sqla/partition_mapping.py b/superset/connectors/sqla/partition_mapping.py
new file mode 100644
index 00000000000..1ead706eac4
--- /dev/null
+++ b/superset/connectors/sqla/partition_mapping.py
@@ -0,0 +1,977 @@
+# 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 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.connectors.sqla.partition_mapping_storage import (
+ ColumnTransform,
+ load_partition_mapping,
+ StoredPartitionMapping,
+)
+from superset.constants import LRU_CACHE_MAX_SIZE, TimeGrain
+from superset.exceptions import SupersetParseError
+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"
+
+#: 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,
+}
+
+
+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)
+
+
+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 {parse_skeleton(transform)}", engine)
+ except SupersetParseError:
+ return None
+
+
+def is_parseable(transform: str | None, engine: str) -> bool:
+ """Whether the transform parses as a single select expression."""
+ if not transform or not transform.strip():
+ return False
+ return _parse_skeleton(transform, engine) is not None
+
+
+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.
+ """
+ if not feature_flag_manager.is_feature_enabled(FEATURE_FLAG):
+ return None
+
+ stored = load_partition_mapping(datasource)
+ partition_column = stored.partition_column
+ 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 = stored.effective_mapped_column(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
+
+ transform_spec = stored.transform_for(mapped_column_name)
+ transform = transform_spec.value_transform
+ # `is_transform_active` rather than a query-time restatement of the same
+ # rules: it *is* `validate_transform` with the messages discarded, so the two
+ # cannot drift. The difference that matters is the non-deterministic
+ # denylist, which a Tier-2-only check omits -- and `CreateDatasetCommand`,
+ # import, and the editor's free-text Extra box all reach this storage without
+ # passing the save-time validation, so a transform calling `now()` would
+ # otherwise be mirrored and its result cached.
+ if not is_transform_active(transform, datasource.database.backend):
+ return None
+
+ if _has_active_advanced_data_type(columns_by_name[mapped_column_name]):
+ # `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=transform_spec.is_monotonic,
+ )
+
+
+def _has_active_advanced_data_type(column: TableColumn) -> bool:
+ 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 build_probe_sql(
+ transform: str,
+ values: list[Any],
+ dialect: Dialect | None = None,
+) -> 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.
+
+ 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)
+
+
+def evaluate_transform(
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+ transform: str,
+ values: list[Any],
+) -> 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.
+ """
+ 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)
+ 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],
+) -> list[Any] | None:
+ try:
+ sql = build_probe_sql(transform, distinct, _dialect_for(database))
+ 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]
+ if len(row) < len(distinct):
+ # The results cannot be aligned back to their inputs; skipping
+ # beats guessing which value produced which column.
+ logger.warning(
+ "Partition transform probe returned %d values for %d inputs",
+ len(row),
+ len(distinct),
+ )
+ return None
+ return [row.iloc[index] for index in range(len(distinct))]
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Partition transform probe failed; queries will not prune",
+ exc_info=True,
+ )
+ return None
+
+
+def _probe_cache_key(
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+ transform: str,
+ values: list[Any],
+) -> str:
+ """
+ Key on everything that can change the answer.
+
+ Note this cache is independent of the chart-data cache: it is keyed on the
+ transform and its inputs, so it is correct to share across every chart on
+ every dataset that happens to use the same transform.
+ """
+ payload = json.dumps(
+ [
+ database.id,
+ database.backend,
+ # The probe asks this connection to evaluate the transform, so the
+ # answer belongs to it. `id` outlives an edit to the URI or to
+ # `extra` (a session timezone, say), which would otherwise serve
+ # values computed against the old environment until the entry
+ # expires. `changed_on` covers the edits that do not show up here,
+ # such as a rotated password, without putting a secret in the key.
+ database.sqlalchemy_uri,
+ database.extra,
+ str(database.changed_on),
+ catalog,
+ schema,
+ transform,
+ [repr(value) for value in values],
+ ],
+ default=repr,
+ )
+ digest = hashlib.md5(payload.encode("utf-8")).hexdigest() # noqa: S324
+ return f"partition_transform_probe:{digest}"
+
+
+def _cache_get(key: str) -> list[Any] | None:
+ try:
+ return cache_manager.cache.get(key)
+ except Exception: # pylint: disable=broad-except # noqa: BLE001
+ return None
+
+
+def _cache_set(key: str, value: list[Any]) -> None:
+ timeout = app.config.get("PARTITION_TRANSFORM_PROBE_CACHE_TIMEOUT", 24 * 60 * 60)
+ try:
+ cache_manager.cache.set(key, value, timeout=timeout)
+ except Exception: # pylint: disable=broad-except
+ logger.warning("Could not cache partition transform probe", exc_info=True)
+
+
+def _dialect_for(database: Database) -> Dialect | None:
+ """
+ The dialect used to render literals in the probe SQL.
+
+ Falls back to SQLAlchemy's default dialect if the database cannot produce
+ one -- the probe is best-effort and a rendering mismatch surfaces as a
+ failed probe, which fails open to no pruning.
+ """
+ try:
+ dialect = database.get_dialect()
+ except Exception: # pylint: disable=broad-except # noqa: BLE001
+ return None
+ return dialect if isinstance(dialect, Dialect) else None
+
+
+def _hashable(value: Any) -> Any:
+ """Values arrive from user filters and are not guaranteed hashable."""
+ try:
+ hash(value)
+ except TypeError:
+ return repr(value)
+ return value
+
+
+@dataclass(frozen=True)
+class MappingValidationIssue:
+ """
+ One problem found with a mapping at save time.
+
+ ``blocking`` issues reject the save (400). The rest save fine and leave the
+ mapping inactive -- the PRD is explicit that a mapping "stays inactive until
+ it parses", so a half-written transform must not cost the owner the rest of
+ their edits.
+ """
+
+ field: str
+ message: str
+ blocking: bool
+
+
+def validate_partition_mapping(
+ mapping: StoredPartitionMapping,
+ *,
+ column_names: set[str],
+ main_dttm_col: str | None,
+ engine: str,
+) -> list[MappingValidationIssue]:
+ """
+ Validate a dataset's partition mapping, in two tiers.
+
+ Tier 1 (``blocking=True``) is structural and safety: the columns have to
+ exist, a column cannot be mapped onto itself, and the transform cannot carry
+ Jinja or call a non-deterministic function. Tier 2 (``blocking=False``) is
+ everything that merely leaves the mapping inactive.
+
+ The Tier-1 transform checks need a successful parse to inspect anything, so
+ an unparseable transform falls through to Tier 2. That leaves them
+ unreachable in exactly the case where it doesn't matter: an unparseable
+ transform is never executed.
+
+ Takes the merged mapping rather than five loose keyword arguments so that
+ every caller -- the save path, the preview endpoint -- validates the same
+ object the store would persist, and none of them has to re-derive which
+ column is the effective one.
+ """
+ partition_column = mapping.partition_column
+ if not partition_column:
+ return []
+
+ partition_mapped_column = mapping.mapped_column
+ transform = mapping.transform_for(
+ mapping.effective_mapped_column(main_dttm_col)
+ ).value_transform
+
+ issues: list[MappingValidationIssue] = []
+
+ if partition_column not in column_names:
+ issues.append(
+ MappingValidationIssue(
+ field="partition_column",
+ message=_(
+ "Partition column %(name)s is not a column on this dataset.",
+ name=partition_column,
+ ),
+ blocking=True,
+ )
+ )
+
+ if partition_mapped_column and partition_mapped_column not in column_names:
+ issues.append(
+ MappingValidationIssue(
+ field="partition_mapped_column",
+ message=_(
+ "Mapped column %(name)s is not a column on this dataset.",
+ name=partition_mapped_column,
+ ),
+ blocking=True,
+ )
+ )
+
+ effective_mapped_column = partition_mapped_column or main_dttm_col
+ if effective_mapped_column and effective_mapped_column == partition_column:
+ issues.append(
+ MappingValidationIssue(
+ field="partition_column",
+ message=_(
+ "The partition column cannot be mapped onto itself. "
+ "%(name)s is both the partition column and the mapped "
+ "column.",
+ name=partition_column,
+ ),
+ blocking=True,
+ )
+ )
+
+ issues.extend(validate_transform(transform, engine))
+ return issues
+
+
+def validate_transform(
+ transform: str | None,
+ engine: str,
+) -> list[MappingValidationIssue]:
+ """Validate the value transform on its own. See `validate_partition_mapping`."""
+ field = "partition_value_transform"
+
+ if contains_jinja(transform):
+ return [
+ MappingValidationIssue(
+ field=field,
+ message=_(
+ "Jinja templating is not supported in a partition value "
+ "transform. The transform is evaluated in a different "
+ "context and at a different time from the chart query, so "
+ "a template would not render the same way."
+ ),
+ blocking=True,
+ )
+ ]
+
+ if not transform or not transform.strip():
+ return [
+ MappingValidationIssue(
+ field=field,
+ message=_(
+ "No value transform is set, so no filter will be mirrored "
+ "onto the partition column."
+ ),
+ blocking=False,
+ )
+ ]
+
+ if not is_parseable(transform, engine):
+ return [
+ MappingValidationIssue(
+ field=field,
+ message=_(
+ "The value transform could not be parsed. The mapping is "
+ "saved but stays inactive until it does."
+ ),
+ blocking=False,
+ )
+ ]
+
+ if not contains_value_placeholder(transform):
+ return [
+ MappingValidationIssue(
+ field=field,
+ message=_(
+ "The value transform must contain the :value placeholder, "
+ "which stands for the filter value being mirrored."
+ ),
+ blocking=False,
+ )
+ ]
+
+ if functions := find_non_deterministic_functions(transform, engine):
+ return [
+ MappingValidationIssue(
+ field=field,
+ message=_(
+ "The value transform calls %(functions)s, whose result "
+ "depends on when and where it runs. The transform is "
+ "evaluated in a separate session and the result is cached, "
+ "so the emitted predicate would freeze a snapshot of that "
+ "moment.",
+ functions=", ".join(sorted(functions)),
+ ),
+ blocking=True,
+ )
+ ]
+
+ return []
+
+
+@lru_cache(maxsize=LRU_CACHE_MAX_SIZE)
+def is_transform_active(transform: str | None, engine: str) -> bool:
+ """
+ Whether Superset will mirror filters through this transform.
+
+ `validate_transform` with the messages discarded, so the Explore indicator
+ cannot advertise a mapping the save path recorded as inactive. Anything
+ cheaper here would be a second, weaker statement of the same rule, free to
+ drift from the one the save path applies.
+
+ Blocking issues count too. Jinja and non-deterministic transforms are
+ rejected on PUT, but `CreateDatasetCommand` and import do not validate the
+ mapping, so one can still reach the database -- and it would never mirror a
+ filter either.
+
+ Memoized because `partition_filter_mapping_summary` is serialized on every
+ Explore and dashboard load and this parses the transform. The key is
+ everything the answer depends on: `validate_transform` reads no session, no
+ locale and no config beyond `SQL_MAX_PARSE_LENGTH`, which bounds the parser
+ rather than changing a verdict within that bound.
+ """
+ return not validate_transform(transform, engine)
+
+
+def preview_partition_mapping(
+ datasource: SqlaTable,
+ *,
+ mapped_column: str,
+ value_transform: str | None,
+ sample_value: str,
+) -> dict[str, Any]:
+ """
+ Evaluate a candidate mapping and describe the predicate it would emit.
+
+ Shares the evaluator -- and therefore the probe cache -- with the query
+ path, so preview and runtime cannot drift and a previewed transform warms
+ the chart path for free.
+
+ Validation runs first and the engine second: a half-typed transform is by
+ definition unparseable, so most of what a text input produces costs zero
+ queries.
+ """
+ partition_column = datasource.partition_column
+ if not partition_column:
+ return {"valid": False, "error": _("No partition column is set.")}
+
+ column_names = {str(column.column_name) for column in datasource.columns}
+ if mapped_column not in column_names:
+ return {
+ "valid": False,
+ "error": _("%(name)s is not a column on this dataset.", name=mapped_column),
+ }
+
+ engine = datasource.database.backend
+ # The candidate, not what is stored: a preview exists to answer "what would
+ # this transform do", so it validates the mapping the editor is holding.
+ candidate = StoredPartitionMapping(
+ partition_column=str(partition_column),
+ mapped_column=mapped_column,
+ column_transforms={mapped_column: ColumnTransform(value_transform)},
+ )
+ for issue in validate_partition_mapping(
+ candidate,
+ column_names=column_names,
+ main_dttm_col=datasource.main_dttm_col,
+ engine=engine,
+ ):
+ return {"valid": False, "error": str(issue.message)}
+
+ evaluated = evaluate_transform(
+ datasource.database,
+ datasource.catalog,
+ datasource.schema,
+ cast(str, value_transform),
+ [sample_value],
+ )
+ if evaluated is None:
+ return {
+ "valid": False,
+ "error": _("The transform could not be evaluated against the database."),
+ }
+
+ return {
+ "valid": True,
+ "emitted_predicate": (
+ f"{datasource.quote_identifier(str(partition_column))} >= "
+ f"{_render_literal(datasource.database, evaluated[0])}"
+ ),
+ }
+
+
+def _render_literal(database: Database, value: Any) -> str:
+ """
+ Render a probed value the way it appears in the generated SQL.
+
+ Compiled by the dialect rather than formatted by hand. A probe returns
+ whatever the warehouse gave it -- an epoch integer, a date, a NULL for an
+ input the transform could not convert -- and only the dialect knows how each
+ of those is written. `str()` would render a NULL as `None` and leave a date
+ unquoted, neither of which is SQL; hand-rolled quote-doubling would only
+ ever have been right for strings.
+ """
+ return str(
+ sa.literal(value).compile(
+ dialect=_dialect_for(database),
+ compile_kwargs={"literal_binds": True},
+ )
+ )
+
+
+def build_mirrored_predicates(
+ datasource: SqlaTable,
+ mapping: PartitionMapping,
+ requests: list[tuple[FilterOperator, Any]],
+) -> list[ColumnElement[Any]]:
+ """
+ Turn collected ``(operator, value)`` mirror requests into predicates.
+
+ ``requests`` is expected to be deduplicated by the caller. Every value is
+ resolved in a single probe round trip -- one per chart query at most -- then
+ emitted as a literal constant, so "View query" shows the reader an ordinary
+ ``WHERE`` clause rather than an inline expression.
+ """
+ if not requests:
+ return []
+
+ partition_column = next(
+ (
+ column
+ for column in datasource.columns
+ if column.column_name == mapping.partition_column
+ ),
+ None,
+ )
+ if partition_column is None:
+ return []
+
+ # Flatten every value that needs probing into one list, remembering how many
+ # each request consumed so the results can be handed back out.
+ flat: list[Any] = []
+ spans: list[tuple[FilterOperator, int, int]] = []
+ for operator, value in requests:
+ values = list(value) if operator == FilterOperator.IN else [value]
+ spans.append((operator, len(flat), len(values)))
+ flat.extend(values)
+
+ evaluated = evaluate_transform(
+ datasource.database,
+ datasource.catalog,
+ datasource.schema,
+ mapping.value_transform,
+ flat,
+ )
+ if evaluated is None:
+ return []
+
+ if not _bounds_are_ordered(evaluated, spans):
+ return []
+
+ sqla_col = datasource.convert_tbl_column_to_sqla_col(partition_column)
+ db_engine_spec = datasource.db_engine_spec
+
+ predicates: list[ColumnElement[Any]] = []
+ for operator, start, length in spans:
+ chunk = evaluated[start : start + length]
+ if any(value is None for value in chunk):
+ continue
+ if operator == FilterOperator.IN:
+ predicates.append(sqla_col.in_(chunk))
+ else:
+ predicates.append(
+ db_engine_spec.handle_comparison_filter(sqla_col, operator, chunk[0])
+ )
+ if not predicates:
+ return []
+
+ # A comparison against a NULL partition value is NULL, so a row parked in a
+ # NULL partition is dropped by the mirror even when the real filter matches
+ # it -- Hive and Impala's default partition, or a transform that returns
+ # NULL for an input it cannot convert. The mirror only has to be no
+ # narrower than the filter it stands in for, so admitting NULL partitions
+ # keeps those rows. Engines still prune; they read one extra partition.
+ return [sa.or_(sa.and_(*predicates), sqla_col.is_(None))]
+
+
+_LOWER_BOUND_OPS = {
+ FilterOperator.GREATER_THAN,
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+}
+_UPPER_BOUND_OPS = {
+ FilterOperator.LESS_THAN,
+ FilterOperator.LESS_THAN_OR_EQUALS,
+}
+
+
+def _bounds_are_ordered(
+ evaluated: list[Any],
+ spans: list[tuple[FilterOperator, int, int]],
+) -> bool:
+ """
+ Backstop for the monotonicity *declaration*: check ``T(lower) <= T(upper)``.
+
+ Both bounds have already been probed, so this costs nothing extra. It is a
+ *necessary* condition, not a sufficient one: it catches an inverted
+ transform, and catches ``hour()`` on any range spanning a day boundary, but
+ not ``hour()`` inside a single day. A cheap sanity check on a claim only the
+ dataset owner can actually make -- not a replacement for the declaration.
+ """
+ lowers = [
+ evaluated[start] for operator, start, _ in spans if operator in _LOWER_BOUND_OPS
+ ]
+ uppers = [
+ evaluated[start] for operator, start, _ in spans if operator in _UPPER_BOUND_OPS
+ ]
+ if not lowers or not uppers:
+ return True
+
+ try:
+ return bool(max(lowers) <= min(uppers))
+ except TypeError:
+ # Probe results arrive through `get_df` as pandas/numpy scalars, which
+ # do not all compare. "Not comparable" is failure, not permission.
+ logger.warning(
+ "Partition transform produced incomparable bounds; not mirroring"
+ )
+ return False
diff --git a/superset/connectors/sqla/partition_mapping_storage.py b/superset/connectors/sqla/partition_mapping_storage.py
new file mode 100644
index 00000000000..146248abba5
--- /dev/null
+++ b/superset/connectors/sqla/partition_mapping_storage.py
@@ -0,0 +1,606 @@
+# 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.
+"""
+Where a dataset's partition filter mapping is persisted.
+
+A mapping is four values -- see `superset.connectors.sqla.partition_mapping` for
+what they mean. This module owns the question of where those four values live,
+and nothing else. Every read goes through `PartitionMappingStore.load` and every
+write through `set_partition_mapping`, so the answer can change without touching
+the model, the API, validation or the query rewriter.
+
+Two backings
+------------
+The feature ships stored as one key in ``tables.extra``
+(`ExtraJsonPartitionMappingStore`) and will move to four real columns
+(`ColumnPartitionMappingStore`) in the migration that adds them. Selected by
+``PARTITION_MAPPING_STORE``.
+
+``tables.extra`` is not merely the option that avoids a migration; it is a
+reasonable home. It is already in `SqlaTable.export_fields` and already listed in
+`ExportDatasetsCommand`'s ``JSON_KEYS``, so a mapping travels through
+export/import with no bundle-format change. SQLAlchemy-Continuum already
+versions the column, so dataset version history and restore work untouched. And
+there is precedent for reading typed settings back out of it:
+``certification``, ``warning_markdown`` and ``timezone`` all live there, the
+first two surfaced through model properties exactly as the mapping now is.
+
+It has to be the *dataset*-level blob, though, never ``table_columns.extra``,
+even for the two values that are per-column after the migration. The dataset
+editor's ``buildExtraJsonObject`` rebuilds every column's and metric's ``extra``
+from ``certification`` and ``warning_markdown`` alone on each save, so anything
+else parked there is destroyed the first time an owner opens the editor.
+Dataset-level ``extra`` is round-tripped verbatim, and its one mutating path
+(``setDatasetCertification``) parses and *merges*, so an unknown key survives.
+
+Reads are cheap; writes are funnelled
+-------------------------------------
+`load` is on the hot path -- `get_extra_cache_keys`, every `get_sqla_query`, and
+once per column in `TableColumn.data` -- so it memoizes on the instance, keyed on
+the raw ``extra`` string, the same self-invalidating trick
+`CertificationMixin.get_extra_dict` uses. (`SqlaTable.extra_dict` is an uncached
+``json.loads``; don't route through it.)
+
+Writes cannot be attribute assignment. Under the extra-JSON store the mapping
+and ``extra`` are the same bytes, so ``table.partition_column = x`` followed by
+``table.extra = <payload>`` loses the first write with no exception and no log --
+and which of the two runs last is a property of marshmallow's field ordering.
+The model's properties are therefore read-only and `set_partition_mapping` is the
+only supported write.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field, replace
+from typing import Any, Mapping, Protocol, TYPE_CHECKING
+
+from flask import current_app as app
+
+from superset.utils import json
+
+if TYPE_CHECKING:
+ from superset.connectors.sqla.models import SqlaTable
+
+logger = logging.getLogger(__name__)
+
+#: Top-level key in ``tables.extra``. Joins ``certification``,
+#: ``warning_markdown``, ``timezone`` and ``disallow_adhoc_metrics``.
+EXTRA_KEY = "partition_filter_mapping"
+
+#: Bounds enforced on read as well as on write. ``DatasetPutSchema`` validates
+#: the typed fields, but the free-text Extra box reaches the same storage without
+#: passing through it, so the store re-checks rather than trusting its contents.
+#: The two lengths are the widths the columns get after the migration.
+MAX_COLUMN_NAME_LENGTH = 250
+MAX_TRANSFORM_LENGTH = 4096
+
+#: Sentinel for "this keyword was not passed", distinct from ``None``, which
+#: means "clear this field".
+UNSET: Any = object()
+
+_MEMO_UNSET: Any = object()
+_MEMO_VALUE_ATTR = "_partition_mapping_memo"
+_MEMO_RAW_ATTR = "_partition_mapping_memo_raw"
+
+
+@dataclass(frozen=True)
+class ColumnTransform:
+ """
+ The two per-column halves of a mapping: the ``:value`` expression and
+ whether it preserves ordering.
+ """
+
+ value_transform: str | None = None
+ is_monotonic: bool = False
+
+ @property
+ def is_empty(self) -> bool:
+ return self == ColumnTransform()
+
+
+@dataclass(frozen=True)
+class StoredPartitionMapping:
+ """
+ A dataset's mapping exactly as persisted -- unresolved and unvalidated.
+
+ ``mapped_column is None`` means "follow ``main_dttm_col``", so re-pointing
+ the dataset's default datetime column moves the mapping with it unless the
+ owner overrides it. Deliberately not resolved at save time: that behaviour
+ then falls out of the model rather than needing code on every write path.
+
+ Transforms are keyed by column name rather than flattened onto the single
+ effective mapped column, because flattening is unsound. An owner who sets a
+ transform on ``A`` (then ``main_dttm_col``) and later re-points
+ ``main_dttm_col`` to ``B`` would, flattened, have ``B``'s values silently
+ mirrored through ``A``'s expression. Keyed, ``B`` has no transform, the
+ mapping goes inactive, and the behaviour matches what four real columns
+ would do -- which is the property that makes the two stores swappable.
+ """
+
+ partition_column: str | None = None
+ mapped_column: str | None = None
+ column_transforms: Mapping[str, ColumnTransform] = field(default_factory=dict)
+
+ @property
+ def is_empty(self) -> bool:
+ return (
+ not self.partition_column
+ and not self.mapped_column
+ and not self.column_transforms
+ )
+
+ def transform_for(self, column_name: str | None) -> ColumnTransform:
+ """
+ The transform declared on one column, or an empty one.
+
+ Empty rather than ``None`` so callers read ``.value_transform`` /
+ ``.is_monotonic`` unconditionally and get the same "not declared"
+ answer a NULL column would give them.
+ """
+ if not column_name:
+ return ColumnTransform()
+ return self.column_transforms.get(column_name, ColumnTransform())
+
+ def effective_mapped_column(self, main_dttm_col: str | None) -> str | None:
+ """The column whose filters are mirrored, override or default."""
+ return self.mapped_column or main_dttm_col
+
+ def with_transform(
+ self, column_name: str, transform: ColumnTransform
+ ) -> StoredPartitionMapping:
+ transforms = dict(self.column_transforms)
+ if transform.is_empty:
+ transforms.pop(column_name, None)
+ else:
+ transforms[column_name] = transform
+ return replace(self, column_transforms=transforms)
+
+
+@dataclass(frozen=True)
+class PartitionMappingPatch:
+ """
+ Only the fields a request actually mentioned.
+
+ Present keys rather than four ``X | None`` attributes, because ``None`` is a
+ meaningful value here: ``{"partition_column": None}`` clears the mapping
+ while an absent key leaves it alone. Collapsing the two is how a partial
+ column payload silently clears a monotonic flag nobody touched -- the exact
+ failure ``DatasetColumnsPutSchema`` avoids by refusing a ``load_default``.
+ """
+
+ fields: Mapping[str, Any] = field(default_factory=dict)
+ column_fields: Mapping[str, Mapping[str, Any]] = field(default_factory=dict)
+
+ @property
+ def is_empty(self) -> bool:
+ return not self.fields and not self.column_fields
+
+ def apply_to(self, mapping: StoredPartitionMapping) -> StoredPartitionMapping:
+ if self.fields:
+ mapping = replace(mapping, **dict(self.fields))
+ for column_name, values in self.column_fields.items():
+ current = mapping.transform_for(column_name)
+ mapping = mapping.with_transform(
+ column_name, replace(current, **dict(values))
+ )
+ return mapping
+
+
+class PartitionMappingStore(Protocol):
+ """
+ Where a dataset's partition mapping lives.
+
+ Exists because the mapping ships before its migration: today it is one key
+ in ``tables.extra``, and a follow-up moves it to four real columns. Both
+ shapes answer the same two questions, so the model facade, the write funnel,
+ validation and the query rewriter are written once against this and never
+ learn which store answered.
+ """
+
+ def load(self, table: SqlaTable) -> StoredPartitionMapping: ...
+
+ def save(self, table: SqlaTable, mapping: StoredPartitionMapping) -> None: ...
+
+
+class ExtraJsonPartitionMappingStore:
+ """One nested object under ``partition_filter_mapping`` in ``tables.extra``."""
+
+ def load(self, table: SqlaTable) -> StoredPartitionMapping:
+ raw = getattr(table, "extra", None)
+ memo_raw = getattr(table, _MEMO_RAW_ATTR, _MEMO_UNSET)
+ if memo_raw is _MEMO_UNSET or memo_raw != raw:
+ # Transient instance attributes, not mapped columns, so they never
+ # reach the database and never confuse Continuum. Safe to hand the
+ # memoized object straight out: it is frozen.
+ object.__setattr__(table, _MEMO_VALUE_ATTR, parse_mapping(raw))
+ object.__setattr__(table, _MEMO_RAW_ATTR, raw)
+ return getattr(table, _MEMO_VALUE_ATTR)
+
+ def save(self, table: SqlaTable, mapping: StoredPartitionMapping) -> None:
+ raw = getattr(table, "extra", None)
+ updated = extra_with_partition_mapping(raw, mapping)
+ if updated is raw:
+ # `extra_with_partition_mapping` hands back the identical object when
+ # the blob already says this. Writing anyway would dirty the column
+ # and mint a Continuum version row for an edit that changed nothing.
+ return
+ # A *new string*, never a mutation of the parsed dict. `decode_extra`
+ # parses fresh and throws its dict away, so mutating one would change
+ # nothing on the model and SQLAlchemy would never see the attribute as
+ # dirty. Assignment is also what invalidates `load`'s memo.
+ table.extra = updated
+
+
+class ColumnPartitionMappingStore:
+ """
+ Reads and writes four real columns. Not yet live.
+
+ Deliberately unimplemented rather than written against attributes that do
+ not exist. It would have to read ``table.partition_column``, which today is
+ the facade property, which calls this store -- unbounded recursion. A
+ property and a ``Column`` of the same name cannot coexist on a model, so
+ this adapter is only coherent once the migration has replaced the
+ properties, and until then it could only be exercised against a hand-built
+ stand-in, which proves nothing about the parts that are actually hard: real
+ columns, Continuum shadow rows, and the FAB metadata
+ (``order_columns``/``search_columns``) that changes meaning the moment these
+ stop being properties.
+
+ The migration that activates it must, in one revision:
+
+ 1. add ``tables.partition_column`` / ``partition_mapped_column`` and
+ ``table_columns.partition_value_transform`` /
+ ``partition_transform_is_monotonic``, plus the four matching columns on
+ the ``*_version`` shadow tables;
+ 2. backfill from ``tables.extra`` and strip the key, leaving a blob it could
+ not parse in place rather than dropping an owner's configuration;
+ 3. delete the facade properties on `SqlaTable` and `TableColumn`, declare
+ the real columns, and add the four names to ``export_fields``;
+ 4. implement this class and flip ``PARTITION_MAPPING_STORE`` to ``columns``.
+
+ The port, the value object and every call site are already right, which is
+ what this indirection buys: the migration is additive, not a rewrite.
+ """
+
+ _MESSAGE = (
+ "PARTITION_MAPPING_STORE='columns' requires the migration that adds the "
+ "partition mapping columns; use 'extra' until it has been applied."
+ )
+
+ def load(self, table: SqlaTable) -> StoredPartitionMapping:
+ raise NotImplementedError(self._MESSAGE)
+
+ def save(self, table: SqlaTable, mapping: StoredPartitionMapping) -> None:
+ raise NotImplementedError(self._MESSAGE)
+
+
+#: A literal table rather than a config-resolved import path: this is a
+#: two-value migration switch with a known end state, not a plugin point.
+STORES: dict[str, type[PartitionMappingStore]] = {
+ "extra": ExtraJsonPartitionMappingStore,
+ "columns": ColumnPartitionMappingStore,
+}
+
+
+def partition_mapping_store() -> PartitionMappingStore:
+ """
+ The configured store.
+
+ Resolved per call, not memoized. ``app.config`` is per-app and tests flip
+ this between cases; construction is a bare ``__init__`` over no state, so a
+ cache would buy nothing and could outlive an app teardown.
+ """
+ name = app.config["PARTITION_MAPPING_STORE"]
+ if name not in STORES:
+ raise ValueError(
+ f"Unknown PARTITION_MAPPING_STORE {name!r}; "
+ f"expected one of {sorted(STORES)}"
+ )
+ return STORES[name]()
+
+
+def load_partition_mapping(table: SqlaTable) -> StoredPartitionMapping:
+ """The dataset's stored mapping. The one read entry point."""
+ return partition_mapping_store().load(table)
+
+
+def set_partition_mapping( # pylint: disable=too-many-arguments
+ table: SqlaTable,
+ *,
+ partition_column: Any = UNSET,
+ mapped_column: Any = UNSET,
+ value_transform: Any = UNSET,
+ is_monotonic: Any = UNSET,
+ column: str | None = None,
+) -> StoredPartitionMapping:
+ """
+ Patch a dataset's mapping. An omitted keyword means "unchanged".
+
+ The one supported write, for the DAO funnel, `update_from_object` and tests.
+ ``value_transform`` and ``is_monotonic`` apply to ``column``, defaulting to
+ the effective mapped column *after* this call's ``mapped_column`` is taken
+ into account -- so setting a mapped column and its transform together does
+ what it reads like.
+ """
+ store = partition_mapping_store()
+ mapping = store.load(table)
+
+ fields: dict[str, Any] = {}
+ if partition_column is not UNSET:
+ fields["partition_column"] = partition_column
+ if mapped_column is not UNSET:
+ fields["mapped_column"] = mapped_column
+
+ column_fields: dict[str, dict[str, Any]] = {}
+ if value_transform is not UNSET or is_monotonic is not UNSET:
+ target = column or StoredPartitionMapping(
+ mapped_column=fields.get("mapped_column", mapping.mapped_column)
+ ).effective_mapped_column(table.main_dttm_col)
+ if not target:
+ raise ValueError(
+ "Cannot set a partition value transform without a column: the "
+ "dataset has no mapped column and no main_dttm_col."
+ )
+ values: dict[str, Any] = {}
+ if value_transform is not UNSET:
+ values["value_transform"] = value_transform
+ if is_monotonic is not UNSET:
+ values["is_monotonic"] = bool(is_monotonic)
+ column_fields[target] = values
+
+ patch = PartitionMappingPatch(fields=fields, column_fields=column_fields)
+ mapping = patch.apply_to(mapping)
+ store.save(table, mapping)
+ return mapping
+
+
+def clear_dangling(
+ mapping: StoredPartitionMapping, surviving_column_names: set[str]
+) -> StoredPartitionMapping:
+ """
+ Drop the parts of a mapping whose columns no longer exist.
+
+ A metadata sync can remove the partition column at the source, which would
+ otherwise leave the dataset pointing at a column that isn't there. The query
+ layer bails out defensively on a dangling mapping, so this is about the
+ dataset's stored state being honest rather than about correctness of the SQL.
+
+ Needed on both write paths: the editor clears the mapping client-side too,
+ but ``override_columns=true`` (an API-driven metadata sync, and SQL Lab's
+ "Save dataset") bypasses the editor entirely, and the upsert path drops every
+ column the payload omits -- so either can take a mapped column out from under
+ the mapping.
+
+ Runs from the DAO's write funnel rather than from ``update_columns``, where
+ a clear would be overwritten moments later by the payload's own ``extra``.
+ """
+ if (
+ mapping.partition_column
+ and mapping.partition_column not in surviving_column_names
+ ):
+ mapping = replace(mapping, partition_column=None, mapped_column=None)
+ elif mapping.mapped_column and mapping.mapped_column not in surviving_column_names:
+ mapping = replace(mapping, mapped_column=None)
+
+ stale = set(mapping.column_transforms) - surviving_column_names
+ for column_name in stale:
+ mapping = mapping.with_transform(column_name, ColumnTransform())
+ return mapping
+
+
+def parse_mapping(raw: Any) -> StoredPartitionMapping:
+ """
+ Read a mapping out of a raw ``extra`` string, tolerating anything.
+
+ ``extra`` is a free-text box in the dataset editor, so every value here is
+ typo- or attacker-supplied. Modelled on `get_dataset_timezone`: a value of
+ the wrong type is "not configured", never propagated.
+
+ ``is_monotonic`` insists on a literal ``True`` rather than coercing with
+ ``bool()``. A truthy ``"no"`` would enable range mirroring, and this is the
+ one field where being wrong drops real rows rather than merely failing to
+ prune.
+ """
+ blob = decode_extra(raw).get(EXTRA_KEY)
+ if not isinstance(blob, dict):
+ return StoredPartitionMapping()
+
+ transforms: dict[str, ColumnTransform] = {}
+ raw_transforms = blob.get("column_transforms")
+ if isinstance(raw_transforms, dict):
+ # Keys are strings by construction: this only ever sees a JSON-decoded
+ # object. The values are whatever someone typed.
+ for column_name, entry in raw_transforms.items():
+ if not isinstance(entry, dict):
+ continue
+ transform = ColumnTransform(
+ value_transform=_as_transform(entry.get("value_transform")),
+ is_monotonic=entry.get("is_monotonic") is True,
+ )
+ if not transform.is_empty:
+ transforms[column_name] = transform
+
+ return StoredPartitionMapping(
+ partition_column=_as_name(blob.get("partition_column")),
+ mapped_column=_as_name(blob.get("mapped_column")),
+ column_transforms=transforms,
+ )
+
+
+def extra_with_partition_mapping(
+ raw: Any, mapping: StoredPartitionMapping
+) -> str | None:
+ """
+ The ``extra`` string that carries this mapping, other keys preserved.
+
+ Pure, and needs no app context, which makes it the way to put a dataset into
+ a given stored state without going through the configured store -- test
+ fixtures today, and the backfill in the migration that moves this to columns.
+
+ Returns ``raw`` *itself* when the blob already says exactly this, so callers
+ can tell "no change" from "same text by coincidence" by identity.
+
+ An ``extra`` that is not decodable JSON cannot be merged into and is replaced
+ rather than blocking the write. That text is already inert -- ``certification``,
+ ``timezone`` and ``warning_markdown`` all read as absent from it -- so the
+ alternative would be to silently drop the mapping instead, which is worse and
+ just as quiet. It is logged, because it is still someone's text.
+ """
+ decoded = decode_extra(raw)
+ if raw and not decoded and not mapping.is_empty:
+ logger.warning(
+ "Replacing an undecodable dataset `extra` in order to store a "
+ "partition filter mapping"
+ )
+ merged = {key: value for key, value in decoded.items() if key != EXTRA_KEY}
+ if not mapping.is_empty:
+ merged[EXTRA_KEY] = _encode_mapping(mapping)
+ if merged == decoded:
+ return raw
+ return json.dumps(merged)
+
+
+def decode_extra(raw: Any) -> dict[str, Any]:
+ """``extra`` as a dict, or an empty one for anything that isn't."""
+ try:
+ decoded = json.loads(raw)
+ except (TypeError, json.JSONDecodeError):
+ return {}
+ return decoded if isinstance(decoded, dict) else {}
+
+
+def _encode_mapping(mapping: StoredPartitionMapping) -> dict[str, Any]:
+ return {
+ "partition_column": mapping.partition_column,
+ "mapped_column": mapping.mapped_column,
+ "column_transforms": {
+ column_name: {
+ "value_transform": transform.value_transform,
+ "is_monotonic": transform.is_monotonic,
+ }
+ for column_name, transform in sorted(mapping.column_transforms.items())
+ },
+ }
+
+
+def _as_name(value: Any) -> str | None:
+ if isinstance(value, str) and value and len(value) <= MAX_COLUMN_NAME_LENGTH:
+ return value
+ return None
+
+
+def _as_transform(value: Any) -> str | None:
+ if isinstance(value, str) and value and len(value) <= MAX_TRANSFORM_LENGTH:
+ return value
+ return None
+
+
+#: Names the dataset API uses for the mapping, mapped to the value object's.
+#: The API keeps the post-migration column names; the value object is free to be
+#: named for what it is, and this is the one place the two meet.
+_TABLE_PAYLOAD_FIELDS = {
+ "partition_column": "partition_column",
+ "partition_mapped_column": "mapped_column",
+}
+_COLUMN_PAYLOAD_FIELDS = {
+ "partition_value_transform": "value_transform",
+ "partition_transform_is_monotonic": "is_monotonic",
+}
+
+
+def extract_partition_mapping_patch(
+ model: SqlaTable | None,
+ attributes: dict[str, Any] | None,
+) -> PartitionMappingPatch | None:
+ """
+ Take the mapping fields out of an API payload, removing them as it goes.
+
+ Stripping is not housekeeping: the model's mapping attributes are read-only
+ properties, so anything left behind reaches a ``setattr`` or a
+ ``TableColumn(**properties)`` and raises. Must run before
+ ``attributes.pop("columns")``.
+ """
+ return _collect_patch(model, attributes, pop=True)
+
+
+def read_partition_mapping_patch(
+ model: SqlaTable | None,
+ attributes: dict[str, Any] | None,
+) -> PartitionMappingPatch | None:
+ """
+ The same patch, without touching the payload.
+
+ For validation, which runs before the write and must leave ``attributes``
+ exactly as the request sent them.
+ """
+ return _collect_patch(model, attributes, pop=False)
+
+
+def _collect_patch(
+ model: SqlaTable | None,
+ attributes: dict[str, Any] | None,
+ *,
+ pop: bool,
+) -> PartitionMappingPatch | None:
+ """
+ Read the mapping fields out of an API payload.
+
+ Returns ``None`` when the payload mentioned none of them, which is how
+ "absent means unchanged" survives all the way up to the caller.
+ """
+ if not attributes:
+ return None
+
+ take = dict.pop if pop else dict.get
+
+ fields: dict[str, Any] = {}
+ for payload_name, field_name in _TABLE_PAYLOAD_FIELDS.items():
+ if payload_name in attributes:
+ fields[field_name] = take(attributes, payload_name)
+
+ column_fields: dict[str, dict[str, Any]] = {}
+ columns_by_id = {column.id: column for column in model.columns} if model else {}
+ for entry in attributes.get("columns") or []:
+ if not isinstance(entry, dict):
+ continue
+ values = {
+ field_name: take(entry, payload_name)
+ for payload_name, field_name in _COLUMN_PAYLOAD_FIELDS.items()
+ if payload_name in entry
+ }
+ if not values:
+ continue
+ if "is_monotonic" in values:
+ # The column would be NOT NULL DEFAULT false after the migration, so
+ # an explicit `null` reads as "not declared" rather than propagating.
+ values["is_monotonic"] = values["is_monotonic"] is True
+ # A partial column payload may carry only `id`, in which case the column
+ # keeps the name it already has.
+ name = entry.get("column_name")
+ if not name and (existing := columns_by_id.get(entry.get("id"))):
+ name = existing.column_name
+ if not name:
+ logger.warning(
+ "Ignoring a partition value transform on a column payload with "
+ "neither a column_name nor a known id"
+ )
+ continue
+ column_fields.setdefault(name, {}).update(values)
+
+ patch = PartitionMappingPatch(fields=fields, column_fields=column_fields)
+ return None if patch.is_empty else patch
diff --git a/superset/daos/dataset.py b/superset/daos/dataset.py
index f571bec5bca..0762b6704ab 100644
--- a/superset/daos/dataset.py
+++ b/superset/daos/dataset.py
@@ -32,6 +32,12 @@ from superset.connectors.sqla.models import (
SqlMetric,
TableColumn,
)
+from superset.connectors.sqla.partition_mapping_storage import (
+ clear_dangling,
+ extract_partition_mapping_patch,
+ partition_mapping_store,
+ PartitionMappingPatch,
+)
from superset.constants import EPOCH_FORMATS
from superset.daos.base import BaseDAO, ColumnOperator, ColumnOperatorEnum
from superset.extensions import db
@@ -441,10 +447,25 @@ class DatasetDAO(BaseDAO[SqlaTable]):
db.session.delete(metric)
attributes = {**(attributes or {}), "changed_on": datetime.now()}
+ # The partition mapping comes out of `attributes` before anything is
+ # applied and goes back in after everything is. The dataset editor's PUT
+ # carries both the typed `partition_column` field and a verbatim copy of
+ # `extra`, and under the extra-JSON store those are the same bytes -- so
+ # whichever `setattr` ran last would win, and which one that is depends
+ # on marshmallow's field ordering. Extracting and re-applying makes the
+ # order a property of this method instead of an accident.
+ #
+ # Stripping is load-bearing beyond ordering: `_upsert_columns` does
+ # `TableColumn(**properties)` with no whitelist, and the model's mapping
+ # attributes are read-only properties, so an unstripped
+ # `partition_value_transform` in a column payload is an AttributeError.
+ mapping_patch = extract_partition_mapping_patch(item, attributes)
+ surviving_column_names: set[str] | None = None
+
if item and attributes:
force_update: bool = False
if "columns" in attributes:
- cls.update_columns(
+ surviving_column_names = cls.update_columns(
item,
attributes.pop("columns"),
override_columns=bool(attributes.get("override_columns")),
@@ -462,7 +483,56 @@ class DatasetDAO(BaseDAO[SqlaTable]):
if force_update:
attributes["changed_on"] = datetime.now()
- return super().update(item, attributes)
+ # `BaseDAO.update` returns `db.session.merge(item)` for a detached item,
+ # so the mapping is written onto whatever came back rather than onto what
+ # went in.
+ updated = super().update(item, attributes)
+
+ if mapping_patch is not None or surviving_column_names is not None:
+ cls.save_partition_mapping(updated, mapping_patch, surviving_column_names)
+ return updated
+
+ @classmethod
+ def create(
+ cls,
+ item: SqlaTable | None = None,
+ attributes: dict[str, Any] | None = None,
+ ) -> SqlaTable:
+ """
+ Creates a Dataset model on the metadata DB.
+
+ Overridden only to route a partition mapping in the payload through the
+ same funnel `update` uses; see the comment there.
+ """
+ mapping_patch = extract_partition_mapping_patch(item, attributes)
+ created = super().create(item, attributes)
+ if mapping_patch is not None:
+ cls.save_partition_mapping(created, mapping_patch, None)
+ return created
+
+ @classmethod
+ def save_partition_mapping(
+ cls,
+ model: SqlaTable,
+ patch: PartitionMappingPatch | None,
+ surviving_column_names: set[str] | None,
+ ) -> None:
+ """
+ The one place a partition filter mapping is written.
+
+ Loads the stored mapping *after* the payload's own `extra` has been
+ applied, so a mapping that arrived through the raw `extra` string is the
+ base the typed fields patch on top of. That gives the two doors into this
+ setting -- the typed API fields and the editor's free-text Extra box --
+ one merge rule instead of two.
+ """
+ store = partition_mapping_store()
+ mapping = store.load(model)
+ if patch is not None:
+ mapping = patch.apply_to(mapping)
+ if surviving_column_names is not None:
+ mapping = clear_dangling(mapping, surviving_column_names)
+ store.save(model, mapping)
@classmethod
def _validate_column_date_formats(
@@ -479,7 +549,7 @@ class DatasetDAO(BaseDAO[SqlaTable]):
@classmethod
def _override_columns(
cls, model: SqlaTable, property_columns: list[dict[str, Any]]
- ) -> None:
+ ) -> set[str]:
"""Replace columns by natural key (``column_name``) — update in place
rather than delete-and-reinsert.
@@ -501,6 +571,8 @@ class DatasetDAO(BaseDAO[SqlaTable]):
columns are preserved. Charts that reference columns by their
``id`` continue to work across a metadata refresh — previously
such references would be invalidated.
+
+ Returns the names of the columns that survive the write.
"""
existing_by_name = {c.column_name: c for c in model.columns}
incoming_by_name = {p["column_name"]: p for p in property_columns}
@@ -512,7 +584,16 @@ class DatasetDAO(BaseDAO[SqlaTable]):
# the old-named row is deleted in the same flush — INSERTs flush
# before DELETEs, so that collides on the PK / UNIQUE(table_id,
# column_name) constraints. ``table_id`` is pinned to *model*.
- protected_keys = ("id", "table_id")
+ # `partition_*` are read-only properties over the mapping store, not
+ # columns, so letting one through would raise on assignment. The DAO's
+ # funnel strips them before this point; this is the backstop for a caller
+ # that reaches `_override_columns` another way.
+ protected_keys = (
+ "id",
+ "table_id",
+ "partition_value_transform",
+ "partition_transform_is_monotonic",
+ )
# Update columns present in both: in-place setattr.
for name, col in existing_by_name.items():
@@ -546,10 +627,17 @@ class DatasetDAO(BaseDAO[SqlaTable]):
}
db.session.add(TableColumn(**{**cleaned, "table_id": model.id}))
+ return set(incoming_by_name)
+
@classmethod
def _upsert_columns(
cls, model: SqlaTable, property_columns: list[dict[str, Any]]
- ) -> None:
+ ) -> set[str]:
+ """
+ Create/update the columns in the payload and delete the rest.
+
+ Returns the names of the columns that survive the write.
+ """
columns_by_id = {column.id: column for column in model.columns}
property_columns_by_id = {
properties["id"]: properties
@@ -557,27 +645,35 @@ class DatasetDAO(BaseDAO[SqlaTable]):
if "id" in properties
}
+ surviving_column_names: set[str] = set()
+
for properties in property_columns:
if "id" not in properties:
db.session.add(TableColumn(**{**properties, "table_id": model.id}))
+ surviving_column_names.add(properties["column_name"])
for properties in property_columns_by_id.values():
col = columns_by_id[properties["id"]]
for key, value in properties.items():
setattr(col, key, value)
+ # A partial update may omit ``column_name``, in which case the
+ # column keeps the name it already had.
+ surviving_column_names.add(col.column_name)
ids_to_keep = property_columns_by_id.keys()
for col in model.columns:
if col.id not in ids_to_keep:
db.session.delete(col)
+ return surviving_column_names
+
@classmethod
def update_columns(
cls,
model: SqlaTable,
property_columns: list[dict[str, Any]],
override_columns: bool = False,
- ) -> None:
+ ) -> set[str]:
"""
Creates/updates and/or deletes a list of columns, based on a
list of Dict.
@@ -589,12 +685,15 @@ class DatasetDAO(BaseDAO[SqlaTable]):
Uses individual ORM operations (not bulk) so that SQLAlchemy-Continuum
can capture each row change in the version history.
+
+ Returns the names of the columns that survive the write, which
+ `save_partition_mapping` needs in order to drop a mapping whose column
+ this call just removed.
"""
cls._validate_column_date_formats(property_columns)
if override_columns:
- cls._override_columns(model, property_columns)
- else:
- cls._upsert_columns(model, property_columns)
+ return cls._override_columns(model, property_columns)
+ return cls._upsert_columns(model, property_columns)
@classmethod
def update_metrics(
diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 5b3e2c9133d..df9592953ea 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -378,6 +378,12 @@ DASHBOARD_DATASET_INACCESSIBLE_FIELDS = (
"main_dttm_col",
"granularity_sqla",
"time_grain_sqla",
+ # Physical column names, same as `columns` above: the partition column is
+ # part of the table's storage layout, which a caller who cannot access the
+ # dataset has no business learning.
+ "partition_column",
+ "partition_mapped_column",
+ "partition_filter_mapping",
)
diff --git a/superset/dashboards/schemas.py b/superset/dashboards/schemas.py
index 7aabb5267e8..d155a11e00a 100644
--- a/superset/dashboards/schemas.py
+++ b/superset/dashboards/schemas.py
@@ -379,6 +379,9 @@ class DashboardDatasetSchema(Schema):
granularity_sqla = fields.List(fields.List(fields.Str()))
normalize_columns = fields.Bool()
always_filter_main_dttm = fields.Bool()
+ partition_column = fields.Str()
+ partition_mapped_column = fields.Str()
+ partition_filter_mapping = fields.Dict(allow_none=True)
# pylint: disable=unused-argument
@post_dump()
diff --git a/superset/datasets/api.py b/superset/datasets/api.py
index 5c276285254..850317d019a 100644
--- a/superset/datasets/api.py
+++ b/superset/datasets/api.py
@@ -23,7 +23,7 @@ from io import BytesIO
from typing import Any, Callable, ClassVar
from zipfile import is_zipfile, ZipFile
-from flask import request, Response
+from flask import current_app as app, request, Response
from flask_appbuilder import permission_name
from flask_appbuilder.api import expose, protect, rison as parse_rison, safe
from flask_appbuilder.api.schemas import get_item_schema
@@ -71,6 +71,7 @@ from superset.commands.purge import (
SoftDeleteBinding,
)
from superset.connectors.sqla.models import SqlaTable
+from superset.connectors.sqla.partition_mapping import preview_partition_mapping
from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
from superset.daos.dashboard import DashboardDAO
from superset.daos.dataset import DatasetDAO
@@ -98,17 +99,21 @@ from superset.datasets.schemas import (
get_related_objects_ids_schema,
GetOrCreateDatasetSchema,
openapi_spec_methods_override,
+ PartitionMappingPreviewSchema,
)
from superset.exceptions import (
OAuth2RedirectError,
+ SupersetSecurityException,
SupersetSyntaxErrorException,
SupersetTemplateException,
SupersetTimeoutException,
)
+from superset.extensions import cache_manager
from superset.jinja_context import BaseTemplateProcessor, get_template_processor
from superset.subjects.filters import FilterRelatedSubjects, subject_type_filter
from superset.utils import json
from superset.utils.core import (
+ get_user_id,
parse_boolean_string,
send_export_zip,
write_zip_entry,
@@ -158,6 +163,51 @@ _DATASET_PURGE_BINDING = SoftDeleteBinding(
)
+#: How long one preview budget lasts, in seconds.
+PREVIEW_RATE_LIMIT_WINDOW = 60
+
+
+def _consume_preview_rate_limit(dataset_id: int) -> bool:
+ """
+ Fixed-window per-user, per-dataset throttle on the preview endpoint.
+
+ Debouncing on the client is a courtesy, not a guard: a held keydown, or a
+ handful of owners with the editor open, becomes sustained load on a
+ production cluster. Returns False once the window's budget is spent.
+
+ Counting is `add` then `inc` rather than read-then-write, which buys two
+ things a `get`/`set` pair cannot. It is atomic where it matters -- on Redis
+ those are `SETNX` and `INCR`, so concurrent previews cannot each read the
+ same sub-limit value and all be let through. And the window is genuinely
+ fixed: only `add` sets a lifetime, so the budget expires a minute after the
+ *first* request rather than a minute after the most recent one, which is
+ what the name promises. (`INCR` leaves the TTL alone; a backend whose `inc`
+ is a read-modify-write may restore its own default lifetime instead, which
+ throttles for longer rather than shorter.)
+
+ Both calls go to the cachelib backend rather than the Flask-Caching wrapper
+ around it, which proxies `add` but not `inc`.
+ """
+ limit = app.config.get("PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT", 30)
+ if not limit:
+ return True
+
+ user_id = get_user_id() or 0
+ key = f"partition_mapping_preview:{user_id}:{dataset_id}"
+ backend = cache_manager.cache.cache
+ try:
+ if backend.add(key, 1, timeout=PREVIEW_RATE_LIMIT_WINDOW):
+ # First request of a fresh window, and the only one that dates it.
+ return True
+ used = backend.inc(key)
+ except Exception: # pylint: disable=broad-except # noqa: BLE001
+ # A cache outage must not take the editor down with it.
+ return True
+ # `inc` reports None when the backend could not count -- a null cache, or a
+ # write that failed. Unenforceable is not the same as spent.
+ return used is None or used <= limit
+
+
class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
datamodel = SQLAInterface(SqlaTable)
@@ -210,6 +260,7 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
"get_or_create_dataset",
"warm_up_cache",
"get_drill_info",
+ "partition_mapping_preview",
"list_versions",
"get_version",
"activity",
@@ -334,6 +385,15 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
"columns.certified_by",
"columns.is_certified",
"columns.warning_markdown",
+ # The partition filter mapping is stored the same way and surfaced the
+ # same way. The two column-level fields are not optional extras here: the
+ # dataset editor seeds itself from this endpoint and then PUTs every
+ # column back, so a transform this payload omits is one the next save
+ # clears.
+ "partition_column",
+ "partition_mapped_column",
+ "columns.partition_value_transform",
+ "columns.partition_transform_is_monotonic",
"metrics.certification_details",
"metrics.certified_by",
"metrics.is_certified",
@@ -370,6 +430,8 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
"description",
"main_dttm_col",
"currency_code_column",
+ "partition_column",
+ "partition_mapped_column",
"normalize_columns",
"always_filter_main_dttm",
"offset",
@@ -446,6 +508,7 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
DatasetPurgeRequestSchema,
DatasetDuplicateSchema,
GetOrCreateDatasetSchema,
+ PartitionMappingPreviewSchema,
VersionListItemSchema,
)
@@ -2075,6 +2138,101 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
entity_concurrency_token(SqlaTable, table.id, table.uuid),
)
+ @expose("/<int:pk>/partition_mapping/preview/", methods=("POST",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs: (
+ f"{self.__class__.__name__}.partition_mapping_preview"
+ ),
+ log_to_statsd=False,
+ )
+ def partition_mapping_preview(self, pk: int) -> Response:
+ """Preview the predicate a partition value transform would emit.
+ ---
+ post:
+ summary: Preview a partition filter mapping
+ description: >-
+ Evaluate a partition value transform at a sample value and return
+ the predicate that would be appended to queries. Parse and
+ denylist checks run before anything reaches the engine, so an
+ unparseable transform costs no query.
+ parameters:
+ - in: path
+ schema:
+ type: integer
+ name: pk
+ description: The dataset ID
+ requestBody:
+ required: true
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/PartitionMappingPreviewSchema'
+ responses:
+ 200:
+ description: Preview result
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ type: object
+ properties:
+ valid:
+ type: boolean
+ emitted_predicate:
+ type: string
+ error:
+ type: string
+ 400:
+ $ref: '#/components/responses/400'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ 429:
+ $ref: '#/components/responses/400'
+ 500:
+ $ref: '#/components/responses/500'
+ """
+ if not is_feature_enabled("PARTITION_FILTER_MAPPING"):
+ return self.response_404()
+
+ dataset = DatasetDAO.find_by_id(pk)
+ if not dataset:
+ return self.response_404()
+ try:
+ security_manager.raise_for_editorship(dataset)
+ except SupersetSecurityException:
+ return self.response_403()
+
+ try:
+ payload = PartitionMappingPreviewSchema().load(request.json)
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+
+ if not _consume_preview_rate_limit(pk):
+ return self.response(
+ 429,
+ message=_(
+ "Too many preview requests for this dataset. "
+ "Wait a moment and try again."
+ ),
+ )
+
+ return self.response(
+ 200,
+ result=preview_partition_mapping(
+ dataset,
+ mapped_column=payload["mapped_column"],
+ value_transform=payload["value_transform"],
+ sample_value=payload["sample_value"],
+ ),
+ )
+
@expose("/<int:pk>/drill_info/", methods=("GET",))
@protect()
@parse_rison(get_drill_info_schema)
diff --git a/superset/datasets/schemas.py b/superset/datasets/schemas.py
index 307f268b584..0ce73f4de18 100644
--- a/superset/datasets/schemas.py
+++ b/superset/datasets/schemas.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
from datetime import datetime
-from typing import Any
+from typing import Any, NoReturn
from uuid import UUID
from dateutil.parser import isoparse
@@ -31,6 +31,11 @@ from marshmallow import (
from marshmallow.validate import Length, OneOf, Range
from superset import security_manager
+from superset.connectors.sqla.partition_mapping_storage import (
+ EXTRA_KEY as PARTITION_MAPPING_EXTRA_KEY,
+ MAX_COLUMN_NAME_LENGTH,
+ MAX_TRANSFORM_LENGTH,
+)
from superset.constants import EPOCH_FORMATS
from superset.exceptions import SupersetMarshmallowValidationError
from superset.models.sql_types import parse_currency_string
@@ -79,6 +84,92 @@ def validate_python_date_format(dt_format: str) -> bool:
return True
+def validate_dataset_extra(value: str | None) -> None:
+ """
+ Structural check on the ``partition_filter_mapping`` key inside ``extra``.
+
+ That key only, and only structurally. ``extra`` is a free-text JSON box in
+ the dataset editor, so the partition mapping has a second door into storage
+ that bypasses the typed fields below -- and a value of the wrong type there is
+ read as "not configured" and silently ignored, which is the failure mode an
+ owner has no way to diagnose. This turns it into an error message.
+
+ Undecodable JSON passes through. ``extra`` has never been validated on this
+ endpoint, so a deployment may already hold a value that is not JSON at all;
+ rejecting it now would fail an owner's next save over a field they did not
+ touch. Nothing else in ``extra`` is checked either -- ``certification``,
+ ``warning_markdown``, ``timezone`` and unknown keys are all legitimate.
+
+ Semantics are not checked here: whether the partition column exists, whether
+ the transform parses, and whether it calls a non-deterministic function all
+ depend on the dataset's columns and its engine, which a schema cannot see.
+ `UpdateDatasetCommand._validate_partition_mapping` merges this value with the
+ typed fields and answers all three.
+ """
+ if not value:
+ return
+ try:
+ decoded = json.loads(value)
+ except (TypeError, json.JSONDecodeError):
+ return
+ if not isinstance(decoded, dict) or PARTITION_MAPPING_EXTRA_KEY not in decoded:
+ return
+
+ blob = decoded[PARTITION_MAPPING_EXTRA_KEY]
+ if not isinstance(blob, dict):
+ _reject("%(key)s in extra must be an object.", key=PARTITION_MAPPING_EXTRA_KEY)
+
+ for key, limit in (
+ ("partition_column", MAX_COLUMN_NAME_LENGTH),
+ ("mapped_column", MAX_COLUMN_NAME_LENGTH),
+ ):
+ _reject_unless_string(blob.get(key), key, limit)
+
+ _validate_column_transforms(blob.get("column_transforms"))
+
+
+def _validate_column_transforms(transforms: Any) -> None:
+ if transforms is None:
+ return
+ if not isinstance(transforms, dict):
+ _reject("column_transforms in extra must be an object.")
+ for column_name, entry in transforms.items():
+ if not isinstance(entry, dict):
+ _reject(
+ "column_transforms.%(name)s in extra must be an object.",
+ name=column_name,
+ )
+ _reject_unless_string(
+ entry.get("value_transform"), "value_transform", MAX_TRANSFORM_LENGTH
+ )
+ if entry.get("is_monotonic") not in (None, True, False):
+ _reject("is_monotonic in extra must be a boolean.")
+
+
+def _reject(message: str, **kwargs: Any) -> NoReturn:
+ """
+ Raise with a rendered message.
+
+ ``str()`` is not decoration: marshmallow special-cases ``str`` and ``dict``
+ and falls back to ``list(messages)`` for anything else, which turns a
+ ``LazyString`` into a list of single characters.
+ """
+ raise ValidationError(str(_(message, **kwargs)))
+
+
+def _reject_unless_string(value: Any, key: str, limit: int) -> None:
+ if value is None:
+ return
+ if not isinstance(value, str):
+ _reject("%(key)s in extra must be a string.", key=key)
+ if len(value) > limit:
+ _reject(
+ "%(key)s in extra must be at most %(limit)d characters.",
+ key=key,
+ limit=limit,
+ )
+
+
class DatasetColumnsPutSchema(Schema):
id = fields.Integer(required=False)
column_name = fields.String(required=True, validate=Length(1, 255))
@@ -101,6 +192,22 @@ class DatasetColumnsPutSchema(Schema):
datetime_format = fields.String(
allow_none=True, validate=[Length(1, 100), validate_python_date_format]
)
+ partition_value_transform = fields.String(
+ allow_none=True,
+ metadata={
+ "description": (
+ "SQL expression containing a :value placeholder. Filters on "
+ "this column are mirrored onto the dataset's partition column "
+ "with the value passed through this transform."
+ )
+ },
+ )
+ # Deliberately no `load_default`: `DatasetDAO.update_columns` applies the
+ # loaded payload field by field onto the stored column, so a default here
+ # would let a partial column payload clear a monotonic flag the request
+ # never mentioned -- and silently stop mirroring range filters. Absent
+ # means "unchanged"; new columns fall back to the model's own default.
+ partition_transform_is_monotonic = fields.Boolean(allow_none=True)
uuid = fields.UUID(allow_none=True)
@@ -179,6 +286,8 @@ class DatasetPostSchema(Schema):
normalize_columns = fields.Boolean(load_default=False)
always_filter_main_dttm = fields.Boolean(load_default=False)
currency_code_column = fields.String(allow_none=True, validate=Length(0, 250))
+ partition_column = fields.String(allow_none=True, validate=Length(0, 250))
+ partition_mapped_column = fields.String(allow_none=True, validate=Length(0, 250))
template_params = fields.String(allow_none=True)
uuid = fields.UUID(allow_none=True)
@@ -194,6 +303,8 @@ class DatasetPutSchema(DiscardIsManagedExternallyMixin, Schema):
description = fields.String(allow_none=True)
main_dttm_col = fields.String(allow_none=True)
currency_code_column = fields.String(allow_none=True, validate=Length(0, 250))
+ partition_column = fields.String(allow_none=True, validate=Length(0, 250))
+ partition_mapped_column = fields.String(allow_none=True, validate=Length(0, 250))
normalize_columns = fields.Boolean(allow_none=True, dump_default=False)
always_filter_main_dttm = fields.Boolean(load_default=False)
offset = fields.Integer(allow_none=True)
@@ -205,7 +316,7 @@ class DatasetPutSchema(DiscardIsManagedExternallyMixin, Schema):
columns = fields.List(fields.Nested(DatasetColumnsPutSchema))
metrics = fields.List(fields.Nested(DatasetMetricsPutSchema))
folders = fields.List(fields.Nested(FolderSchema), required=False)
- extra = fields.String(allow_none=True)
+ extra = fields.String(allow_none=True, validate=validate_dataset_extra)
external_url = fields.String(allow_none=True)
uuid = fields.UUID(allow_none=True)
@@ -469,6 +580,9 @@ class ImportV1DatasetSchema(Schema):
external_url = fields.String(allow_none=True)
normalize_columns = fields.Boolean(load_default=False)
always_filter_main_dttm = fields.Boolean(load_default=False)
+ # No partition mapping fields: it rides inside `extra`, which is already a
+ # `fields.Dict` here and already round-trips through export/import, so
+ # declaring them would export the same values twice.
folders = fields.List(fields.Nested(FolderSchema), required=False, allow_none=True)
# data_file is used by the example loading system to reference Parquet files
data_file = fields.String(allow_none=True, load_default=None)
@@ -496,6 +610,38 @@ class GetOrCreateDatasetSchema(Schema):
always_filter_main_dttm = fields.Boolean(load_default=False)
+class PartitionMappingPreviewSchema(Schema):
+ """
+ Payload for the dataset editor's partition mapping preview panel.
+
+ Every field is bounded. The endpoint parses `value_transform` with sqlglot
+ and then evaluates it against the warehouse, so an unbounded string is
+ parser time and warehouse time an owner can spend at will; the bounds keep
+ a malformed or oversized payload a 400 rather than work.
+ """
+
+ mapped_column = fields.String(
+ required=True,
+ # Matches the `String(250)` the mapping columns are stored in.
+ validate=Length(1, 250),
+ metadata={"description": "Column whose filters would be mirrored"},
+ )
+ value_transform = fields.String(
+ required=True,
+ allow_none=True,
+ # The stored column is `Text`, so this bounds the *request*, not the
+ # feature: a transform is one expression around `:value`, and 1024
+ # characters is far past anything that reads as one.
+ validate=Length(1, 1024),
+ metadata={"description": "SQL expression containing a :value placeholder"},
+ )
+ sample_value = fields.String(
+ required=True,
+ validate=Length(1, 250),
+ metadata={"description": "Value to evaluate the transform at"},
+ )
+
+
class DatasetCacheWarmUpRequestSchema(Schema):
db_name = fields.String(
required=True,
diff --git a/superset/mcp_service/common/schema_discovery.py b/superset/mcp_service/common/schema_discovery.py
index 9e1ee130a8f..36788a2b637 100644
--- a/superset/mcp_service/common/schema_discovery.py
+++ b/superset/mcp_service/common/schema_discovery.py
@@ -365,6 +365,25 @@ DATASET_SORTABLE_COLUMNS = [
# UUID search separately and converts it to an exact ``uuid`` filter.
DATASET_SEARCH_COLUMNS = ["table_name", "description", "schema", "sql"]
DATASET_EXTRA_COLUMNS: dict[str, ColumnMetadata] = {
+ # Partition filter mapping. Listed here rather than in
+ # `_COLUMN_DESCRIPTIONS` because `get_columns_from_model` walks
+ # `mapper.columns`, and these are read-only properties over the mapping
+ # store, not columns -- a description there would never be reached.
+ "partition_column": ColumnMetadata(
+ name="partition_column",
+ description="Physical column the engine partitions on",
+ type="str",
+ is_default=False,
+ ),
+ "partition_mapped_column": ColumnMetadata(
+ name="partition_mapped_column",
+ description=(
+ "Column whose filters are mirrored onto the partition column; "
+ "defaults to the main datetime column"
+ ),
+ type="str",
+ is_default=False,
+ ),
"database_name": ColumnMetadata(
name="database_name",
description="Database connection name",
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
index 47e10c8b708..5820cd07628 100644
--- a/superset/models/helpers.py
+++ b/superset/models/helpers.py
@@ -50,6 +50,7 @@ import pandas as pd
import pytz
import sqlalchemy as sa
import yaml
+from dateutil.relativedelta import relativedelta
from flask import current_app as app, g
from flask_appbuilder import Model
from flask_appbuilder.models.decorators import renders
@@ -89,6 +90,12 @@ from superset.common.utils.time_range_utils import (
get_since_until_from_query_object,
get_since_until_from_time_range,
)
+from superset.connectors.sqla.partition_mapping import (
+ build_mirrored_predicates,
+ grain_bucket_width,
+ PartitionMapping,
+ resolve_partition_mapping,
+)
from superset.constants import (
CacheRegion,
EMPTY_STRING,
@@ -170,7 +177,7 @@ class ValidationResultDict(TypedDict):
if TYPE_CHECKING:
from superset.common.query_object import QueryObject
- from superset.connectors.sqla.models import SqlMetric, TableColumn
+ from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn
from superset.db_engine_specs import BaseEngineSpec
from superset.models.core import Database
@@ -4288,36 +4295,153 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
return f"""'{dttm.strftime("%Y-%m-%d %H:%M:%S.%f")}'"""
- def get_time_filter( # pylint: disable=too-many-arguments # noqa: C901
+ def _collect_partition_mirror_range(
self,
- time_col: "TableColumn",
+ mapping: Optional["PartitionMapping"],
+ sink: list[tuple[utils.FilterOperator, Any]],
+ column_name: str,
start_dttm: Optional[sa.DateTime],
end_dttm: Optional[sa.DateTime],
- time_grain: Optional[str] = None,
- label: Optional[str] = "__time",
- template_processor: Optional[BaseTemplateProcessor] = None,
- ) -> Optional[ColumnElement]:
- col = (
- time_col.get_timestamp_expression(
- time_grain=time_grain,
- label=label,
- template_processor=template_processor,
- )
- if time_grain
- else self.convert_tbl_column_to_sqla_col(
- time_col, label=label, template_processor=template_processor
- )
+ widen_bounds_by: Optional[relativedelta] = None,
+ ) -> None:
+ """
+ Record a time range for mirroring onto the partition column.
+
+ The bounds are adjusted first, by the same helper `get_time_filter` uses:
+ the mirrored predicate has to describe the same instants as the
+ timestamp predicate it stands in for, or the pruning is wrong by exactly
+ the dataset's timezone offset -- silently.
+
+ Either bound may be `None` for an open-ended range, in which case only
+ the bound that exists is mirrored.
+
+ `widen_bounds_by` is one grain bucket width, passed when the filter this
+ stands in for compares a *truncated* column. The real predicate then
+ keeps rows the raw bounds exclude, and widening both ends by one bucket
+ is the smallest range guaranteed to contain all of them -- see the "Time
+ grains" section of `superset.connectors.sqla.partition_mapping`. Note
+ the widened upper bound is `<=` rather than `<`: `ts < until + width`
+ only gives `T(ts) <= T(until + width)` for a transform that is monotonic
+ but not strictly so, such as `unix_timestamp` on a sub-second column.
+
+ A widened range can also be wide enough to fail `_bounds_are_ordered`
+ against a tighter filter on the same column. That fails open -- no
+ mirroring, no pruning, no dropped rows.
+ """
+ if mapping is None or column_name != mapping.mapped_column:
+ return
+ if not mapping.mirrors(utils.FilterOperator.TEMPORAL_RANGE):
+ return
+
+ # Resolve the column so the offset adjustment sees the same type
+ # `get_time_filter` does -- an offset the column's type cannot represent
+ # is dropped, and the mirrored bounds have to agree with the originals.
+ mapped_col = next(
+ (col for col in self.columns if col.column_name == column_name), None
)
+ start_dttm, end_dttm = self.adjust_time_bounds(start_dttm, end_dttm, mapped_col)
+
+ # Adjust first, then widen. `adjust_time_bounds` moves the naive UI
+ # bounds into the frame the column is *stored* in, and the grain
+ # truncation in the real predicate happens in that same frame, so the
+ # bucket width has to be added there too. The two commute for the
+ # hour-offset branch but not for the `ZoneInfo` one, where widening
+ # across a DST boundary first lands an hour out.
+ upper_operator = utils.FilterOperator.LESS_THAN
+ if widen_bounds_by is not None:
+ if start_dttm is not None:
+ start_dttm -= widen_bounds_by
+ if end_dttm is not None:
+ end_dttm += widen_bounds_by
+ upper_operator = utils.FilterOperator.LESS_THAN_OR_EQUALS
+
+ if start_dttm is not None:
+ sink.append((utils.FilterOperator.GREATER_THAN_OR_EQUALS, start_dttm))
+ if end_dttm is not None:
+ sink.append((upper_operator, end_dttm))
+
+ def _collect_partition_mirror_filter(
+ self,
+ mapping: Optional["PartitionMapping"],
+ sink: list[tuple[utils.FilterOperator, Any]],
+ column_name: str,
+ operator: utils.FilterOperator,
+ value: Any,
+ ) -> None:
+ """
+ Record a structured filter for mirroring onto the partition column.
+
+ Only operators the mapping declares safe are recorded; see the operator
+ matrix in `superset.connectors.sqla.partition_mapping`.
+ """
+ if mapping is None or column_name != mapping.mapped_column:
+ return
+ if not mapping.mirrors(operator):
+ return
+
+ if operator == utils.FilterOperator.IN:
+ if not isinstance(value, (list, tuple)) or not value:
+ return
+ if any(item is None for item in value):
+ # A `None` in the list widens the real predicate to
+ # `col IS NULL OR col IN (...)`. Mirroring only the non-null
+ # members would be *narrower* than the original filter and
+ # would drop rows it keeps.
+ return
+ sink.append((operator, tuple(value)))
+ elif value is not None:
+ sink.append((operator, value))
+
+ def _build_partition_mirror_predicates(
+ self,
+ mapping: "PartitionMapping",
+ requests: list[tuple[utils.FilterOperator, Any]],
+ ) -> list[ColumnElement]:
+ """
+ Resolve collected mirror requests into predicates on the partition column.
+
+ Requests are deduplicated first: a chart can reach both collection points
+ for the same column -- a `granularity` time filter *and* a
+ `TEMPORAL_RANGE` ad-hoc filter on the same column is a routine Explore
+ configuration -- and emitting the predicate twice is harmless SQL but
+ makes "View query" surprising.
+ """
+ if not requests:
+ return []
- # Resolve dataset-level time-boundary adjustments. A configured
- # `extra.timezone` (an IANA name, DST-aware) takes precedence: naive UI
- # boundaries are interpreted in that zone and converted to UTC for
- # comparison with UTC-stored data. When no timezone is configured, fall
- # back to the legacy "Hour Offset" field instead: displayed values are
- # shifted by +offset hours (see normalize_df / DateColumn in
- # superset.utils.core), but the time filter compares raw stored values, so
- # bounds are shifted by -offset to stay consistent with what's displayed
- # (#104810).
+ deduped: list[tuple[utils.FilterOperator, Any]] = []
+ seen: set[Any] = set()
+ for operator, value in requests:
+ key = (operator, value if isinstance(value, Hashable) else repr(value))
+ if key not in seen:
+ seen.add(key)
+ deduped.append((operator, value))
+
+ return build_mirrored_predicates(cast("SqlaTable", self), mapping, deduped)
+
+ def adjust_time_bounds(
+ self,
+ start_dttm: Optional[sa.DateTime],
+ end_dttm: Optional[sa.DateTime],
+ time_col: Optional["TableColumn"] = None,
+ ) -> tuple[Optional[sa.DateTime], Optional[sa.DateTime]]:
+ """
+ Apply the dataset's time-boundary adjustments to a pair of bounds.
+
+ A configured `extra.timezone` (an IANA name, DST-aware) takes
+ precedence: naive UI boundaries are interpreted in that zone and
+ converted to UTC for comparison with UTC-stored data. When no timezone
+ is configured, fall back to the legacy "Hour Offset" field instead:
+ displayed values are shifted by +offset hours (see normalize_df /
+ DateColumn in superset.utils.core), but the time filter compares raw
+ stored values, so bounds are shifted by -offset to stay consistent with
+ what's displayed (#104810).
+
+ Extracted so callers that need the *adjusted* bounds for something other
+ than building the clause -- partition filter mirroring resolves them
+ against the engine -- see exactly the instants the clause compares
+ against, rather than a copy of this logic that can drift.
+ """
dataset_timezone = self.get_dataset_timezone()
if dataset_timezone and (start_dttm or end_dttm):
@@ -4350,7 +4474,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
if not dataset_timezone and (offset_hours := getattr(self, "offset", 0) or 0):
offset_hours = get_effective_hours_offset(
self.db_engine_spec,
- time_col.type,
+ time_col.type if time_col is not None else None,
offset_hours,
db_extra=self.db_extra,
)
@@ -4359,6 +4483,31 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
if end_dttm is not None:
end_dttm = end_dttm - timedelta(hours=offset_hours)
+ return start_dttm, end_dttm
+
+ def get_time_filter( # pylint: disable=too-many-arguments
+ self,
+ time_col: "TableColumn",
+ start_dttm: Optional[sa.DateTime],
+ end_dttm: Optional[sa.DateTime],
+ time_grain: Optional[str] = None,
+ label: Optional[str] = "__time",
+ template_processor: Optional[BaseTemplateProcessor] = None,
+ ) -> Optional[ColumnElement]:
+ col = (
+ time_col.get_timestamp_expression(
+ time_grain=time_grain,
+ label=label,
+ template_processor=template_processor,
+ )
+ if time_grain
+ else self.convert_tbl_column_to_sqla_col(
+ time_col, label=label, template_processor=template_processor
+ )
+ )
+
+ start_dttm, end_dttm = self.adjust_time_bounds(start_dttm, end_dttm, time_col)
+
l = [] # noqa: E741
if start_dttm:
l.append(
@@ -5000,6 +5149,13 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
time_filters = []
+ # Partition filter mapping: mirror requests are collected here as
+ # `(operator, value)` pairs and resolved in a single probe round trip
+ # after the filter loop, rather than emitted inline at each of the eight
+ # append sites. `None` when the dataset has no usable mapping.
+ partition_mapping = resolve_partition_mapping(cast("SqlaTable", self))
+ partition_mirror: list[tuple[utils.FilterOperator, Any]] = []
+
# Process FROM clause early to populate removed_filters from virtual dataset
# templates before we decide whether to add time filters
tbl, cte = self.get_from_clause(template_processor)
@@ -5039,6 +5195,16 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
)
if _main_dttm_filter is not None:
time_filters.append(_main_dttm_filter)
+ # This block filters `main_dttm_col`, a *different* column
+ # from `dttm_col`. Checking only `dttm_col` below would miss
+ # it when the mapping tracks the main datetime column.
+ self._collect_partition_mirror_range(
+ partition_mapping,
+ partition_mirror,
+ self.main_dttm_col,
+ from_dttm,
+ to_dttm,
+ )
# Check if time filter should be skipped because it was handled in template.
# Check both the actual column name and __timestamp alias
@@ -5056,6 +5222,13 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
)
if time_filter_column is not None:
time_filters.append(time_filter_column)
+ self._collect_partition_mirror_range(
+ partition_mapping,
+ partition_mirror,
+ dttm_col.column_name,
+ from_dttm,
+ to_dttm,
+ )
# Gate on `groupby_all_columns` rather than the raw dimensions: it is the
# real GROUP BY signal and also captures the timeseries time bucket. A
@@ -5315,6 +5488,39 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
db_extra=self.db_extra,
)
+ # Mirror onto the partition column. This single site covers
+ # ad-hoc filters, dashboard native filters and cross-filters,
+ # because they all arrive as entries in `filter`. `eq` rather
+ # than the raw `val`: it is the value the real predicate uses.
+ # `TEMPORAL_RANGE` is collected in its own branch below, where
+ # the range has been resolved into a pair of bounds.
+ #
+ # A grain makes the real predicate compare the *truncated*
+ # column, so the raw value no longer describes the rows it
+ # matches: drill-to-detail sends `==` on a bucket start, which
+ # every row in the bucket satisfies once truncated. Mirroring it
+ # raw would keep only the bucket's first instant.
+ #
+ # Grained ranges do mirror, by widening (see the TEMPORAL_RANGE
+ # branch below). The same trick is deferred rather than unsafe
+ # here: `==` on a bucket would have to become a two-sided range,
+ # which needs a monotonic transform -- `EQUALS` mirrors without
+ # one today -- and a parse back into a datetime; and a grained
+ # `IN` would need a union of buckets, which this sink cannot
+ # express because its entries are AND-ed.
+ if (
+ col_obj is not None
+ and op != utils.FilterOperator.TEMPORAL_RANGE
+ and not filter_grain
+ ):
+ self._collect_partition_mirror_filter(
+ partition_mapping,
+ partition_mirror,
+ col_obj.column_name,
+ op,
+ eq,
+ )
+
# Get ADVANCED_DATA_TYPES from config when needed
ADVANCED_DATA_TYPES = app.config.get("ADVANCED_DATA_TYPES", {}) # noqa: N806
@@ -5592,6 +5798,29 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
)
if _temporal_filter is not None:
target_clause_list.append(_temporal_filter)
+ # A grained range truncates the column before
+ # comparing, so a row in the final partial bucket
+ # satisfies `DATE_TRUNC(...) < until` while the raw
+ # upper bound excludes it. Which direction a grain
+ # rounds is not knowable from the duration -- "week
+ # ending Saturday" rounds forward, Ocient rounds to
+ # nearest -- but the displacement is bounded by the
+ # grain's own bucket width either way, so widening
+ # both bounds by one bucket is never narrower than
+ # the real predicate. A grain whose SQL an operator
+ # supplied has no width we know, and still skips.
+ widen_by = grain_bucket_width(
+ flt_grain, db_engine_spec.engine
+ )
+ if not flt_grain or widen_by is not None:
+ self._collect_partition_mirror_range(
+ partition_mapping,
+ partition_mirror,
+ col_obj.column_name,
+ _since,
+ _until,
+ widen_bounds_by=widen_by,
+ )
else:
raise QueryObjectValidationError(
_("Invalid filter operation type: %(op)s", op=op)
@@ -5600,6 +5829,11 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
# col_obj is None and sqla_col is None - column not found!
# Silently skip - this can happen for removed columns or invalid filters
pass
+ if partition_mapping is not None:
+ where_clause_and += self._build_partition_mirror_predicates(
+ partition_mapping,
+ partition_mirror,
+ )
where_clause_and += self.get_sqla_row_level_filters(template_processor)
if extras:
where = extras.get("where")
diff --git a/superset/sql/parse.py b/superset/sql/parse.py
index 667b818e94c..ff975481ee3 100644
--- a/superset/sql/parse.py
+++ b/superset/sql/parse.py
@@ -1043,6 +1043,44 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
}
)
+ def get_niladic_functions(self) -> set[str]:
+ """
+ Names of functions called with no arguments.
+
+ Some functions mean something entirely different with an empty argument
+ list: on Hive and Impala ``unix_timestamp()`` is the current time while
+ ``unix_timestamp(x)`` is a pure conversion. Callers that care about
+ determinism have to tell those apart by arity, so a name-based check
+ like ``check_functions_present`` is not enough.
+ """
+ niladic: set[str] = set()
+ for function in self._parsed.find_all(exp.Func):
+ sql_name = function.sql_name()
+ name = function.name.upper() if sql_name == "ANONYMOUS" else sql_name
+ if not self._function_args(function):
+ niladic.add(name.upper())
+ return niladic
+
+ @staticmethod
+ def _function_args(function: exp.Func) -> list[Any]:
+ """
+ The arguments a function node was called with.
+
+ `exp.Anonymous` keeps the function *name* in `this` and the arguments in
+ `expressions`, while a named node like `exp.Lower` keeps its single
+ argument in `this` -- so the two shapes have to be read differently or
+ every anonymous call looks like it takes one argument.
+ """
+ if isinstance(function, exp.Anonymous):
+ return list(function.expressions or [])
+
+ args: list[Any] = []
+ for value in function.args.values():
+ if value is None:
+ continue
+ args.extend(value if isinstance(value, list) else [value])
+ return args
+
def __init__(
self,
statement: str | None = None,
diff --git a/superset/superset_typing.py b/superset/superset_typing.py
index 56769707a85..78a0bf6f46e 100644
--- a/superset/superset_typing.py
+++ b/superset/superset_typing.py
@@ -292,6 +292,9 @@ class ExplorableData(TypedDict, total=False):
time_grain_sqla: Available time grains
main_dttm_col: Main datetime column
currency_code_column: Column containing currency codes for dynamic formatting
+ partition_column: Physical column the engine partitions on
+ partition_mapped_column: Explicit override for the mirrored column
+ partition_filter_mapping: Summary of the active mapping, or None
fetch_values_predicate: Predicate for fetching filter values
template_params: Template parameters for Jinja
is_sqllab_view: Whether this is a SQL Lab view
@@ -351,6 +354,12 @@ class ExplorableData(TypedDict, total=False):
extra: str | None
always_filter_main_dttm: bool
normalize_columns: bool
+ partition_column: str | None
+ partition_mapped_column: str | None
+ # Self-contained summary for the Explore indicator. Kept separate from
+ # `columns` because `data_for_slices` prunes columns no chart references,
+ # and the partition column is typically referenced by none of them.
+ partition_filter_mapping: dict[str, Any] | None
rls_filters: list[dict[str, Any]]
# Set by datasources that cannot return raw row samples (e.g. semantic
# views, which only expose pre-defined metrics and dimensions).
diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py
index 6a6ff7e8a23..f0e83b28454 100644
--- a/tests/integration_tests/datasets/api_tests.py
+++ b/tests/integration_tests/datasets/api_tests.py
@@ -802,6 +802,7 @@ class TestDatasetApi(SupersetTestCase):
"can_get_or_create_dataset",
"can_warm_up_cache",
"can_get_drill_info",
+ "can_partition_mapping_preview",
}
def test_create_dataset_item(self):
diff --git a/tests/unit_tests/commands/dataset/update_test.py b/tests/unit_tests/commands/dataset/update_test.py
index 86bd6158107..9f2b8b81a9c 100644
--- a/tests/unit_tests/commands/dataset/update_test.py
+++ b/tests/unit_tests/commands/dataset/update_test.py
@@ -36,6 +36,11 @@ from superset.commands.dataset.update import (
UpdateDatasetCommand,
validate_folders,
)
+from superset.connectors.sqla.partition_mapping_storage import (
+ ColumnTransform,
+ extra_with_partition_mapping,
+ StoredPartitionMapping,
+)
from superset.datasets.schemas import FolderSchema
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetSecurityException
@@ -95,6 +100,7 @@ def test_update_dataset_sql_authorized_schema(mocker: MockerFixture) -> None:
mock_dataset.schema = "public"
mock_dataset.table_name = "test_table"
mock_dataset.editors = [] # No editors to avoid computation issues
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_database
@@ -140,6 +146,7 @@ def test_update_dataset_sql_unauthorized_schema(mocker: MockerFixture) -> None:
mock_dataset.schema = "public"
mock_dataset.table_name = "test_table"
mock_dataset.editors = [] # No editors to avoid computation issues
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_database
@@ -203,6 +210,7 @@ def test_update_dataset_database_id_change_checks_new_database_access(
mock_dataset.schema = "public"
mock_dataset.table_name = "test_table"
mock_dataset.editors = [] # No editors to avoid computation issues
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_new_database
@@ -259,6 +267,7 @@ def test_update_dataset_database_id_change_allowed_with_access(
mock_dataset.schema = "public"
mock_dataset.table_name = "test_table"
mock_dataset.editors = [] # No editors to avoid computation issues
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_new_database
@@ -308,6 +317,7 @@ def test_update_dataset_physical_repoint_requires_table_access(
mock_dataset.table_name = "allowed_table"
mock_dataset.sql = None # physical dataset
mock_dataset.editors = []
+ mock_dataset.extra = None
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.validate_update_uniqueness.return_value = True
@@ -501,6 +511,7 @@ def test_update_dataset_rejects_malicious_expression(
mock_dataset.database = mock_database
mock_dataset.catalog = "catalog"
mock_dataset.schema = None
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_database
mock_dataset_dao.validate_update_uniqueness.return_value = True
@@ -550,6 +561,7 @@ def test_update_dataset_accepts_benign_expression(mocker: MockerFixture) -> None
mock_dataset.database = mock_database
mock_dataset.catalog = "catalog"
mock_dataset.schema = None
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_database
mock_dataset_dao.validate_update_uniqueness.return_value = True
@@ -591,6 +603,7 @@ def test_update_dataset_accepts_jinja_expression(mocker: MockerFixture) -> None:
mock_dataset.database = mock_database
mock_dataset.catalog = "catalog"
mock_dataset.schema = None
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_database
mock_dataset_dao.validate_update_uniqueness.return_value = True
@@ -1352,6 +1365,7 @@ def test_update_dataset_rejects_malicious_fetch_values_predicate(
mock_dataset.database = mock_database
mock_dataset.catalog = "catalog"
mock_dataset.schema = None
+ mock_dataset.extra = None # No partition filter mapping
mock_dataset_dao.find_by_id.return_value = mock_dataset
mock_dataset_dao.get_database_by_id.return_value = mock_database
mock_dataset_dao.validate_update_uniqueness.return_value = True
@@ -1366,3 +1380,130 @@ def test_update_dataset_rejects_malicious_fetch_values_predicate(
and "fetch_values_predicate" in (exc.field_name or "")
for exc in excinfo.value._exceptions
)
+
+
+def _mapping_command(
+ mocker: MockerFixture,
+ transform: str | None,
+ *,
+ properties: dict[str, Any] | None = None,
+) -> UpdateDatasetCommand:
+ """
+ A command whose stored dataset maps `event_time` onto `dt_epoch`.
+
+ The mapping is set up as stored state -- an `extra` blob -- rather than as
+ mocked attributes, because the command reads it back through the store. That
+ is the point: validation runs against exactly the mapping the DAO will
+ persist, so a mock of the model's properties would be testing nothing.
+ """
+ mapped_column = mocker.MagicMock()
+ mapped_column.column_name = "event_time"
+ partition_column = mocker.MagicMock()
+ partition_column.column_name = "dt_epoch"
+
+ mock_dataset = mocker.MagicMock(is_managed_externally=False)
+ mock_dataset.database.backend = "sqlite"
+ mock_dataset.catalog = None
+ mock_dataset.schema = "main"
+ mock_dataset.columns = [mapped_column, partition_column]
+ mock_dataset.main_dttm_col = "event_time"
+ mock_dataset.extra = extra_with_partition_mapping(
+ None,
+ StoredPartitionMapping(
+ partition_column="dt_epoch",
+ column_transforms=(
+ {"event_time": ColumnTransform(transform)} if transform else {}
+ ),
+ ),
+ )
+
+ command = UpdateDatasetCommand(1, properties or {})
+ command._model = mock_dataset
+ return command
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_an_unparseable_transform_does_not_block_the_save(
+ mocker: MockerFixture,
+) -> None:
+ """
+ An unparseable transform is a Tier-2 issue: the mapping saves and stays
+ inactive. `validate_stored_expression` rejects anything it cannot parse, so
+ running it here would turn that into a blocking error and cost the owner the
+ rest of their edits.
+ """
+ gate = mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value")
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert exceptions == []
+ gate.assert_not_called()
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_a_parseable_transform_still_goes_through_the_stored_expression_gate(
+ mocker: MockerFixture,
+) -> None:
+ """The parser gate that governs every other stored expression still runs."""
+ gate = mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value)")
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert exceptions == []
+ gate.assert_called_once()
+ assert ":value" not in gate.call_args.args[-1]
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_a_jinja_transform_is_still_rejected(mocker: MockerFixture) -> None:
+ """
+ Skipping the gate for unparseable transforms is not a hole for templating:
+ Jinja is a blocking issue of its own, reported before the gate is reached.
+ """
+ mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp({{ current_user_id() }})")
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert [exc.field_name for exc in exceptions] == ["partition_value_transform"]
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=True)
+def test_a_self_mapping_is_rejected_while_the_feature_is_on(
+ mocker: MockerFixture,
+) -> None:
+ """A column cannot stand in for itself: that is a Tier-1 blocking issue."""
+ mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value)")
+ command._properties["partition_mapped_column"] = "dt_epoch"
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert [exc.field_name for exc in exceptions] == ["partition_column"]
+
+
+@with_feature_flags(PARTITION_FILTER_MAPPING=False)
+def test_no_mapping_validation_runs_while_the_feature_is_off(
+ mocker: MockerFixture,
+) -> None:
+ """
+ With the flag off nothing mirrors, so a stored mapping can never be
+ consumed. Rejecting the save over it would hand the owner a validation
+ error they have no way to act on -- and every path that reads a mapping is
+ gated the same way.
+ """
+ gate = mocker.patch("superset.commands.dataset.update.validate_stored_expression")
+ command = _mapping_command(mocker, "unix_timestamp(:value)")
+ command._properties["partition_mapped_column"] = "dt_epoch"
+
+ exceptions: list[ValidationError] = []
+ command._validate_partition_mapping(exceptions)
+
+ assert exceptions == []
+ gate.assert_not_called()
diff --git a/tests/unit_tests/connectors/sqla/partition_mapping_storage_test.py b/tests/unit_tests/connectors/sqla/partition_mapping_storage_test.py
new file mode 100644
index 00000000000..314b81d03f8
--- /dev/null
+++ b/tests/unit_tests/connectors/sqla/partition_mapping_storage_test.py
@@ -0,0 +1,544 @@
+# 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.
+"""
+Where a partition filter mapping is persisted.
+
+The feature ships stored in ``tables.extra`` and moves to four real columns in a
+later migration, so these tests are about the seam rather than the feature: that
+the blob survives its neighbours, that a malformed one degrades to "not
+configured" instead of propagating, that a write dirties the column exactly when
+it should, and that the store is the only way in.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+import sqlalchemy as sa
+from flask import Flask
+
+from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.connectors.sqla.partition_mapping_storage import (
+ ColumnPartitionMappingStore,
+ ColumnTransform,
+ EXTRA_KEY,
+ extra_with_partition_mapping,
+ ExtraJsonPartitionMappingStore,
+ load_partition_mapping,
+ MAX_COLUMN_NAME_LENGTH,
+ MAX_TRANSFORM_LENGTH,
+ parse_mapping,
+ partition_mapping_store,
+ set_partition_mapping,
+ StoredPartitionMapping,
+)
+from superset.models.core import Database
+from superset.utils import json
+
+MAPPING = StoredPartitionMapping(
+ partition_column="dt_epoch",
+ mapped_column=None,
+ column_transforms={
+ "event_time": ColumnTransform("unix_timestamp(:value)", is_monotonic=True)
+ },
+)
+
+
+def _table(extra: str | None = None) -> SqlaTable:
+ database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
+ return SqlaTable(
+ table_name="web_events",
+ database=database,
+ main_dttm_col="event_time",
+ extra=extra,
+ columns=[
+ TableColumn(column_name="event_time", is_dttm=True, type="TIMESTAMP"),
+ TableColumn(column_name="dt_epoch", type="BIGINT"),
+ ],
+ )
+
+
+# ---------------------------------------------------------------------------
+# Round-tripping
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "mapping",
+ [
+ StoredPartitionMapping(),
+ StoredPartitionMapping(partition_column="dt_epoch"),
+ StoredPartitionMapping(partition_column="dt_epoch", mapped_column="event_time"),
+ MAPPING,
+ StoredPartitionMapping(
+ partition_column="dt_epoch",
+ column_transforms={
+ "event_time": ColumnTransform("unix_timestamp(:value)"),
+ "country": ColumnTransform("lower(:value)", is_monotonic=False),
+ },
+ ),
+ ],
+)
+def test_the_store_round_trips_a_mapping(
+ app: Flask, mapping: StoredPartitionMapping
+) -> None:
+ table = _table()
+
+ with app.app_context():
+ store = partition_mapping_store()
+ store.save(table, mapping)
+ assert store.load(table) == mapping
+
+
+def test_a_write_is_a_new_string_not_a_mutation(app: Flask) -> None:
+ """
+ `decode_extra` parses fresh and throws its dict away, so mutating a parsed
+ copy would change nothing on the model -- SQLAlchemy would never see the
+ attribute as dirty, and neither would Continuum.
+ """
+ table = _table()
+
+ with app.app_context():
+ partition_mapping_store().save(table, MAPPING)
+
+ assert sa.inspect(table).attrs.extra.history.has_changes()
+ assert json.loads(table.extra)[EXTRA_KEY]["partition_column"] == "dt_epoch"
+
+
+def test_a_no_op_write_does_not_dirty_extra(app: Flask) -> None:
+ """A version row for an edit that changed nothing is noise in the history."""
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+ sa.inspect(table).attrs.extra.history # noqa: B018
+ before = table.extra
+
+ with app.app_context():
+ partition_mapping_store().save(table, MAPPING)
+
+ assert table.extra is before
+
+
+def test_clearing_a_mapping_removes_the_key(app: Flask) -> None:
+ """
+ Rather than leaving ``"partition_filter_mapping": null`` behind. Datasets
+ that never used the feature, and datasets that stopped, should look the same
+ in the blob.
+ """
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+
+ with app.app_context():
+ partition_mapping_store().save(table, StoredPartitionMapping())
+
+ assert EXTRA_KEY not in json.loads(table.extra)
+
+
+# ---------------------------------------------------------------------------
+# Defensive reads: ``extra`` is a free-text box
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "raw",
+ [
+ None,
+ "",
+ "not json",
+ "[]",
+ "3",
+ '"a string"',
+ "{}",
+ '{"partition_filter_mapping": null}',
+ '{"partition_filter_mapping": "nope"}',
+ '{"partition_filter_mapping": []}',
+ '{"partition_filter_mapping": {"partition_column": 42}}',
+ '{"partition_filter_mapping": {"partition_column": ""}}',
+ '{"partition_filter_mapping": {"partition_column": {"a": 1}}}',
+ ],
+)
+def test_an_unusable_blob_reads_as_not_configured(raw: str | None) -> None:
+ """
+ Modelled on `get_dataset_timezone`: a value of the wrong type is "not
+ configured", never propagated. Every one of these is something an owner can
+ type into the editor's Extra box.
+ """
+ assert parse_mapping(raw) == StoredPartitionMapping()
+
+
+def test_an_over_long_column_name_is_not_configured() -> None:
+ """
+ The typed API fields are length-validated; the Extra box is not, and these
+ are the widths the columns get after the migration.
+ """
+ over = "x" * (MAX_COLUMN_NAME_LENGTH + 1)
+ blob = {EXTRA_KEY: {"partition_column": over}}
+ assert parse_mapping(json.dumps(blob)).partition_column is None
+
+ at_limit = "x" * MAX_COLUMN_NAME_LENGTH
+ blob = {EXTRA_KEY: {"partition_column": at_limit}}
+ assert parse_mapping(json.dumps(blob)).partition_column == at_limit
+
+
+def test_an_over_long_transform_is_not_configured() -> None:
+ """An unbounded string would otherwise reach sqlglot on every query build."""
+ blob = {
+ EXTRA_KEY: {
+ "partition_column": "dt_epoch",
+ "column_transforms": {
+ "event_time": {"value_transform": "x" * (MAX_TRANSFORM_LENGTH + 1)}
+ },
+ }
+ }
+ mapping = parse_mapping(json.dumps(blob))
+ assert mapping.partition_column == "dt_epoch"
+ assert mapping.transform_for("event_time").value_transform is None
+
+
+@pytest.mark.parametrize("value", ["true", "no", 1, 0, [], {}, None])
+def test_only_a_literal_true_declares_monotonicity(value: Any) -> None:
+ """
+ The one field where being wrong drops real rows rather than merely failing to
+ prune: a truthy ``"no"`` coerced with ``bool()`` would enable range
+ mirroring through a transform nobody declared order-preserving.
+ """
+ blob = {
+ EXTRA_KEY: {
+ "partition_column": "dt_epoch",
+ "column_transforms": {
+ "event_time": {
+ "value_transform": "unix_timestamp(:value)",
+ "is_monotonic": value,
+ }
+ },
+ }
+ }
+ assert parse_mapping(json.dumps(blob)).transform_for("event_time").is_monotonic is (
+ False
+ )
+
+
+def test_a_malformed_transform_entry_is_skipped_not_fatal() -> None:
+ """One bad column must not cost the owner the rest of the mapping."""
+ blob = {
+ EXTRA_KEY: {
+ "partition_column": "dt_epoch",
+ "column_transforms": {
+ "event_time": {"value_transform": "unix_timestamp(:value)"},
+ "country": "not an object",
+ "region_key": [],
+ "dt_epoch": None,
+ },
+ }
+ }
+ mapping = parse_mapping(json.dumps(blob))
+ assert set(mapping.column_transforms) == {"event_time"}
+
+
+def test_an_undecodable_extra_is_replaced_and_logged(app: Flask) -> None:
+ """
+ It cannot be merged into, and the alternative -- silently dropping the
+ mapping -- is worse and just as quiet. The text is already inert: nothing
+ reads `certification` or `timezone` out of it either.
+ """
+ table = _table("not json at all")
+
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping_storage.logger"
+ ) as logger:
+ partition_mapping_store().save(table, MAPPING)
+
+ assert load_partition_mapping.__module__ # keep the import honest
+ assert logger.warning.called
+ assert json.loads(table.extra)[EXTRA_KEY]["partition_column"] == "dt_epoch"
+
+
+# ---------------------------------------------------------------------------
+# The memo
+# ---------------------------------------------------------------------------
+
+
+def test_the_read_is_memoized_per_instance(app: Flask) -> None:
+ """
+ `load` is on the hot path -- `get_extra_cache_keys`, every `get_sqla_query`,
+ and once per column in `TableColumn.data`.
+ """
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping_storage.parse_mapping",
+ side_effect=parse_mapping,
+ ) as spy:
+ for _ in range(5):
+ assert load_partition_mapping(table) == MAPPING
+
+ assert spy.call_count == 1
+
+
+def test_the_memo_invalidates_when_extra_changes(app: Flask) -> None:
+ """Keyed on the raw string, so an assignment anywhere invalidates it."""
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+
+ with app.app_context():
+ assert load_partition_mapping(table).partition_column == "dt_epoch"
+ table.extra = json.dumps({"timezone": "Europe/Berlin"})
+ assert load_partition_mapping(table) == StoredPartitionMapping()
+
+
+def test_the_memo_is_not_a_mapped_attribute() -> None:
+ """It must never reach the database, or confuse Continuum's diffing."""
+ names = set(sa.inspect(SqlaTable).attrs.keys())
+ assert not {name for name in names if name.startswith("_partition_mapping")}
+
+
+# ---------------------------------------------------------------------------
+# The model facade
+# ---------------------------------------------------------------------------
+
+
+def test_the_facade_reads_through_the_store(app: Flask) -> None:
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+
+ with app.app_context():
+ assert table.partition_column == "dt_epoch"
+ # `None` means "follow main_dttm_col", which is not resolved on read.
+ assert table.partition_mapped_column is None
+
+ event_time, dt_epoch = table.columns
+ assert event_time.partition_value_transform == "unix_timestamp(:value)"
+ assert event_time.partition_transform_is_monotonic is True
+ assert dt_epoch.partition_value_transform is None
+ assert dt_epoch.partition_transform_is_monotonic is False
+
+
+def test_a_detached_column_reads_as_not_declared(app: Flask) -> None:
+ """
+ `DatasetDAO._override_columns` constructs `TableColumn(table_id=...)` with no
+ relationship, so an unattached column is a normal state. Reporting "not
+ declared" is the safe direction: mirroring stops rather than starting wrongly.
+ """
+ orphan = TableColumn(column_name="event_time")
+
+ with app.app_context():
+ assert orphan.partition_value_transform is None
+ assert orphan.partition_transform_is_monotonic is False
+
+
+@pytest.mark.parametrize(
+ "attribute",
+ ["partition_column", "partition_mapped_column"],
+)
+def test_the_dataset_attributes_are_read_only(app: Flask, attribute: str) -> None:
+ """
+ Assignment must raise rather than vanish. Under the extra-JSON store a
+ setter's write is lost the moment `extra` is assigned afterwards, with no
+ exception and no log -- and which of the two runs last depends on
+ marshmallow's field ordering. This is the guard against reintroducing setters.
+ """
+ table = _table()
+ with app.app_context():
+ with pytest.raises(AttributeError):
+ setattr(table, attribute, "dt_epoch")
+
+
+@pytest.mark.parametrize(
+ "attribute",
+ ["partition_value_transform", "partition_transform_is_monotonic"],
+)
+def test_the_column_attributes_are_read_only(app: Flask, attribute: str) -> None:
+ column = TableColumn(column_name="event_time")
+ with app.app_context():
+ with pytest.raises(AttributeError):
+ setattr(column, attribute, "x")
+
+
+def test_the_column_attributes_are_not_constructor_kwargs() -> None:
+ """
+ `DatasetDAO._upsert_columns` does `TableColumn(**properties)` with no
+ whitelist, which is why the DAO's funnel strips these from the payload before
+ it gets there.
+ """
+ with pytest.raises((AttributeError, TypeError)):
+ TableColumn(column_name="event_time", partition_value_transform="x")
+
+
+# ---------------------------------------------------------------------------
+# The write API
+# ---------------------------------------------------------------------------
+
+
+def test_set_partition_mapping_patches_rather_than_replaces(app: Flask) -> None:
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+
+ with app.app_context():
+ set_partition_mapping(table, partition_column="other_epoch")
+
+ assert table.partition_column == "other_epoch"
+ # The transform, and the mapped-column override, are untouched.
+ assert table.columns[0].partition_value_transform == "unix_timestamp(:value)"
+ assert table.partition_mapped_column is None
+
+
+def test_set_partition_mapping_targets_the_effective_mapped_column(
+ app: Flask,
+) -> None:
+ """
+ Setting a mapped column and its transform in one call has to do what it reads
+ like, so the target is resolved after this call's `mapped_column`, not before.
+ """
+ table = _table()
+
+ with app.app_context():
+ set_partition_mapping(
+ table,
+ partition_column="dt_epoch",
+ mapped_column="event_time",
+ value_transform="unix_timestamp(:value)",
+ )
+
+ assert table.columns[0].partition_value_transform == "unix_timestamp(:value)"
+
+
+def test_set_partition_mapping_can_target_a_named_column(app: Flask) -> None:
+ """A transform may be declared on a column that is not (yet) the mapped one."""
+ table = _table()
+
+ with app.app_context():
+ set_partition_mapping(table, value_transform="lower(:value)", column="dt_epoch")
+ mapping = load_partition_mapping(table)
+
+ assert mapping.transform_for("dt_epoch").value_transform == "lower(:value)"
+ assert mapping.transform_for("event_time").value_transform is None
+
+
+def test_set_partition_mapping_refuses_a_transform_with_no_column(app: Flask) -> None:
+ table = _table()
+ table.main_dttm_col = None
+
+ with app.app_context():
+ with pytest.raises(ValueError, match="no mapped column"):
+ set_partition_mapping(table, value_transform="unix_timestamp(:value)")
+
+
+def test_an_omitted_keyword_means_unchanged(app: Flask) -> None:
+ table = _table(extra_with_partition_mapping(None, MAPPING))
+
+ with app.app_context():
+ set_partition_mapping(table, is_monotonic=False)
+ mapping = load_partition_mapping(table)
+
+ assert mapping.partition_column == "dt_epoch"
+ assert (
+ mapping.transform_for("event_time").value_transform == "unix_timestamp(:value)"
+ )
+ assert mapping.transform_for("event_time").is_monotonic is False
+
+
+# ---------------------------------------------------------------------------
+# The swap point
+# ---------------------------------------------------------------------------
+
+
+def test_the_factory_resolves_the_configured_store(app: Flask) -> None:
+ with app.app_context():
+ assert isinstance(partition_mapping_store(), ExtraJsonPartitionMappingStore)
+
+
+def test_the_column_store_names_the_migration_it_needs(app: Flask) -> None:
+ """
+ Deliberately unimplemented rather than written against columns that do not
+ exist: it would read `table.partition_column`, which today is the facade
+ property, which calls the store. The error has to say what to do about it.
+ """
+ app.config["PARTITION_MAPPING_STORE"] = "columns"
+ try:
+ with app.app_context():
+ store = partition_mapping_store()
+ assert isinstance(store, ColumnPartitionMappingStore)
+ for call in (
+ lambda: store.load(_table()),
+ lambda: store.save(_table(), MAPPING),
+ ):
+ with pytest.raises(NotImplementedError, match="migration"):
+ call()
+ finally:
+ app.config["PARTITION_MAPPING_STORE"] = "extra"
+
+
+def test_an_unknown_store_is_a_configuration_error(app: Flask) -> None:
+ app.config["PARTITION_MAPPING_STORE"] = "bogus"
+ try:
+ with app.app_context():
+ with pytest.raises(ValueError, match="PARTITION_MAPPING_STORE"):
+ partition_mapping_store()
+ finally:
+ app.config["PARTITION_MAPPING_STORE"] = "extra"
+
+
+def test_the_facade_and_the_dao_read_whichever_store_is_configured(
+ app: Flask,
+) -> None:
+ """
+ The claim the indirection exists to support: swap the store and the model,
+ the DAO and the query rewriter follow, with no other change. Proved against a
+ dict-backed stand-in rather than the column store, which cannot be
+ implemented before its migration.
+ """
+ from superset.connectors.sqla import partition_mapping_storage as storage
+
+ class FakeStore:
+ def __init__(self) -> None:
+ self.mappings: dict[int, StoredPartitionMapping] = {}
+
+ def load(self, table: SqlaTable) -> StoredPartitionMapping:
+ return self.mappings.get(id(table), StoredPartitionMapping())
+
+ def save(self, table: SqlaTable, mapping: StoredPartitionMapping) -> None:
+ self.mappings[id(table)] = mapping
+
+ fake = FakeStore()
+ app.config["PARTITION_MAPPING_STORE"] = "fake"
+ storage.STORES["fake"] = lambda: fake # type: ignore[assignment]
+ try:
+ table = _table()
+ with app.app_context():
+ set_partition_mapping(
+ table,
+ partition_column="dt_epoch",
+ mapped_column="event_time",
+ value_transform="unix_timestamp(:value)",
+ is_monotonic=True,
+ )
+
+ # Nothing was written to `extra` -- the store, not the column, is the
+ # source of truth.
+ assert table.extra is None
+ assert table.partition_column == "dt_epoch"
+ assert table.partition_mapped_column == "event_time"
+ assert (
+ table.columns[0].partition_value_transform == "unix_timestamp(:value)"
+ )
+ assert table.columns[0].partition_transform_is_monotonic is True
+ assert table.partition_filter_mapping_summary == {
+ "partition_column": "dt_epoch",
+ "mapped_column": "event_time",
+ "active": True,
+ }
+ finally:
+ del storage.STORES["fake"]
+ app.config["PARTITION_MAPPING_STORE"] = "extra"
diff --git a/tests/unit_tests/connectors/sqla/partition_mapping_test.py b/tests/unit_tests/connectors/sqla/partition_mapping_test.py
new file mode 100644
index 00000000000..7cdc7656324
--- /dev/null
+++ b/tests/unit_tests/connectors/sqla/partition_mapping_test.py
@@ -0,0 +1,984 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for ``superset.connectors.sqla.partition_mapping``."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, cast
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+import pytest
+from dateutil.relativedelta import relativedelta
+from flask import Flask
+
+from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.connectors.sqla.partition_mapping import (
+ _probe_cache_key,
+ contains_jinja,
+ contains_value_placeholder,
+ evaluate_transform,
+ find_non_deterministic_functions,
+ grain_bucket_width,
+ GRAIN_BUCKET_WIDTHS,
+ is_transform_active,
+ MappingValidationIssue,
+ MIRRORABLE_ALWAYS,
+ MIRRORABLE_IF_MONOTONIC,
+ mirrorable_operators,
+ resolve_partition_mapping,
+ validate_partition_mapping,
+ validate_transform,
+)
+from superset.connectors.sqla.partition_mapping_storage import (
+ ColumnTransform,
+ EXTRA_KEY,
+ extra_with_partition_mapping,
+ StoredPartitionMapping,
+)
+from superset.constants import TimeGrain
+from superset.models.core import Database
+from superset.utils import json
+from superset.utils.core import FilterOperator
+
+
+@pytest.fixture(autouse=True)
+def real_probe_cache(app: Flask) -> Any:
+ """
+ The test app runs a null cache, which would make every cache assertion here
+ vacuously pass. Swap in a real in-memory one for the duration.
+ """
+ from flask_caching import Cache
+
+ from superset.extensions import cache_manager
+
+ cache = Cache(config={"CACHE_TYPE": "SimpleCache", "CACHE_DEFAULT_TIMEOUT": 300})
+ cache.init_app(app)
+ original = cache_manager._cache # noqa: SLF001
+ cache_manager._cache = cache # noqa: SLF001
+ yield
+ cache_manager._cache = original # noqa: SLF001
+
+
+@pytest.fixture(autouse=True)
+def enable_partition_filter_mapping(app: Flask) -> Any:
+ """The feature ships off; every test here exercises it on."""
+ original = app.config["DEFAULT_FEATURE_FLAGS"].get("PARTITION_FILTER_MAPPING")
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = True
+ yield
+ if original is None:
+ del app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"]
+ else:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = original
+
+
+def _table(**kwargs: Any) -> SqlaTable:
+ database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
+ defaults: dict[str, Any] = {
+ "table_name": "web_events",
+ "database": database,
+ "main_dttm_col": "event_time",
+ "columns": [
+ TableColumn(column_name="event_time", is_dttm=True, type="TIMESTAMP"),
+ TableColumn(column_name="dt_epoch", type="BIGINT"),
+ TableColumn(column_name="country", type="VARCHAR"),
+ TableColumn(column_name="region_key", type="VARCHAR"),
+ ],
+ }
+ defaults.update(kwargs)
+ return SqlaTable(**defaults)
+
+
+def _mapped_table(
+ *,
+ partition_column: str = "dt_epoch",
+ mapped_column: str = "event_time",
+ transform: str | None = "unix_timestamp(:value)",
+ monotonic: bool = True,
+ partition_mapped_column: str | None = None,
+ main_dttm_col: str | None = "event_time",
+) -> SqlaTable:
+ table = _table(main_dttm_col=main_dttm_col)
+ return _store_mapping(
+ table,
+ partition_column=partition_column,
+ mapped_column=partition_mapped_column,
+ transforms={
+ mapped_column: ColumnTransform(
+ value_transform=transform, is_monotonic=monotonic
+ )
+ },
+ )
+
+
+def _store_mapping(
+ table: SqlaTable,
+ *,
+ partition_column: str | None,
+ mapped_column: str | None = None,
+ transforms: dict[str, ColumnTransform] | None = None,
+) -> SqlaTable:
+ """
+ Persist a mapping onto a dataset, no app context required.
+
+ The model's four mapping attributes are read-only properties over the
+ configured store, so a test sets up stored state rather than assigning to
+ them. `extra_with_partition_mapping` is the store's own serializer, which is
+ pure, so this works in a factory that has no app context.
+ """
+ table.extra = extra_with_partition_mapping(
+ table.extra,
+ StoredPartitionMapping(
+ partition_column=partition_column,
+ mapped_column=mapped_column,
+ column_transforms=transforms or {},
+ ),
+ )
+ return table
+
+
+# ---------------------------------------------------------------------------
+# §2 — operator safety matrix
+# ---------------------------------------------------------------------------
+
+
+def test_equality_and_in_are_always_mirrorable() -> None:
+ """``=`` and ``IN`` are safe for any function ``T``."""
+ assert MIRRORABLE_ALWAYS == {FilterOperator.EQUALS, FilterOperator.IN}
+
+
+def test_range_operators_require_a_monotonic_transform() -> None:
+ assert MIRRORABLE_IF_MONOTONIC == {
+ FilterOperator.GREATER_THAN,
+ FilterOperator.GREATER_THAN_OR_EQUALS,
+ FilterOperator.LESS_THAN,
+ FilterOperator.LESS_THAN_OR_EQUALS,
+ FilterOperator.TEMPORAL_RANGE,
+ }
+
+
+def test_mirrorable_operators_excludes_ranges_when_not_monotonic() -> None:
+ assert mirrorable_operators(is_monotonic=False) == MIRRORABLE_ALWAYS
+
+
+def test_mirrorable_operators_includes_ranges_when_monotonic() -> None:
+ assert mirrorable_operators(is_monotonic=True) == (
+ MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC
+ )
+
+
+@pytest.mark.parametrize(
+ "operator",
+ [
+ FilterOperator.NOT_EQUALS,
+ FilterOperator.NOT_IN,
+ FilterOperator.LIKE,
+ FilterOperator.ILIKE,
+ FilterOperator.NOT_LIKE,
+ FilterOperator.NOT_ILIKE,
+ FilterOperator.IS_NULL,
+ FilterOperator.IS_NOT_NULL,
+ FilterOperator.IS_TRUE,
+ FilterOperator.IS_FALSE,
+ ],
+)
+def test_negations_and_pattern_matches_are_never_mirrorable(
+ operator: FilterOperator,
+) -> None:
+ """
+ ``T`` is not injective, so ``col != v`` does **not** imply ``T(col) != T(v)``:
+ mirroring it would drop rows the original filter keeps.
+ """
+ assert operator not in mirrorable_operators(is_monotonic=True)
+
+
+# ---------------------------------------------------------------------------
+# Time grain bucket widths
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("grain", list(TimeGrain))
+def test_every_builtin_grain_has_a_known_bucket_width(
+ app: Flask, grain: TimeGrain
+) -> None:
+ """
+ The widening is only sound for a grain whose bucket width we know, so a
+ grain added upstream without a width here would silently stop mirroring.
+ This is the test that notices.
+ """
+ with app.app_context():
+ assert grain_bucket_width(grain.value, "sqlite") is not None
+
+
+def test_a_calendar_grain_widens_by_a_calendar_month(app: Flask) -> None:
+ """A month is not thirty days, and the month it lands in decides its width."""
+ with app.app_context():
+ width = grain_bucket_width(TimeGrain.MONTH.value, "sqlite")
+
+ assert width is not None
+ assert datetime(2026, 1, 31) + width == datetime(2026, 2, 28)
+ assert datetime(2026, 2, 1) + width == datetime(2026, 3, 1)
+
+
+def test_a_fractional_grain_has_a_width(app: Flask) -> None:
+ """
+ ``PT0.5H`` and ``P0.25Y`` are why the widths are written out rather than
+ parsed -- and the anchored week grains are intervals ``parse_duration``
+ rejects outright.
+ """
+ with app.app_context():
+ assert grain_bucket_width(TimeGrain.HALF_HOUR.value, "sqlite") == relativedelta(
+ minutes=30
+ )
+ assert grain_bucket_width(
+ TimeGrain.WEEK_ENDING_SATURDAY.value, "sqlite"
+ ) == relativedelta(days=7)
+
+
+def test_an_unknown_grain_has_no_known_width(app: Flask) -> None:
+ with app.app_context():
+ assert grain_bucket_width("P13X", "sqlite") is None
+ assert grain_bucket_width(None, "sqlite") is None
+
+
+def test_an_addon_grain_has_no_known_width(app: Flask) -> None:
+ """An operator-declared grain carries whatever duration string they typed."""
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDONS"] = {"PT2S": "2 second"}
+ try:
+ assert grain_bucket_width("PT2S", "sqlite") is None
+ finally:
+ app.config["TIME_GRAIN_ADDONS"] = {}
+
+
+def test_a_builtin_grain_an_operator_redefined_has_no_known_width(
+ app: Flask,
+) -> None:
+ """
+ ``TIME_GRAIN_ADDON_EXPRESSIONS`` replaces a built-in grain's SQL for one
+ engine, so its duration no longer bounds the displacement *there* -- but it
+ still does on every other engine.
+ """
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {
+ "sqlite": {"P1D": "DATETIME({col}, 'start of year')"}
+ }
+ try:
+ assert grain_bucket_width(TimeGrain.DAY.value, "sqlite") is None
+ assert grain_bucket_width(TimeGrain.DAY.value, "hive") is not None
+ finally:
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {}
+
+
+def test_the_width_table_is_keyed_on_grain_durations(app: Flask) -> None:
+ """Keys are the ISO duration a filter carries, not the enum member name."""
+ assert GRAIN_BUCKET_WIDTHS[TimeGrain.DAY] == relativedelta(days=1)
+ assert set(GRAIN_BUCKET_WIDTHS) == {grain.value for grain in TimeGrain}
+
+
+# ---------------------------------------------------------------------------
+# §4.1 — resolving the effective mapping, and the defensive bail-outs
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_returns_the_mapping_when_everything_lines_up(app: Flask) -> None:
+ with app.app_context():
+ mapping = resolve_partition_mapping(_mapped_table())
+
+ assert mapping is not None
+ assert mapping.partition_column == "dt_epoch"
+ assert mapping.mapped_column == "event_time"
+ assert mapping.value_transform == "unix_timestamp(:value)"
+ assert mapping.is_monotonic is True
+
+
+def test_effective_mapped_column_follows_main_dttm_col(app: Flask) -> None:
+ """No explicit override: the mapping follows the default datetime column."""
+ with app.app_context():
+ mapping = resolve_partition_mapping(_mapped_table())
+
+ assert mapping is not None
+ assert mapping.mapped_column == "event_time"
+
+
+def test_explicit_override_wins_over_main_dttm_col(app: Flask) -> None:
+ table = _mapped_table(
+ mapped_column="country",
+ transform="lower(:value)",
+ monotonic=False,
+ partition_mapped_column="country",
+ )
+ _store_mapping(
+ table,
+ partition_column="region_key",
+ mapped_column="country",
+ transforms={"country": ColumnTransform("lower(:value)", is_monotonic=False)},
+ )
+
+ with app.app_context():
+ mapping = resolve_partition_mapping(table)
+
+ assert mapping is not None
+ assert mapping.mapped_column == "country"
+ assert mapping.partition_column == "region_key"
+ assert mapping.is_monotonic is False
+
+
+def test_resolve_returns_none_when_the_feature_flag_is_off(app: Flask) -> None:
+ table = _mapped_table()
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ return_value=False,
+ ):
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_without_a_partition_column(app: Flask) -> None:
+ table = _mapped_table()
+ _store_mapping(table, partition_column=None)
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_partition_column_no_longer_exists(
+ app: Flask,
+) -> None:
+ """A column sync can drop the physical partition column out from under us."""
+ table = _mapped_table()
+ _store_mapping(
+ table,
+ partition_column="dt_epoch_gone",
+ transforms={"event_time": ColumnTransform("unix_timestamp(:value)", True)},
+ )
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_mapped_column_no_longer_exists(
+ app: Flask,
+) -> None:
+ table = _mapped_table(main_dttm_col="vanished")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_on_self_mapping(app: Flask) -> None:
+ """
+ ``partition_column == effective mapped column`` mirrors a column onto itself.
+ Save-time validation rejects it, but rows predating that validation exist.
+ """
+ table = _mapped_table(
+ partition_column="event_time",
+ mapped_column="event_time",
+ )
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+@pytest.mark.parametrize("transform", [None, "", " "])
+def test_resolve_returns_none_without_a_transform(
+ app: Flask, transform: str | None
+) -> None:
+ table = _mapped_table(transform=transform)
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_transform_lacks_the_placeholder(
+ app: Flask,
+) -> None:
+ table = _mapped_table(transform="unix_timestamp(event_time)")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_for_an_unparseable_transform(app: Flask) -> None:
+ table = _mapped_table(transform="unix_timestamp(:value")
+ with app.app_context():
+ assert resolve_partition_mapping(table) is None
+
+
+def test_resolve_returns_none_when_the_mapped_column_has_an_advanced_data_type(
+ app: Flask,
+) -> None:
+ """
+ ``translate_filter`` builds its own predicate shape from *translated* values,
+ so the ``(operator, value)`` pair the operator matrix reasons about does not
+ exist. Mirroring anyway would silently apply the wrong values (§4.1).
+ """
+ table = _mapped_table()
+ for column in table.columns:
+ if column.column_name == "event_time":
+ column.advanced_data_type = "port"
+
+ with app.app_context():
+ with patch.dict(app.config["ADVANCED_DATA_TYPES"], {"port": MagicMock()}):
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ side_effect=lambda flag: flag
+ in {"PARTITION_FILTER_MAPPING", "ENABLE_ADVANCED_DATA_TYPES"},
+ ):
+ assert resolve_partition_mapping(table) is None
+
+
+def test_advanced_data_type_does_not_block_when_the_flag_is_off(app: Flask) -> None:
+ """An inert ``advanced_data_type`` is not a reason to skip mirroring."""
+ table = _mapped_table()
+ for column in table.columns:
+ if column.column_name == "event_time":
+ column.advanced_data_type = "port"
+
+ with app.app_context():
+ with patch(
+ "superset.connectors.sqla.partition_mapping.feature_flag_manager."
+ "is_feature_enabled",
+ side_effect=lambda flag: flag == "PARTITION_FILTER_MAPPING",
+ ):
+ assert resolve_partition_mapping(table) is not None
+
+
+# ---------------------------------------------------------------------------
+# §5 — transform inspection helpers
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "transform,expected",
+ [
+ ("unix_timestamp(:value)", True),
+ ("lower(:value)", True),
+ ("CAST(:value AS BIGINT)", True),
+ ("unix_timestamp(event_time)", False),
+ (":values", False),
+ ("", False),
+ (None, False),
+ ],
+)
+def test_contains_value_placeholder(transform: str | None, expected: bool) -> None:
+ assert contains_value_placeholder(transform) is expected
+
+
+@pytest.mark.parametrize(
+ "transform,expected",
+ [
+ ("unix_timestamp(:value)", False),
+ ("{{ current_username() }}", True),
+ ("lower({% if x %}:value{% endif %})", True),
+ ("lower(:value) -- {# comment #}", True),
+ ],
+)
+def test_contains_jinja(transform: str, expected: bool) -> None:
+ assert contains_jinja(transform) is expected
+
+
+@pytest.mark.parametrize(
+ "transform",
+ [
+ "unix_timestamp(:value)",
+ "lower(:value)",
+ "CAST(:value AS BIGINT)",
+ "date_format(:value, 'yyyyMMdd')",
+ ],
+)
+def test_pure_transforms_report_no_non_deterministic_functions(
+ transform: str,
+) -> None:
+ assert find_non_deterministic_functions(transform, "hive") == set()
+
+
+@pytest.mark.parametrize(
+ "transform,expected_name",
+ [
+ ("date_diff(:value, now())", "NOW"),
+ ("CAST(:value AS BIGINT) + rand()", "RAND"),
+ ("CAST(:value AS DATE) - current_date", "CURRENT_DATE"),
+ ],
+)
+def test_non_deterministic_functions_are_reported(
+ transform: str, expected_name: str
+) -> None:
+ """
+ The probe runs at a different moment and in a different session from the
+ chart query, and its result is cached, so anything time- or
+ randomness-dependent freezes a snapshot of probe time into the predicate.
+ """
+ assert expected_name in find_non_deterministic_functions(transform, "hive")
+
+
+def test_niladic_unix_timestamp_is_rejected_but_the_unary_form_is_not() -> None:
+ """
+ On Hive/Impala ``unix_timestamp()`` means "now" while ``unix_timestamp(x)``
+ -- the canonical temporal transform -- is pure. The distinction is the whole
+ reason this check inspects arity rather than just the name.
+ """
+ assert find_non_deterministic_functions("unix_timestamp(:value)", "hive") == set()
+ assert find_non_deterministic_functions(
+ "unix_timestamp(:value) - unix_timestamp()", "hive"
+ )
+
+
+# ---------------------------------------------------------------------------
+# §4.2 — probe evaluation
+# ---------------------------------------------------------------------------
+
+
+def _database_returning(values: list[Any]) -> Database:
+ """A real ``Database`` (so the dialect is real) with only ``get_df`` stubbed."""
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=pd.DataFrame(
+ [values], columns=[f"v{i}" for i in range(len(values))]
+ )
+ )
+ return database
+
+
+def _probe(database: Database) -> MagicMock:
+ """
+ The stubbed ``get_df``, typed for call assertions.
+
+ ``Database.get_df`` is a real method, so mypy reads its declared signature
+ rather than the mock that replaced it and rejects ``.call_args`` and
+ friends.
+ """
+ return cast(MagicMock, database.get_df)
+
+
+def test_evaluate_transform_returns_one_value_per_input(app: Flask) -> None:
+ database = _database_returning([1767225600, 1769904000])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database,
+ None,
+ "default",
+ "unix_timestamp(:value)",
+ ["2026-01-01 00:00:00", "2026-02-01 00:00:00"],
+ )
+
+ assert result == [1767225600, 1769904000]
+ _probe(database).assert_called_once()
+
+
+def test_evaluate_transform_binds_values_rather_than_interpolating(
+ app: Flask,
+) -> None:
+ """
+ Filter values are attacker-controlled (a Gamma user picks them), so they must
+ be bound and escaped, never pasted into the probe SQL.
+ """
+ database = _database_returning(["o''brien"])
+
+ with app.app_context():
+ evaluate_transform(database, None, None, "lower(:value)", ["O'Brien"])
+
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert "O'Brien" not in sql
+ assert "O''Brien" in sql
+
+
+def test_evaluate_transform_dedupes_repeated_values(app: Flask) -> None:
+ """Three inputs, two distinct: the probe evaluates the transform twice."""
+ database = _database_returning(["us", "us"])
+
+ with app.app_context():
+ result = evaluate_transform(
+ database, None, None, "lower(:value)", ["US", "US", "us"]
+ )
+
+ assert result == ["us", "us", "us"]
+ sql = _probe(database).call_args.kwargs["sql"]
+ assert sql.count("lower") == 2
+
+
+def test_evaluate_transform_fails_open_on_a_short_result_row(app: Flask) -> None:
+ """A row narrower than the probe asked for means the results cannot be
+ aligned back to their inputs; pruning is skipped rather than guessed at."""
+ database = _database_returning(["us"])
+
+ with app.app_context():
+ assert (
+ evaluate_transform(database, None, None, "lower(:value)", ["US", "CA"])
+ is None
+ )
+
+
+def test_evaluate_transform_pins_catalog_and_schema(app: Flask) -> None:
+ """
+ The probe runs in a different session from the chart query; pinning the
+ catalog and schema keeps session settings as close as the pool allows.
+ """
+ database = _database_returning([1])
+
+ with app.app_context():
+ evaluate_transform(database, "prod", "analytics", "lower(:value)", ["x"])
+
+ kwargs = _probe(database).call_args.kwargs
+ assert kwargs["catalog"] == "prod"
+ assert kwargs["schema"] == "analytics"
+
+
+def test_evaluate_transform_fails_open_when_the_probe_raises(app: Flask) -> None:
+ """A wedged engine must not break the chart — it just stops pruning."""
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ side_effect=RuntimeError("connection reset")
+ )
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)", ["x"]) is None
+
+
+def test_evaluate_transform_fails_open_on_an_empty_result(app: Flask) -> None:
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ return_value=pd.DataFrame()
+ )
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)", ["x"]) is None
+
+
+def test_evaluate_transform_returns_none_for_no_values(app: Flask) -> None:
+ database = _database_returning([])
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)", []) is None
+
+ _probe(database).assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# §5 — save-time validation, in two tiers
+# ---------------------------------------------------------------------------
+
+
+def _issues(
+ *,
+ partition_column: str | None = "dt_epoch",
+ partition_mapped_column: str | None = None,
+ transform: str | None = "unix_timestamp(:value)",
+ main_dttm_col: str | None = "event_time",
+ column_names: set[str] | None = None,
+ engine: str = "hive",
+) -> list[MappingValidationIssue]:
+ """
+ Validate one mapping described in the API's vocabulary.
+
+ Keeps the four payload-shaped keywords the cases below read best with, and
+ assembles the value object `validate_partition_mapping` takes. The transform
+ is attached to the effective mapped column, which is what the store would
+ hold.
+ """
+ mapping = StoredPartitionMapping(
+ partition_column=partition_column,
+ mapped_column=partition_mapped_column,
+ )
+ effective = mapping.effective_mapped_column(main_dttm_col)
+ if effective and transform is not None:
+ mapping = mapping.with_transform(effective, ColumnTransform(transform))
+ return validate_partition_mapping(
+ mapping,
+ column_names=(
+ column_names
+ if column_names is not None
+ else {"event_time", "dt_epoch", "country", "region_key"}
+ ),
+ main_dttm_col=main_dttm_col,
+ engine=engine,
+ )
+
+
+def _blocking(issues: list[MappingValidationIssue]) -> list[MappingValidationIssue]:
+ return [issue for issue in issues if issue.blocking]
+
+
+def _warnings(issues: list[MappingValidationIssue]) -> list[MappingValidationIssue]:
+ return [issue for issue in issues if not issue.blocking]
+
+
+def test_a_well_formed_mapping_raises_nothing() -> None:
+ assert _issues() == []
+
+
+def test_no_partition_column_means_nothing_to_validate() -> None:
+ assert _issues(partition_column=None, transform=None) == []
+
+
+# Tier 1 — blocks the save
+
+
+def test_an_unknown_partition_column_blocks_the_save() -> None:
+ issues = _issues(partition_column="nope")
+ assert len(_blocking(issues)) == 1
+ assert issues[0].field == "partition_column"
+
+
+def test_an_unknown_mapped_column_override_blocks_the_save() -> None:
+ issues = _issues(partition_mapped_column="nope")
+ assert len(_blocking(issues)) == 1
+ assert issues[0].field == "partition_mapped_column"
+
+
+def test_an_explicit_self_mapping_blocks_the_save() -> None:
+ issues = _blocking(_issues(partition_mapped_column="dt_epoch"))
+ assert len(issues) == 1
+ assert "itself" in issues[0].message
+
+
+def test_an_implicit_self_mapping_blocks_the_save() -> None:
+ """
+ Checking only the explicit override misses the case an owner actually hits:
+ setting ``partition_column`` to the column that is *already* the default
+ datetime column, with no override in play.
+ """
+ issues = _blocking(
+ _issues(partition_column="event_time", main_dttm_col="event_time")
+ )
+ assert len(issues) == 1
+ assert "itself" in issues[0].message
+
+
+def test_jinja_in_the_transform_blocks_the_save() -> None:
+ """
+ The probe would render the template in a different context at a different
+ time from the chart query, so v1 disallows it outright.
+ """
+ issues = _blocking(_issues(transform="unix_timestamp('{{ ds }}' , :value)"))
+ assert len(issues) == 1
+ assert "Jinja" in issues[0].message
+
+
+@pytest.mark.parametrize(
+ "transform",
+ [
+ "unix_timestamp(:value) - unix_timestamp()",
+ "date_diff(:value, now())",
+ "CAST(:value AS BIGINT) + rand()",
+ ],
+)
+def test_a_non_deterministic_transform_blocks_the_save(transform: str) -> None:
+ issues = _blocking(_issues(transform=transform))
+ assert len(issues) == 1
+ assert issues[0].field == "partition_value_transform"
+
+
+# Tier 2 — saves, but the mapping stays inactive
+
+
+def test_an_unparseable_transform_saves_with_a_warning() -> None:
+ """The PRD is explicit: a bad transform still saves, it just stays inactive."""
+ issues = _issues(transform="unix_timestamp(:value")
+ assert _blocking(issues) == []
+ assert len(_warnings(issues)) == 1
+
+
+def test_a_transform_without_the_placeholder_saves_with_a_warning() -> None:
+ issues = _issues(transform="unix_timestamp(event_time)")
+ assert _blocking(issues) == []
+ assert len(_warnings(issues)) == 1
+
+
+def test_a_missing_transform_saves_with_a_warning() -> None:
+ issues = _issues(transform=None)
+ assert _blocking(issues) == []
+ assert len(_warnings(issues)) == 1
+
+
+def test_an_unparseable_transform_skips_the_checks_that_need_a_parse() -> None:
+ """
+ The Jinja and non-determinism checks require a successful parse. When there
+ is nothing to inspect, fall through to a warning rather than reporting a
+ blocking error the owner cannot act on.
+ """
+ issues = _issues(transform="now(:value")
+ assert _blocking(issues) == []
+
+
+# ---------------------------------------------------------------------------
+# §4.2 — the probe cache
+# ---------------------------------------------------------------------------
+
+
+def test_a_repeated_probe_is_served_from_cache(app: Flask) -> None:
+ """
+ Day-aligned ranges ("Last month") repeat constantly across charts, so the
+ hit rate is what keeps the added round trip off the hot path.
+ """
+ database = _database_returning([1767225600])
+
+ with app.app_context():
+ first = evaluate_transform(
+ database, None, None, "unix_timestamp(:value)", ["2026-01-01"]
+ )
+ second = evaluate_transform(
+ database, None, None, "unix_timestamp(:value)", ["2026-01-01"]
+ )
+
+ assert first == second == [1767225600]
+ _probe(database).assert_called_once()
+
+
+def test_the_cache_key_includes_the_transform(app: Flask) -> None:
+ """Editing the transform must not serve the old transform's results."""
+ database = _database_returning([1])
+
+ with app.app_context():
+ evaluate_transform(database, None, None, "unix_timestamp(:value)", ["x"])
+ evaluate_transform(database, None, None, "lower(:value)", ["x"])
+
+ assert _probe(database).call_count == 2
+
+
+def test_the_cache_key_includes_the_values(app: Flask) -> None:
+ database = _database_returning([1])
+
+ with app.app_context():
+ evaluate_transform(database, None, None, "lower(:value)", ["x"])
+ evaluate_transform(database, None, None, "lower(:value)", ["y"])
+
+ assert _probe(database).call_count == 2
+
+
+def test_the_cache_key_includes_the_catalog_and_schema(app: Flask) -> None:
+ """
+ The transform is evaluated against a pinned catalog and schema; the same
+ expression can resolve differently under a different one.
+ """
+ database = _database_returning([1])
+
+ with app.app_context():
+ evaluate_transform(database, "prod", "analytics", "lower(:value)", ["x"])
+ evaluate_transform(database, "prod", "staging", "lower(:value)", ["x"])
+
+ assert _probe(database).call_count == 2
+
+
+def test_a_failed_probe_is_not_cached(app: Flask) -> None:
+ """Caching a failure would keep a transient engine blip pruning-free for a day."""
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.get_df = MagicMock( # type: ignore[method-assign]
+ side_effect=[
+ RuntimeError("connection reset"),
+ pd.DataFrame([[42]], columns=["v0"]),
+ ]
+ )
+
+ with app.app_context():
+ assert evaluate_transform(database, None, None, "lower(:value)", ["x"]) is None
+ assert evaluate_transform(database, None, None, "lower(:value)", ["x"]) == [42]
+
+
+def test_an_undeclared_monotonic_flag_reads_as_not_declared(app: Flask) -> None:
+ """
+ A mapping whose stored blob omits (or mistypes) the flag must fail closed:
+ ranges stop mirroring rather than mirroring through a transform nobody
+ declared order-preserving. `parse_mapping` insists on a literal `true` for
+ exactly this reason -- a truthy `"no"` would enable range mirroring.
+ """
+ table = _mapped_table()
+ table.extra = json.dumps(
+ {
+ EXTRA_KEY: {
+ "partition_column": "dt_epoch",
+ "column_transforms": {
+ "event_time": {
+ "value_transform": "unix_timestamp(:value)",
+ "is_monotonic": "no",
+ }
+ },
+ }
+ }
+ )
+
+ with app.app_context():
+ mapping = resolve_partition_mapping(table)
+
+ assert mapping is not None
+ assert mapping.is_monotonic is False
+ assert not mapping.mirrors(FilterOperator.TEMPORAL_RANGE)
+ assert mapping.mirrors(FilterOperator.EQUALS)
+
+
+def test_probe_cache_key_tracks_the_connection(app: Flask) -> None:
+ """
+ The probe asks a specific database to evaluate the transform, so the answer
+ belongs to that connection. ``database.id`` survives an edit to the URI or
+ to ``extra`` (a session timezone, say), which would otherwise serve values
+ computed against the old environment for the whole cache timeout -- and a
+ wrong transformed bound prunes away rows the real filter keeps.
+ """
+ database = Database(database_name="probe_db", sqlalchemy_uri="sqlite://")
+ database.id = 1
+
+ with app.app_context():
+ before = _probe_cache_key(database, None, None, "lower(:value)", ["US"])
+ database.sqlalchemy_uri = "postgresql://host/db"
+ after_uri = _probe_cache_key(database, None, None, "lower(:value)", ["US"])
+ database.extra = '{"engine_params": {"connect_args": {"timezone": "UTC"}}}'
+ after_extra = _probe_cache_key(database, None, None, "lower(:value)", ["US"])
+
+ assert before != after_uri
+ assert after_uri != after_extra
+
+
+# ---------------------------------------------------------------------------
+# §6 — the read-side predicate the Explore indicator asks
+# ---------------------------------------------------------------------------
+
+#: One transform per branch of ``validate_transform``, blocking and not.
+INACTIVE_TRANSFORMS = [
+ None,
+ "",
+ " ",
+ "unix_timestamp(:value",
+ "unix_timestamp(event_time)",
+ "unix_timestamp('{{ ds }}', :value)",
+ "date_diff(:value, now())",
+]
+
+
+def test_a_well_formed_transform_is_active() -> None:
+ assert is_transform_active("unix_timestamp(:value)", "hive") is True
+
+
+@pytest.mark.parametrize("transform", INACTIVE_TRANSFORMS)
+def test_every_transform_issue_leaves_it_inactive(transform: str | None) -> None:
+ """
+ Blocking issues count as much as the rest. A Jinja or non-deterministic
+ transform is rejected on PUT, but create and import do not validate the
+ mapping, so one can still be read back -- and it mirrors nothing either.
+ """
+ assert is_transform_active(transform, "hive") is False
+
+
+@pytest.mark.parametrize(
+ "transform",
+ ["unix_timestamp(:value)", "CAST(:value AS BIGINT)", *INACTIVE_TRANSFORMS],
+)
+def test_activity_is_exactly_the_absence_of_issues(transform: str | None) -> None:
+ """
+ The indicator and the save path have to answer the same question. Pinning
+ the equivalence is what stops the two from drifting apart again.
+ """
+ assert is_transform_active(transform, "hive") is (
+ validate_transform(transform, "hive") == []
+ )
diff --git a/tests/unit_tests/dao/dataset_test.py b/tests/unit_tests/dao/dataset_test.py
index a5ffc3c7ea3..31279076fb2 100644
--- a/tests/unit_tests/dao/dataset_test.py
+++ b/tests/unit_tests/dao/dataset_test.py
@@ -24,9 +24,11 @@ import pytest
from freezegun import freeze_time
from sqlalchemy.orm.session import Session
+from superset.connectors.sqla.partition_mapping_storage import set_partition_mapping
from superset.daos.base import BaseDAO
from superset.daos.dataset import DatasetDAO
from superset.sql.parse import Table
+from superset.utils import json
def test_validate_update_uniqueness(session: Session) -> None:
@@ -423,3 +425,198 @@ def test_override_columns_rename_flushes_delete_before_insert(
db.session.flush()
cols = db.session.query(TableColumn).filter_by(table_id=table.id).all()
assert [c.column_name for c in cols] == ["new"]
+
+
+def _mapped_table(session: Session, name: str) -> Any:
+ """A saved dataset with two columns and a live partition mapping."""
+ from superset import db
+ from superset.connectors.sqla.models import SqlaTable, TableColumn
+ from superset.models.core import Database
+
+ SqlaTable.metadata.create_all(session.get_bind())
+ database = Database(database_name=f"{name}_db", sqlalchemy_uri="sqlite://")
+ table = SqlaTable(table_name=name, schema="main", database=database)
+ table.columns = [
+ TableColumn(column_name="event_time"),
+ TableColumn(column_name="dt_epoch"),
+ ]
+ db.session.add_all([database, table])
+ db.session.flush()
+ set_partition_mapping(
+ table,
+ partition_column="dt_epoch",
+ mapped_column="event_time",
+ value_transform="unix_timestamp(:value)",
+ is_monotonic=True,
+ )
+ db.session.flush()
+ return table
+
+
+def test_a_sync_that_drops_the_partition_column_clears_the_mapping(
+ session: Session,
+) -> None:
+ """The upsert path deletes every column the payload omits — including, when
+ a metadata sync drops it, the dataset's partition column. The mapping has to
+ be cleared with it, exactly as on the ``override_columns=true`` path."""
+ from superset import db
+
+ table = _mapped_table(session, "pm_t")
+ event_time_id = next(c.id for c in table.columns if c.column_name == "event_time")
+
+ # The payload keeps only "event_time", so "dt_epoch" is deleted.
+ DatasetDAO.update(
+ table,
+ {"columns": [{"id": event_time_id, "column_name": "event_time"}]},
+ )
+ db.session.flush()
+
+ assert table.partition_column is None
+ assert table.partition_mapped_column is None
+
+
+def test_a_sync_that_keeps_both_columns_leaves_the_mapping_alone(
+ session: Session,
+) -> None:
+ from superset import db
+
+ table = _mapped_table(session, "pm_t2")
+ ids = {c.column_name: c.id for c in table.columns}
+
+ DatasetDAO.update(
+ table,
+ {
+ "columns": [
+ {"id": ids["event_time"], "verbose_name": "Event time"},
+ {"id": ids["dt_epoch"]},
+ ]
+ },
+ )
+ db.session.flush()
+
+ assert table.partition_column == "dt_epoch"
+ assert table.partition_mapped_column == "event_time"
+ assert table.columns[0].partition_value_transform == "unix_timestamp(:value)"
+
+
+def test_a_typed_field_beats_a_stale_extra_in_the_same_payload(
+ session: Session,
+) -> None:
+ """
+ The regression the DAO's extract-then-apply funnel exists to prevent.
+
+ The dataset editor GETs a dataset, holds `extra` as an opaque string, and
+ PUTs it back verbatim alongside the typed `partition_column` field. Under the
+ extra-JSON store those are the same bytes, so a `setattr` loop that happens
+ to apply `extra` last silently discards the typed field -- and which order
+ that is depends on marshmallow's field ordering, not on anything visible
+ here.
+ """
+ from superset import db
+
+ table = _mapped_table(session, "pm_t3")
+ stale_extra = json.dumps({"timezone": "Europe/Berlin"})
+
+ DatasetDAO.update(
+ table,
+ {"partition_column": "dt_epoch", "extra": stale_extra},
+ )
+ db.session.flush()
+
+ assert table.partition_column == "dt_epoch"
+ # The payload's other `extra` keys survive too; the funnel patches, it does
+ # not replace.
+ assert json.loads(table.extra)["timezone"] == "Europe/Berlin"
+
+
+def test_a_mapping_written_through_extra_alone_lands(session: Session) -> None:
+ """
+ The second door: an owner typing into the editor's free-text Extra box.
+
+ No typed field in the payload at all, so the funnel has no patch to apply and
+ must leave what `extra` carried intact.
+ """
+ from superset import db
+
+ table = _mapped_table(session, "pm_t4")
+
+ DatasetDAO.update(
+ table,
+ {
+ "extra": json.dumps(
+ {
+ "partition_filter_mapping": {
+ "partition_column": "dt_epoch",
+ "mapped_column": "event_time",
+ }
+ }
+ )
+ },
+ )
+ db.session.flush()
+
+ assert table.partition_column == "dt_epoch"
+ assert table.partition_mapped_column == "event_time"
+
+
+def test_an_absent_field_leaves_the_mapping_unchanged(session: Session) -> None:
+ """ "Absent" must mean "unchanged", never "cleared"."""
+ from superset import db
+
+ table = _mapped_table(session, "pm_t5")
+
+ DatasetDAO.update(table, {"description": "unrelated edit"})
+ db.session.flush()
+
+ assert table.partition_column == "dt_epoch"
+ assert table.partition_mapped_column == "event_time"
+ assert table.columns[0].partition_transform_is_monotonic is True
+
+
+def test_a_partial_column_payload_keeps_the_monotonic_flag(session: Session) -> None:
+ """
+ A column payload that names only `id` and one unrelated field must not clear
+ a flag it never mentioned. This is why `DatasetColumnsPutSchema` refuses a
+ `load_default` for it and why the patch carries present keys only.
+ """
+ from superset import db
+
+ table = _mapped_table(session, "pm_t6")
+ ids = {c.column_name: c.id for c in table.columns}
+
+ DatasetDAO.update(
+ table,
+ {
+ "columns": [
+ {"id": ids["event_time"], "verbose_name": "Event time"},
+ {"id": ids["dt_epoch"]},
+ ]
+ },
+ )
+ db.session.flush()
+
+ event_time = next(c for c in table.columns if c.column_name == "event_time")
+ assert event_time.partition_value_transform == "unix_timestamp(:value)"
+ assert event_time.partition_transform_is_monotonic is True
+
+
+def test_the_mapping_attributes_are_read_only(session: Session) -> None:
+ """
+ Assignment has to raise rather than vanish.
+
+ Under the extra-JSON store a setter's write is lost the moment `extra` is
+ assigned afterwards, with no exception and no log. Read-only properties turn
+ that silent loss into an immediate error, so this is the guard against
+ someone reintroducing setters later.
+ """
+ table = _mapped_table(session, "pm_t7")
+ column = table.columns[0]
+
+ for target, attribute in (
+ (table, "partition_column"),
+ (table, "partition_mapped_column"),
+ (column, "partition_value_transform"),
+ (column, "partition_transform_is_monotonic"),
+ ):
+ with pytest.raises(AttributeError):
+ setattr(target, attribute, "x")
diff --git a/tests/unit_tests/dashboards/datasets_test.py b/tests/unit_tests/dashboards/datasets_test.py
index b1a426e543b..9c8d54213f6 100644
--- a/tests/unit_tests/dashboards/datasets_test.py
+++ b/tests/unit_tests/dashboards/datasets_test.py
@@ -330,6 +330,9 @@ def test_dashboard_table_serialization_includes_capabilities_and_parent() -> Non
"health_check_message",
"always_filter_main_dttm",
"normalize_columns",
+ "partition_column",
+ "partition_mapped_column",
+ "partition_filter_mapping",
"column_types",
"column_names",
}
diff --git a/tests/unit_tests/datasets/partition_mapping_preview_test.py b/tests/unit_tests/datasets/partition_mapping_preview_test.py
new file mode 100644
index 00000000000..510e30198a2
--- /dev/null
+++ b/tests/unit_tests/datasets/partition_mapping_preview_test.py
@@ -0,0 +1,406 @@
+# 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.
+"""
+The partition mapping preview endpoint.
+
+This route fires a real warehouse query from a text input in the dataset
+editor, so the order of its guards is the point: parse and denylist checks run
+*before* anything reaches the engine, which is what keeps a half-typed
+expression from costing a query at all.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+from flask import Flask
+from sqlalchemy.orm.session import Session
+
+from superset import db
+from superset.connectors.sqla.partition_mapping_storage import set_partition_mapping
+
+PROBE = "superset.connectors.sqla.partition_mapping.evaluate_transform"
+
+
+@pytest.fixture(autouse=True)
+def enable_partition_filter_mapping(app: Flask) -> Any:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = True
+ yield
+ del app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"]
+
+
+@pytest.fixture(autouse=True)
+def real_cache(app: Flask) -> Any:
+ """
+ The test app runs a null cache, which would make the rate limiter a no-op.
+ """
+ from flask_caching import Cache
+
+ from superset.extensions import cache_manager
+
+ cache = Cache(config={"CACHE_TYPE": "SimpleCache", "CACHE_DEFAULT_TIMEOUT": 300})
+ cache.init_app(app)
+ original = cache_manager._cache # noqa: SLF001
+ cache_manager._cache = cache # noqa: SLF001
+ yield
+ cache_manager._cache = original # noqa: SLF001
+
+
+@pytest.fixture(autouse=True)
+def allow_editorship(mocker: Any) -> Any:
+ """
+ The route gates on per-object editorship on top of ``@protect()``. The unit
+ test app has no real roles, so grant it here rather than in every test.
+ """
+ mocker.patch(
+ "superset.datasets.api.security_manager.raise_for_editorship",
+ )
+
+
+@pytest.fixture
+def dataset(session: Session) -> Any:
+ from superset.connectors.sqla.models import SqlaTable, TableColumn
+ from superset.models.core import Database
+
+ SqlaTable.metadata.create_all(db.session.get_bind())
+ database = Database(database_name="my_db", sqlalchemy_uri="sqlite://")
+ table = SqlaTable(
+ table_name="web_events",
+ database=database,
+ main_dttm_col="event_time",
+ columns=[
+ TableColumn(column_name="event_time", is_dttm=True, type="TIMESTAMP"),
+ TableColumn(column_name="dt_epoch", type="BIGINT"),
+ ],
+ )
+ # Through the real write API: this fixture has an app context, so it can
+ # exercise the configured store rather than hand-building its blob.
+ set_partition_mapping(table, partition_column="dt_epoch")
+ db.session.add(table)
+ db.session.flush()
+ return table
+
+
+def test_preview_returns_the_emitted_predicate(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ with patch(PROBE, return_value=[1768435200]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15 00:00:00",
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json["result"] == {
+ "valid": True,
+ "emitted_predicate": "dt_epoch >= 1768435200",
+ }
+
+
+def test_preview_reports_a_parse_error_without_touching_the_engine(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ Validate first, probe second. A half-typed transform is by definition
+ unparseable, which is most of the traffic a debounced text input produces.
+ """
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value",
+ "sample_value": "2026-01-15 00:00:00",
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json["result"]["valid"] is False
+ assert response.json["result"]["error"]
+
+
+def test_preview_rejects_a_non_deterministic_transform(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp()",
+ "sample_value": "2026-01-15 00:00:00",
+ },
+ )
+
+ assert response.json["result"]["valid"] is False
+
+
+def test_preview_reports_an_unknown_mapped_column(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "nope",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15 00:00:00",
+ },
+ )
+
+ assert response.json["result"]["valid"] is False
+
+
+def test_preview_reports_a_failed_probe_rather_than_erroring(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ with patch(PROBE, return_value=None):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15 00:00:00",
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json["result"]["valid"] is False
+
+
+def test_preview_404s_for_an_unknown_dataset(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ response = client.post(
+ "/api/v1/dataset/99999/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15",
+ },
+ )
+ assert response.status_code == 404
+
+
+def test_preview_is_gated_on_the_feature_flag(
+ app: Flask, client: Any, full_api_access: None, dataset: Any
+) -> None:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = False
+
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15",
+ },
+ )
+ assert response.status_code == 404
+
+
+def test_preview_rejects_an_invalid_payload(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={"nonsense": True},
+ )
+ assert response.status_code == 400
+
+
+def test_preview_is_rate_limited_per_user_and_dataset(
+ app: Flask, client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ Debouncing is a client-side courtesy, not a guard: a held keydown, or a few
+ owners with the editor open, is sustained load on a production cluster.
+ """
+ app.config["PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT"] = 2
+
+ payload = {
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15",
+ }
+ with patch(PROBE, return_value=[1]):
+ statuses = [
+ client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={**payload, "sample_value": f"2026-01-{day:02d}"},
+ ).status_code
+ for day in range(1, 5)
+ ]
+
+ assert statuses[:2] == [200, 200]
+ assert 429 in statuses[2:]
+
+
+def test_the_budget_is_spent_exactly_once_per_request(
+ app: Flask, client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ Counting is `add` then `inc`, so the first request of a window and every
+ one after it cost the same single unit. A read-then-write pair would let
+ concurrent previews all observe the same sub-limit value and all through.
+ """
+ app.config["PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT"] = 3
+
+ payload = {
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15",
+ }
+ with patch(PROBE, return_value=[1]):
+ statuses = [
+ client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={**payload, "sample_value": f"2026-01-{day:02d}"},
+ ).status_code
+ for day in range(1, 6)
+ ]
+
+ assert statuses == [200, 200, 200, 429, 429]
+
+
+def test_the_window_is_dated_by_its_first_request_not_its_last(
+ app: Flask, client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ Only `add` sets a lifetime. Re-stamping the key on every increment would
+ make this a sliding window, where sustained typing keeps the budget spent
+ indefinitely instead of recovering a minute after the burst began.
+ """
+ from superset.datasets.api import PREVIEW_RATE_LIMIT_WINDOW
+ from superset.extensions import cache_manager
+
+ app.config["PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT"] = 5
+
+ payload = {
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15",
+ }
+ timeouts: list[Any] = []
+ backend = cache_manager.cache.cache
+ original_add = backend.add
+
+ def record(key: str, value: Any, timeout: Any = None) -> Any:
+ timeouts.append(timeout)
+ return original_add(key, value, timeout=timeout)
+
+ with patch.object(backend, "add", side_effect=record):
+ with patch(PROBE, return_value=[1]):
+ for day in range(1, 4):
+ client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={**payload, "sample_value": f"2026-01-{day:02d}"},
+ )
+
+ # Three requests, three `add` attempts, but only the first one takes -- and
+ # every later one is a no-op that leaves the original expiry alone.
+ assert timeouts == [PREVIEW_RATE_LIMIT_WINDOW] * 3
+
+
+def test_a_cache_that_cannot_count_does_not_lock_the_editor_out(
+ app: Flask, client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """Unenforceable is not the same as spent."""
+ from superset.extensions import cache_manager
+
+ app.config["PARTITION_TRANSFORM_PREVIEW_RATE_LIMIT"] = 1
+
+ payload = {
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "2026-01-15",
+ }
+ backend = cache_manager.cache.cache
+ with patch.object(backend, "add", return_value=False):
+ with patch.object(backend, "inc", return_value=None):
+ with patch(PROBE, return_value=[1]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json=payload,
+ )
+
+ assert response.status_code == 200
+
+
+def test_a_probed_null_renders_as_sql_null(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ A transform returns NULL for an input it cannot convert. `str()` would put
+ the Python repr `None` in the predicate, which is not SQL.
+ """
+ with patch(PROBE, return_value=[None]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_value": "not a date",
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json["result"]["emitted_predicate"] == "dt_epoch >= NULL"
+
+
+def test_a_probed_string_is_quoted_and_escaped_by_the_dialect(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """Escaping is the dialect's job, not a hand-rolled `replace`."""
+ with patch(PROBE, return_value=["o'hara"]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "lower(:value)",
+ "sample_value": "O'Hara",
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json["result"]["emitted_predicate"] == "dt_epoch >= 'o''hara'"
+
+
+def test_an_over_long_transform_is_rejected_before_the_engine(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ The transform is parsed and then run against the warehouse, so an unbounded
+ string is parser time and warehouse time an owner can spend at will.
+ """
+ with patch(PROBE) as probe:
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)" + " " * 2000,
+ "sample_value": "2026-01-15",
+ },
+ )
+
+ assert response.status_code == 400
+ probe.assert_not_called()
diff --git a/tests/unit_tests/datasets/partition_mapping_serialization_test.py b/tests/unit_tests/datasets/partition_mapping_serialization_test.py
new file mode 100644
index 00000000000..86a99bffde7
--- /dev/null
+++ b/tests/unit_tests/datasets/partition_mapping_serialization_test.py
@@ -0,0 +1,504 @@
+# 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.
+"""
+The partition mapping has to survive every layer it passes through.
+
+``always_filter_main_dttm`` is the template these follow: a dataset-level setting
+that names a column and appears in the ORM, the ``data`` payload, the API schemas
+and the frontend types. A field missing from any one of them is dropped silently,
+which is exactly the failure mode these tests exist to catch.
+
+Export and import get a shorter list than they would with real columns: the
+mapping lives inside ``tables.extra``, which is already in ``export_fields`` and
+already one of `ExportDatasetsCommand`'s ``JSON_KEYS``, so it travels with no
+bundle-format change. What has to be proved instead is that it comes back out
+the other side, and that nothing else in the blob is lost on the way.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+from flask import Flask
+
+from superset.commands.dataset.export import ExportDatasetsCommand, JSON_KEYS
+from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.connectors.sqla.partition_mapping_storage import (
+ clear_dangling,
+ ColumnTransform,
+ EXTRA_KEY,
+ extra_with_partition_mapping,
+ parse_mapping,
+ StoredPartitionMapping,
+)
+from superset.datasets.schemas import (
+ DatasetColumnsPutSchema,
+ DatasetPutSchema,
+ ImportV1DatasetSchema,
+)
+from superset.exceptions import SupersetMarshmallowValidationError
+from superset.models.core import Database
+from superset.utils import json
+
+DATASET_FIELDS = ["partition_column", "partition_mapped_column"]
+COLUMN_FIELDS = ["partition_value_transform", "partition_transform_is_monotonic"]
+
+MAPPING = StoredPartitionMapping(
+ partition_column="dt_epoch",
+ mapped_column=None,
+ column_transforms={
+ "event_time": ColumnTransform("unix_timestamp(:value)", is_monotonic=True)
+ },
+)
+
+
+@pytest.fixture(autouse=True)
+def enable_partition_filter_mapping(app: Flask) -> Any:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = True
+ yield
+ del app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"]
+
+
+def _table(
+ mapping: StoredPartitionMapping = MAPPING, extra: str | None = None
+) -> SqlaTable:
+ database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
+ table = SqlaTable(
+ table_name="web_events",
+ database=database,
+ main_dttm_col="event_time",
+ columns=[
+ TableColumn(column_name="event_time", is_dttm=True, type="TIMESTAMP"),
+ TableColumn(column_name="dt_epoch", type="BIGINT"),
+ ],
+ )
+ # Stored state, not attribute assignment: the model's four mapping attributes
+ # are read-only properties over the configured store.
+ table.extra = extra_with_partition_mapping(extra, mapping)
+ return table
+
+
+# ---------------------------------------------------------------------------
+# The Explore and dashboard payloads
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("field", DATASET_FIELDS)
+def test_dataset_fields_reach_the_explore_payload(app: Flask, field: str) -> None:
+ with app.app_context():
+ data = _table().data
+ assert field in data
+
+
+@pytest.mark.parametrize("field", COLUMN_FIELDS)
+def test_column_fields_reach_the_explore_payload(app: Flask, field: str) -> None:
+ with app.app_context():
+ data = _table().data
+ assert field in data["columns"][0]
+
+
+def test_the_transform_is_reported_only_on_the_mapped_column(app: Flask) -> None:
+ """
+ Each column reports what its own row would say after the migration, so the
+ partition column must not claim the mapped column's transform. Scoping this
+ is also what stops `resolve_partition_mapping` reading a transform off the
+ wrong column.
+ """
+ with app.app_context():
+ columns = {column["column_name"]: column for column in _table().data["columns"]}
+
+ assert (
+ columns["event_time"]["partition_value_transform"] == "unix_timestamp(:value)"
+ )
+ assert columns["event_time"]["partition_transform_is_monotonic"] is True
+ assert columns["dt_epoch"]["partition_value_transform"] is None
+ assert columns["dt_epoch"]["partition_transform_is_monotonic"] is False
+
+
+def test_a_transform_does_not_follow_a_re_pointed_main_dttm_col(app: Flask) -> None:
+ """
+ The reason transforms are keyed by column name rather than flattened onto the
+ mapping. An owner who declares a transform on ``event_time`` and later
+ re-points ``main_dttm_col`` at another column must not have that column's
+ values mirrored through ``event_time``'s expression -- the mapping goes
+ inactive instead, which is what four real columns would do.
+ """
+ table = _table()
+ table.main_dttm_col = "dt_epoch"
+
+ with app.app_context():
+ summary = table.data["partition_filter_mapping"]
+
+ # `dt_epoch` is now both the partition column and the effective mapped one,
+ # and it carries no transform of its own.
+ assert summary is not None
+ assert summary["active"] is False
+
+
+def test_the_mapping_summary_survives_dashboard_payload_pruning(app: Flask) -> None:
+ """
+ ``data_for_slices`` prunes columns no chart references, and the partition
+ column is typically referenced by none of them. The Explore indicator
+ therefore reads a self-contained dataset-level dict rather than looking the
+ column up inside ``datasource.columns``.
+ """
+ with app.app_context():
+ data = _table().data_for_slices([])
+
+ assert data["partition_filter_mapping"] == {
+ "partition_column": "dt_epoch",
+ "mapped_column": "event_time",
+ "active": True,
+ }
+
+
+def test_the_mapping_summary_reports_inactive_without_a_transform(app: Flask) -> None:
+ table = _table(StoredPartitionMapping(partition_column="dt_epoch"))
+
+ with app.app_context():
+ summary = table.data["partition_filter_mapping"]
+
+ assert summary is not None
+ assert summary["active"] is False
+
+
+@pytest.mark.parametrize(
+ "transform",
+ [
+ # Missing the placeholder, so there is no filter value to substitute.
+ "unix_timestamp(event_time)",
+ # Unparseable, so nothing can be emitted from it.
+ "unix_timestamp(:value",
+ # Blocking on PUT, but create, import and the free-text Extra box do not
+ # go through that validation, so it can still be read back.
+ "{{ current_username() }}",
+ # Non-deterministic: rejected on PUT, and inactive wherever it lands.
+ "unix_timestamp(now())",
+ ],
+)
+def test_the_mapping_summary_reports_inactive_for_an_invalid_transform(
+ app: Flask,
+ transform: str,
+) -> None:
+ """
+ A transform that fails validation is saved inactive on purpose. The indicator
+ has to say the same thing, or it advertises a mapping that will never mirror
+ a filter.
+ """
+ table = _table(
+ StoredPartitionMapping(
+ partition_column="dt_epoch",
+ column_transforms={"event_time": ColumnTransform(transform)},
+ )
+ )
+
+ with app.app_context():
+ summary = table.data["partition_filter_mapping"]
+
+ assert summary is not None
+ assert summary["active"] is False
+
+
+def test_the_mapping_summary_still_names_the_columns_when_inactive(app: Flask) -> None:
+ """
+ An inactive mapping is exactly the one worth naming: the editor has to tell
+ the owner which columns to fix.
+ """
+ table = _table(
+ StoredPartitionMapping(
+ partition_column="dt_epoch",
+ column_transforms={
+ "event_time": ColumnTransform("unix_timestamp(event_time)")
+ },
+ )
+ )
+
+ with app.app_context():
+ summary = table.data["partition_filter_mapping"]
+
+ assert summary == {
+ "partition_column": "dt_epoch",
+ "mapped_column": "event_time",
+ "active": False,
+ }
+
+
+def test_there_is_no_mapping_summary_without_a_partition_column(app: Flask) -> None:
+ table = _table(StoredPartitionMapping())
+
+ with app.app_context():
+ assert table.data["partition_filter_mapping"] is None
+
+
+# ---------------------------------------------------------------------------
+# The API schemas
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("field", DATASET_FIELDS)
+def test_put_schema_accepts_the_dataset_fields(field: str) -> None:
+ loaded = DatasetPutSchema().load({field: "dt_epoch"})
+ assert loaded[field] == "dt_epoch"
+
+
+def test_put_schema_accepts_the_column_fields() -> None:
+ loaded = DatasetColumnsPutSchema().load(
+ {
+ "column_name": "event_time",
+ "partition_value_transform": "unix_timestamp(:value)",
+ "partition_transform_is_monotonic": True,
+ }
+ )
+ assert loaded["partition_value_transform"] == "unix_timestamp(:value)"
+ assert loaded["partition_transform_is_monotonic"] is True
+
+
+def test_put_schema_leaves_out_a_monotonic_flag_the_payload_omits() -> None:
+ """
+ The DAO's patch carries present keys only, so a default here would let a
+ partial column update clear a monotonic flag the request never mentioned --
+ and silently stop mirroring range filters.
+ """
+ loaded = DatasetColumnsPutSchema().load({"column_name": "event_time"})
+ assert "partition_transform_is_monotonic" not in loaded
+
+
+def test_put_schema_allows_clearing_the_mapping() -> None:
+ """Removing a mapping is a null, not an omission."""
+ loaded = DatasetPutSchema().load({"partition_column": None})
+ assert loaded["partition_column"] is None
+
+
+@pytest.mark.parametrize("field", DATASET_FIELDS + COLUMN_FIELDS)
+def test_the_api_exposes_the_mapping(field: str) -> None:
+ """
+ The two column fields are not optional here. The dataset editor seeds itself
+ from `GET /api/v1/dataset/<pk>` and then PUTs every column back, so a
+ transform this payload omits is one the next save clears.
+ """
+ from superset.datasets.api import DatasetRestApi
+
+ names = set(DatasetRestApi.show_columns)
+ assert field in names or f"columns.{field}" in names
+
+
+@pytest.mark.parametrize("field", DATASET_FIELDS)
+def test_the_api_accepts_the_dataset_fields(field: str) -> None:
+ from superset.datasets.api import DatasetRestApi
+
+ assert field in DatasetRestApi.edit_columns
+
+
+def test_the_mapping_is_hidden_from_an_inaccessible_dashboard_dataset() -> None:
+ """
+ ``partition_column`` is a physical column name, the same class of detail as
+ ``columns`` and ``sql``, so a caller who cannot access the dataset must not
+ receive it in the dashboard payload.
+ """
+ from superset.dashboards.api import DASHBOARD_DATASET_INACCESSIBLE_FIELDS
+
+ for field in (*DATASET_FIELDS, "partition_filter_mapping"):
+ assert field in DASHBOARD_DATASET_INACCESSIBLE_FIELDS
+
+
+# ---------------------------------------------------------------------------
+# Export and import, which carry the mapping inside ``extra``
+# ---------------------------------------------------------------------------
+
+
+def test_export_writes_the_mapping_as_nested_yaml(app: Flask) -> None:
+ """
+ No export-format change was needed: ``extra`` is already in
+ ``SqlaTable.export_fields`` and already one of ``JSON_KEYS``, so the command
+ parses it into a nested mapping in the YAML for free.
+ """
+ assert "extra" in SqlaTable.export_fields
+ assert "extra" in JSON_KEYS
+
+ with app.app_context():
+ content = ExportDatasetsCommand._file_content(_table())
+
+ assert EXTRA_KEY in content
+ assert "unix_timestamp(:value)" in content
+
+
+def test_the_mapping_round_trips_through_the_import_schema(app: Flask) -> None:
+ """
+ The bundle carries a nested ``extra`` mapping; `fix_extra` also accepts the
+ string form, and `import_dataset` re-encodes it before
+ `import_from_dict` writes it. Either way the mapping comes back.
+ """
+ with app.app_context():
+ exported = _table().extra
+
+ for payload_extra in (json.loads(exported), exported):
+ loaded = ImportV1DatasetSchema().load(
+ {
+ "table_name": "web_events",
+ "uuid": "00000000-0000-0000-0000-000000000001",
+ "database_uuid": "00000000-0000-0000-0000-000000000002",
+ "version": "1.0.0",
+ "extra": payload_extra,
+ }
+ )
+ assert parse_mapping(json.dumps(loaded["extra"])) == MAPPING
+
+
+def test_a_bundle_predating_the_flag_does_not_claim_monotonicity() -> None:
+ """
+ The flag gates range mirroring, so a blob that omits it -- or carries
+ something merely truthy -- must not silently claim the transform preserves
+ ordering.
+ """
+ for value in (None, "true", 1, "no"):
+ blob: dict[str, Any] = {
+ "partition_column": "dt_epoch",
+ "column_transforms": {
+ "event_time": {"value_transform": "unix_timestamp(:value)"}
+ },
+ }
+ if value is not None:
+ blob["column_transforms"]["event_time"]["is_monotonic"] = value
+ mapping = parse_mapping(json.dumps({EXTRA_KEY: blob}))
+ assert mapping.transform_for("event_time").is_monotonic is False
+
+
+def test_writing_the_mapping_preserves_the_rest_of_extra() -> None:
+ """
+ ``extra`` is shared: ``certification``, ``warning_markdown``, ``timezone``
+ and ``disallow_adhoc_metrics`` all live in the same blob, and the dataset
+ editor round-trips the whole string verbatim. A mapping write merges.
+ """
+ neighbours = {
+ "certification": {"certified_by": "Data Eng", "details": "reviewed"},
+ "warning_markdown": "**beware**",
+ "timezone": "Europe/Berlin",
+ "disallow_adhoc_metrics": True,
+ }
+ written = extra_with_partition_mapping(json.dumps(neighbours), MAPPING)
+ decoded = json.loads(str(written))
+
+ assert {key: decoded[key] for key in neighbours} == neighbours
+ assert parse_mapping(written) == MAPPING
+
+ # And clearing it takes only the one key with it.
+ cleared = json.loads(
+ str(extra_with_partition_mapping(written, StoredPartitionMapping()))
+ )
+ assert cleared == neighbours
+
+
+# ---------------------------------------------------------------------------
+# A column sync can pull the partition column out from under the mapping
+# ---------------------------------------------------------------------------
+
+
+def test_a_sync_that_removes_the_partition_column_clears_the_mapping() -> None:
+ """
+ An API-driven ``override_columns=true`` sync must not leave a dangling
+ mapping. This is the authoritative path -- the client-side sync clears the
+ mapping too, but a caller can bypass the editor entirely.
+ """
+ cleared = clear_dangling(MAPPING, {"event_time"})
+
+ assert cleared.partition_column is None
+ assert cleared.mapped_column is None
+
+
+def test_a_sync_that_removes_the_mapped_column_clears_only_the_override() -> None:
+ """
+ The partition column is still real, so the designation survives; the mapping
+ falls back to "no mapped column" and goes inactive until one is chosen.
+ """
+ mapping = StoredPartitionMapping(
+ partition_column="dt_epoch",
+ mapped_column="event_time",
+ column_transforms=dict(MAPPING.column_transforms),
+ )
+ cleared = clear_dangling(mapping, {"dt_epoch"})
+
+ assert cleared.partition_column == "dt_epoch"
+ assert cleared.mapped_column is None
+
+
+def test_a_sync_drops_transforms_for_columns_that_are_gone() -> None:
+ """A transform keyed on a deleted column would linger forever otherwise."""
+ cleared = clear_dangling(MAPPING, {"dt_epoch"})
+
+ assert cleared.column_transforms == {}
+
+
+def test_a_sync_that_keeps_both_columns_leaves_the_mapping_alone() -> None:
+ assert clear_dangling(MAPPING, {"event_time", "dt_epoch"}) == MAPPING
+
+
+# ---------------------------------------------------------------------------
+# The second door: a mapping typed into the free-text Extra box
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "blob",
+ [
+ 42,
+ "a string",
+ [],
+ {"partition_column": 7},
+ {"mapped_column": ["event_time"]},
+ {"partition_column": "x" * 251},
+ {"column_transforms": "nope"},
+ {"column_transforms": {"event_time": "nope"}},
+ {"column_transforms": {"event_time": {"value_transform": 7}}},
+ {"column_transforms": {"event_time": {"is_monotonic": "yes"}}},
+ ],
+)
+def test_a_mistyped_mapping_in_extra_is_rejected(blob: Any) -> None:
+ """
+ `parse_mapping` degrades any of these to "not configured", which an owner has
+ no way to diagnose. `DatasetPutSchema.extra` turns the silence into a message.
+ """
+ with pytest.raises(SupersetMarshmallowValidationError) as excinfo:
+ DatasetPutSchema().load({"extra": json.dumps({EXTRA_KEY: blob})})
+
+ # One readable sentence, not a list of single characters -- which is what
+ # marshmallow makes of a `LazyString` passed straight to `ValidationError`.
+ message = excinfo.value.error.extra["messages"]["extra"][0]
+ assert isinstance(message, str)
+ assert message.endswith(".")
+
+
+@pytest.mark.parametrize(
+ "extra",
+ [
+ # Never validated on this endpoint before, so a deployment may already
+ # hold one of these. Failing an owner's next save over a field they did
+ # not touch would be a regression.
+ "not json at all",
+ "[1, 2, 3]",
+ # Everything else in the blob is somebody else's business.
+ '{"certification": {"certified_by": "Data Eng"}, "unknown_key": 1}',
+ '{"timezone": "Europe/Berlin"}',
+ ],
+)
+def test_extra_validation_is_scoped_to_the_mapping_key(extra: str) -> None:
+ assert DatasetPutSchema().load({"extra": extra})["extra"] == extra
+
+
+def test_a_well_formed_mapping_in_extra_is_accepted() -> None:
+ extra = extra_with_partition_mapping('{"timezone": "Europe/Berlin"}', MAPPING)
+ assert DatasetPutSchema().load({"extra": extra})["extra"] == extra
diff --git a/tests/unit_tests/models/partition_mirroring_test.py b/tests/unit_tests/models/partition_mirroring_test.py
new file mode 100644
index 00000000000..897ec8e43bf
--- /dev/null
+++ b/tests/unit_tests/models/partition_mirroring_test.py
@@ -0,0 +1,908 @@
+# 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 mirroring inside ``ExploreMixin.get_sqla_query``.
+
+The probe that resolves ``T(v)`` against the engine is stubbed throughout; what
+these tests pin down is *which* predicates get mirrored and with what values,
+which is where the correctness argument lives.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+from flask import Flask
+
+from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn
+from superset.connectors.sqla.partition_mapping_storage import (
+ ColumnTransform,
+ extra_with_partition_mapping,
+ StoredPartitionMapping,
+)
+from superset.models.core import Database
+from superset.utils import json
+from superset.utils.core import FilterOperator
+
+PROBE = "superset.connectors.sqla.partition_mapping.evaluate_transform"
+
+
+@pytest.fixture(autouse=True)
+def enable_partition_filter_mapping(app: Flask) -> Any:
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = True
+ yield
+ del app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"]
+
+
+def _table(
+ *,
+ transform: str = "unix_timestamp(:value)",
+ monotonic: bool = True,
+ partition_column: str | None = "dt_epoch",
+ mapped_column: str = "event_time",
+ partition_mapped_column: str | None = None,
+ main_dttm_col: str | None = "event_time",
+) -> SqlaTable:
+ database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
+ columns = [
+ TableColumn(column_name="event_time", is_dttm=True, type="TIMESTAMP"),
+ TableColumn(column_name="other_time", is_dttm=True, type="TIMESTAMP"),
+ TableColumn(column_name="dt_epoch", type="BIGINT"),
+ TableColumn(column_name="country", type="VARCHAR"),
+ TableColumn(column_name="region_key", type="VARCHAR"),
+ ]
+ table = SqlaTable(
+ table_name="web_events",
+ database=database,
+ schema=None,
+ main_dttm_col=main_dttm_col,
+ columns=columns,
+ metrics=[SqlMetric(metric_name="hits", expression="COUNT(*)")],
+ )
+ # The mapping is persisted, not assigned: the model's four attributes are
+ # read-only properties over the configured store. `extra_with_partition_mapping`
+ # is the store's own serializer and needs no app context, which the factory
+ # does not have.
+ table.extra = extra_with_partition_mapping(
+ table.extra,
+ StoredPartitionMapping(
+ partition_column=partition_column,
+ mapped_column=partition_mapped_column,
+ column_transforms={
+ mapped_column: ColumnTransform(
+ value_transform=transform, is_monotonic=monotonic
+ )
+ },
+ ),
+ )
+ return table
+
+
+def _set_timezone(table: SqlaTable, timezone: str) -> None:
+ """
+ Add a dataset timezone without dropping the mapping.
+
+ Assigning a fresh `extra` string would: under the extra-JSON store the
+ mapping lives in the same blob, next to `timezone`, `certification` and the
+ rest. Merging is what every writer has to do, including the frontend.
+ """
+ extra = json.loads(table.extra) if table.extra else {}
+ extra["timezone"] = timezone
+ table.extra = json.dumps(extra)
+
+
+def _query(table: SqlaTable, **kwargs: Any) -> str:
+ defaults: dict[str, Any] = {
+ "columns": ["country"],
+ "metrics": [],
+ "orderby": [],
+ "extras": {},
+ "filter": [],
+ "granularity": None,
+ "is_timeseries": False,
+ }
+ defaults.update(kwargs)
+ result = table.get_sqla_query(**defaults)
+ return str(
+ result.sqla_query.compile(compile_kwargs={"literal_binds": True})
+ ).replace("\n", " ")
+
+
+# ---------------------------------------------------------------------------
+# The safe operators
+# ---------------------------------------------------------------------------
+
+
+def test_equality_filter_mirrors_onto_the_partition_column(app: Flask) -> None:
+ table = _table(
+ transform="lower(:value)",
+ monotonic=False,
+ mapped_column="country",
+ partition_mapped_column="country",
+ partition_column="region_key",
+ )
+
+ with app.app_context():
+ with patch(PROBE, return_value=["us"]):
+ sql = _query(
+ table,
+ filter=[
+ {"col": "country", "op": FilterOperator.EQUALS.value, "val": "US"}
+ ],
+ )
+
+ assert "region_key = 'us'" in sql
+ assert "country = 'US'" in sql
+
+
+def test_in_filter_mirrors_element_wise(app: Flask) -> None:
+ table = _table(
+ transform="lower(:value)",
+ monotonic=False,
+ mapped_column="country",
+ partition_mapped_column="country",
+ partition_column="region_key",
+ )
+
+ with app.app_context():
+ with patch(PROBE, return_value=["us", "ca"]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "country",
+ "op": FilterOperator.IN.value,
+ "val": ["US", "CA"],
+ }
+ ],
+ )
+
+ assert "region_key IN ('us', 'ca')" in sql
+
+
+def test_time_range_mirrors_both_bounds(app: Flask) -> None:
+ """
+ The Explore time range is the most important operator in the feature, and it
+ is a range operator -- so it only mirrors on a declared-monotonic transform.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch >= 1767225600" in sql
+ assert "dt_epoch < 1769904000" in sql
+
+
+def test_temporal_range_filter_mirrors_both_bounds(app: Flask) -> None:
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-01 : 2026-02-01",
+ }
+ ],
+ )
+
+ assert "dt_epoch >= 1767225600" in sql
+ assert "dt_epoch < 1769904000" in sql
+
+
+@pytest.mark.parametrize(
+ "operator,expected",
+ [
+ (FilterOperator.GREATER_THAN, "dt_epoch > 1767225600"),
+ (FilterOperator.GREATER_THAN_OR_EQUALS, "dt_epoch >= 1767225600"),
+ (FilterOperator.LESS_THAN, "dt_epoch < 1767225600"),
+ (FilterOperator.LESS_THAN_OR_EQUALS, "dt_epoch <= 1767225600"),
+ ],
+)
+def test_range_operators_mirror_with_the_same_direction(
+ app: Flask, operator: FilterOperator, expected: str
+) -> None:
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": operator.value,
+ "val": "2026-01-01 00:00:00",
+ }
+ ],
+ )
+
+ assert expected in sql
+
+
+# ---------------------------------------------------------------------------
+# The unsafe operators — nothing is emitted
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "flt",
+ [
+ {"col": "country", "op": FilterOperator.NOT_EQUALS.value, "val": "US"},
+ {"col": "country", "op": FilterOperator.NOT_IN.value, "val": ["US"]},
+ {"col": "country", "op": FilterOperator.LIKE.value, "val": "U%"},
+ {"col": "country", "op": FilterOperator.ILIKE.value, "val": "U%"},
+ {"col": "country", "op": FilterOperator.NOT_LIKE.value, "val": "U%"},
+ {"col": "country", "op": FilterOperator.IS_NULL.value},
+ {"col": "country", "op": FilterOperator.IS_NOT_NULL.value},
+ ],
+)
+def test_unsafe_operators_emit_nothing(app: Flask, flt: dict[str, Any]) -> None:
+ """
+ ``T`` need not be injective, so ``country != 'US'`` does not imply
+ ``region_key != 'us'`` -- mirroring it would drop rows whose ``country`` is
+ already lowercase, rows the original filter keeps.
+ """
+ table = _table(
+ transform="lower(:value)",
+ monotonic=False,
+ mapped_column="country",
+ partition_mapped_column="country",
+ partition_column="region_key",
+ )
+
+ with app.app_context():
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ sql = _query(table, filter=[flt])
+
+ assert "region_key" not in sql
+
+
+def test_ranges_do_not_mirror_when_the_transform_is_not_order_preserving(
+ app: Flask,
+) -> None:
+ """
+ ``hour(:value)`` is a perfectly reasonable partition transform on a
+ ``TIMESTAMP`` column and it is not monotonic, so a time range must not
+ mirror through it.
+ """
+ table = _table(transform="hour(:value)", monotonic=False)
+
+ with app.app_context():
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_equality_still_mirrors_when_the_transform_is_not_order_preserving(
+ app: Flask,
+) -> None:
+ table = _table(transform="hour(:value)", monotonic=False)
+
+ with app.app_context():
+ with patch(PROBE, return_value=[13]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.EQUALS.value,
+ "val": "2026-01-01 13:00:00",
+ }
+ ],
+ )
+
+ assert "dt_epoch = 13" in sql
+
+
+# ---------------------------------------------------------------------------
+# Bail-outs and edge cases
+# ---------------------------------------------------------------------------
+
+
+def test_nothing_mirrors_when_the_feature_flag_is_off(app: Flask) -> None:
+ table = _table()
+ app.config["DEFAULT_FEATURE_FLAGS"]["PARTITION_FILTER_MAPPING"] = False
+
+ with app.app_context():
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_an_open_ended_range_mirrors_only_the_bound_it_has(app: Flask) -> None:
+ """``from_dttm``/``to_dttm`` are ``None`` for open-ended ranges."""
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600]) as probe:
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=None,
+ )
+
+ assert probe.call_args.args[-1] == [datetime(2026, 1, 1)]
+ assert "dt_epoch >= 1767225600" in sql
+ assert "dt_epoch <" not in sql
+
+
+def test_no_filter_range_mirrors_nothing(app: Flask) -> None:
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=None,
+ to_dttm=None,
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_the_same_bound_is_not_mirrored_twice(app: Flask) -> None:
+ """
+ A ``granularity`` time filter *and* a ``TEMPORAL_RANGE`` ad-hoc filter on the
+ same column is a routine Explore configuration. Emitting the predicate twice
+ is harmless SQL but makes "View query" surprising.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-01 : 2026-02-01",
+ }
+ ],
+ )
+
+ assert sql.count("dt_epoch >= 1767225600") == 1
+ assert sql.count("dt_epoch < 1769904000") == 1
+
+
+def test_always_filter_main_dttm_mirrors_the_main_column_too(app: Flask) -> None:
+ """
+ With ``always_filter_main_dttm`` the query also filters ``main_dttm_col``,
+ which is a *different* column from the one the chart grouped by. That filter
+ is the one the mapping tracks, so it has to mirror.
+ """
+ table = _table()
+ table.always_filter_main_dttm = True
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]):
+ sql = _query(
+ table,
+ granularity="other_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch >= 1767225600" in sql
+ assert "dt_epoch < 1769904000" in sql
+
+
+def test_a_failing_probe_leaves_the_query_correct_and_unpruned(app: Flask) -> None:
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=None):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch" not in sql
+ assert "event_time" in sql
+
+
+def test_an_inverted_transform_emits_nothing(app: Flask) -> None:
+ """
+ ``T(lower) <= T(upper)`` is a nearly-free runtime backstop for the
+ monotonicity *declaration*: it catches inverted transforms, and catches
+ ``hour()`` on any range spanning a day boundary. Necessary, not sufficient.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1769904000, 1767225600]):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_incomparable_probe_results_emit_nothing(app: Flask) -> None:
+ """Probe results come back as pandas scalars; not every pair compares."""
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[object(), object()]):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_self_mapping_is_skipped_defensively(app: Flask) -> None:
+ """Save-time validation rejects this, but older rows can carry it."""
+ table = _table(partition_column="event_time")
+
+ with app.app_context():
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ sql = _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ )
+
+ assert sql.count("event_time >=") == 1
+
+
+def test_a_filter_on_an_unmapped_column_mirrors_nothing(app: Flask) -> None:
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, side_effect=AssertionError("probe must not run")):
+ sql = _query(
+ table,
+ filter=[
+ {"col": "country", "op": FilterOperator.EQUALS.value, "val": "US"}
+ ],
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_the_probe_receives_timezone_adjusted_bounds(app: Flask) -> None:
+ """
+ ``get_time_filter`` shifts the bounds by the dataset's timezone before
+ building the clause. Probing the *raw* bounds would produce epoch bounds
+ describing a different instant than the timestamp bounds they mirror --
+ wrong by exactly the offset, silently.
+ """
+ table = _table()
+ _set_timezone(table, "Europe/Berlin")
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1, 2]) as probe:
+ _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 9),
+ to_dttm=datetime(2026, 1, 10),
+ )
+
+ # Berlin is UTC+1 in January, so local midnight is 23:00 the day before.
+ assert probe.call_args.args[-1] == [
+ datetime(2026, 1, 8, 23, 0),
+ datetime(2026, 1, 9, 23, 0),
+ ]
+
+
+def test_the_probe_receives_hour_offset_adjusted_bounds(app: Flask) -> None:
+ """The legacy ``offset`` field shifts bounds too, and must reach the probe."""
+ table = _table()
+ table.offset = 5
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1, 2]) as probe:
+ _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 9, 12),
+ to_dttm=datetime(2026, 1, 10, 12),
+ )
+
+ assert probe.call_args.args[-1] == [
+ datetime(2026, 1, 9, 7),
+ datetime(2026, 1, 10, 7),
+ ]
+
+
+def test_one_probe_round_trip_per_query(app: Flask) -> None:
+ """Mirror requests are collected and resolved together, not one at a time."""
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]) as probe:
+ _query(
+ table,
+ granularity="event_time",
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-01 : 2026-02-01",
+ }
+ ],
+ )
+
+ probe.assert_called_once()
+
+
+def test_mirrored_predicates_reach_the_series_limit_subquery(app: Flask) -> None:
+ """
+ The mirrored predicate is appended to ``where_clause_and``, which the
+ series-limit subquery reuses -- so pruning applies there too.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]):
+ sql = _query(
+ table,
+ columns=["country"],
+ metrics=["hits"],
+ granularity="event_time",
+ is_timeseries=True,
+ from_dttm=datetime(2026, 1, 1),
+ to_dttm=datetime(2026, 2, 1),
+ timeseries_limit=5,
+ timeseries_limit_metric="hits",
+ )
+
+ assert sql.count("dt_epoch >= 1767225600") >= 2
+
+
+def test_extra_cache_keys_include_the_mapping(app: Flask) -> None:
+ """
+ The mapping changes the SQL a cached chart result came from, so it has to
+ participate in the chart-data cache key or a mapping fix leaves stale pruned
+ results behind.
+ """
+ table = _table()
+
+ with app.app_context():
+ keys = table.get_extra_cache_keys({})
+
+ assert any("dt_epoch" in str(key) for key in keys)
+
+
+def test_extra_cache_keys_are_unchanged_without_a_mapping(app: Flask) -> None:
+ """Cache keys must not churn for the entire installed base."""
+ table = _table(partition_column=None)
+
+ with app.app_context():
+ assert table.get_extra_cache_keys({}) == []
+
+
+# ---------------------------------------------------------------------------
+# Time grains
+# ---------------------------------------------------------------------------
+
+
+def test_a_grain_bearing_equality_filter_does_not_mirror(app: Flask) -> None:
+ """
+ Drill-to-detail sends the clicked bucket as `==` plus the chart's grain, and
+ the real predicate compares the *truncated* column. Mirroring the raw bucket
+ start would keep only the rows at the bucket's first instant and silently
+ drop the rest of the bucket.
+
+ Deferred rather than unsafe: `==` on a bucket could become a two-sided
+ range, but that needs a monotonic transform -- `EQUALS` mirrors without one
+ today -- and a parse back into a datetime.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.EQUALS.value,
+ "val": "2026-01-05",
+ "grain": "P1W",
+ }
+ ],
+ )
+
+ assert "dt_epoch" not in sql
+
+
+def test_a_grain_bearing_temporal_range_widens_both_bounds(app: Flask) -> None:
+ """
+ A grained range compares the truncated column, so the raw bounds it carries
+ do not describe the rows it keeps. Widening both ends by one bucket is the
+ smallest range guaranteed to contain all of them -- both ends, because which
+ way a grain rounds is not knowable from its duration.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1, 2]) as probe:
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-08 : 2026-02-01",
+ "grain": "P1D",
+ }
+ ],
+ )
+
+ assert probe.call_args.args[-1] == [
+ datetime(2026, 1, 7),
+ datetime(2026, 2, 2),
+ ]
+ assert "dt_epoch" in sql
+ # The widened upper bound is inclusive: `ts < until + width` only gives
+ # `T(ts) <= T(until + width)` for a transform that is monotonic but not
+ # strictly so.
+ assert "dt_epoch <= 2" in sql
+ assert "dt_epoch >= 1" in sql
+
+
+def test_a_grain_bearing_range_widens_a_month_across_the_boundary(
+ app: Flask,
+) -> None:
+ """A month is not thirty days: the widening has to be calendar arithmetic."""
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1, 2]) as probe:
+ _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-31 : 2026-03-31",
+ "grain": "P1M",
+ }
+ ],
+ )
+
+ assert probe.call_args.args[-1] == [
+ datetime(2025, 12, 31),
+ datetime(2026, 4, 30),
+ ]
+
+
+def test_a_forward_rounding_grain_still_mirrors(app: Flask) -> None:
+ """
+ ``WEEK_ENDING_SATURDAY`` rounds *forward* on Hive and Presto, which is
+ exactly why the widening is symmetric: a floor-only version would drop up to
+ a week of rows off the bottom of the range.
+ """
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1, 2]) as probe:
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-08 : 2026-02-01",
+ "grain": "P1W/1970-01-03T00:00:00Z",
+ }
+ ],
+ )
+
+ assert probe.call_args.args[-1] == [
+ datetime(2026, 1, 1),
+ datetime(2026, 2, 8),
+ ]
+ assert "dt_epoch" in sql
+
+
+def test_an_addon_grain_still_does_not_mirror(app: Flask) -> None:
+ """An operator-supplied grain has no width Superset can claim to know."""
+ table = _table()
+
+ with app.app_context():
+ # An operator declares both halves: the label and the SQL that builds it.
+ app.config["TIME_GRAIN_ADDONS"] = {"PT2S": "2 second"}
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {
+ "sqlite": {"PT2S": "DATETIME({col}, 'start of day')"}
+ }
+ try:
+ with patch(PROBE, return_value=[1, 2]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-01 : 2026-02-01",
+ "grain": "PT2S",
+ }
+ ],
+ )
+ finally:
+ app.config["TIME_GRAIN_ADDONS"] = {}
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {}
+
+ assert "dt_epoch" not in sql
+
+
+def test_a_builtin_grain_an_operator_redefined_does_not_mirror(
+ app: Flask,
+) -> None:
+ """
+ ``TIME_GRAIN_ADDON_EXPRESSIONS`` can replace a built-in grain's SQL for one
+ engine, so ``P1D`` there may have nothing to do with a day. The duration
+ alone no longer bounds the displacement.
+ """
+ table = _table()
+
+ with app.app_context():
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {
+ "sqlite": {"P1D": "DATETIME({col}, 'start of year')"}
+ }
+ try:
+ with patch(PROBE, return_value=[1, 2]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-01 : 2026-02-01",
+ "grain": "P1D",
+ }
+ ],
+ )
+ finally:
+ app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {}
+
+ assert "dt_epoch" not in sql
+
+
+def test_widening_is_applied_after_the_timezone_adjustment(app: Flask) -> None:
+ """
+ The grain truncation happens in the frame the column is *stored* in, so the
+ bucket width has to be added there too. Widening first and converting after
+ lands an hour out across a DST boundary.
+ """
+ table = _table()
+ _set_timezone(table, "Europe/Berlin")
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1, 2]) as probe:
+ _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-09 : 2026-01-10",
+ "grain": "P1D",
+ }
+ ],
+ )
+
+ # Berlin is UTC+1 in January: adjust to 23:00 the day before, then widen.
+ assert probe.call_args.args[-1] == [
+ datetime(2026, 1, 7, 23, 0),
+ datetime(2026, 1, 10, 23, 0),
+ ]
+
+
+def test_an_ungrained_filter_still_mirrors(app: Flask) -> None:
+ """The grain guard must not cost pruning for the ordinary case."""
+ table = _table()
+
+ with app.app_context():
+ with patch(PROBE, return_value=[1767225600, 1769904000]):
+ sql = _query(
+ table,
+ filter=[
+ {
+ "col": "event_time",
+ "op": FilterOperator.TEMPORAL_RANGE.value,
+ "val": "2026-01-01 : 2026-02-01",
+ }
+ ],
+ )
+
+ assert "dt_epoch" in sql
+ # The widened `<=` must not leak into the common path, which mirrors an
+ # exclusive upper bound exactly.
+ assert "dt_epoch < 1769904000" in sql
+ assert "dt_epoch <= " not in sql
+
+
+# ---------------------------------------------------------------------------
+# NULL partition values
+# ---------------------------------------------------------------------------
+
+
+def test_rows_in_a_null_partition_survive_mirroring(app: Flask) -> None:
+ """
+ A mirrored comparison against NULL is NULL, so a row whose partition value
+ is NULL is dropped by the mirror even when the real filter matches it. Hive
+ and Impala park such rows in the default partition, and a transform that
+ returns NULL for an input it cannot convert produces them on any engine.
+ The mirror only has to be *no narrower* than the real filter, so widening it
+ to admit NULL partitions keeps those rows while still pruning.
+ """
+ table = _table(
+ transform="lower(:value)",
+ monotonic=False,
+ mapped_column="country",
+ partition_mapped_column="country",
+ partition_column="region_key",
+ )
+
+ with app.app_context():
+ with patch(PROBE, return_value=["us"]):
+ sql = _query(
+ table,
+ filter=[
+ {"col": "country", "op": FilterOperator.EQUALS.value, "val": "US"}
+ ],
+ )
+
+ assert "region_key = 'us'" in sql
+ assert "region_key IS NULL" in sql
+ assert "OR" in sql
diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py
index 616cb802905..82a520bce2f 100644
--- a/tests/unit_tests/sql/parse_tests.py
+++ b/tests/unit_tests/sql/parse_tests.py
@@ -6966,3 +6966,38 @@ def test_rls_returns_whether_applied(
)
is expected
)
+
+
+@pytest.mark.parametrize(
+ "sql, engine, expected",
+ [
+ # Hive's parser resolves the zero-argument form to CURRENT_TIMESTAMP,
+ # which is what it actually means, so that is the name reported.
+ ("SELECT unix_timestamp()", "hive", {"CURRENT_TIMESTAMP"}),
+ ("SELECT unix_timestamp( )", "hive", {"CURRENT_TIMESTAMP"}),
+ ("SELECT unix_timestamp(ds) - unix_timestamp()", "hive", {"CURRENT_TIMESTAMP"}),
+ # Dialects that do not special-case it report the name as written.
+ ("SELECT unix_timestamp()", "sqlite", {"UNIX_TIMESTAMP"}),
+ ("SELECT unix_timestamp(ds)", "hive", set()),
+ ("SELECT unix_timestamp(ds)", "sqlite", set()),
+ ("SELECT lower(country)", "hive", set()),
+ ("SELECT * FROM some_table", "hive", set()),
+ # A named node whose arguments span all three shapes sqlglot uses: a
+ # scalar (`this`), a list (`expressions`), and unset optional slots
+ # left as `None`. Reading only `this` would count this as niladic.
+ ("SELECT coalesce(a, b)", "hive", set()),
+ # `expressions` is the only argument here, so the list branch is what
+ # decides whether the call looks niladic at all.
+ ("SELECT concat(a, b)", "hive", set()),
+ ],
+)
+def test_get_niladic_functions(sql: str, engine: str, expected: set[str]) -> None:
+ """
+ Check the `get_niladic_functions` method.
+
+ Some functions mean something entirely different with no arguments -- on
+ Hive and Impala `unix_timestamp()` is the current time while
+ `unix_timestamp(x)` is a pure conversion -- so callers that care about
+ determinism need to distinguish the two by arity, not by name.
+ """
+ assert SQLStatement(sql, engine).get_niladic_functions() == expected