sadpandajoe commented on code in PR #43870:
URL: https://github.com/apache/superset/pull/43870#discussion_r4337086623
##########
superset/datasets/schemas.py:
##########
@@ -128,6 +133,23 @@ class DatasetColumnsPutSchema(ResolveIsCertifiedMixin,
Schema):
datetime_format = fields.String(
allow_none=True, validate=[Length(1, 100), validate_python_date_format]
)
+ partition_value_transform = fields.String(
+ allow_none=True,
+ validate=Length(1, MAX_TRANSFORM_LENGTH),
+ metadata={
+ "description": (
+ "SQL expression containing a :value placeholder. Filters on "
+ "this column are mirrored onto the dataset's partition column "
+ "with the value passed through this transform."
+ )
+ },
+ )
+ # Deliberately no `load_default`: `DatasetDAO.update_columns` applies the
+ # loaded payload field by field onto the stored column, so a default here
+ # would let a partial column payload clear a monotonic flag the request
+ # never mentioned -- and silently stop mirroring range filters. Absent
+ # means "unchanged"; new columns fall back to the model's own default.
+ partition_transform_is_monotonic = fields.Boolean(allow_none=True)
Review Comment:
A PUT that replaces `partition_value_transform` but omits this field keeps
the old stored `True`, so the new transform is silently treated as monotonic.
For example, switching `p1 = n` / `:value` to `p2 = abs(n)` / `abs(:value)` and
then filtering `n >= -1` adds `p2 >= 1`, which drops the `n = 0` rows with no
error. The import path resets the declaration when it replaces a transform;
should the REST update reset it to False (or require it) whenever the transform
changes?
##########
superset/views/datasource/utils.py:
##########
@@ -217,12 +232,85 @@ def get_samples( # pylint: disable=too-many-arguments
sample_data["page"] = page
sample_data["per_page"] = per_page
- sample_data["total_count"] = count_star_data["data"][0]["COUNT(*)"]
+ if skip_count:
+ _set_lower_bound_total(
+ sample_data, row_offset, row_limit, probed=not via_cursor
+ )
+ else:
+ sample_data["total_count"] = count_star_data["data"][0]["COUNT(*)"]
+ sample_data["total_count_is_lower_bound"] = False
return sample_data
except (IndexError, KeyError) as exc:
raise DatasetSamplesFailedError from exc
+def _count_star_query_context(
+ datasource: Datasource,
+ payload: SamplesPayloadSchema | None,
+ form_data: dict[str, Any] | None,
+ force: bool,
+) -> QueryContext:
+ count_star_metric = {
+ "metrics": [
+ {
+ "expressionType": "SQL",
+ "sqlExpression": "COUNT(*)",
+ "label": "COUNT(*)",
+ }
+ ]
+ }
+ return QueryContextFactory().create(
+ datasource={
+ "type": datasource.type,
+ "id": datasource.id,
+ },
+ queries=[{**payload, **count_star_metric} if payload else
count_star_metric],
+ form_data=form_data,
+ result_type=ChartDataResultType.FULL,
+ force=force,
+ )
+
+
+def _get_count_star_data(
+ count_star_instance: QueryContext | None,
+) -> dict[str, Any]:
+ """Run the count query, or return an empty payload when it was skipped."""
+ if count_star_instance is None:
+ return {}
+
+ count_star_instance.raise_for_access()
+ count_star_data = count_star_instance.get_payload()["queries"][0]
+ if count_star_data.get("status") == QueryStatus.FAILED:
+ raise DatasetSamplesFailedError(count_star_data.get("error"))
+ return count_star_data
+
+
+def _set_lower_bound_total(
+ sample_data: dict[str, Any], row_offset: int, row_limit: int, probed: bool
+) -> None:
+ """
+ Report the total as "the rows up to and including this page".
+
+ The rows query asked for one probe row past the page. If it came back there
+ is at least one more page, so the probe is trimmed and the total flagged as
+ a lower bound; if it did not, this page is the last one and the total is
+ exact. The cursor path pages with an exact page size and has no probe row
+ (``probed=False``), so a full page is taken to mean there may be more.
+ """
+ data = sample_data.get("data") or []
+ page_size = row_limit or app.config["SAMPLES_ROW_LIMIT"]
+ if probed:
+ has_more = len(data) > page_size
+ if has_more:
+ data = data[:page_size]
+ sample_data["data"] = data
+ sample_data["rowcount"] = len(data)
+ else:
+ has_more = len(data) >= page_size
+ sample_data["total_count"] = row_offset + len(data)
Review Comment:
When the requested page is past the end, `data` is empty, so this reports
`total_count = row_offset` with `total_count_is_lower_bound = False`. For a
51-row dataset, requesting page 3 at 50 per page returns an exact total of 100.
An empty page doesn't establish a count, and this can also happen when rows
shrink mid-pagination. Should an empty page with a non-zero offset be reported
as a lower bound instead?
##########
superset-frontend/src/components/Datasource/components/DatasourceEditor/DatasourceEditor.tsx:
##########
@@ -1257,6 +1386,178 @@ function DatasourceEditor({
[],
);
+ // Which column's row expand to open, for the "map a different column" links.
+ // Consumed by the Columns tab, which scrolls the row into view and expands
it.
+ //
+ // Carries a nonce because the request is an event, not a state: clicking the
+ // same link twice -- after collapsing the row by hand in between -- asks for
+ // the same column name, and a bare string would make the second
+ // `setColumnToReveal` a no-op. React would bail out of the render, nothing
+ // downstream would see a change, and the row would stay shut.
+ const [columnToReveal, setColumnToReveal] = useState<{
+ name: string;
+ nonce: number;
+ } | null>(null);
+
+ // What the mapped column's own preview said, relayed up from its row expand
+ // so the dataset-level banner cannot claim a mirror that panel has already
+ // refused. Reported rather than lifted: the transform the preview judges is
+ // the debounced local one that panel owns. `null` means no verdict yet.
+ const [previewMirrors, setPreviewMirrors] = useState<boolean | null>(null);
Review Comment:
`previewMirrors` is only ever written by the mapped row's effect, and
nothing resets it when the mapping moves. Preview a transform that fails
against the current partition target (say `upper(:value)` on a numeric
partition column), collapse the row, then pick a compatible partition column:
the row is unmounted so no new verdict is reported, and the banner keeps saying
filters won't mirror. Should `handlePartitionColumnChange` (and the other
mapping moves) reset this to `null`?
##########
tests/unit_tests/datasets/partition_mapping_preview_test.py:
##########
@@ -0,0 +1,1342 @@
+# 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 MagicMock, 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"
+
+
[email protected](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"]
+
+
[email protected](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
+
+
[email protected](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",
+ )
+
+
[email protected]
+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"),
+ # A second, text-typed partition key. A transform that answers with
+ # text cannot be previewed against `dt_epoch` -- the predicate it
+ # would emit is `dt_epoch = 'us'`, which the database refuses, and
+ # `probed_value_type_error` says so before the preview can call it
+ # valid. Tests whose transform answers with text point here.
+ TableColumn(column_name="region_key", type="VARCHAR"),
+ # And a temporal one, for a transform that answers with a date or a
+ # timestamp. `region_key` will not do: the type gate holds a text
+ # key to text results, because not every engine accepts anything
+ # else in a text comparison and one of them compares it as never
+ # equal rather than refusing it.
+ TableColumn(column_name="part_ts", type="TIMESTAMP"),
+ # A text *mapped* column. A temporal one will not stand in for it:
+ # a sample that is not an ISO instant declines on a column the
+ # engine compares more coarsely than the value carries, which is
+ # the right answer there and noise in a test about anything else.
+ TableColumn(column_name="country", type="VARCHAR"),
+ ],
+ )
+ 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_values": ["2026-01-15 00:00:00"],
+ },
+ )
+
+ assert response.status_code == 200
+ assert response.json["result"] == {
+ "valid": True,
+ "sample_input": "event_time == '2026-01-15 00:00:00'",
+ "emitted_predicate": "dt_epoch = 1768435200 OR dt_epoch IS NULL",
+ }
+
+
+def test_preview_mirrors_a_range_when_the_transform_is_monotonic(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ The case the whole feature exists for: an Explore time range bound.
+ """
+ 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_values": ["2026-01-15 00:00:00"],
+ "operator": ">=",
+ "is_monotonic": True,
+ },
+ )
+
+ assert response.json["result"] == {
+ "valid": True,
+ "sample_input": "event_time >= '2026-01-15 00:00:00'",
+ "emitted_predicate": "dt_epoch >= 1768435200 OR dt_epoch IS NULL",
+ }
+
+
+def test_preview_shows_a_strict_bound_mirroring_non_strictly(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ The preview and the query path agree only because both go through
+ `build_mirrored_predicates`. A preview rendering `<` while the chart
emitted
+ `<=` would be the editor claiming something the SQL does not do.
+
+ `sample_input` describes the owner's filter and `emitted_predicate` the
+ mirror, so the two carrying different operators is the behaviour under test
+ rather than a defect.
+ """
+ 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_values": ["2026-01-15 00:00:00"],
+ "operator": "<",
+ "is_monotonic": True,
+ },
+ )
+
+ assert response.json["result"] == {
+ "valid": True,
+ "sample_input": "event_time < '2026-01-15 00:00:00'",
+ "emitted_predicate": "dt_epoch <= 1768435200 OR dt_epoch IS NULL",
+ }
+
+
+def test_preview_refuses_a_range_when_the_transform_is_not_monotonic(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ Refusing here rather than showing a predicate is the point: the query path
+ would not mirror this filter either, and a preview that implied otherwise
+ would be the one thing worse than no preview.
+ """
+ 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_values": ["2026-01-15 00:00:00"],
+ "operator": ">=",
+ "is_monotonic": False,
+ },
+ )
+
+ result = response.json["result"]
+ assert result["valid"] is False
+ assert "preserves ordering" in result["error"]
+
+
+def test_preview_names_the_comparison_when_text_equality_cannot_mirror(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ On an engine whose string comparison is not byte-exact, `=` and `IN` on a
+ text column are withdrawn whatever the owner declares about ordering. The
+ reason has to say so, or it sends them to a checkbox that changes nothing.
+ """
+ with patch.object(
+ dataset.database.db_engine_spec, "binary_string_comparison", False
+ ):
+ 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": "country",
+ "partition_column": "region_key",
+ "value_transform": "lower(:value)",
+ "sample_values": ["US"],
+ "operator": "IN",
+ "is_monotonic": True,
+ },
+ )
+
+ result = response.json["result"]
+ assert result["valid"] is False
+ assert result["reason"] == "operator"
+ assert "preserves ordering" not in result["error"]
+ assert "case or trailing spaces" in result["error"]
+
+
+def test_preview_mirrors_in_element_wise(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ Wireframe 1h: a non-temporal mapping previews as an `IN`, not as a `>=`
+ against the first value.
+ """
+ with patch(PROBE, return_value=["us", "ca"]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "country",
+ "partition_column": "region_key",
+ "value_transform": "lower(:value)",
+ "sample_values": ["US", "CA"],
+ "operator": "IN",
+ },
+ )
+
+ assert response.json["result"] == {
+ "valid": True,
+ "sample_input": "country IN ('US', 'CA')",
+ "emitted_predicate": "region_key IN ('us', 'ca') OR region_key IS
NULL",
+ }
+
+
+def test_preview_rejects_an_operator_that_can_never_mirror(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_values": ["2026-01-15 00:00:00"],
+ "operator": "NOT IN",
+ },
+ )
+
+ assert response.status_code == 400
+
+
+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_values": ["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_values": ["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_values": ["2026-01-15 00:00:00"],
+ },
+ )
+
+ assert response.json["result"]["valid"] is False
+
+
+def test_preview_refuses_a_transform_whose_result_the_key_cannot_hold(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ The one verdict a preview must never give: "Valid" for a mapping that makes
+ every chart fail. `cast(:value as text) || 'x'` evaluates perfectly, so the
+ probe succeeds -- and the preview went on to render the predicate the chart
+ would carry, `dt_epoch >= '2026-01-15 00:00:00x'` against a BIGINT, which
+ Postgres answers with a 400.
+
+ Its own `reason`, because "the database would not evaluate this" and "it
+ evaluated and answered with the wrong kind of thing" are different problems
+ and only the owner can say which transform they meant.
+ """
+ with patch(PROBE, return_value=["2026-01-15 00:00:00x"]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "value_transform": "cast(:value as text) || 'x'",
+ "sample_values": ["2026-01-15 00:00:00"],
+ },
+ )
+
+ assert response.status_code == 200
+ result = response.json["result"]
+ assert result["valid"] is False
+ assert result["reason"] == "type"
+ assert "dt_epoch" in result["error"]
+ assert result["sample_input"] == "event_time == '2026-01-15 00:00:00'"
+
+
+def test_preview_refuses_a_number_against_a_text_partition_key(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """
+ The mirror image, and the one that used to be exempt. A text partition key
+ was held to no type at all, on the belief that every engine takes a number
+ in a text comparison -- but Postgres refuses `character varying = integer`,
+ Trino `varchar = bigint` and BigQuery `STRING = INT64`, so the preview said
+ "Valid" for a mapping whose every chart 400s.
+
+ SQLite is the reason this is a decline rather than a coercion: it raises
+ nothing and compares a number against text as never equal, so the mirror
+ would silently drop every row the filter keeps.
+ """
+ with patch(PROBE, return_value=[1768435200]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "partition_column": "region_key",
+ "value_transform": "unix_timestamp(:value)",
+ "sample_values": ["2026-01-15 00:00:00"],
+ },
+ )
+
+ assert response.status_code == 200
+ result = response.json["result"]
+ assert result["valid"] is False
+ assert result["reason"] == "type"
+ assert "region_key" in result["error"]
+ assert "string" in result["error"]
+
+
+def test_preview_mirrors_text_against_a_text_partition_key(
+ client: Any, full_api_access: None, dataset: Any
+) -> None:
+ """The other half: a text key still takes a transform that answers text."""
+ with patch(PROBE, return_value=["20260115"]):
+ response = client.post(
+ f"/api/v1/dataset/{dataset.id}/partition_mapping/preview/",
+ json={
+ "mapped_column": "event_time",
+ "partition_column": "region_key",
+ "value_transform": "to_char(:value, 'YYYYMMDD')",
+ "sample_values": ["2026-01-15 00:00:00"],
+ },
+ )
+
+ assert response.status_code == 200
+ result = response.json["result"]
+ assert result["valid"] is True
+ assert "'20260115'" in result["emitted_predicate"]
+
+
+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_values": ["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_values": ["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_values": ["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(
Review Comment:
This test is named per user and dataset, but it uses a single caller and one
dataset, so removing either part of the limiter key keeps it green. If the key
dropped the user, one owner exhausting the budget would 429 every other owner's
editor on that dataset. Could it exhaust user A / dataset X, then assert user B
on X and user A on another dataset still get 200 while A on X gets 429?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
