From dfd3685c11e5561c80ca937367c77c4d4da092b3 Mon Sep 17 00:00:00 2001 From: Oliver Newland Date: Mon, 19 Oct 2026 15:11:04 -0700 Subject: [PATCH] feat(clickhouse): Skip past-retention partitions in attach-partitions attach-partitions had an upper date bound (--max-future-days) but no lower one, so it attached source partitions that were already past retention, including garbage weeks such as 90-19700101. TTL deletes those rows on the destination anyway, and attaching them recreates the destination's block_numbers znode, which prune-partitions then forgets again. On ClickHouse 25.8 an attach into a partition whose znode was forgotten can also fail with ZNONODE when the replica still has the path cached. Skip partitions that prune-partitions would classify as dead_past: past retention plus --retention-margin-days (default 14). Both commands share classify_partition, so they agree on what is dead. The filter only applies when the source table is partitioned by (retention_days, toMonday(...)), since other keys' IDs can't be read as a retention and a week, and the partition key is only read when some candidate looks expired. --include-expired restores the old behaviour. A single --partition-id that is past retention is refused unless --include-expired is passed. Co-Authored-By: Claude --- snuba/cli/attach_partitions.py | 50 +++++- snuba/clickhouse/partition_management.py | 92 +++++++++- snuba/clickhouse/prune_partitions.py | 11 +- tests/cli/test_attach_partitions.py | 63 +++++++ tests/clickhouse/test_partition_management.py | 166 ++++++++++++++++++ 5 files changed, 376 insertions(+), 6 deletions(-) create mode 100644 tests/cli/test_attach_partitions.py diff --git a/snuba/cli/attach_partitions.py b/snuba/cli/attach_partitions.py index 4a7b68561b..8f1ea6d842 100644 --- a/snuba/cli/attach_partitions.py +++ b/snuba/cli/attach_partitions.py @@ -1,6 +1,6 @@ import logging from collections.abc import Sequence -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta import click @@ -10,10 +10,12 @@ attach_partition_from_table, attach_partitions_from_table, build_health_check, + filter_expired_partitions, get_partition_boundaries, is_partition_too_far_in_future, ) from snuba.clickhouse.pool import ClickhousePool +from snuba.clickhouse.prune_partitions import DEFAULT_RETENTION_MARGIN from snuba.clusters.cluster import ( ClickhouseClientSettings, ClickhouseNode, @@ -82,6 +84,24 @@ "cutoff is kept. Pass 0 to attach future-dated partitions regardless." ), ) +@click.option( + "--retention-margin-days", + type=click.IntRange(min=0), + default=DEFAULT_RETENTION_MARGIN.days, + show_default=True, + help=( + "Skip partitions past retention by more than this many days, using the " + "same rule as prune-partitions. Only applies to tables partitioned by " + "(retention_days, toMonday(timestamp)). TTL would delete their rows " + "anyway, and attaching them recreates block_numbers znodes that " + "prune-partitions forgets." + ), +) +@click.option( + "--include-expired/--skip-expired", + default=False, + help="Attach partitions past retention instead of skipping them.", +) @click.option("--log-level", help="Logging level to use.") def attach_partitions( source_table: str, @@ -97,14 +117,17 @@ def attach_partitions( partition_id: str | None, health_check_query: str | None, max_future_days: int, + retention_margin_days: int, + include_expired: bool, log_level: str | None, ) -> None: """Attach SOURCE_TABLE partitions to DESTINATION_TABLE one at a time. Table names are unqualified and must be in the database used by STORAGE. - In discovery mode, existing destination partitions and partitions dated too - far in the future are skipped. A specific partition ID is attached directly. - The source data is not removed. + In discovery mode, existing destination partitions, partitions dated too + far in the future and partitions past retention are skipped. A specific + partition ID is attached directly, unless it is too far in the future or + past retention. The source data is not removed. """ setup_logging(log_level) setup_sentry() @@ -116,6 +139,7 @@ def attach_partitions( raise click.UsageError("--max-future-days must not be negative") max_future_age = timedelta(days=max_future_days) if max_future_days > 0 else None + retention_margin = None if include_expired else timedelta(days=retention_margin_days) try: storage = get_storage(StorageKey(storage_name)) @@ -174,6 +198,18 @@ def attach_partitions( f"Partition {partition_id} begins more than {max_future_days} day(s) " "in the future. Pass --max-future-days 0 to attach it anyway." ) + if retention_margin is not None and not filter_expired_partitions( + connection, + database, + source_table, + [partition_id], + today=datetime.now(UTC).date(), + retention_margin=retention_margin, + ): + raise click.UsageError( + f"Partition {partition_id} is past retention by more than " + f"{retention_margin_days} day(s). Pass --include-expired to attach it anyway." + ) if execute: build_health_check( connection, @@ -198,6 +234,7 @@ def attach_partitions( destination_table, health_check_query=health_check_query, max_future_age=max_future_age, + retention_margin=retention_margin, dry_run=not execute, on_partition_attached=lambda attached_partition_id: logger.info( "Attached partition %s from %s to %s", @@ -211,6 +248,11 @@ def attach_partitions( boundary, max_future_days, ), + on_partition_expired=lambda expired_partition_id: logger.info( + "Skipping partition %s past retention by more than %d day(s)", + expired_partition_id, + retention_margin_days, + ), ) except PartitionBoundaryError as error: raise click.UsageError(str(error)) from error diff --git a/snuba/clickhouse/partition_management.py b/snuba/clickhouse/partition_management.py index 215f06794d..d02723458e 100644 --- a/snuba/clickhouse/partition_management.py +++ b/snuba/clickhouse/partition_management.py @@ -1,15 +1,23 @@ import logging import re from collections.abc import Callable, Sequence -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta from snuba.clickhouse.escaping import escape_identifier from snuba.clickhouse.pool import ClickhousePool +from snuba.clickhouse.prune_partitions import ( + DEFAULT_RETENTION_MARGIN, + Category, + classify_partition, + is_weekly_retention_partition_key, +) +from snuba.clickhouse.prune_partitions import today as utc_today logger = logging.getLogger("snuba.clickhouse.partition_management") HealthCheck = Callable[[str], None] PartitionAttached = Callable[[str], None] +PartitionExpired = Callable[[str], None] PartitionSkipped = Callable[[str, str], None] PARTITION_START_PARAM = "partition_start" @@ -256,6 +264,69 @@ def filter_future_partitions( return kept +def get_partition_key(clickhouse: ClickhousePool, database: str, table: str) -> str | None: + """Return a table's ``partition_key`` as rendered by ``system.tables``.""" + response = clickhouse.execute( + "SELECT partition_key FROM system.tables WHERE database = %(database)s AND name = %(table)s", + {"database": database, "table": table}, + ) + if not response.results: + return None + return str(response.results[0][0]) + + +def filter_expired_partitions( + clickhouse: ClickhousePool, + database: str, + table: str, + partition_ids: Sequence[str], + *, + today: date, + retention_margin: timedelta = DEFAULT_RETENTION_MARGIN, + on_partition_expired: PartitionExpired | None = None, +) -> Sequence[str]: + """ + Drop partitions that are past retention plus ``retention_margin``, using + the same rule ``prune-partitions`` uses to decide a partition is dead. + + TTL would delete the attached rows anyway. Attaching them also recreates + the destination's ``block_numbers`` znode, which ``prune-partitions`` then + forgets again, and on ClickHouse 25.8 an attach into a partition whose + znode was forgotten can fail with ZNONODE. + + Only tables partitioned by ``(retention_days, toMonday(timestamp))`` are + filtered, because other keys' partition IDs can't be read as a retention and + a week. ``table`` must be the source table; ATTACH PARTITION FROM requires + the destination to share its partition key. The key is only read when some + partition ID looks expired. + """ + expired = { + partition_id + for partition_id in partition_ids + if classify_partition(partition_id, today, retention_margin) == Category.DEAD_PAST + } + if not expired: + return list(partition_ids) + + partition_key = get_partition_key(clickhouse, database, table) + if partition_key is None or not is_weekly_retention_partition_key(partition_key): + return list(partition_ids) + + kept = [] + for partition_id in partition_ids: + if partition_id in expired: + logger.info( + "Skipping partition %s: past retention by more than %s", + partition_id, + retention_margin, + ) + if on_partition_expired is not None: + on_partition_expired(partition_id) + continue + kept.append(partition_id) + return kept + + def get_active_partition_ids( clickhouse: ClickhousePool, database: str, table: str ) -> Sequence[str]: @@ -347,7 +418,9 @@ def attach_partitions_from_table( health_check_query: str | None = None, on_partition_attached: PartitionAttached | None = None, on_partition_skipped: PartitionSkipped | None = None, + on_partition_expired: PartitionExpired | None = None, max_future_age: timedelta | None = DEFAULT_MAX_FUTURE_PARTITION_AGE, + retention_margin: timedelta | None = DEFAULT_RETENTION_MARGIN, now: datetime | None = None, dry_run: bool = False, ) -> Sequence[str]: @@ -368,6 +441,10 @@ def attach_partitions_from_table( skipped, which keeps implausibly future-dated partitions out of the destination. Pass ``max_future_age=None`` to attach them regardless. + Partitions past retention by more than ``retention_margin`` are skipped too, + as described in ``filter_expired_partitions``. Pass + ``retention_margin=None`` to attach them regardless. + ClickHouse's ``ATTACH PARTITION ... FROM`` leaves the source partition in place. Re-running this function is safe because active destination partition IDs are skipped. @@ -385,6 +462,19 @@ def attach_partitions_from_table( if partition_id not in destination_partition_ids ] + if retention_margin is not None: + pending_partition_ids = list( + filter_expired_partitions( + clickhouse, + database, + source_table, + pending_partition_ids, + today=now.date() if now is not None else utc_today(), + retention_margin=retention_margin, + on_partition_expired=on_partition_expired, + ) + ) + # Boundaries resolve hashed partition IDs, and are read once because the # date filter and a scoped health check query both need them. The date # filter only needs them when a pending partition ID is not itself a diff --git a/snuba/clickhouse/prune_partitions.py b/snuba/clickhouse/prune_partitions.py index 92ac677024..f189977b0d 100644 --- a/snuba/clickhouse/prune_partitions.py +++ b/snuba/clickhouse/prune_partitions.py @@ -116,6 +116,15 @@ def parse_partition_id(partition_id: str) -> tuple[int, date] | None: return int(match.group(1)), week_start +def is_weekly_retention_partition_key(partition_key: str) -> bool: + """ + Report whether a ``system.tables`` ``partition_key`` is + ``(, toMonday())``, the only key whose partition + IDs ``classify_partition`` can read as a retention and a week. + """ + return _PARTITION_KEY.match(partition_key) is not None + + def classify_partition( partition_id: str, today: date, @@ -163,7 +172,7 @@ def check_table_settings(clickhouse: ClickhousePool, database: str, table: str) engine_full, partition_key = result.results[0] if not all(flag in engine_full for flag in _BLOCK_FLAGS): raise SkipTable("block number and offset columns are not both enabled") - if not _PARTITION_KEY.match(partition_key): + if not is_weekly_retention_partition_key(partition_key): raise SkipTable(f"unsupported partition key {partition_key!r}") diff --git a/tests/cli/test_attach_partitions.py b/tests/cli/test_attach_partitions.py new file mode 100644 index 0000000000..c736b685dd --- /dev/null +++ b/tests/cli/test_attach_partitions.py @@ -0,0 +1,63 @@ +from collections.abc import Mapping +from typing import Any +from unittest.mock import Mock, patch + +from click.testing import CliRunner + +from snuba.cli.attach_partitions import attach_partitions +from snuba.clickhouse.pool import ClickhousePool, ClickhouseResult + +EXPIRED = "90-19700101" + + +def make_connection() -> Mock: + def execute(query: str, params: Mapping[str, Any] | None = None) -> ClickhouseResult: + if "FROM system.tables" in query: + return ClickhouseResult([("(retention_days, toMonday(timestamp))",)]) + if "FROM system.parts" in query: + return ClickhouseResult([(EXPIRED, "(90,'1970-01-01')")]) + raise AssertionError(f"unexpected query: {query}") + + connection = Mock(spec=ClickhousePool) + connection.execute.side_effect = execute + return connection + + +def invoke(connection: Mock, *args: str) -> Any: + with patch("snuba.cli.attach_partitions.build_pool", return_value=connection): + return CliRunner().invoke( + attach_partitions, + [ + "eap_items_1_local", + "eap_items_2_local", + "--storage", + "eap_items", + "--clickhouse-host", + "localhost", + "--clickhouse-port", + "8123", + "--partition-id", + EXPIRED, + *args, + ], + ) + + +def test_refuses_a_single_expired_partition() -> None: + connection = make_connection() + + result = invoke(connection, "--execute") + + assert result.exit_code == 2, result.output + assert "past retention" in result.output + assert "--include-expired" in result.output + assert not [c for c in connection.execute.call_args_list if "ATTACH" in c.args[0]] + + +def test_include_expired_allows_a_single_expired_partition() -> None: + connection = make_connection() + + result = invoke(connection, "--include-expired") + + assert result.exit_code == 0, result.output + assert not [c for c in connection.execute.call_args_list if "system.tables" in c.args[0]] diff --git a/tests/clickhouse/test_partition_management.py b/tests/clickhouse/test_partition_management.py index 74847c650a..b5bd8a32a7 100644 --- a/tests/clickhouse/test_partition_management.py +++ b/tests/clickhouse/test_partition_management.py @@ -10,6 +10,7 @@ attach_partition_from_table, attach_partitions_from_table, build_health_check, + filter_expired_partitions, get_partition_boundaries, is_partition_too_far_in_future, parse_partition_start, @@ -203,6 +204,8 @@ def test_resolves_boundaries_from_the_source_table_for_hashed_partition_ids() -> "source", "destination", health_check_query="SELECT count() FROM destination WHERE timestamp >= %(partition_start)s", + # These 2024 partitions are long past retention; expiry is tested separately. + retention_margin=None, ) assert result == ["90-20240212"] @@ -302,6 +305,8 @@ def test_retention_days_distinguishes_partitions_sharing_a_boundary() -> None: "SELECT count() FROM destination " "WHERE timestamp >= %(partition_start)s AND retention_days = %(retention_days)s" ), + # These 2024 partitions are long past retention; expiry is tested separately. + retention_margin=None, ) assert result == ["30-20240212", "90-20240212"] @@ -523,3 +528,164 @@ def test_dry_run_reports_the_date_filtered_partitions() -> None: for call_args in clickhouse.execute.call_args_list if "ATTACH PARTITION" in call_args.args[0] ] + + +WEEKLY_RETENTION_KEY = "(retention_days, toMonday(timestamp))" +# 2026-08-31 + 7 days + 30 days retention + 14 days margin = 2026-10-21. +EXPIRY_NOW = datetime(2026, 10, 22) + + +def make_eap_clickhouse( + source_partitions: list[str], + destination_partitions: list[str], + partition_key: str = WEEKLY_RETENTION_KEY, +) -> Mock: + clickhouse = Mock(spec=ClickhousePool) + clickhouse.execute.side_effect = [ + ClickhouseResult([(partition,) for partition in source_partitions]), + ClickhouseResult([(partition,) for partition in destination_partitions]), + ClickhouseResult([(partition_key,)]), + *[ClickhouseResult() for _ in source_partitions], + ] + return clickhouse + + +def test_skips_partitions_past_retention() -> None: + clickhouse = make_eap_clickhouse( + ["30-20260831", "30-20260907", "90-19700101", "90-20260831"], [] + ) + expired = Mock(spec=Callable[[str], None]) + + result = attach_partitions_from_table( + clickhouse, + "default", + "source", + "destination", + health_check=Mock(spec=Callable[[str], None]), + on_partition_expired=expired, + max_future_age=None, + now=EXPIRY_NOW, + ) + + # 30-20260907 and the 90 day week are still within retention plus margin. + assert result == ["30-20260907", "90-20260831"] + assert expired.call_args_list == [call("30-20260831"), call("90-19700101")] + assert clickhouse.execute.call_args_list[2] == call( + "SELECT partition_key FROM system.tables WHERE database = %(database)s AND name = %(table)s", + {"database": "default", "table": "source"}, + ) + + +def test_expiry_uses_the_retention_margin() -> None: + clickhouse = make_eap_clickhouse(["30-20260831"], []) + + result = attach_partitions_from_table( + clickhouse, + "default", + "source", + "destination", + health_check=Mock(spec=Callable[[str], None]), + retention_margin=timedelta(days=30), + max_future_age=None, + now=EXPIRY_NOW, + ) + + assert result == ["30-20260831"] + + +def test_attaches_expired_partitions_when_the_filter_is_disabled() -> None: + clickhouse = make_clickhouse(["30-20260831", "90-19700101"], []) + + result = attach_partitions_from_table( + clickhouse, + "default", + "source", + "destination", + health_check=Mock(spec=Callable[[str], None]), + retention_margin=None, + max_future_age=None, + now=EXPIRY_NOW, + ) + + assert result == ["30-20260831", "90-19700101"] + assert not [ + call_args + for call_args in clickhouse.execute.call_args_list + if "system.tables" in call_args.args[0] + ] + + +def test_keeps_expired_looking_partitions_of_other_partition_keys() -> None: + # A monthly key renders the same ID shape, but the ID can't be read as a week. + clickhouse = make_eap_clickhouse( + ["30-20260801"], [], partition_key="(retention_days, toStartOfMonth(timestamp))" + ) + + result = attach_partitions_from_table( + clickhouse, + "default", + "source", + "destination", + health_check=Mock(spec=Callable[[str], None]), + max_future_age=None, + now=EXPIRY_NOW, + ) + + assert result == ["30-20260801"] + + +def test_does_not_read_the_partition_key_without_expired_candidates() -> None: + clickhouse = make_clickhouse(["30-20261012", "20240213"], []) + + result = attach_partitions_from_table( + clickhouse, + "default", + "source", + "destination", + health_check=Mock(spec=Callable[[str], None]), + max_future_age=None, + now=EXPIRY_NOW, + ) + + assert result == ["30-20261012", "20240213"] + assert not [ + call_args + for call_args in clickhouse.execute.call_args_list + if "system.tables" in call_args.args[0] + ] + + +def test_dry_run_reports_the_expiry_filtered_partitions() -> None: + clickhouse = make_eap_clickhouse(["30-20260831", "30-20260907"], []) + + result = attach_partitions_from_table( + clickhouse, + "default", + "source", + "destination", + max_future_age=None, + now=EXPIRY_NOW, + dry_run=True, + ) + + assert result == ["30-20260907"] + assert not [ + call_args + for call_args in clickhouse.execute.call_args_list + if "ATTACH PARTITION" in call_args.args[0] + ] + + +def test_filter_expired_partitions_skips_missing_tables() -> None: + clickhouse = Mock(spec=ClickhousePool) + clickhouse.execute.return_value = ClickhouseResult([]) + + result = filter_expired_partitions( + clickhouse, + "default", + "missing", + ["90-19700101"], + today=EXPIRY_NOW.date(), + ) + + assert result == ["90-19700101"]