sadpandajoe commented on code in PR #45029:
URL: https://github.com/apache/superset/pull/45029#discussion_r4210255121
##########
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:
The docstring says a failed or blocked root doesn't consume the budget, but
the fixture only returns a plain `purged=False` result: no `blocked_reason`,
and `_purge_one` never raises. If the guard were changed to count `purged +
blocked + failures` against `max_per_run`, this test would still pass, and a
persistent blocked root at the lowest ID would starve every eligible root
behind it on each run. Could you drive `_purge_model` with a blocked result, an
exception, then two successes under cap `2` and assert `(2, 0, 1, 1)`?
##########
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:
The scheduled-entrypoint tests cover the defaults (`max_per_run=1000,
dry_run=False`) and invalid values, but no test sets a valid
`VERSION_HISTORY_PRUNE_DRY_RUN=True` or a non-default cap through
`prune_old_versions()`. If the wrapper hardcoded either argument, every test
here would still pass, and an operator's preview run would delete data. Could
you add one case with `dry_run=True` and cap `2` that asserts both values reach
`_prune_old_versions_impl`?
--
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]
