mirror of
https://github.com/apache/superset.git
synced 2026-09-09 00:34:49 +00:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
20fd771e43 | ||
|
|
c7a6548b4a | ||
|
|
3d7615aa1c | ||
|
|
6f5fb95215 |
@@ -0,0 +1,161 @@
|
||||
---
|
||||
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
|
||||
|
||||
In the dataset editor's **Settings** tab, pick a **Partition column**. By default the
|
||||
mapping follows the dataset's default datetime column, so re-pointing that column moves the
|
||||
mapping with it; set an explicit override if you want it pinned to a different column.
|
||||
|
||||
On the mapped column's row, set the **value transform**: a SQL expression containing a
|
||||
`:value` placeholder, which stands for the filter value being mirrored.
|
||||
|
||||
| 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 you check **Transform preserves ordering**.
|
||||
|
||||
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.
|
||||
|
||||
When the box is unchecked, `=` 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 editor's preview panel runs a real query |
|
||||
|
||||
Mappings travel with the dataset in import/export.
|
||||
Vendored
+6
@@ -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": "Mirror filters on a dataset's business column onto its physical partition column, so engines that require an explicit partition predicate can prune."
|
||||
},
|
||||
{
|
||||
"name": "PRESTO_EXPAND_DATA",
|
||||
"default": false,
|
||||
|
||||
@@ -97,9 +97,21 @@ 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;
|
||||
}
|
||||
|
||||
export interface PartitionFilterMapping {
|
||||
partition_column: string;
|
||||
mapped_column: string | null;
|
||||
active: boolean;
|
||||
}
|
||||
|
||||
export interface ControlPanelState {
|
||||
slice: {
|
||||
slice_id: number;
|
||||
|
||||
@@ -128,6 +128,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:
|
||||
@@ -171,6 +173,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),
|
||||
}),
|
||||
|
||||
+37
-2
@@ -93,7 +93,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,
|
||||
@@ -151,6 +155,8 @@ interface Column {
|
||||
certified_by?: string;
|
||||
certification_details?: string;
|
||||
is_certified?: boolean;
|
||||
partition_value_transform?: string | null;
|
||||
partition_transform_is_monotonic?: boolean;
|
||||
}
|
||||
|
||||
interface Database {
|
||||
@@ -190,6 +196,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[];
|
||||
@@ -1216,6 +1224,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);
|
||||
}
|
||||
@@ -1239,7 +1268,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
|
||||
|
||||
@@ -38,6 +38,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 {
|
||||
@@ -222,6 +229,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,
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
@@ -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 = {
|
||||
@@ -78,6 +82,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;
|
||||
|
||||
@@ -46,6 +46,10 @@ from superset.commands.dataset.exceptions import (
|
||||
)
|
||||
from superset.commands.utils import compute_subjects
|
||||
from superset.connectors.sqla.models import SqlaTable, validate_stored_expression
|
||||
from superset.connectors.sqla.partition_mapping import (
|
||||
parse_skeleton,
|
||||
validate_partition_mapping,
|
||||
)
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
from superset.datasets.schemas import FolderSchema
|
||||
from superset.exceptions import (
|
||||
@@ -281,6 +285,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:
|
||||
@@ -389,6 +395,102 @@ 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.
|
||||
"""
|
||||
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}
|
||||
)
|
||||
|
||||
partition_column = self._properties.get(
|
||||
"partition_column", self._model.partition_column
|
||||
)
|
||||
partition_mapped_column = self._properties.get(
|
||||
"partition_mapped_column", self._model.partition_mapped_column
|
||||
)
|
||||
main_dttm_col = self._properties.get("main_dttm_col", self._model.main_dttm_col)
|
||||
if not 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)
|
||||
|
||||
effective_mapped_column = partition_mapped_column or main_dttm_col
|
||||
transform = self._effective_transform(columns, effective_mapped_column)
|
||||
|
||||
for issue in validate_partition_mapping(
|
||||
column_names=column_names,
|
||||
partition_column=partition_column,
|
||||
partition_mapped_column=partition_mapped_column,
|
||||
main_dttm_col=main_dttm_col,
|
||||
transform=transform,
|
||||
engine=database.backend,
|
||||
):
|
||||
if issue.blocking:
|
||||
exceptions.append(
|
||||
ValidationError(str(issue.message), field_name=issue.field)
|
||||
)
|
||||
|
||||
if transform:
|
||||
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 _effective_transform(
|
||||
self,
|
||||
columns: list[dict[str, Any]] | None,
|
||||
mapped_column: str | None,
|
||||
) -> str | None:
|
||||
"""
|
||||
The value transform on the effective mapped column.
|
||||
|
||||
Reads from the payload when the request carries columns, and from the
|
||||
persisted model otherwise -- a PUT that changes only `partition_column`
|
||||
still has to be validated against the transform already stored.
|
||||
"""
|
||||
if not mapped_column:
|
||||
return None
|
||||
if columns is not None:
|
||||
for column in columns:
|
||||
if column.get("column_name") == mapped_column:
|
||||
return column.get("partition_value_transform")
|
||||
return None
|
||||
for existing in self._model.columns:
|
||||
if existing.column_name == mapped_column:
|
||||
return existing.partition_value_transform
|
||||
return None
|
||||
|
||||
def _validate_fetch_values_predicate(
|
||||
self,
|
||||
predicate: str,
|
||||
|
||||
@@ -400,6 +400,17 @@ 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"
|
||||
# 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"
|
||||
|
||||
@@ -719,6 +730,10 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = {
|
||||
# Try to optimize SQL queries — for now only predicate pushdown is supported
|
||||
# @lifecycle: development
|
||||
"OPTIMIZE_SQL": False,
|
||||
# Mirror filters on a dataset's business column onto its physical partition
|
||||
# column, so engines that require an explicit partition predicate can prune.
|
||||
# @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
|
||||
|
||||
@@ -71,6 +71,7 @@ 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 resolve_partition_mapping
|
||||
from superset.connectors.sqla.utils import (
|
||||
get_columns_description,
|
||||
get_physical_table_metadata,
|
||||
@@ -1060,6 +1061,18 @@ class TableColumn(AuditMixinNullable, ImportExportMixin, CertificationMixin, Mod
|
||||
python_date_format = Column(String(255))
|
||||
datetime_format = Column(String(100))
|
||||
extra = Column(Text)
|
||||
# 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>`.
|
||||
partition_value_transform = Column(Text)
|
||||
# Whether the transform preserves ordering. Range operators (and time
|
||||
# ranges) are only mirrored when it does; see the operator matrix in
|
||||
# `superset.connectors.sqla.partition_mapping`. Non-null with a default
|
||||
# rather than a nullable tri-state, matching `normalize_columns`.
|
||||
partition_transform_is_monotonic = Column(
|
||||
Boolean, nullable=False, default=False, server_default=sa.false()
|
||||
)
|
||||
|
||||
table: Mapped["SqlaTable"] = relationship(
|
||||
"SqlaTable",
|
||||
@@ -1082,6 +1095,8 @@ class TableColumn(AuditMixinNullable, ImportExportMixin, CertificationMixin, Mod
|
||||
"python_date_format",
|
||||
"datetime_format",
|
||||
"extra",
|
||||
"partition_value_transform",
|
||||
"partition_transform_is_monotonic",
|
||||
]
|
||||
|
||||
update_from_object_fields = [s for s in export_fields if s not in ("table_id",)]
|
||||
@@ -1308,6 +1323,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)}
|
||||
@@ -1537,6 +1554,13 @@ class SqlaTable(
|
||||
normalize_columns = Column(Boolean, default=False)
|
||||
always_filter_main_dttm = Column(Boolean, default=False)
|
||||
folders = Column(JSON, nullable=True)
|
||||
# Physical column the engine partitions on. Filters on the effective mapped
|
||||
# column are mirrored onto it so the engine can prune partitions.
|
||||
partition_column = Column(String(250))
|
||||
# Explicit override for the column whose filters are mirrored. NULL means
|
||||
# "follow `main_dttm_col`", so re-pointing the default datetime column moves
|
||||
# the mapping with it.
|
||||
partition_mapped_column = Column(String(250))
|
||||
|
||||
baselink = "tablemodelview"
|
||||
|
||||
@@ -1560,6 +1584,8 @@ class SqlaTable(
|
||||
"normalize_columns",
|
||||
"always_filter_main_dttm",
|
||||
"folders",
|
||||
"partition_column",
|
||||
"partition_mapped_column",
|
||||
]
|
||||
update_from_object_fields = [f for f in export_fields if f != "database_id"]
|
||||
export_parent = "database"
|
||||
@@ -1787,8 +1813,44 @@ 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_
|
||||
|
||||
@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 derived from cheap signals only. This property is serialized
|
||||
on every chart and dashboard load, so parsing the transform here would
|
||||
put a per-request cost on a hot path for a value that only changes on
|
||||
save.
|
||||
"""
|
||||
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 (mapped_column.partition_value_transform or "").strip()
|
||||
)
|
||||
return {
|
||||
"partition_column": self.partition_column,
|
||||
"mapped_column": mapped_column_name,
|
||||
"active": active,
|
||||
}
|
||||
|
||||
@property
|
||||
def extra_dict(self) -> dict[str, Any]:
|
||||
try:
|
||||
@@ -2410,6 +2472,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
|
||||
|
||||
@@ -0,0 +1,827 @@
|
||||
# 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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, cast, TYPE_CHECKING
|
||||
|
||||
import sqlalchemy as sa
|
||||
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.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)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
partition_column = getattr(datasource, "partition_column", None)
|
||||
if not partition_column:
|
||||
return None
|
||||
|
||||
columns_by_name = {column.column_name: column for column in datasource.columns}
|
||||
if partition_column not in columns_by_name:
|
||||
# The partition column was dropped by a column sync or at the source.
|
||||
return None
|
||||
|
||||
mapped_column_name = (
|
||||
getattr(datasource, "partition_mapped_column", None) or datasource.main_dttm_col
|
||||
)
|
||||
if not mapped_column_name or mapped_column_name not in columns_by_name:
|
||||
return None
|
||||
|
||||
if mapped_column_name == partition_column:
|
||||
# Self-mapping: the mirrored predicate would duplicate the original.
|
||||
return None
|
||||
|
||||
mapped_column = columns_by_name[mapped_column_name]
|
||||
transform = getattr(mapped_column, "partition_value_transform", None)
|
||||
if not _transform_is_usable(transform, datasource.database.backend):
|
||||
return None
|
||||
|
||||
if _has_active_advanced_data_type(mapped_column):
|
||||
# `translate_filter` builds its own predicate shape from *translated*
|
||||
# values, so the `(operator, value)` pair the operator matrix reasons
|
||||
# about does not exist and mirroring would apply the wrong values.
|
||||
return None
|
||||
|
||||
return PartitionMapping(
|
||||
partition_column=str(partition_column),
|
||||
mapped_column=str(mapped_column_name),
|
||||
value_transform=cast(str, transform),
|
||||
is_monotonic=bool(
|
||||
getattr(mapped_column, "partition_transform_is_monotonic", False)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _transform_is_usable(transform: str | None, engine: str) -> bool:
|
||||
"""
|
||||
Whether the transform is safe to evaluate and mirror through.
|
||||
|
||||
Mirrors the Tier-2 half of `validate_partition_mapping` plus the Jinja
|
||||
block, so a mapping saved before a check existed -- or one whose engine
|
||||
changed underneath it -- is still skipped at query time.
|
||||
"""
|
||||
if not transform or not transform.strip():
|
||||
return False
|
||||
if not contains_value_placeholder(transform):
|
||||
return False
|
||||
if contains_jinja(transform):
|
||||
return False
|
||||
return is_parseable(transform, engine)
|
||||
|
||||
|
||||
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,
|
||||
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( # pylint: disable=too-many-arguments
|
||||
*,
|
||||
column_names: set[str],
|
||||
partition_column: str | None,
|
||||
partition_mapped_column: str | None,
|
||||
main_dttm_col: str | None,
|
||||
transform: 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.
|
||||
"""
|
||||
if not partition_column:
|
||||
return []
|
||||
|
||||
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 []
|
||||
|
||||
|
||||
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
|
||||
for issue in validate_partition_mapping(
|
||||
column_names=column_names,
|
||||
partition_column=str(partition_column),
|
||||
partition_mapped_column=mapped_column,
|
||||
main_dttm_col=datasource.main_dttm_col,
|
||||
transform=value_transform,
|
||||
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"{partition_column} >= {_render_literal(evaluated[0])}"),
|
||||
}
|
||||
|
||||
|
||||
def _render_literal(value: Any) -> str:
|
||||
"""Render a probed value the way it appears in the generated SQL."""
|
||||
if isinstance(value, str):
|
||||
escaped = value.replace("'", "''")
|
||||
return f"'{escaped}'"
|
||||
return str(value)
|
||||
|
||||
|
||||
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])
|
||||
)
|
||||
return predicates
|
||||
|
||||
|
||||
_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
|
||||
@@ -448,6 +448,37 @@ class DatasetDAO(BaseDAO[SqlaTable]):
|
||||
"python_date_format is an invalid date/timestamp format."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def clear_dangling_partition_mapping(
|
||||
model: SqlaTable, surviving_column_names: set[str]
|
||||
) -> None:
|
||||
"""
|
||||
Drop parts of the partition filter 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.
|
||||
|
||||
Called from the backend `override_columns=true` path, which is the
|
||||
authoritative one: the editor clears the mapping client-side too, but an
|
||||
API-driven sync bypasses the editor entirely.
|
||||
"""
|
||||
if (
|
||||
model.partition_column
|
||||
and model.partition_column not in surviving_column_names
|
||||
):
|
||||
model.partition_column = None
|
||||
model.partition_mapped_column = None
|
||||
return
|
||||
|
||||
if (
|
||||
model.partition_mapped_column
|
||||
and model.partition_mapped_column not in surviving_column_names
|
||||
):
|
||||
model.partition_mapped_column = None
|
||||
|
||||
@classmethod
|
||||
def _override_columns(
|
||||
cls, model: SqlaTable, property_columns: list[dict[str, Any]]
|
||||
@@ -518,6 +549,8 @@ class DatasetDAO(BaseDAO[SqlaTable]):
|
||||
}
|
||||
db.session.add(TableColumn(**{**cleaned, "table_id": model.id}))
|
||||
|
||||
cls.clear_dangling_partition_mapping(model, set(incoming_by_name))
|
||||
|
||||
@classmethod
|
||||
def _upsert_columns(
|
||||
cls, model: SqlaTable, property_columns: list[dict[str, Any]]
|
||||
|
||||
@@ -366,6 +366,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()
|
||||
|
||||
+137
-3
@@ -23,13 +23,13 @@ from io import BytesIO
|
||||
from typing import Any, Callable
|
||||
from zipfile import is_zipfile, ZipFile
|
||||
|
||||
from flask import request, Response, send_file
|
||||
from flask import current_app as app, request, Response, send_file
|
||||
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
|
||||
from flask_appbuilder.const import API_RESULT_RES_KEY, API_SELECT_COLUMNS_RIS_KEY
|
||||
from flask_appbuilder.models.sqla.interface import SQLAInterface
|
||||
from flask_babel import ngettext
|
||||
from flask_babel import gettext as __, ngettext
|
||||
from jinja2.exceptions import TemplateError
|
||||
from marshmallow import ValidationError
|
||||
from sqlalchemy.orm.exc import MultipleResultsFound
|
||||
@@ -61,6 +61,7 @@ from superset.commands.importers.exceptions import NoValidFilesFoundError
|
||||
from superset.commands.importers.v1.utils import get_contents_from_bundle
|
||||
from superset.commands.purge import PurgeArchivedCommand, 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
|
||||
@@ -85,15 +86,22 @@ from superset.datasets.schemas import (
|
||||
get_export_ids_schema,
|
||||
GetOrCreateDatasetSchema,
|
||||
openapi_spec_methods_override,
|
||||
PartitionMappingPreviewSchema,
|
||||
)
|
||||
from superset.exceptions import (
|
||||
SupersetSecurityException,
|
||||
SupersetSyntaxErrorException,
|
||||
SupersetTemplateException,
|
||||
)
|
||||
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 parse_boolean_string, sanitize_cookie_token
|
||||
from superset.utils.core import (
|
||||
get_user_id,
|
||||
parse_boolean_string,
|
||||
sanitize_cookie_token,
|
||||
)
|
||||
from superset.versioning.api_helpers import (
|
||||
current_entity_etag_uuid,
|
||||
current_entity_version_info,
|
||||
@@ -128,6 +136,31 @@ _DATASET_PURGE_BINDING = SoftDeleteBinding(
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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}"
|
||||
try:
|
||||
used = cache_manager.cache.get(key) or 0
|
||||
if used >= limit:
|
||||
return False
|
||||
cache_manager.cache.set(key, used + 1, timeout=60)
|
||||
except Exception: # pylint: disable=broad-except # noqa: BLE001
|
||||
# A cache outage must not take the editor down with it.
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
|
||||
datamodel = SQLAInterface(SqlaTable)
|
||||
base_filters = [["id", DatasourceFilter, lambda: []]]
|
||||
@@ -162,6 +195,7 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
|
||||
"get_or_create_dataset",
|
||||
"warm_up_cache",
|
||||
"get_drill_info",
|
||||
"partition_mapping_preview",
|
||||
"list_versions",
|
||||
"get_version",
|
||||
"activity",
|
||||
@@ -219,6 +253,8 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
|
||||
"description",
|
||||
"main_dttm_col",
|
||||
"currency_code_column",
|
||||
"partition_column",
|
||||
"partition_mapped_column",
|
||||
"normalize_columns",
|
||||
"always_filter_main_dttm",
|
||||
"offset",
|
||||
@@ -310,6 +346,8 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
|
||||
"description",
|
||||
"main_dttm_col",
|
||||
"currency_code_column",
|
||||
"partition_column",
|
||||
"partition_mapped_column",
|
||||
"normalize_columns",
|
||||
"always_filter_main_dttm",
|
||||
"offset",
|
||||
@@ -383,6 +421,7 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
|
||||
DatasetRelatedObjectsResponse,
|
||||
DatasetDuplicateSchema,
|
||||
GetOrCreateDatasetSchema,
|
||||
PartitionMappingPreviewSchema,
|
||||
VersionListItemSchema,
|
||||
)
|
||||
|
||||
@@ -1699,6 +1738,101 @@ class DatasetRestApi(SoftDeleteApiMixin, BaseSupersetModelRestApi):
|
||||
current_entity_etag_uuid(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)
|
||||
|
||||
@@ -100,6 +100,17 @@ 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."
|
||||
)
|
||||
},
|
||||
)
|
||||
partition_transform_is_monotonic = fields.Boolean(load_default=False)
|
||||
uuid = fields.UUID(allow_none=True)
|
||||
|
||||
|
||||
@@ -178,6 +189,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)
|
||||
|
||||
@@ -193,6 +206,8 @@ class DatasetPutSchema(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)
|
||||
@@ -285,6 +300,10 @@ class ImportV1ColumnSchema(Schema):
|
||||
description = fields.String(allow_none=True)
|
||||
python_date_format = fields.String(allow_none=True)
|
||||
datetime_format = fields.String(allow_none=True)
|
||||
partition_value_transform = fields.String(allow_none=True)
|
||||
# Bundles predating the field must not claim their transform preserves
|
||||
# ordering, which would silently enable range mirroring on import.
|
||||
partition_transform_is_monotonic = fields.Boolean(load_default=False)
|
||||
uuid = fields.UUID(allow_none=True)
|
||||
|
||||
|
||||
@@ -407,6 +426,8 @@ 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)
|
||||
partition_column = fields.String(allow_none=True)
|
||||
partition_mapped_column = fields.String(allow_none=True)
|
||||
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)
|
||||
@@ -434,6 +455,24 @@ 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."""
|
||||
|
||||
mapped_column = fields.String(
|
||||
required=True,
|
||||
metadata={"description": "Column whose filters would be mirrored"},
|
||||
)
|
||||
value_transform = fields.String(
|
||||
required=True,
|
||||
allow_none=True,
|
||||
metadata={"description": "SQL expression containing a :value placeholder"},
|
||||
)
|
||||
sample_value = fields.String(
|
||||
required=True,
|
||||
metadata={"description": "Value to evaluate the transform at"},
|
||||
)
|
||||
|
||||
|
||||
class DatasetCacheWarmUpRequestSchema(Schema):
|
||||
db_name = fields.String(
|
||||
required=True,
|
||||
|
||||
@@ -160,6 +160,11 @@ _COLUMN_DESCRIPTIONS: dict[str, str] = {
|
||||
"filter_select_enabled": "Whether filter select is enabled",
|
||||
"normalize_columns": "Whether to normalize column names",
|
||||
"always_filter_main_dttm": "Whether to always filter on the main datetime column",
|
||||
"partition_column": "Physical column the engine partitions on",
|
||||
"partition_mapped_column": (
|
||||
"Column whose filters are mirrored onto the partition column; "
|
||||
"defaults to the main datetime column"
|
||||
),
|
||||
"fetch_values_predicate": "SQL predicate for fetching filter values",
|
||||
"default_endpoint": "Default endpoint URL",
|
||||
"offset": "Row offset for queries",
|
||||
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
# 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.
|
||||
"""add partition filter mapping
|
||||
|
||||
Adds the four columns behind the ``PARTITION_FILTER_MAPPING`` feature:
|
||||
|
||||
- ``tables.partition_column`` -- the physical column the engine partitions on
|
||||
- ``tables.partition_mapped_column`` -- explicit override for the column whose
|
||||
filters are mirrored; NULL means "follow ``main_dttm_col``"
|
||||
- ``table_columns.partition_value_transform`` -- the ``:value`` expression
|
||||
- ``table_columns.partition_transform_is_monotonic`` -- gates range operators
|
||||
|
||||
The Continuum shadow tables get the same columns so dataset version history and
|
||||
restore keep working.
|
||||
|
||||
Revision ID: a7f3c2e91d84
|
||||
Revises: 1072de5ed955
|
||||
Create Date: 2026-08-31 22:30:00.000000
|
||||
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from superset.migrations.shared.utils import add_columns, drop_columns
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "a7f3c2e91d84"
|
||||
down_revision = "1072de5ed955"
|
||||
|
||||
|
||||
def upgrade():
|
||||
add_columns(
|
||||
"tables",
|
||||
sa.Column("partition_column", sa.String(250), nullable=True),
|
||||
sa.Column("partition_mapped_column", sa.String(250), nullable=True),
|
||||
)
|
||||
add_columns(
|
||||
"table_columns",
|
||||
sa.Column("partition_value_transform", sa.Text(), nullable=True),
|
||||
# Non-null with a default rather than a nullable tri-state: a nullable
|
||||
# boolean invites `if x:` bugs where `None` and `False` need
|
||||
# distinguishing and don't get it. Matches `normalize_columns`.
|
||||
sa.Column(
|
||||
"partition_transform_is_monotonic",
|
||||
sa.Boolean(),
|
||||
nullable=False,
|
||||
server_default=sa.false(),
|
||||
),
|
||||
)
|
||||
|
||||
# Shadow tables are nullable throughout -- a version row records the state
|
||||
# of the columns that changed, so every column has to tolerate NULL.
|
||||
add_columns(
|
||||
"tables_version",
|
||||
sa.Column("partition_column", sa.String(250), nullable=True),
|
||||
sa.Column("partition_mapped_column", sa.String(250), nullable=True),
|
||||
)
|
||||
add_columns(
|
||||
"table_columns_version",
|
||||
sa.Column("partition_value_transform", sa.Text(), nullable=True),
|
||||
sa.Column("partition_transform_is_monotonic", sa.Boolean(), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
drop_columns(
|
||||
"table_columns_version",
|
||||
"partition_value_transform",
|
||||
"partition_transform_is_monotonic",
|
||||
)
|
||||
drop_columns(
|
||||
"tables_version",
|
||||
"partition_column",
|
||||
"partition_mapped_column",
|
||||
)
|
||||
drop_columns(
|
||||
"table_columns",
|
||||
"partition_value_transform",
|
||||
"partition_transform_is_monotonic",
|
||||
)
|
||||
drop_columns(
|
||||
"tables",
|
||||
"partition_column",
|
||||
"partition_mapped_column",
|
||||
)
|
||||
+188
-26
@@ -88,6 +88,11 @@ 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,
|
||||
PartitionMapping,
|
||||
resolve_partition_mapping,
|
||||
)
|
||||
from superset.constants import (
|
||||
CacheRegion,
|
||||
EMPTY_STRING,
|
||||
@@ -3563,36 +3568,117 @@ 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
|
||||
)
|
||||
)
|
||||
) -> None:
|
||||
"""
|
||||
Record a time range for mirroring onto the partition column.
|
||||
|
||||
# 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).
|
||||
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.
|
||||
"""
|
||||
if mapping is None or column_name != mapping.mapped_column:
|
||||
return
|
||||
if not mapping.mirrors(utils.FilterOperator.TEMPORAL_RANGE):
|
||||
return
|
||||
|
||||
start_dttm, end_dttm = self.adjust_time_bounds(start_dttm, end_dttm)
|
||||
if start_dttm is not None:
|
||||
sink.append((utils.FilterOperator.GREATER_THAN_OR_EQUALS, start_dttm))
|
||||
if end_dttm is not None:
|
||||
sink.append((utils.FilterOperator.LESS_THAN, 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 []
|
||||
|
||||
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(self, mapping, deduped)
|
||||
|
||||
def adjust_time_bounds(
|
||||
self,
|
||||
start_dttm: Optional[sa.DateTime],
|
||||
end_dttm: Optional[sa.DateTime],
|
||||
) -> 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):
|
||||
@@ -3628,6 +3714,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)
|
||||
|
||||
l = [] # noqa: E741
|
||||
if start_dttm:
|
||||
l.append(
|
||||
@@ -4207,6 +4318,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(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)
|
||||
@@ -4243,6 +4361,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
|
||||
@@ -4260,6 +4388,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
|
||||
@@ -4461,6 +4596,21 @@ 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.
|
||||
if col_obj is not None and op != utils.FilterOperator.TEMPORAL_RANGE:
|
||||
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
|
||||
|
||||
@@ -4606,6 +4756,13 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
)
|
||||
if _temporal_filter is not None:
|
||||
target_clause_list.append(_temporal_filter)
|
||||
self._collect_partition_mirror_range(
|
||||
partition_mapping,
|
||||
partition_mirror,
|
||||
col_obj.column_name,
|
||||
_since,
|
||||
_until,
|
||||
)
|
||||
else:
|
||||
raise QueryObjectValidationError(
|
||||
_("Invalid filter operation type: %(op)s", op=op)
|
||||
@@ -4614,6 +4771,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")
|
||||
|
||||
@@ -594,6 +594,44 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
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 check_tables_present(
|
||||
self, tables: set[str], default_schema: str | None = None
|
||||
) -> bool:
|
||||
|
||||
@@ -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
|
||||
@@ -345,6 +348,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).
|
||||
|
||||
@@ -92,6 +92,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.partition_column = 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
|
||||
@@ -137,6 +138,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.partition_column = 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
|
||||
@@ -200,6 +202,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.partition_column = 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
|
||||
@@ -256,6 +259,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.partition_column = 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
|
||||
@@ -305,6 +309,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.partition_column = None
|
||||
|
||||
mock_dataset_dao.find_by_id.return_value = mock_dataset
|
||||
mock_dataset_dao.validate_update_uniqueness.return_value = True
|
||||
@@ -454,6 +459,7 @@ def test_update_dataset_rejects_malicious_expression(
|
||||
mock_dataset.database = mock_database
|
||||
mock_dataset.catalog = "catalog"
|
||||
mock_dataset.schema = None
|
||||
mock_dataset.partition_column = 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
|
||||
@@ -499,6 +505,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.partition_column = 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
|
||||
@@ -540,6 +547,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.partition_column = 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
|
||||
@@ -1301,6 +1309,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.partition_column = 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
|
||||
|
||||
@@ -0,0 +1,720 @@
|
||||
# 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 typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from superset.connectors.sqla.models import SqlaTable, TableColumn
|
||||
from superset.connectors.sqla.partition_mapping import (
|
||||
contains_jinja,
|
||||
contains_value_placeholder,
|
||||
evaluate_transform,
|
||||
find_non_deterministic_functions,
|
||||
MappingValidationIssue,
|
||||
MIRRORABLE_ALWAYS,
|
||||
MIRRORABLE_IF_MONOTONIC,
|
||||
mirrorable_operators,
|
||||
resolve_partition_mapping,
|
||||
validate_partition_mapping,
|
||||
)
|
||||
from superset.models.core import Database
|
||||
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",
|
||||
partition_value_transform=None,
|
||||
),
|
||||
TableColumn(column_name="country", type="VARCHAR"),
|
||||
TableColumn(column_name="region_key", type="VARCHAR"),
|
||||
],
|
||||
}
|
||||
defaults.update(kwargs)
|
||||
return SqlaTable(**defaults)
|
||||
|
||||
|
||||
def _mapped_table(
|
||||
*,
|
||||
partition_column: str = "dt_epoch",
|
||||
mapped_column: str = "event_time",
|
||||
transform: str | None = "unix_timestamp(:value)",
|
||||
monotonic: bool = True,
|
||||
partition_mapped_column: str | None = None,
|
||||
main_dttm_col: str | None = "event_time",
|
||||
) -> SqlaTable:
|
||||
table = _table(main_dttm_col=main_dttm_col)
|
||||
table.partition_column = partition_column
|
||||
table.partition_mapped_column = partition_mapped_column
|
||||
for column in table.columns:
|
||||
if column.column_name == mapped_column:
|
||||
column.partition_value_transform = transform
|
||||
column.partition_transform_is_monotonic = monotonic
|
||||
return table
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §2 — operator safety matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_equality_and_in_are_always_mirrorable() -> None:
|
||||
"""``=`` and ``IN`` are safe for any function ``T``."""
|
||||
assert MIRRORABLE_ALWAYS == {FilterOperator.EQUALS, FilterOperator.IN}
|
||||
|
||||
|
||||
def test_range_operators_require_a_monotonic_transform() -> None:
|
||||
assert MIRRORABLE_IF_MONOTONIC == {
|
||||
FilterOperator.GREATER_THAN,
|
||||
FilterOperator.GREATER_THAN_OR_EQUALS,
|
||||
FilterOperator.LESS_THAN,
|
||||
FilterOperator.LESS_THAN_OR_EQUALS,
|
||||
FilterOperator.TEMPORAL_RANGE,
|
||||
}
|
||||
|
||||
|
||||
def test_mirrorable_operators_excludes_ranges_when_not_monotonic() -> None:
|
||||
assert mirrorable_operators(is_monotonic=False) == MIRRORABLE_ALWAYS
|
||||
|
||||
|
||||
def test_mirrorable_operators_includes_ranges_when_monotonic() -> None:
|
||||
assert mirrorable_operators(is_monotonic=True) == (
|
||||
MIRRORABLE_ALWAYS | MIRRORABLE_IF_MONOTONIC
|
||||
)
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §4.1 — resolving the effective mapping, and the defensive bail-outs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resolve_returns_the_mapping_when_everything_lines_up(app: Flask) -> None:
|
||||
with app.app_context():
|
||||
mapping = resolve_partition_mapping(_mapped_table())
|
||||
|
||||
assert mapping is not None
|
||||
assert mapping.partition_column == "dt_epoch"
|
||||
assert mapping.mapped_column == "event_time"
|
||||
assert mapping.value_transform == "unix_timestamp(:value)"
|
||||
assert mapping.is_monotonic is True
|
||||
|
||||
|
||||
def test_effective_mapped_column_follows_main_dttm_col(app: Flask) -> None:
|
||||
"""No explicit override: the mapping follows the default datetime column."""
|
||||
with app.app_context():
|
||||
mapping = resolve_partition_mapping(_mapped_table())
|
||||
|
||||
assert mapping is not None
|
||||
assert mapping.mapped_column == "event_time"
|
||||
|
||||
|
||||
def test_explicit_override_wins_over_main_dttm_col(app: Flask) -> None:
|
||||
table = _mapped_table(
|
||||
mapped_column="country",
|
||||
transform="lower(:value)",
|
||||
monotonic=False,
|
||||
partition_mapped_column="country",
|
||||
)
|
||||
table.partition_column = "region_key"
|
||||
|
||||
with app.app_context():
|
||||
mapping = resolve_partition_mapping(table)
|
||||
|
||||
assert mapping is not None
|
||||
assert mapping.mapped_column == "country"
|
||||
assert mapping.partition_column == "region_key"
|
||||
assert mapping.is_monotonic is False
|
||||
|
||||
|
||||
def test_resolve_returns_none_when_the_feature_flag_is_off(app: Flask) -> None:
|
||||
table = _mapped_table()
|
||||
with app.app_context():
|
||||
with patch(
|
||||
"superset.connectors.sqla.partition_mapping.feature_flag_manager."
|
||||
"is_feature_enabled",
|
||||
return_value=False,
|
||||
):
|
||||
assert resolve_partition_mapping(table) is None
|
||||
|
||||
|
||||
def test_resolve_returns_none_without_a_partition_column(app: Flask) -> None:
|
||||
table = _mapped_table()
|
||||
table.partition_column = None
|
||||
with app.app_context():
|
||||
assert resolve_partition_mapping(table) is None
|
||||
|
||||
|
||||
def test_resolve_returns_none_when_partition_column_no_longer_exists(
|
||||
app: Flask,
|
||||
) -> None:
|
||||
"""A column sync can drop the physical partition column out from under us."""
|
||||
table = _mapped_table()
|
||||
table.partition_column = "dt_epoch_gone"
|
||||
with app.app_context():
|
||||
assert resolve_partition_mapping(table) is None
|
||||
|
||||
|
||||
def test_resolve_returns_none_when_the_mapped_column_no_longer_exists(
|
||||
app: Flask,
|
||||
) -> None:
|
||||
table = _mapped_table(main_dttm_col="vanished")
|
||||
with app.app_context():
|
||||
assert resolve_partition_mapping(table) is None
|
||||
|
||||
|
||||
def test_resolve_returns_none_on_self_mapping(app: Flask) -> None:
|
||||
"""
|
||||
``partition_column == effective mapped column`` mirrors a column onto itself.
|
||||
Save-time validation rejects it, but rows predating that validation exist.
|
||||
"""
|
||||
table = _mapped_table(
|
||||
partition_column="event_time",
|
||||
mapped_column="event_time",
|
||||
)
|
||||
with app.app_context():
|
||||
assert resolve_partition_mapping(table) is None
|
||||
|
||||
|
||||
@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 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]
|
||||
database.get_df.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 = database.get_df.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 = database.get_df.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 = database.get_df.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
|
||||
|
||||
database.get_df.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §5 — save-time validation, in two tiers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _issues(**kwargs: Any) -> list[MappingValidationIssue]:
|
||||
defaults: dict[str, Any] = {
|
||||
"column_names": {"event_time", "dt_epoch", "country", "region_key"},
|
||||
"partition_column": "dt_epoch",
|
||||
"partition_mapped_column": None,
|
||||
"main_dttm_col": "event_time",
|
||||
"transform": "unix_timestamp(:value)",
|
||||
"engine": "hive",
|
||||
}
|
||||
defaults.update(kwargs)
|
||||
return validate_partition_mapping(**defaults)
|
||||
|
||||
|
||||
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]
|
||||
database.get_df.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 database.get_df.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 database.get_df.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 database.get_df.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]
|
||||
@@ -216,6 +216,8 @@ folders:
|
||||
- uuid: 00000000-0000-0000-0000-000000000005
|
||||
type: column
|
||||
name: profit
|
||||
partition_column: null
|
||||
partition_mapped_column: null
|
||||
uuid: {payload["uuid"]}
|
||||
metrics:
|
||||
- metric_name: cnt
|
||||
@@ -244,6 +246,8 @@ columns:
|
||||
datetime_format: null
|
||||
extra:
|
||||
certified_by: User
|
||||
partition_value_transform: null
|
||||
partition_transform_is_monotonic: false
|
||||
uuid: 00000000-0000-0000-0000-000000000005
|
||||
- column_name: ds
|
||||
verbose_name: null
|
||||
@@ -258,6 +262,8 @@ columns:
|
||||
python_date_format: null
|
||||
datetime_format: null
|
||||
extra: null
|
||||
partition_value_transform: null
|
||||
partition_transform_is_monotonic: false
|
||||
uuid: 00000000-0000-0000-0000-000000000006
|
||||
- column_name: user_id
|
||||
verbose_name: null
|
||||
@@ -272,6 +278,8 @@ columns:
|
||||
python_date_format: null
|
||||
datetime_format: null
|
||||
extra: null
|
||||
partition_value_transform: null
|
||||
partition_transform_is_monotonic: false
|
||||
uuid: 00000000-0000-0000-0000-000000000007
|
||||
- column_name: revenue
|
||||
verbose_name: null
|
||||
@@ -286,6 +294,8 @@ columns:
|
||||
python_date_format: null
|
||||
datetime_format: null
|
||||
extra: null
|
||||
partition_value_transform: null
|
||||
partition_transform_is_monotonic: false
|
||||
uuid: 00000000-0000-0000-0000-000000000008
|
||||
- column_name: expenses
|
||||
verbose_name: null
|
||||
@@ -300,6 +310,8 @@ columns:
|
||||
python_date_format: null
|
||||
datetime_format: null
|
||||
extra: null
|
||||
partition_value_transform: null
|
||||
partition_transform_is_monotonic: false
|
||||
uuid: 00000000-0000-0000-0000-000000000009
|
||||
version: 1.0.0
|
||||
database_uuid: {database.uuid}
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
# 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
|
||||
|
||||
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"),
|
||||
],
|
||||
)
|
||||
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:]
|
||||
@@ -0,0 +1,245 @@
|
||||
# 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 export fields, 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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from superset.connectors.sqla.models import SqlaTable, TableColumn
|
||||
from superset.datasets.schemas import (
|
||||
DatasetColumnsPutSchema,
|
||||
DatasetPutSchema,
|
||||
ImportV1ColumnSchema,
|
||||
ImportV1DatasetSchema,
|
||||
)
|
||||
from superset.models.core import Database
|
||||
|
||||
DATASET_FIELDS = ["partition_column", "partition_mapped_column"]
|
||||
COLUMN_FIELDS = ["partition_value_transform", "partition_transform_is_monotonic"]
|
||||
|
||||
|
||||
@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() -> SqlaTable:
|
||||
database = Database(database_name="test_db", sqlalchemy_uri="sqlite://")
|
||||
column = TableColumn(column_name="event_time", is_dttm=True, type="TIMESTAMP")
|
||||
column.partition_value_transform = "unix_timestamp(:value)"
|
||||
column.partition_transform_is_monotonic = True
|
||||
table = SqlaTable(
|
||||
table_name="web_events",
|
||||
database=database,
|
||||
main_dttm_col="event_time",
|
||||
columns=[column, TableColumn(column_name="dt_epoch", type="BIGINT")],
|
||||
)
|
||||
table.partition_column = "dt_epoch"
|
||||
table.partition_mapped_column = None
|
||||
return table
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", DATASET_FIELDS)
|
||||
def test_dataset_fields_are_exported(field: str) -> None:
|
||||
"""
|
||||
Being in ``export_fields`` is what makes the mapping travel in dataset YAML,
|
||||
and what makes ``update_from_object`` write it back on import.
|
||||
"""
|
||||
assert field in SqlaTable.export_fields
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", COLUMN_FIELDS)
|
||||
def test_column_fields_are_exported(field: str) -> None:
|
||||
assert field in TableColumn.export_fields
|
||||
|
||||
|
||||
@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_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()
|
||||
table.columns[0].partition_value_transform = None
|
||||
|
||||
with app.app_context():
|
||||
assert table.data["partition_filter_mapping"]["active"] is False
|
||||
|
||||
|
||||
def test_there_is_no_mapping_summary_without_a_partition_column(
|
||||
app: Flask,
|
||||
) -> None:
|
||||
table = _table()
|
||||
table.partition_column = None
|
||||
|
||||
with app.app_context():
|
||||
assert table.data["partition_filter_mapping"] is None
|
||||
|
||||
|
||||
@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_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
|
||||
|
||||
|
||||
def test_import_schema_round_trips_the_mapping() -> None:
|
||||
loaded = ImportV1DatasetSchema().load(
|
||||
{
|
||||
"table_name": "web_events",
|
||||
"uuid": "00000000-0000-0000-0000-000000000001",
|
||||
"database_uuid": "00000000-0000-0000-0000-000000000002",
|
||||
"version": "1.0.0",
|
||||
"partition_column": "dt_epoch",
|
||||
"partition_mapped_column": "event_time",
|
||||
}
|
||||
)
|
||||
assert loaded["partition_column"] == "dt_epoch"
|
||||
assert loaded["partition_mapped_column"] == "event_time"
|
||||
|
||||
|
||||
def test_import_column_schema_round_trips_the_transform() -> None:
|
||||
loaded = ImportV1ColumnSchema().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_import_column_schema_defaults_the_monotonic_flag_to_false() -> None:
|
||||
"""
|
||||
The flag gates range mirroring. A dataset imported from a bundle that
|
||||
predates the field must not silently claim its transform preserves ordering.
|
||||
"""
|
||||
loaded = ImportV1ColumnSchema().load({"column_name": "event_time"})
|
||||
assert loaded["partition_transform_is_monotonic"] is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", DATASET_FIELDS)
|
||||
def test_the_api_exposes_and_accepts_the_dataset_fields(field: str) -> None:
|
||||
from superset.datasets.api import DatasetRestApi
|
||||
|
||||
assert field in DatasetRestApi.show_select_columns
|
||||
assert field in DatasetRestApi.edit_columns
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §10 — 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.
|
||||
"""
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
table = _table()
|
||||
DatasetDAO.clear_dangling_partition_mapping(table, {"event_time"})
|
||||
|
||||
assert table.partition_column is None
|
||||
assert table.partition_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.
|
||||
"""
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
table = _table()
|
||||
table.partition_mapped_column = "event_time"
|
||||
|
||||
DatasetDAO.clear_dangling_partition_mapping(table, {"dt_epoch"})
|
||||
|
||||
assert table.partition_column == "dt_epoch"
|
||||
assert table.partition_mapped_column is None
|
||||
|
||||
|
||||
def test_a_sync_that_keeps_both_columns_leaves_the_mapping_alone() -> None:
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
table = _table()
|
||||
DatasetDAO.clear_dangling_partition_mapping(table, {"event_time", "dt_epoch"})
|
||||
|
||||
assert table.partition_column == "dt_epoch"
|
||||
@@ -0,0 +1,601 @@
|
||||
# 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.models.core import Database
|
||||
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(*)")],
|
||||
)
|
||||
table.partition_column = partition_column
|
||||
table.partition_mapped_column = partition_mapped_column
|
||||
for column in columns:
|
||||
if column.column_name == mapped_column:
|
||||
column.partition_value_transform = transform
|
||||
column.partition_transform_is_monotonic = monotonic
|
||||
return table
|
||||
|
||||
|
||||
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()
|
||||
table.extra = '{"timezone": "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({}) == []
|
||||
@@ -5471,3 +5471,31 @@ def test_has_aggregate(expression: str, expected: bool) -> None:
|
||||
function sqlglot can't model.
|
||||
"""
|
||||
assert has_aggregate(expression) 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()),
|
||||
],
|
||||
)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user