sadpandajoe commented on code in PR #44657: URL: https://github.com/apache/superset/pull/44657#discussion_r4242803187
########## tests/unit_tests/migrations/test_tag_pivot_tables_with_restored_aggregation.py: ########## @@ -0,0 +1,283 @@ +# 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. +"""Tests for migration +``141b8ada7731_tag_pivot_tables_with_restored_aggregation``. + +Covers the helper (_has_legacy_aggregate_function), the full upgrade() path +across multiple slices (tagging, idempotency, and that unaffected slices are +left alone -- and that query_context is never touched, see the module +docstring for why), and that downgrade() strips the default "Metric" value. +""" + +from __future__ import annotations + +from importlib import import_module +from unittest.mock import patch + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session + +from superset.tags.models import ObjectType, TagType +from superset.utils import json + +migration = import_module( + "superset.migrations.versions." + "2026-09-25_00-00_141b8ada7731_tag_pivot_tables_with_restored_" +) + +Slice = migration.Slice +Tag = migration.Tag +TaggedObject = migration.TaggedObject +_has_legacy_aggregate_function = migration._has_legacy_aggregate_function +_VIZ_TYPE = migration._VIZ_TYPE +_FIELD = migration._FIELD +_LEGACY_AGGREGATE_FUNCTIONS = migration._LEGACY_AGGREGATE_FUNCTIONS +LEGACY_AGGREGATION_TAG = migration.LEGACY_AGGREGATION_TAG + + [email protected] +def engine() -> Engine: + engine = create_engine("sqlite:///:memory:") + migration.Base.metadata.create_all(engine) + return engine + + +# --------------------------------------------------------------------------- +# _has_legacy_aggregate_function +# --------------------------------------------------------------------------- + + [email protected]("value", sorted(_LEGACY_AGGREGATE_FUNCTIONS)) +def test_has_legacy_aggregate_function_true_for_every_recognized_value( + value: str, +) -> None: + slc = Slice(params=json.dumps({"viz_type": _VIZ_TYPE, _FIELD: value})) + assert _has_legacy_aggregate_function(slc) + + +def test_has_legacy_aggregate_function_false_for_unrecognized_value() -> None: + params = json.dumps({"viz_type": _VIZ_TYPE, _FIELD: "Metric"}) + assert not _has_legacy_aggregate_function(Slice(params=params)) + + +def test_has_legacy_aggregate_function_false_when_field_absent() -> None: + params = json.dumps({"viz_type": _VIZ_TYPE, "colTotals": True}) + assert not _has_legacy_aggregate_function(Slice(params=params)) + + [email protected]("malformed_value", [[], {}, 1, None]) +def test_has_legacy_aggregate_function_false_when_field_not_a_string( + malformed_value: object, +) -> None: + """A malformed non-string aggregateFunction value (e.g. `[]`/`{}` from an + unrelated historical bug) must be skipped, not raise -- `in` on a + frozenset requires a hashable left operand, so an unguarded membership + check would TypeError on a list/dict and abort `superset db upgrade` + partway through paginated_update's batches.""" + params = json.dumps({"viz_type": _VIZ_TYPE, _FIELD: malformed_value}) + assert not _has_legacy_aggregate_function(Slice(params=params)) + + +def test_has_legacy_aggregate_function_false_when_params_empty() -> None: + assert not _has_legacy_aggregate_function(Slice(params=None)) + + [email protected]("malformed_params", ["[]", "null", "1", '"a string"']) +def test_has_legacy_aggregate_function_false_when_params_not_dict( + malformed_params: str, +) -> None: + """A historically malformed non-dict params value must be skipped, not + raise, so one bad row can't abort `superset db upgrade` partway through + paginated_update's batches.""" + assert not _has_legacy_aggregate_function(Slice(params=malformed_params)) + + +def test_has_legacy_aggregate_function_false_on_invalid_json() -> None: + assert not _has_legacy_aggregate_function(Slice(params="not-json")) + + +# --------------------------------------------------------------------------- +# Full upgrade() integration test +# --------------------------------------------------------------------------- + + +def _run_upgrade(engine: Engine) -> None: + # Simulate the Alembic transaction the same way the real migration runs: + # bind the session to an explicit connection so paginated_update's + # internal commits are visible once the outer transaction closes. + with engine.begin() as conn: + upgrade_session = Session(bind=conn) + with ( + patch.object(migration, "op") as mock_op, + patch.object(migration, "db") as mock_db, + ): + mock_op.get_bind.return_value = conn + mock_db.Session.return_value = upgrade_session + migration.upgrade() + + +def test_upgrade_tags_and_invalidates_only_affected_pivot_tables( + engine: Engine, +) -> None: + with Session(engine) as seed: + seed.add_all( + [ + # A legacy aggregateFunction value -- must be tagged. + Slice( Review Comment: The module docstring promises the upgrade leaves `query_context` untouched, but none of the seeded upgrade charts set one, and the assertions only check tagging. If a later edit to `upgrade()` started clearing or rewriting `query_context`, these tests would keep passing while cache warm-up and reports start rejecting those charts. Could you seed a legacy-`aggregateFunction` chart with a non-null `query_context` here and assert it is byte-identical after `_run_upgrade`? -- 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]
