diff --git a/clickhouse_alembic/__init__.py b/clickhouse_alembic/__init__.py index 364307f..818590b 100644 --- a/clickhouse_alembic/__init__.py +++ b/clickhouse_alembic/__init__.py @@ -13,12 +13,22 @@ # Lazy imports to avoid import errors before dependencies are created def __getattr__(name: str) -> Any: - if name in ("read_sql", "get_db", "create_dictionary"): - from clickhouse_alembic.helpers import create_dictionary, get_db, read_sql + if name in ("read_sql", "get_db", "create_dictionary", "on_cluster", "get_cluster"): + from clickhouse_alembic.helpers import ( + create_dictionary, + get_cluster, + get_db, + on_cluster, + read_sql, + ) - return {"read_sql": read_sql, "get_db": get_db, "create_dictionary": create_dictionary}[ - name - ] + return { + "read_sql": read_sql, + "get_db": get_db, + "create_dictionary": create_dictionary, + "on_cluster": on_cluster, + "get_cluster": get_cluster, + }[name] elif name == "get_env_config": from clickhouse_alembic.config import get_env_config @@ -40,6 +50,8 @@ def __getattr__(name: str) -> Any: "get_db", "get_env_config", "create_dictionary", + "on_cluster", + "get_cluster", "get_secret", "SSMSecretNotFoundError", "SSMJsonKeyError", diff --git a/clickhouse_alembic/cli.py b/clickhouse_alembic/cli.py index 0bc7464..0219bdd 100644 --- a/clickhouse_alembic/cli.py +++ b/clickhouse_alembic/cli.py @@ -333,12 +333,16 @@ def history(environment: str) -> None: @click.option( "--dict", "-d", "dict_name", help="Create SQL file for dictionary (e.g., --dict regions)" ) +@click.option( + "--exchange", is_flag=True, help="Generate EXCHANGE TABLES scaffold (requires --table)" +) def new( environment: str, name: str, table_name: str | None, view_name: str | None, dict_name: str | None, + exchange: bool, ) -> None: """Create a new migration. @@ -347,13 +351,32 @@ def new( Use --table, --view, or --dict with the object name to create a SQL history file: + \b ch-migrate new dev add_status_column --table logs + + Use --exchange with --table to generate a zero-downtime EXCHANGE TABLES scaffold: + + \b + ch-migrate new dev alter_users --table users --exchange """ + if exchange and not table_name: + click.echo("Error: --exchange requires --table", err=True) + sys.exit(1) + result = _run_alembic(environment, ["revision", "-m", name], exit_on_complete=False) if result is None or result.returncode != 0: sys.exit(1 if result is None else result.returncode) + revision = _extract_revision_from_output(result.stdout) + if not revision: + click.echo("Warning: Could not extract revision ID, SQL file not created", err=True) + sys.exit(0) + + if exchange: + _create_exchange_scaffold(environment, table_name, revision, result.stdout) + sys.exit(0) + # Determine object type and name from options object_name: str | None = None object_type: str | None = None @@ -366,13 +389,9 @@ def new( # If object specified, create SQL file if object_name and object_type: - revision = _extract_revision_from_output(result.stdout) - if revision: - sql_path = _create_sql_file(object_name, object_type, revision) - if sql_path: - click.echo(f" Created {sql_path.relative_to(Path.cwd())}") - else: - click.echo("Warning: Could not extract revision ID, SQL file not created", err=True) + sql_path = _create_sql_file(object_name, object_type, revision) + if sql_path: + click.echo(f" Created {sql_path.relative_to(Path.cwd())}") sys.exit(0) @@ -429,6 +448,64 @@ def _create_sql_file(name: str, object_type: str, revision: str) -> Path | None: return sql_file +def _create_exchange_scaffold( + environment: str, table_name: str, revision: str, alembic_stdout: str +) -> None: + """Create EXCHANGE TABLES migration scaffold. + + Rewrites the alembic-generated migration with the EXCHANGE pattern + and creates a SQL history file for the shadow table. + """ + from clickhouse_alembic.scaffold import ( + fetch_current_ddl, + find_dependent_dictionaries, + generate_exchange_sql, + rewrite_migration_file, + ) + + config_path = Path.cwd() / "config.yaml" + + # Try to connect to live DB for current DDL and dict detection + current_ddl: str | None = None + dict_names: list[str] = [] + try: + env_config = get_env_config(environment, config_path) + current_ddl = fetch_current_ddl(env_config, table_name) + if current_ddl: + click.echo(f" Fetched current DDL for {table_name}") + dict_names = find_dependent_dictionaries(env_config, table_name) + if dict_names: + click.echo(f" Detected dependent dictionaries: {', '.join(dict_names)}") + except Exception: + click.echo(" Note: Could not connect to DB; using placeholder DDL", err=True) + + # Create SQL history file with shadow table DDL + sql_content = generate_exchange_sql(table_name, current_ddl) + sql_path = _create_sql_file(table_name, "table", revision) + if sql_path: + sql_path.write_text(sql_content) + click.echo(f" Created {sql_path.relative_to(Path.cwd())}") + + # Rewrite the migration .py with EXCHANGE pattern + migration_path = _find_migration_file(alembic_stdout) + if migration_path: + rel_sql = str(sql_path.relative_to(Path.cwd() / "migrations" / "sql")) + rewrite_migration_file(migration_path, table_name, rel_sql, dict_names or None) + click.echo(f" Rewrote {migration_path.name} with EXCHANGE TABLES pattern") + else: + click.echo("Warning: Could not locate migration file to rewrite", err=True) + + +def _find_migration_file(alembic_stdout: str) -> Path | None: + """Find the migration .py file path from alembic output.""" + match = re.search(r"Generating (.+?\.py)", alembic_stdout, re.DOTALL) + if not match: + return None + file_path = re.sub(r"\n\s*", "", match.group(1)) + path = Path(file_path) + return path if path.exists() else None + + @main.command() @click.argument("environment") @click.option("--onto", default=None, help="Target revision to rebase onto (skips auto-detection)") diff --git a/clickhouse_alembic/env.py b/clickhouse_alembic/env.py index c3466da..510449f 100644 --- a/clickhouse_alembic/env.py +++ b/clickhouse_alembic/env.py @@ -92,6 +92,10 @@ def get_sqlalchemy_url() -> str: DATABASE_NAME = env_config["database"] os.environ["CH_DATABASE"] = DATABASE_NAME +# Export cluster name if configured (for ON CLUSTER support) +if env_config.get("cluster"): + os.environ["CH_CLUSTER"] = env_config["cluster"] + def bootstrap_version_table(connection: Connection) -> None: """ diff --git a/clickhouse_alembic/helpers.py b/clickhouse_alembic/helpers.py index e695754..def0364 100644 --- a/clickhouse_alembic/helpers.py +++ b/clickhouse_alembic/helpers.py @@ -49,6 +49,38 @@ def get_db() -> str: return os.environ.get("CH_DATABASE", "default") +def get_cluster() -> str | None: + """ + Get the cluster name from environment. + + Returns: + Cluster name from CH_CLUSTER env var, or None if not set + """ + return os.environ.get("CH_CLUSTER") or None + + +def on_cluster() -> str: + """ + Get the ON CLUSTER clause for use in DDL statements. + + This is an opt-in template variable. Not all DDL supports ON CLUSTER + equally — dictionary creation, some ALTER operations, and system queries + have version-dependent ON CLUSTER support. Use this explicitly in + statements where ON CLUSTER is appropriate. + + Returns: + "ON CLUSTER cluster_name" if cluster is configured, empty string otherwise + + Example: + >>> read_sql("tables/users.sql", db=get_db(), on_cluster=on_cluster()) + # In SQL: CREATE TABLE {db}.users {on_cluster} (...) + """ + cluster = get_cluster() + if cluster: + return f"ON CLUSTER {cluster}" + return "" + + def get_config_value(key: str) -> str | None: """ Get a configuration value from environment. diff --git a/clickhouse_alembic/plans/2026-03-05_architectural_review_DRY-13.md b/clickhouse_alembic/plans/2026-03-05_architectural_review_DRY-13.md new file mode 100644 index 0000000..e2825b7 --- /dev/null +++ b/clickhouse_alembic/plans/2026-03-05_architectural_review_DRY-13.md @@ -0,0 +1,223 @@ +# Architectural Review: ch-migrate Feature Proposals + +**Reviewer:** Principal Architect +**Date:** 2026-03-05 +**Design under review:** DRY-11 feature proposals (12 proposed features for ch-migrate) +**Related issues:** [DRY-11](/issues/DRY-11) (feature proposals), [DRY-12](/issues/DRY-12) (design), [DRY-13](/issues/DRY-13) (this review) + +--- + +## Summary + +The feature proposal identifies the right market gap: no tool today handles the full ClickHouse lifecycle in a single Python package. The competitive analysis is accurate and the 12 proposed features are individually reasonable. However, the proposal treats ch-migrate as a monolithic CLI where features are bolted on independently. Several features (diff, snapshot, dependency graph, linting) share a common need to introspect a live ClickHouse database -- yet the proposal designs each as a standalone command. This review recommends introducing an **introspection layer** as a foundational module before building analysis features on top, and reorders priorities accordingly. We also flag architectural constraints in the current codebase that the proposal doesn't address. + +--- + +## Key Recommendations + +### 1. [Blocker] Introduce an introspection module before building analysis features + +**Concern:** Features 1 (diff), 2 (dependency graph), 3 (linting, runtime checks), and 6 (snapshot) all need to query `system.tables`, `system.columns`, `SHOW CREATE TABLE`, and `system.dictionaries`. Without a shared introspection layer, each feature will independently implement system table queries, DDL parsing, and object model construction. + +**Recommendation:** Create `clickhouse_alembic/introspect.py` with: +- `get_live_schema(client, database) -> Schema` -- query system tables, return structured object model +- `parse_create_statement(ddl: str) -> ObjectDefinition` -- normalize ClickHouse DDL for comparison +- `get_dependencies(client, database) -> DependencyGraph` -- query `system.tables` for MV source/target relationships + +This module becomes the foundation for diff, snapshot, lint, and dependency features. Without it, we build the same plumbing four times. + +**Effort:** Medium (but it's amortized across 4 features) + +### 2. [Blocker] Address the Alembic subprocess boundary before adding hooks + +**Concern:** The current `_run_alembic()` in `cli.py:37-92` delegates to Alembic via subprocess. This is a deliberate design choice (isolation, pyenv compat) but it means ch-migrate cannot intercept individual migration execution. Features 3 (lint pre-checks) and 11 (execution hooks) need to run logic before/after each migration, not just before/after the entire `alembic upgrade` command. + +**Recommendation:** For linting, static analysis of SQL files can work without touching the Alembic boundary. But runtime checks (row counts, MV dependency validation) and hooks require either: +- **Option A:** Add an Alembic plugin (`EnvironmentContext.configure` hooks in `env.py`) that calls back into ch-migrate +- **Option B:** Replace subprocess orchestration with in-process Alembic API calls for commands that need hooks + +Option A is simpler and preserves the subprocess isolation for `up`/`down`. We'd extend the generated `env.py` to import and call a hook registry. + +### 3. [Concern] ON CLUSTER support needs a guard list, not blind appending + +**Concern:** The proposal suggests adding `ON CLUSTER {cluster}` config and automatically appending it to DDL. Not all DDL supports ON CLUSTER equally -- dictionary creation, some ALTER operations, and system queries have version-dependent ON CLUSTER support. Blindly appending will produce invalid SQL on older ClickHouse versions. + +**Recommendation:** Implement as a template variable `{on_cluster}` that resolves to `ON CLUSTER {cluster_name}` or empty string. Migration authors use it explicitly in their SQL: +```sql +CREATE TABLE {db}.users {on_cluster} (...) +``` +This is opt-in per statement rather than automatic, avoiding silent breakage. + +### 4. [Concern] PG-to-CH helper should be a separate package, not in ch-migrate core + +**Concern:** Feature 8 (Postgres-to-ClickHouse migration helper) introduces `psycopg2` as a runtime dependency, Postgres schema introspection, type mapping tables, and ENGINE/ORDER BY heuristics. This is a substantial body of code with a different concern (ETL/migration) than schema versioning. + +**Recommendation:** Build as `ch-migrate-from-pg` -- a separate package that depends on ch-migrate for output (generating migration files) but keeps Postgres concerns isolated. This avoids bloating ch-migrate's dependency surface and allows independent release cycles. + +### 5. [Suggestion] Reorder priorities around the introspection layer + +**Current proposed order:** EXCHANGE scaffolding -> Linting -> GitHub Actions -> Snapshot -> Diff -> Dependencies -> ON CLUSTER -> PG helper + +**Recommended order:** + +| Phase | Feature | Rationale | +|-------|---------|-----------| +| 1a | EXCHANGE TABLES scaffolding | Quick win, no new architecture needed | +| 1b | ON CLUSTER template variable | Config change + template variable, small scope | +| 2 | Introspection module | Foundation for 3 subsequent features | +| 3a | Schema snapshot | First consumer of introspection, validates the layer | +| 3b | Schema diff | Second consumer, builds on snapshot | +| 3c | MV dependency graph | Third consumer, highest differentiation | +| 4 | Migration linting (static) | File-based analysis, no DB connection needed | +| 5 | Execution hooks | Requires env.py refactoring | +| 6 | GitHub Actions | Packaging concern, parallel track | +| Separate | PG-to-CH helper | Separate package | + +--- + +## Detailed Findings + +### Feature 1: Schema Diff / Drift Detection + +**Severity: Concern** + +The proposal shows clean CLI output but underestimates the DDL normalization problem. ClickHouse's `SHOW CREATE TABLE` output includes: +- Codec declarations (`CODEC(LZ4HC(9))`) +- TTL expressions with complex date arithmetic +- MATERIALIZED/ALIAS column expressions +- Settings clauses (`SETTINGS index_granularity = 8192`) +- Engine-specific parameters that vary by version + +Comparing raw DDL strings will produce false positives. We need a structured comparison that normalizes whitespace, ordering of columns, and default settings that ClickHouse adds implicitly. + +**Recommendation:** The introspection module should parse DDL into a structured `TableDefinition` (columns, engine, order_by, partition_by, ttl, settings) and compare field-by-field. String-level diff is a fallback for objects we can't parse, not the primary comparison. + +### Feature 2: MV Dependency Graph + +**Severity: Suggestion** + +Good feature, genuine differentiator. Two implementation notes: + +1. **Source of truth is `system.tables`**, not migration files. MVs created outside ch-migrate (manually, by CDC tools) still need to appear in the dependency graph. The proposal implies parsing migration SQL, but live DB introspection is more reliable. + +2. **ClickHouse MVs are insert triggers, not views.** Dropping the source table doesn't break the MV definition -- it breaks data flow. The dependency graph should distinguish between "schema dependency" (MV references table in SELECT) and "data flow dependency" (MV triggers on INSERT to source). The proposal conflates these. + +### Feature 3: Migration Linting + +**Severity: Suggestion** + +The proposal mixes static analysis (can run without DB connection) and runtime checks (needs row counts, MV state). These should be two modes: + +- `ch-migrate lint` -- static analysis of SQL files (no DB needed, fast, CI-friendly) +- `ch-migrate lint dev` -- static + runtime analysis with live DB (needs connection) + +Static rules: missing IF EXISTS, reserved words, destructive DDL detection +Runtime rules: large table mutation warnings, MV dependency checks, ON CLUSTER validation + +This separation matters for CI/CD: static linting runs on every PR without credentials. + +### Feature 4: EXCHANGE TABLES Scaffolding + +**Severity: Note -- on track** + +This is the right first feature. It fits naturally into the existing `_create_sql_file()` pattern at `cli.py:401-429`. Implementation is straightforward: add `--exchange` flag to `ch-migrate new` that generates a multi-step SQL template instead of an empty file. No new architecture needed. + +One addition: the scaffold should include the `SYSTEM RELOAD DICTIONARY` statement when the table being exchanged is a dictionary source. The current `create_dictionary()` helper in `helpers.py:65-106` already parses source tables -- reuse that logic. + +### Feature 5: GitHub Actions + +**Severity: Note** + +This is packaging and CI, not a code feature. It can proceed in parallel on a separate track without touching ch-migrate's architecture. The action itself is a thin wrapper around `ch-migrate bootstrap` + `ch-migrate up` + `ch-migrate lint`. + +Worth noting: the action should use ClickHouse's official Docker image (`clickhouse/clickhouse-server`) and the test should validate both Cloud-compatible (SharedMergeTree) and self-hosted (ReplicatedMergeTree) modes. + +### Feature 6: Schema Snapshot + +**Severity: Suggestion** + +Good feature, natural complement to diff. But the proposal doesn't address: + +1. **Objects created outside ch-migrate** (PeerDB tables, manual DDL, system objects). Snapshot should have a `--filter` or `--exclude` option. + +2. **Snapshot format.** The proposal puts snapshots in `migrations/sql/snapshot/`. This conflicts with the existing object-centric history structure. Snapshots are point-in-time state, not migration history. Recommend `migrations/sql/snapshots//` to avoid confusion. + +### Feature 7: ON CLUSTER Support + +See Key Recommendation #3 above. Template variable approach over automatic appending. + +### Feature 8: PG-to-CH Migration Helper + +See Key Recommendation #4 above. Separate package. + +### Features 9-12 (Declarative Mode, TTL Management, Hooks, TUI) + +These are correctly categorized as future/nice-to-have. Declarative mode (feature 9) is a very large scope change that would require rethinking the Alembic foundation. Only pursue if there's strong user demand. + +Hooks (feature 11) is the most likely to be needed soon -- dictionary reloads after migrations are a common pattern in the Metopio codebase. + +--- + +## Cross-Cutting Concerns + +### Scalability + +The current codebase (2K lines, 10 modules) is well-structured for a single-developer tool. Adding 6-8 features could triple the codebase to 6-8K lines. The proposed features cluster into natural modules: + +| Module | Features | Estimated size | +|--------|----------|---------------| +| `introspect.py` | Shared by diff, snapshot, deps | ~400 lines | +| `diff.py` | Diff + snapshot | ~300 lines | +| `deps.py` | Dependency graph | ~200 lines | +| `lint.py` | Static + runtime linting | ~300 lines | +| `scaffold.py` | EXCHANGE TABLES templates | ~150 lines | + +This keeps each module focused and under 500 lines -- consistent with the current architecture style. + +### Operational Burden + +The introspection-based features (diff, snapshot, deps) add a new failure mode: stale or incomplete system table data. ClickHouse's `system.tables` is eventually consistent in clustered deployments. The introspection module should document this and add a `--node` flag for targeting specific cluster nodes if needed. + +### Security + +No new security concerns. The introspection module uses the existing `migration_user` connection which already has `SELECT` on system tables (granted in `bootstrap.py`). No new credentials or elevated permissions needed. + +### Incremental Deliverability + +The phased approach (EXCHANGE scaffolding -> introspection -> analysis features) allows shipping value at each phase. EXCHANGE scaffolding and ON CLUSTER support ship standalone. Introspection enables the next three features without being user-visible itself. + +### Continuity with Existing Architecture + +The proposed features follow existing patterns: +- CLI commands via Click (consistent with `cli.py`) +- Config-driven behavior via `config.yaml` (consistent with `config.py`) +- SQL file generation (consistent with `_create_sql_file()`) +- Rich formatted output (consistent with `display.py`) + +The one pattern break is the introspection module, which introduces live DB introspection as a first-class concept. Currently, only `status` and `history` read from the DB, and they only touch `alembic_version`. The introspection module queries system tables, which is a new category of DB interaction. + +--- + +## Positive Patterns + +1. **Competitive analysis is thorough and accurate.** The gap identification (no Python tool does schema diffing + MV dependency tracking) is correct and well-argued. + +2. **EXCHANGE TABLES scaffolding as first feature is the right call.** Quick win, teaches best practices through tooling, builds on existing patterns. + +3. **The priority split into tiers is sensible.** Tier 1 features are genuinely high-impact. Tier 3 features are correctly deferred. + +4. **The "lead magnet for consulting agency" framing** is a good product lens. Features 1 (diff), 3 (lint), and 8 (PG helper) have the strongest lead magnet potential. + +5. **The proposal acknowledges that declarative mode is too large** for the current phase. Good restraint. + +--- + +## Open Questions + +1. **[OPEN] Test strategy:** The proposal doesn't mention how to test introspection and diffing features. Do we spin up a ClickHouse container in CI, or mock system tables? This affects the GitHub Actions feature too. + +2. **[OPEN] Plugin architecture:** Should ch-migrate support third-party plugins (e.g., community linting rules, custom scaffold templates)? The proposal hints at extensibility but doesn't commit. A plugin interface adds complexity but enables community contributions. + +3. **[OPEN] Versioning and compatibility:** Features like diff and snapshot need to handle ClickHouse version differences (e.g., SharedMergeTree only exists in Cloud). Should the introspection module detect CH version and adapt, or should users declare their target version in config? + +4. **[OPEN] Existing users:** How do existing ch-migrate users adopt new features? Is there a migration path for projects that already have SQL history files but no snapshot baseline? diff --git a/clickhouse_alembic/scaffold.py b/clickhouse_alembic/scaffold.py new file mode 100644 index 0000000..4014155 --- /dev/null +++ b/clickhouse_alembic/scaffold.py @@ -0,0 +1,246 @@ +"""EXCHANGE TABLES migration scaffold generation.""" + +from __future__ import annotations + +import re +from datetime import datetime +from pathlib import Path +from typing import Any + + +def fetch_current_ddl(env_config: dict[str, Any], table_name: str) -> str | None: + """Fetch current CREATE TABLE statement from the live database. + + Args: + env_config: Environment config dict from get_env_config(). + table_name: Table name to inspect. + + Returns: + DDL string or None if connection fails or table doesn't exist. + """ + from clickhouse_alembic.connection import get_client + + try: + client = get_client(env_config) + db = env_config["database"] + result = client.query(f"SHOW CREATE TABLE {db}.{table_name}") + if result.result_rows: + return result.result_rows[0][0] + except Exception: + return None + return None + + +def find_dependent_dictionaries( + env_config: dict[str, Any], table_name: str +) -> list[str]: + """Find dictionaries that use this table as a source. + + Queries system.dictionaries to find any dictionary whose source + references the given table. + + Args: + env_config: Environment config dict from get_env_config(). + table_name: Table name to check. + + Returns: + List of dictionary names that depend on this table. + """ + from clickhouse_alembic.connection import get_client + + try: + client = get_client(env_config) + db = env_config["database"] + result = client.query( + "SELECT name FROM system.dictionaries " + "WHERE database = {db:String} AND source LIKE {pattern:String}", + parameters={"db": db, "pattern": f"%{table_name}%"}, + ) + return [row[0] for row in result.result_rows] + except Exception: + return [] + + +def _make_shadow_ddl(ddl: str, table_name: str) -> str: + """Transform a CREATE TABLE statement into a shadow table version. + + Replaces the table name with _shadow and adds IF NOT EXISTS. + """ + # Replace table name (handles db.table and just table patterns) + shadow = re.sub( + rf"(CREATE\s+TABLE\s+)(\S+\.)?{re.escape(table_name)}\b", + rf"\g<1>\g<2>{table_name}_shadow", + ddl, + count=1, + flags=re.IGNORECASE, + ) + # Add IF NOT EXISTS if not present + if "IF NOT EXISTS" not in shadow.upper(): + shadow = re.sub( + r"(CREATE\s+TABLE\s+)", + r"\1IF NOT EXISTS ", + shadow, + count=1, + flags=re.IGNORECASE, + ) + return shadow + + +def generate_exchange_sql( + table_name: str, current_ddl: str | None = None +) -> str: + """Generate the SQL file content for an EXCHANGE TABLES migration. + + This creates the shadow table DDL that the user should modify with + their desired schema changes before running the migration. + + Args: + table_name: Name of the table being altered. + current_ddl: Current CREATE TABLE DDL from the database, if available. + + Returns: + SQL file content for the shadow table creation. + """ + if current_ddl: + shadow_ddl = _make_shadow_ddl(current_ddl, table_name) + return ( + f"-- Shadow table for EXCHANGE TABLES migration\n" + f"-- Modify this schema with your desired changes.\n" + f"--\n" + f"-- Original DDL fetched from live database.\n" + f"-- The migration will:\n" + f"-- 1. Create this shadow table\n" + f"-- 2. Copy data from {table_name} into it\n" + f"-- 3. Atomically swap via EXCHANGE TABLES\n" + f"-- 4. Drop the old table\n\n" + f"{shadow_ddl}\n" + ) + + # Placeholder when no live DDL is available + return ( + f"-- Shadow table for EXCHANGE TABLES migration\n" + f"-- Replace this placeholder with your desired schema.\n" + f"--\n" + f"-- TIP: Run `clickhouse-client --query 'SHOW CREATE TABLE {{db}}.{table_name}'`\n" + f"-- to get the current schema, then modify it here.\n\n" + f"CREATE TABLE IF NOT EXISTS {{db}}.{table_name}_shadow\n" + f"(\n" + f" -- TODO: Define columns here\n" + f")\n" + f"ENGINE = MergeTree\n" + f"ORDER BY tuple()\n" + ) + + +def generate_exchange_migration( + revision: str, + down_revision: str | None, + message: str, + table_name: str, + sql_path: str, + dict_names: list[str] | None = None, +) -> str: + """Generate migration .py content with the EXCHANGE TABLES pattern. + + Args: + revision: Alembic revision ID. + down_revision: Previous revision ID. + message: Migration description. + table_name: Table being exchanged. + sql_path: Relative path to the SQL history file (from migrations/sql/). + dict_names: Dictionaries to reload after exchange, if any. + + Returns: + Complete migration .py file content. + """ + down_repr = repr(down_revision) + now = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + dict_lines = "" + if dict_names: + dict_lines = "\n # Reload dependent dictionaries\n" + for d in dict_names: + dict_lines += f' op.execute("SYSTEM RELOAD DICTIONARY {{db}}.{d}")\n' + + return f'''"""{message} + +Revision ID: {revision} +Revises: {down_revision or "None"} +Create Date: {now} + +EXCHANGE TABLES migration for: {table_name} +Steps: CREATE shadow -> INSERT SELECT -> EXCHANGE -> DROP +""" + +from alembic import op + +from clickhouse_alembic import get_db, read_sql + +# revision identifiers +revision = {repr(revision)} +down_revision = {down_repr} +branch_labels = None +depends_on = None + + +def upgrade() -> None: + db = get_db() + + # 1. Create shadow table with new schema + op.execute(read_sql("{sql_path}", db=db)) + + # 2. Copy data from original table into shadow + # NOTE: Modify the SELECT if columns changed (added/removed/renamed) + op.execute(f"INSERT INTO {{db}}.{table_name}_shadow SELECT * FROM {{db}}.{table_name}") + + # 3. Atomically swap tables + op.execute(f"EXCHANGE TABLES {{db}}.{table_name} AND {{db}}.{table_name}_shadow") + + # 4. Drop old table (now named {table_name}_shadow) + op.execute(f"DROP TABLE IF EXISTS {{db}}.{table_name}_shadow") +{dict_lines} + +def downgrade() -> None: + raise NotImplementedError( + "EXCHANGE TABLES migrations cannot be automatically reversed. " + "Create a new forward migration to restore the previous schema." + ) +''' + + +def rewrite_migration_file( + migration_path: Path, + table_name: str, + sql_path: str, + dict_names: list[str] | None = None, +) -> None: + """Rewrite an alembic-generated migration file with EXCHANGE pattern. + + Reads the revision and down_revision from the existing file, then + overwrites it with the EXCHANGE TABLES template. + + Args: + migration_path: Path to the generated migration .py file. + table_name: Table being exchanged. + sql_path: Relative path to the SQL history file. + dict_names: Dictionaries to reload after exchange, if any. + """ + content = migration_path.read_text() + + rev_match = re.search(r'revision\s*=\s*["\'](\w+)["\']', content) + down_match = re.search(r'down_revision\s*=\s*["\'](\w+)["\']', content) + msg_match = re.search(r'^"""(.+?)$', content, re.MULTILINE) + + revision = rev_match.group(1) if rev_match else "UNKNOWN" + down_revision = down_match.group(1) if down_match else None + message = msg_match.group(1) if msg_match else table_name + + new_content = generate_exchange_migration( + revision=revision, + down_revision=down_revision, + message=message, + table_name=table_name, + sql_path=sql_path, + dict_names=dict_names, + ) + migration_path.write_text(new_content) diff --git a/clickhouse_alembic/templates/project/config.yaml.template b/clickhouse_alembic/templates/project/config.yaml.template index 8e1c923..e363fa8 100644 --- a/clickhouse_alembic/templates/project/config.yaml.template +++ b/clickhouse_alembic/templates/project/config.yaml.template @@ -9,6 +9,8 @@ defaults: port: 8443 secure: true admin_user: default + # Optional: cluster name for ON CLUSTER DDL (self-hosted deployments) + # cluster: my_cluster # Optional: dict_reader for dictionary sources # dict_reader_name: dict_reader # Optional: mcp_user for read-only MCP tool access diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 2aa2341..15b3edc 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -4,7 +4,13 @@ import pytest -from clickhouse_alembic.helpers import _parse_source_table, get_db, read_sql +from clickhouse_alembic.helpers import ( + _parse_source_table, + get_cluster, + get_db, + on_cluster, + read_sql, +) class TestReadSql: @@ -49,6 +55,52 @@ def test_returns_default_when_not_set(self, monkeypatch): assert get_db() == "default" +class TestGetCluster: + def test_returns_cluster_from_env(self, monkeypatch): + monkeypatch.setenv("CH_CLUSTER", "my_cluster") + assert get_cluster() == "my_cluster" + + def test_returns_none_when_not_set(self, monkeypatch): + monkeypatch.delenv("CH_CLUSTER", raising=False) + assert get_cluster() is None + + def test_returns_none_for_empty_string(self, monkeypatch): + monkeypatch.setenv("CH_CLUSTER", "") + assert get_cluster() is None + + +class TestOnCluster: + def test_returns_on_cluster_clause(self, monkeypatch): + monkeypatch.setenv("CH_CLUSTER", "my_cluster") + assert on_cluster() == "ON CLUSTER my_cluster" + + def test_returns_empty_string_when_not_set(self, monkeypatch): + monkeypatch.delenv("CH_CLUSTER", raising=False) + assert on_cluster() == "" + + def test_works_in_sql_template(self, tmp_path, monkeypatch): + sql_dir = tmp_path / "migrations" / "sql" + sql_dir.mkdir(parents=True) + sql_file = sql_dir / "test.sql" + sql_file.write_text("CREATE TABLE {db}.users {on_cluster} (id UInt64)") + monkeypatch.chdir(tmp_path) + + monkeypatch.setenv("CH_CLUSTER", "default") + result = read_sql("test.sql", db="mydb", on_cluster=on_cluster()) + assert result == "CREATE TABLE mydb.users ON CLUSTER default (id UInt64)" + + def test_no_cluster_in_sql_template(self, tmp_path, monkeypatch): + sql_dir = tmp_path / "migrations" / "sql" + sql_dir.mkdir(parents=True) + sql_file = sql_dir / "test.sql" + sql_file.write_text("CREATE TABLE {db}.users {on_cluster} (id UInt64)") + monkeypatch.chdir(tmp_path) + + monkeypatch.delenv("CH_CLUSTER", raising=False) + result = read_sql("test.sql", db="mydb", on_cluster=on_cluster()) + assert result == "CREATE TABLE mydb.users (id UInt64)" + + class TestParseSourceTable: def test_parses_table_pattern(self): sql = """ diff --git a/tests/test_scaffold.py b/tests/test_scaffold.py new file mode 100644 index 0000000..692e3d5 --- /dev/null +++ b/tests/test_scaffold.py @@ -0,0 +1,157 @@ +"""Tests for EXCHANGE TABLES scaffold generation.""" + +from pathlib import Path + +import pytest + +from clickhouse_alembic.scaffold import ( + _make_shadow_ddl, + generate_exchange_migration, + generate_exchange_sql, + rewrite_migration_file, +) + + +class TestMakeShadowDdl: + def test_renames_table(self): + ddl = "CREATE TABLE mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + result = _make_shadow_ddl(ddl, "users") + assert "mydb.users_shadow" in result + assert "mydb.users " not in result + + def test_adds_if_not_exists(self): + ddl = "CREATE TABLE mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + result = _make_shadow_ddl(ddl, "users") + assert "IF NOT EXISTS" in result + + def test_preserves_existing_if_not_exists(self): + ddl = "CREATE TABLE IF NOT EXISTS mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + result = _make_shadow_ddl(ddl, "users") + assert result.count("IF NOT EXISTS") == 1 + + def test_handles_no_database_prefix(self): + ddl = "CREATE TABLE users (id UInt64) ENGINE = MergeTree ORDER BY id" + result = _make_shadow_ddl(ddl, "users") + assert "users_shadow" in result + + +class TestGenerateExchangeSql: + def test_with_current_ddl(self): + ddl = "CREATE TABLE mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + result = generate_exchange_sql("users", ddl) + assert "users_shadow" in result + assert "Shadow table for EXCHANGE TABLES migration" in result + assert "Original DDL fetched from live database" in result + + def test_without_current_ddl(self): + result = generate_exchange_sql("users") + assert "users_shadow" in result + assert "TODO: Define columns here" in result + assert "SHOW CREATE TABLE {db}.users" in result + + def test_placeholder_uses_db_variable(self): + result = generate_exchange_sql("users") + assert "{db}.users_shadow" in result + + +class TestGenerateExchangeMigration: + def test_generates_valid_migration(self): + content = generate_exchange_migration( + revision="abc123", + down_revision="def456", + message="alter_users", + table_name="users", + sql_path="history/tables/users/2024_01_01_0000_abc123.sql", + ) + assert "revision = 'abc123'" in content + assert "down_revision = 'def456'" in content + assert "EXCHANGE TABLES" in content + assert "users_shadow" in content + assert "INSERT INTO {db}.users_shadow SELECT * FROM {db}.users" in content + assert "DROP TABLE IF EXISTS {db}.users_shadow" in content + + def test_includes_dict_reload(self): + content = generate_exchange_migration( + revision="abc123", + down_revision="def456", + message="alter_users", + table_name="users", + sql_path="history/tables/users/001_abc123.sql", + dict_names=["dict_users", "dict_user_roles"], + ) + assert "SYSTEM RELOAD DICTIONARY {db}.dict_users" in content + assert "SYSTEM RELOAD DICTIONARY {db}.dict_user_roles" in content + + def test_no_dict_reload_when_empty(self): + content = generate_exchange_migration( + revision="abc123", + down_revision=None, + message="alter_users", + table_name="users", + sql_path="history/tables/users/001_abc123.sql", + ) + assert "SYSTEM RELOAD DICTIONARY" not in content + + def test_downgrade_raises(self): + content = generate_exchange_migration( + revision="abc123", + down_revision="def456", + message="alter_users", + table_name="users", + sql_path="history/tables/users/001_abc123.sql", + ) + assert "NotImplementedError" in content + + def test_none_down_revision(self): + content = generate_exchange_migration( + revision="abc123", + down_revision=None, + message="first_migration", + table_name="users", + sql_path="history/tables/users/001_abc123.sql", + ) + assert "down_revision = None" in content + + +class TestRewriteMigrationFile: + def test_rewrites_migration(self, tmp_path: Path): + migration = tmp_path / "001_abc123.py" + migration.write_text( + '"""alter_users\n\n' + "Revision ID: abc123\n" + "Revises: def456\n" + '"""\n\n' + "revision = 'abc123'\n" + "down_revision = 'def456'\n" + "\n\ndef upgrade():\n pass\n\ndef downgrade():\n pass\n" + ) + + rewrite_migration_file( + migration, + table_name="users", + sql_path="history/tables/users/001_abc123.sql", + ) + + content = migration.read_text() + assert "EXCHANGE TABLES" in content + assert "users_shadow" in content + assert "revision = 'abc123'" in content + assert "down_revision = 'def456'" in content + + def test_rewrites_with_dict_names(self, tmp_path: Path): + migration = tmp_path / "001_abc123.py" + migration.write_text( + '"""alter_users\n\n"""\n\n' + "revision = 'abc123'\n" + "down_revision = 'def456'\n" + ) + + rewrite_migration_file( + migration, + table_name="users", + sql_path="history/tables/users/001_abc123.sql", + dict_names=["dict_users"], + ) + + content = migration.read_text() + assert "SYSTEM RELOAD DICTIONARY {db}.dict_users" in content