feat(clickhouse): add partition attachment tooling - #8357
Conversation
The health check for partition attachment was a fixed SELECT 1, so it could only tell whether ClickHouse was reachable. Allow an operator to supply a query instead, and bind the partition being attached into it. Derive the boundary from system.parts.partition rather than the partition ID. A tuple key such as (retention_days, toMonday(timestamp)) renders as (90,'2024-02-12'), which exposes the boundary even when the partition ID is hashed, as it is once a String column joins the key. Bind retention_days alongside the boundary. Partitions that differ only by retention share a boundary, so partition_start alone cannot identify one of them, and a check scoped only by time reads rows belonging to a partition that was already attached. Fail before running a query that references a parameter the partition does not carry, so a stray placeholder cannot quietly pass every check.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 09e204c. Configure here.
StorageKey accepts any string without validating it, so an unknown storage reaches the registry and fails on a dict lookup. That raises KeyError, which the handler did not catch, so a typo in --storage exited with a traceback instead of the intended message. Reported by Cursor Bugbot.
| from snuba.clickhouse.partition_management import ( | ||
| PartitionBoundaryError, | ||
| attach_partition_from_table, | ||
| attach_partitions_from_table, | ||
| build_health_check, | ||
| get_partition_boundaries, | ||
| ) | ||
| from snuba.clickhouse.pool import ClickhousePool | ||
| from snuba.clusters.cluster import ClickhouseNode, build_pool |
There was a problem hiding this comment.
can we move this to the top of the file?
There was a problem hiding this comment.
Done in 2347ef8. Moved all three (partition_management, ClickhousePool, ClickhouseNode/build_pool) to the top, merging the cluster names into the existing import.
Safe to hoist here: SnubaCLI.get_command calls initialize_snuba() before compiling the command module, and the file already imported from snuba.clusters.cluster and snuba.datasets.storages.factory at module level. Verified --help, a live dry run, and snuba --help for the other commands, so no import cycle.
![sentry[bot]](https://avatars.githubusercontent.com/in/12637?s=60&v=4){kind=small}
The command imported its dependencies inside the function body. The CLI loader initializes snuba before compiling command code, so a module level import is already safe here, and the top-level import block covered the same packages.
| database, | ||
| source_table, | ||
| destination_table, | ||
| health_check_query=health_check_query, | ||
| dry_run=not execute, | ||
| on_partition_attached=lambda attached_partition_id: click.echo( | ||
| f"Attached partition {attached_partition_id}" | ||
| ), | ||
| ) | ||
| except PartitionBoundaryError as error: | ||
| raise click.UsageError(str(error)) from error |
There was a problem hiding this comment.
Bug: The attach-partitions CLI command doesn't handle UnhealthyError, causing a full traceback to be displayed to the user when a health check fails during execution.
Severity: MEDIUM
Suggested Fix
Update the try/except block in snuba/cli/attach_partitions.py to also catch UnhealthyError. Convert the caught exception into a click.UsageError or click.ClickException to provide a clean error message to the user, following the existing pattern for CLI error handling in the codebase.
Prompt for AI Agent
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.
Location: snuba/cli/attach_partitions.py#L128-L159
Potential issue: The `try/except` block in the `attach-partitions` CLI command is
designed to catch errors during partition attachment. However, it only catches
`PartitionBoundaryError`. When a health check fails, the `run_health_check_query`
function raises an `UnhealthyError`. This exception is not caught by the CLI's exception
handler, causing a full Python traceback to be displayed to the user instead of a clean,
user-friendly error message. This can happen when running the command with the
`--execute` flag and a health check fails.
2 participants
Summary
snuba attach-partitionscommand with optional--partition-idtargeting--health-check-queryHealth check
--health-check-querylets an operator supply a custom gate:
The query runs before every attach, gating the first operation and each one
following an attach. A failure aborts immediately, so a bad destination stops
the run rather than absorbing every remaining partition. Unhealthy means the
query raised, returned no rows, or returned a falsy first value.
Bound parameters
Two parameters are substituted so the check can scope to the partition about
to be attached:
%(partition_start)s— start of the partition's time boundary%(retention_days)s— bound when the partition key carries oneBoth are needed for EAP tables.
eap_itemspartitions by(retention_days, toMonday(timestamp)), so30-20240212and90-20240212share the boundary
2024-02-12. A check scoped only by time cannot tell themapart: it reads rows from a partition that was already attached, and reports a
false unhealthy on a healthy destination. Scoping on
retention_daysas wellidentifies a single partition.
Referencing a parameter the partition does not carry (for example
%(retention_days)son a date-only key) fails before the query runs, listingwhat is available, so a stray placeholder cannot quietly pass every check.
Risk
Low. There is currently no entrypoint into this code path:
snuba.clickhouse.partition_managementoutside its own tests
snuba attach-partitionsis picked up by CLI auto-discovery but is notreferenced by any GoCD pipeline, GitHub workflow, Makefile target, cron, or
manual job, so it only runs when a human invokes it
SELECT 1, so--health-check-queryis opt-in--dry-run;--executeis required to writeATTACH PARTITION ... FROMleaves the source intact, and active destinationpartition IDs are skipped, so re-running is safe
Testing
Verified end to end against a local ClickHouse using tables with the real
(retention_days, toMonday(timestamp))key, plus a String-key table withhashed partition IDs: boundaries derived correctly in both, all partitions
attached, re-run was a no-op, and an unhealthy query aborted before any attach
with the destination left empty.