mikebridge commented on code in PR #45029:
URL: https://github.com/apache/superset/pull/45029#discussion_r4212222051
##########
tests/unit_tests/tasks/test_deletion_retention.py:
##########
@@ -573,4 +576,214 @@ def test_standalone_window_bounds_reach_safe_purge_cutoff(
assert result == {"skipped": 1}
purge.assert_not_called()
else:
- purge.assert_called_once_with(Slice, now - timedelta(days=expected),
False)
+ purge.assert_called_once_with(
+ Slice, now - timedelta(days=expected), False, max_per_run=None
+ )
+
+
+def test_purge_cap_counts_committed_roots_across_batches(app_context: None) ->
None:
Review Comment:
Added a cap-2 regression test that drives `_purge_model` with one blocked
root, one root that raises, then two successful roots. It asserts the counts
are `(2, 0, 1, 1)` and that all four IDs were attempted, so blocked and failed
attempts do not consume the purge budget. The test is in
https://github.com/apache/superset/commit/a54556a2965e70577d12fb97af7b0994f7a3b27b.
##########
tests/unit_tests/tasks/test_version_history_retention.py:
##########
@@ -407,9 +410,353 @@ def test_immediate_cutoff_and_invalid_skip(stats:
MagicMock, days: int) -> None:
days
)
if days == -1:
- run_pass.assert_called_once_with(now, [], 0)
+ run_pass.assert_called_once_with(now, [], 0, max_prune=None)
assert result["cutoff"] == now.isoformat()
else:
assert result == {"skipped": 1}
tables.assert_not_called()
run_pass.assert_not_called()
+
+
+def test_prune_cap_accounts_only_for_committed_passes(stats: MagicMock) ->
None:
+ """A retried pass uses the same allowance and cannot overshoot the cap."""
+ tables: version_history_retention.ShadowTables = (
+ version_history_retention.ShadowTables(
+ parent=[], child=[], m2m=None, transaction=MagicMock()
+ )
+ )
+ first: dict[str, int] = {
+ "candidate_count": 1000,
+ "max_candidate_id": 1000,
+ "pruned_transactions": 2,
+ }
+ second: dict[str, int] = {
+ "candidate_count": 1000,
+ "max_candidate_id": 2000,
+ "pruned_transactions": 1,
+ }
+ run_pass: MagicMock
+ with (
+ patch.object(
+ version_history_retention, "_resolve_shadow_tables",
return_value=tables
+ ),
+ patch.object(
+ version_history_retention,
+ "_run_pass_with_retry",
+ side_effect=[(first, 1), (second, 0)],
+ ) as run_pass,
+ patch.object(
+ version_history_retention, "_probe_prunable", return_value=(4,
False)
+ ),
+ ):
+ result: dict[str, Any] =
version_history_retention._prune_old_versions_impl(
+ retention_days=30, max_per_run=3
+ )
+
+ assert result["pruned_transactions"] == 3
+ assert result["retried"] == 1
+ assert result["cap_reached"] is True
+ assert result["remaining_eligible"] == 4
+ assert result["remaining_count_complete"] is False
+ stats.gauge.assert_any_call(
+ "superset.versioning.retention.remaining_count_complete", 0
+ )
+ assert run_pass.call_args_list[0].kwargs["max_prune"] == 3
+ assert run_pass.call_args_list[1].kwargs["max_prune"] == 1
+
+
+def test_prune_cap_probes_after_preserved_candidate_windows(
+ stats: MagicMock,
+) -> None:
+ """The remainder probe skips windows already scanned before the cap."""
+ tables: version_history_retention.ShadowTables = (
+ version_history_retention.ShadowTables(
+ parent=[], child=[], m2m=None, transaction=MagicMock()
+ )
+ )
+ batch_size: int = version_history_retention._MAX_PRUNE_BATCH
+ windows: dict[int, version_history_retention._PruneWindow] = {
+ 0: version_history_retention._PruneWindow([], batch_size, 1000),
+ 1000: version_history_retention._PruneWindow([1500], batch_size, 2000),
+ }
+
+ def resolve_window(
+ _conn: sa.engine.Connection,
+ _cutoff: datetime,
+ _shadow_tables: list[sa.Table],
+ after_id: int,
+ _limit: int,
+ ) -> version_history_retention._PruneWindow:
+ """Return the candidate window at the requested scan cursor."""
+ return windows[after_id]
+
+ resolve: MagicMock
+ with (
+ patch.object(
+ version_history_retention, "_resolve_shadow_tables",
return_value=tables
+ ),
+ patch.object(
+ version_history_retention,
+ "_run_pass_with_retry",
+ side_effect=[
+ ({"candidate_count": batch_size, "max_candidate_id": 1000}, 0),
+ (
+ {
+ "candidate_count": batch_size,
+ "max_candidate_id": 2000,
+ "pruned_transactions": 1,
+ },
+ 0,
+ ),
+ ],
+ ),
+ patch.object(version_history_retention, "db"),
+ patch.object(
+ version_history_retention,
+ "_resolve_prune_window",
+ side_effect=resolve_window,
+ ) as resolve,
+ ):
+ result: dict[str, Any] =
version_history_retention._prune_old_versions_impl(
+ retention_days=30, max_per_run=1
+ )
+
+ assert result["pruned_transactions"] == 1
+ assert result["remaining_eligible"] == 1
+ assert result["remaining_count_complete"] is False
+ assert resolve.call_args.args[3] == 1000
+ stats.gauge.assert_any_call(
+ "superset.versioning.retention.remaining_eligible_at_least", 1
+ )
+
+
[email protected]("cap", [None, 0, 3])
+def test_prune_dry_run_counts_all_eligible_without_writes(
+ stats: MagicMock, cap: int | None
+) -> None:
+ """A dry run reports the whole backlog and capped-run estimate."""
+ tables: version_history_retention.ShadowTables = (
+ version_history_retention.ShadowTables(
+ parent=[], child=[], m2m=None, transaction=MagicMock()
+ )
+ )
+ run_pass: MagicMock
+ with (
+ patch.object(
+ version_history_retention, "_resolve_shadow_tables",
return_value=tables
+ ),
+ patch.object(version_history_retention, "_count_prunable",
return_value=7),
+ patch.object(version_history_retention, "_run_pass_with_retry") as
run_pass,
+ ):
+ result: dict[str, Any] =
version_history_retention._prune_old_versions_impl(
+ 30, max_per_run=cap, dry_run=True
+ )
+
+ assert result["eligible_backlog"] == 7
+ assert result["estimated_capped_runs"] == (3 if cap == 3 else 1)
+ run_pass.assert_not_called()
+ stats.gauge.assert_not_called()
+
+
[email protected]("helper_name", ["_count_prunable", "_probe_prunable"])
+def test_remainder_reads_use_default_isolation(helper_name: str) -> None:
+ """Backlog measurement does not request the delete pass's isolation."""
+ tables: version_history_retention.ShadowTables = (
+ version_history_retention.ShadowTables(
+ parent=[], child=[], m2m=None, transaction=MagicMock()
+ )
+ )
+ engine: MagicMock = MagicMock()
+ window: version_history_retention._PruneWindow = (
+ version_history_retention._PruneWindow(
+ prunable=[], candidate_count=0, max_candidate_id=0
+ )
+ )
+ with (
+ patch.object(version_history_retention, "db") as mock_db,
+ patch.object(
+ version_history_retention, "_resolve_prune_window",
return_value=window
+ ),
+ ):
+ mock_db.engine = engine
+ result: int | tuple[int, bool] = getattr(
+ version_history_retention, helper_name
+ )(datetime(2026, 1, 1), tables)
+
+ assert result in (0, (0, True))
+ engine.connect.return_value.execution_options.assert_not_called()
+
+
+def test_dry_run_count_accumulates_windows_and_closes_read_transactions() ->
None:
+ """A paged backlog count sums windows and releases each read snapshot."""
+ tables: version_history_retention.ShadowTables = (
+ version_history_retention.ShadowTables(
+ parent=[], child=[], m2m=None, transaction=MagicMock()
+ )
+ )
+ engine: MagicMock = MagicMock()
+ full_window: version_history_retention._PruneWindow = (
+ version_history_retention._PruneWindow(
+ prunable=[1, 2],
+ candidate_count=version_history_retention._MAX_PRUNE_BATCH,
+ max_candidate_id=1000,
+ )
+ )
+ final_window: version_history_retention._PruneWindow = (
+ version_history_retention._PruneWindow(
+ prunable=[1001], candidate_count=1, max_candidate_id=1001
+ )
+ )
+ mock_db: MagicMock
+ resolve: MagicMock
+ with (
+ patch.object(version_history_retention, "db") as mock_db,
+ patch.object(
+ version_history_retention,
+ "_resolve_prune_window",
+ side_effect=[full_window, final_window],
+ ) as resolve,
+ ):
+ mock_db.engine = engine
+ count: int = version_history_retention._count_prunable(
+ datetime(2026, 1, 1), tables
+ )
+
+ assert count == 3
+ assert resolve.call_args_list[0].args[3] == 0
+ assert resolve.call_args_list[1].args[3] == full_window.max_candidate_id
+ assert engine.connect.call_count == 2
+ assert engine.connect.return_value.__exit__.call_count == 2
+ first_exit: int = engine.mock_calls.index(call.connect().__exit__(None,
None, None))
+ second_connect: int = engine.mock_calls.index(call.connect(), first_exit +
1)
+ assert first_exit < second_connect
+
+
[email protected]("invalid", [-1, True, "3", 2.5])
+def test_prune_rejects_invalid_cap_before_work(invalid: object) -> None:
+ """Malformed budgets fail closed before resolving any shadow table."""
+ with (
+ patch.object(version_history_retention, "_resolve_shadow_tables") as
resolve,
+ pytest.raises(ValueError, match="VERSION_HISTORY_PRUNE_MAX"),
+ ):
+ version_history_retention._prune_old_versions_impl(
+ 30, max_per_run=cast(int | None, invalid)
+ )
+ resolve.assert_not_called()
+
+
+def test_prune_retry_reuses_the_same_transaction_budget(stats: MagicMock) ->
None:
+ """A rolled-back SERIALIZABLE attempt consumes no transaction allowance."""
+ tables: version_history_retention.ShadowTables = (
+ version_history_retention.ShadowTables(
+ parent=[], child=[], m2m=None, transaction=MagicMock()
+ )
+ )
+ cutoff: datetime = datetime(2026, 1, 1)
+ run_pass: MagicMock
+ with (
+ patch.object(
+ version_history_retention,
+ "_run_prune_pass",
+ side_effect=[
+ OperationalError("SELECT 1", {}, Exception("serialization
conflict")),
+ {"pruned_transactions": 2},
+ ],
+ ) as run_pass,
+ patch.object(version_history_retention.time, "sleep"),
+ ):
+ result: tuple[dict[str, Any], int] = (
+ version_history_retention._run_pass_with_retry(
+ cutoff, tables, after_id=4, max_prune=2
+ )
+ )
+
+ assert result == ({"pruned_transactions": 2}, 1)
+ assert run_pass.call_args_list == [
+ call(cutoff, tables, 4, max_prune=2),
+ call(cutoff, tables, 4, max_prune=2),
+ ]
+ stats.incr.assert_called_once_with("superset.versioning.retention.retried")
+
+
+def test_scheduled_prune_rejects_invalid_cap_before_work(stats: MagicMock) ->
None:
+ """A malformed scheduled budget cannot enter the prune implementation."""
+ app: Flask = Flask(__name__)
+ app.config.update(
+ VERSION_HISTORY_RETENTION_DAYS=30,
+ VERSION_HISTORY_PRUNE_MAX_TRANSACTIONS_PER_RUN=-1,
+ )
+ prune: MagicMock
+ with (
+ app.app_context(),
+ patch.object(version_history_retention, "_prune_old_versions_impl") as
prune,
+ ):
+ result: dict[str, Any] = version_history_retention.prune_old_versions()
+
+ assert result == {"skipped_invalid_cap": 1}
+ prune.assert_not_called()
+ stats.incr.assert_called_once_with(
+ "superset.versioning.retention.skipped_invalid_cap"
+ )
+
+
[email protected]("invalid", [None, 0, 1, "false", [], {}])
+def test_scheduled_prune_skips_invalid_dry_run_before_work(
Review Comment:
Added a scheduled-task test that sets `VERSION_HISTORY_PRUNE_DRY_RUN=True`
and a cap of 2 and calls `prune_old_versions()`. It asserts
`_prune_old_versions_impl` receives both `dry_run=True` and `max_per_run=2`, so
a wrapper that hardcodes either value now fails. The test is in
https://github.com/apache/superset/commit/a54556a2965e70577d12fb97af7b0994f7a3b27b.
--
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]
