feat(clickhouse): Skip past-retention partitions in attach-partitions

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

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

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
+    ``(<retention column>, toMonday(<column>))``, 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}")
 
 

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]]