From b5b0a67124a16c205c8419e5fedd2517332f750f Mon Sep 17 00:00:00 2001 From: Daniel Young Date: Thu, 5 Mar 2026 17:38:26 -0500 Subject: [PATCH 1/2] [lint] Add migration linting with static and runtime rules Implements `ch-migrate lint` (static-only, CI-friendly) and `ch-migrate lint ` (static + runtime with live DB). Static rules: DestructiveChangeRule, IdempotencyRule, ReservedWordRule, MissingOnClusterRule. Runtime rules: LargeTableMutationRule, MVDependencyRule. Rules are configurable (warn/error/off) via config.yaml lint section. Exit code 0 = clean, 1 = errors found. Solves: DRY-22 Phase 4 migration linting Tests: 44 tests covering all rules, config, and integration --- clickhouse_alembic/cli.py | 62 ++++ clickhouse_alembic/display.py | 55 ++++ clickhouse_alembic/lint.py | 548 ++++++++++++++++++++++++++++++++++ tests/test_lint.py | 458 ++++++++++++++++++++++++++++ 4 files changed, 1123 insertions(+) create mode 100644 clickhouse_alembic/lint.py create mode 100644 tests/test_lint.py diff --git a/clickhouse_alembic/cli.py b/clickhouse_alembic/cli.py index a1d4242..d890068 100644 --- a/clickhouse_alembic/cli.py +++ b/clickhouse_alembic/cli.py @@ -568,5 +568,67 @@ def skill(target: str) -> None: click.echo(f"Installed skill to {skill_dst}") +@main.command() +@click.argument("environment", required=False, default=None) +def lint(environment: str | None) -> None: + """Lint pending migration files. + + Without an environment argument, runs static analysis only (no DB connection, + no credentials needed, CI-friendly). + + With an environment argument, runs static + runtime analysis (connects to the + live database for row counts and dependency checks). + + \b + Examples: + ch-migrate lint # Static only (CI-friendly) + ch-migrate lint dev # Static + runtime (needs DB) + """ + from clickhouse_alembic.config import load_config + from clickhouse_alembic.display import render_lint_report + from clickhouse_alembic.lint import LintConfig, lint_migrations + + versions_dir = Path.cwd() / "migrations" / "versions" + if not versions_dir.exists(): + click.echo("Error: migrations/versions/ not found", err=True) + sys.exit(1) + + config_path = Path.cwd() / "config.yaml" + lint_config = LintConfig() + if config_path.exists(): + try: + raw_config = load_config(config_path) + lint_config = LintConfig.from_config(raw_config) + except Exception: + pass + + client = None + database = None + + if environment: + try: + env_config = get_env_config(environment, config_path) + database = env_config["database"] + + from clickhouse_alembic.connection import get_client + + client = get_client(env_config) + except Exception as e: + click.echo(f"Warning: Could not connect to {environment}: {e}", err=True) + click.echo("Falling back to static-only analysis.", err=True) + click.echo() + + report = lint_migrations( + versions_dir, + config=lint_config, + client=client, + database=database, + ) + + render_lint_report(report, runtime=environment is not None) + + sys.exit(1 if report.has_errors else 0) + + if __name__ == "__main__": main() diff --git a/clickhouse_alembic/display.py b/clickhouse_alembic/display.py index 5f73d9d..0dccca3 100644 --- a/clickhouse_alembic/display.py +++ b/clickhouse_alembic/display.py @@ -201,3 +201,58 @@ def render_status( border_style="blue", ) console.print(panel) + + +def render_lint_report( + report: Any, + *, + runtime: bool = False, + console: Console | None = None, +) -> None: + """Render lint results as a Rich table. + + Args: + report: LintReport with results. + runtime: Whether runtime rules were included. + console: Optional Console for testability. + """ + console = console or Console() + + if not report.results: + mode = "static + runtime" if runtime else "static" + console.print(f"[green]No lint issues found ({mode} analysis).[/green]") + return + + table = Table(title="Lint Results", show_lines=False) + table.add_column("Severity", width=8) + table.add_column("Rule", width=22) + table.add_column("File", width=30) + table.add_column("Line", width=5, justify="right") + table.add_column("Message") + + for r in report.results: + if r.severity.value == "error": + sev_style = "bold red" + sev_text = "ERROR" + else: + sev_style = "yellow" + sev_text = "WARN" + + table.add_row( + Text(sev_text, style=sev_style), + r.rule, + r.file or "", + str(r.line) if r.line else "", + r.message, + ) + + console.print(table) + console.print() + + summary_parts = [] + if report.error_count: + summary_parts.append(f"[bold red]{report.error_count} error(s)[/bold red]") + if report.warning_count: + summary_parts.append(f"[yellow]{report.warning_count} warning(s)[/yellow]") + + console.print(f" {', '.join(summary_parts)}") diff --git a/clickhouse_alembic/lint.py b/clickhouse_alembic/lint.py new file mode 100644 index 0000000..ddd0969 --- /dev/null +++ b/clickhouse_alembic/lint.py @@ -0,0 +1,548 @@ +"""Migration linting: static and runtime analysis rules for ch-migrate.""" + +from __future__ import annotations + +import re +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any + +from clickhouse_alembic.rebase import RevisionGraph, build_revision_graph, parse_migration + + +# --------------------------------------------------------------------------- +# Data models +# --------------------------------------------------------------------------- + + +class Severity(str, Enum): + ERROR = "error" + WARN = "warn" + OFF = "off" + + +@dataclass +class LintResult: + rule: str + message: str + severity: Severity + file: str | None = None + line: int | None = None + + +@dataclass +class LintReport: + results: list[LintResult] = field(default_factory=list) + + @property + def has_errors(self) -> bool: + return any(r.severity == Severity.ERROR for r in self.results) + + @property + def error_count(self) -> int: + return sum(1 for r in self.results if r.severity == Severity.ERROR) + + @property + def warning_count(self) -> int: + return sum(1 for r in self.results if r.severity == Severity.WARN) + + +@dataclass +class LintConfig: + """Lint configuration loaded from config.yaml.""" + + large_table_threshold: int = 100_000_000 + rules: dict[str, Severity] = field(default_factory=dict) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> LintConfig: + lint_section = config.get("lint", {}) + if not lint_section: + return cls() + + threshold = lint_section.get("large_table_threshold", 100_000_000) + rules_raw = lint_section.get("rules", {}) + rules = {} + for name, level in rules_raw.items(): + try: + rules[name] = Severity(level) + except ValueError: + pass + + return cls(large_table_threshold=threshold, rules=rules) + + +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +class LintRule(ABC): + """Base class for lint rules. + + Subclasses implement `check()` which receives migration SQL and context, + returning a list of LintResult. Each rule has a `name` used for config lookup. + """ + + name: str = "" + default_severity: Severity = Severity.WARN + requires_db: bool = False + + def get_severity(self, config: LintConfig) -> Severity: + return config.rules.get(self.name, self.default_severity) + + @abstractmethod + def check( + self, + sql: str, + *, + file_path: str | None = None, + config: LintConfig | None = None, + client: Any | None = None, + database: str | None = None, + graph: RevisionGraph | None = None, + ) -> list[LintResult]: + ... + + +# --------------------------------------------------------------------------- +# ClickHouse reserved words +# --------------------------------------------------------------------------- + +# Subset of CH reserved words that commonly collide with column names. +# Full list is version-dependent; these are the most common traps. +_CH_RESERVED_WORDS = frozenset({ + "add", "after", "alias", "all", "alter", "and", "anti", "any", "array", + "as", "asc", "attach", "between", "both", "by", "case", "cast", "check", + "cluster", "collate", "column", "comment", "constraint", "create", + "cross", "cube", "current", "database", "databases", "date", "day", + "default", "delete", "desc", "describe", "detach", "dictionaries", + "dictionary", "distinct", "distributed", "drop", "else", "end", "engine", + "events", "except", "exists", "explain", "expression", "extract", "fetch", + "final", "first", "flush", "following", "for", "format", "from", "full", + "function", "global", "granularity", "group", "having", "hour", "if", + "ilike", "in", "index", "inject", "inner", "insert", "interval", "into", + "is", "join", "key", "kill", "last", "layout", "leading", "left", "like", + "limit", "live", "local", "logs", "materialize", "materialized", "max", + "merges", "min", "minute", "modify", "month", "move", "mutation", "no", + "not", "null", "nulls", "offset", "on", "optimize", "or", "order", + "outer", "outfile", "over", "partition", "populate", "preceding", + "primary", "prewhere", "projection", "quarter", "range", "reload", + "remove", "rename", "replace", "right", "rollup", "row", "rows", + "sample", "second", "select", "semi", "set", "settings", "show", + "source", "start", "stop", "system", "table", "tables", "temporary", + "test", "then", "ties", "timestamp", "to", "top", "totals", "trailing", + "trim", "truncate", "type", "unbounded", "union", "update", "use", + "using", "uuid", "values", "view", "volume", "watch", "week", "when", + "where", "window", "with", "year", +}) + + +# --------------------------------------------------------------------------- +# Static rules (no DB connection needed) +# --------------------------------------------------------------------------- + + +class DestructiveChangeRule(LintRule): + """Flags DROP TABLE and DROP COLUMN statements.""" + + name = "destructive_changes" + default_severity = Severity.WARN + + _RE_DROP_TABLE = re.compile( + r"\bDROP\s+TABLE\b", re.IGNORECASE + ) + _RE_DROP_COLUMN = re.compile( + r"\bDROP\s+COLUMN\b", re.IGNORECASE + ) + + def check(self, sql: str, **kwargs: Any) -> list[LintResult]: + config = kwargs.get("config") or LintConfig() + severity = self.get_severity(config) + if severity == Severity.OFF: + return [] + + results: list[LintResult] = [] + file_path = kwargs.get("file_path") + + for match in self._RE_DROP_TABLE.finditer(sql): + line = sql[:match.start()].count("\n") + 1 + results.append(LintResult( + rule=self.name, + message="DROP TABLE is destructive and irreversible", + severity=severity, + file=file_path, + line=line, + )) + + for match in self._RE_DROP_COLUMN.finditer(sql): + line = sql[:match.start()].count("\n") + 1 + results.append(LintResult( + rule=self.name, + message="DROP COLUMN is destructive and irreversible", + severity=severity, + file=file_path, + line=line, + )) + + return results + + +class IdempotencyRule(LintRule): + """Flags CREATE/DROP without IF EXISTS / IF NOT EXISTS.""" + + name = "idempotency" + default_severity = Severity.WARN + + _RE_CREATE_NO_IF = re.compile( + r"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW|MATERIALIZED\s+VIEW|DICTIONARY)\s+" + r"(?!IF\s+NOT\s+EXISTS\b)", + re.IGNORECASE, + ) + _RE_DROP_NO_IF = re.compile( + r"\bDROP\s+(?:TABLE|VIEW|DICTIONARY)\s+(?!IF\s+EXISTS\b)", + re.IGNORECASE, + ) + + def check(self, sql: str, **kwargs: Any) -> list[LintResult]: + config = kwargs.get("config") or LintConfig() + severity = self.get_severity(config) + if severity == Severity.OFF: + return [] + + results: list[LintResult] = [] + file_path = kwargs.get("file_path") + + for match in self._RE_CREATE_NO_IF.finditer(sql): + # Skip CREATE OR REPLACE (already idempotent) + matched_text = match.group(0) + if re.search(r"OR\s+REPLACE", matched_text, re.IGNORECASE): + continue + line = sql[:match.start()].count("\n") + 1 + results.append(LintResult( + rule=self.name, + message="CREATE without IF NOT EXISTS is not idempotent", + severity=severity, + file=file_path, + line=line, + )) + + for match in self._RE_DROP_NO_IF.finditer(sql): + line = sql[:match.start()].count("\n") + 1 + results.append(LintResult( + rule=self.name, + message="DROP without IF EXISTS is not idempotent", + severity=severity, + file=file_path, + line=line, + )) + + return results + + +class ReservedWordRule(LintRule): + """Flags column names that are ClickHouse reserved words.""" + + name = "reserved_words" + default_severity = Severity.WARN + + _RE_COLUMN_DEF = re.compile( + r"^\s+`?(\w+)`?\s+(?:Nullable|UInt|Int|Float|String|Date|Array|Tuple|Map|Bool|Enum)", + re.IGNORECASE | re.MULTILINE, + ) + + def check(self, sql: str, **kwargs: Any) -> list[LintResult]: + config = kwargs.get("config") or LintConfig() + severity = self.get_severity(config) + if severity == Severity.OFF: + return [] + + results: list[LintResult] = [] + file_path = kwargs.get("file_path") + + for match in self._RE_COLUMN_DEF.finditer(sql): + col_name = match.group(1) + if col_name.lower() in _CH_RESERVED_WORDS: + line = sql[:match.start()].count("\n") + 1 + results.append(LintResult( + rule=self.name, + message=f"Column '{col_name}' is a ClickHouse reserved word", + severity=severity, + file=file_path, + line=line, + )) + + return results + + +class MissingOnClusterRule(LintRule): + """Flags DDL without {on_cluster} when cluster is configured.""" + + name = "missing_on_cluster" + default_severity = Severity.OFF # Off by default — only relevant for clustered setups + + _RE_DDL = re.compile( + r"\b(CREATE|ALTER|DROP)\s+(?:OR\s+REPLACE\s+)?" + r"(?:TABLE|VIEW|MATERIALIZED\s+VIEW|DICTIONARY)\b", + re.IGNORECASE, + ) + + def check(self, sql: str, **kwargs: Any) -> list[LintResult]: + config = kwargs.get("config") or LintConfig() + severity = self.get_severity(config) + if severity == Severity.OFF: + return [] + + results: list[LintResult] = [] + file_path = kwargs.get("file_path") + + for match in self._RE_DDL.finditer(sql): + # Check if ON CLUSTER or {on_cluster} appears nearby + rest = sql[match.end():match.end() + 200] + if not re.search(r"(?:ON\s+CLUSTER|{on_cluster})", rest, re.IGNORECASE): + line = sql[:match.start()].count("\n") + 1 + stmt_type = match.group(0).strip() + results.append(LintResult( + rule=self.name, + message=f"{stmt_type} without ON CLUSTER or {{on_cluster}} placeholder", + severity=severity, + file=file_path, + line=line, + )) + + return results + + +# --------------------------------------------------------------------------- +# Runtime rules (require DB connection) +# --------------------------------------------------------------------------- + + +class LargeTableMutationRule(LintRule): + """Flags ALTER on tables above a configurable row threshold.""" + + name = "large_table_mutation" + default_severity = Severity.WARN + requires_db = True + + _RE_ALTER_TABLE = re.compile( + r"\bALTER\s+TABLE\s+(?:`?(\w+)`?\.)?`?(\w+)`?", + re.IGNORECASE, + ) + + def check(self, sql: str, **kwargs: Any) -> list[LintResult]: + config = kwargs.get("config") or LintConfig() + severity = self.get_severity(config) + if severity == Severity.OFF: + return [] + + client = kwargs.get("client") + database = kwargs.get("database") + if not client or not database: + return [] + + results: list[LintResult] = [] + file_path = kwargs.get("file_path") + threshold = config.large_table_threshold + + for match in self._RE_ALTER_TABLE.finditer(sql): + db = match.group(1) or database + table_name = match.group(2) + try: + result = client.query( + "SELECT count() FROM system.parts " + "WHERE database = {db:String} AND table = {tbl:String} AND active", + parameters={"db": db, "tbl": table_name}, + ) + if result.result_rows: + row_count = result.result_rows[0][0] + if row_count > threshold: + line = sql[:match.start()].count("\n") + 1 + results.append(LintResult( + rule=self.name, + message=( + f"ALTER on '{table_name}' which has {row_count:,} parts " + f"(threshold: {threshold:,})" + ), + severity=severity, + file=file_path, + line=line, + )) + except Exception: + pass + + return results + + +class MVDependencyRule(LintRule): + """Flags operations on tables that have materialized view dependencies.""" + + name = "mv_dependency" + default_severity = Severity.WARN + requires_db = True + + _RE_DROP_TABLE = re.compile( + r"\bDROP\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:`?(\w+)`?\.)?`?(\w+)`?", + re.IGNORECASE, + ) + _RE_ALTER_TABLE = re.compile( + r"\bALTER\s+TABLE\s+(?:`?(\w+)`?\.)?`?(\w+)`?", + re.IGNORECASE, + ) + + def check(self, sql: str, **kwargs: Any) -> list[LintResult]: + config = kwargs.get("config") or LintConfig() + severity = self.get_severity(config) + if severity == Severity.OFF: + return [] + + client = kwargs.get("client") + database = kwargs.get("database") + if not client or not database: + return [] + + results: list[LintResult] = [] + file_path = kwargs.get("file_path") + + from clickhouse_alembic.introspect import get_dependencies + + try: + dep_graph = get_dependencies(client, database) + except Exception: + return [] + + tables_to_check: list[tuple[str, re.Match[str]]] = [] + for match in self._RE_DROP_TABLE.finditer(sql): + tables_to_check.append((match.group(2), match)) + for match in self._RE_ALTER_TABLE.finditer(sql): + tables_to_check.append((match.group(2), match)) + + for table_name, match in tables_to_check: + affected = dep_graph.affected_by_drop(table_name) + if affected: + mv_names = [ + n.name for n in affected if n.obj_type == "materialized_view" + ] + dict_names = [ + n.name for n in affected if n.obj_type == "dictionary" + ] + if mv_names or dict_names: + line = sql[:match.start()].count("\n") + 1 + deps = [] + if mv_names: + deps.append(f"MVs: {', '.join(mv_names)}") + if dict_names: + deps.append(f"Dicts: {', '.join(dict_names)}") + results.append(LintResult( + rule=self.name, + message=( + f"'{table_name}' has dependent objects: {'; '.join(deps)}" + ), + severity=severity, + file=file_path, + line=line, + )) + + return results + + +# --------------------------------------------------------------------------- +# Rule registry +# --------------------------------------------------------------------------- + +STATIC_RULES: list[LintRule] = [ + DestructiveChangeRule(), + IdempotencyRule(), + ReservedWordRule(), + MissingOnClusterRule(), +] + +RUNTIME_RULES: list[LintRule] = [ + LargeTableMutationRule(), + MVDependencyRule(), +] + +ALL_RULES: list[LintRule] = STATIC_RULES + RUNTIME_RULES + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def _read_migration_sql(migration_path: Path) -> str: + """Read SQL from a migration file's upgrade function. + + Extracts SQL string literals from read_sql() calls and op.execute() calls. + Falls back to reading any associated .sql files. + """ + content = migration_path.read_text() + sql_parts: list[str] = [] + + # Extract string literals from op.execute(...) calls + for match in re.finditer(r'op\.execute\(\s*(?:f?"""(.*?)"""|f?"([^"]*)")', content, re.DOTALL): + sql_parts.append(match.group(1) or match.group(2) or "") + + # Extract paths from read_sql(...) calls and try to read them + for match in re.finditer(r'read_sql\(\s*["\']([^"\']+)["\']', content): + sql_path = migration_path.parent.parent / "sql" / match.group(1) + if sql_path.exists(): + sql_parts.append(sql_path.read_text()) + + return "\n".join(sql_parts) if sql_parts else content + + +def lint_migrations( + versions_dir: Path, + *, + config: LintConfig | None = None, + client: Any | None = None, + database: str | None = None, +) -> LintReport: + """Run lint rules against pending migration files. + + Args: + versions_dir: Path to migrations/versions/ directory. + config: Lint configuration. Defaults to LintConfig(). + client: Optional clickhouse-connect client for runtime rules. + database: Database name for runtime rules. + + Returns: + LintReport with all findings. + """ + if config is None: + config = LintConfig() + + graph = build_revision_graph(versions_dir) + report = LintReport() + + rules = list(STATIC_RULES) + if client is not None: + rules.extend(RUNTIME_RULES) + + for migration in graph.migrations.values(): + sql = _read_migration_sql(migration.path) + if not sql.strip(): + continue + + file_path = str(migration.path.name) + for rule in rules: + severity = rule.get_severity(config) + if severity == Severity.OFF: + continue + if rule.requires_db and client is None: + continue + + findings = rule.check( + sql, + file_path=file_path, + config=config, + client=client, + database=database, + graph=graph, + ) + report.results.extend(findings) + + return report diff --git a/tests/test_lint.py b/tests/test_lint.py new file mode 100644 index 0000000..d120466 --- /dev/null +++ b/tests/test_lint.py @@ -0,0 +1,458 @@ +"""Tests for migration linting rules.""" + +from __future__ import annotations + +import textwrap +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from clickhouse_alembic.lint import ( + ALL_RULES, + RUNTIME_RULES, + STATIC_RULES, + DestructiveChangeRule, + IdempotencyRule, + LargeTableMutationRule, + LintConfig, + LintReport, + LintResult, + MissingOnClusterRule, + MVDependencyRule, + ReservedWordRule, + Severity, + lint_migrations, +) + + +# --------------------------------------------------------------------------- +# LintConfig tests +# --------------------------------------------------------------------------- + + +class TestLintConfig: + def test_defaults(self): + config = LintConfig() + assert config.large_table_threshold == 100_000_000 + assert config.rules == {} + + def test_from_config_empty(self): + config = LintConfig.from_config({}) + assert config.large_table_threshold == 100_000_000 + + def test_from_config_with_values(self): + config = LintConfig.from_config({ + "lint": { + "large_table_threshold": 50_000, + "rules": { + "destructive_changes": "error", + "missing_on_cluster": "off", + }, + } + }) + assert config.large_table_threshold == 50_000 + assert config.rules["destructive_changes"] == Severity.ERROR + assert config.rules["missing_on_cluster"] == Severity.OFF + + def test_from_config_invalid_severity_ignored(self): + config = LintConfig.from_config({ + "lint": { + "rules": {"destructive_changes": "invalid_value"}, + } + }) + assert "destructive_changes" not in config.rules + + +# --------------------------------------------------------------------------- +# LintReport tests +# --------------------------------------------------------------------------- + + +class TestLintReport: + def test_empty_report(self): + report = LintReport() + assert not report.has_errors + assert report.error_count == 0 + assert report.warning_count == 0 + + def test_report_with_errors(self): + report = LintReport(results=[ + LintResult(rule="test", message="bad", severity=Severity.ERROR), + LintResult(rule="test", message="meh", severity=Severity.WARN), + ]) + assert report.has_errors + assert report.error_count == 1 + assert report.warning_count == 1 + + +# --------------------------------------------------------------------------- +# DestructiveChangeRule tests +# --------------------------------------------------------------------------- + + +class TestDestructiveChangeRule: + def test_flags_drop_table(self): + sql = "DROP TABLE mydb.users" + results = DestructiveChangeRule().check(sql) + assert len(results) == 1 + assert "DROP TABLE" in results[0].message + assert results[0].severity == Severity.WARN + + def test_flags_drop_column(self): + sql = "ALTER TABLE mydb.users DROP COLUMN email" + results = DestructiveChangeRule().check(sql) + assert len(results) == 1 + assert "DROP COLUMN" in results[0].message + + def test_no_flags_on_safe_sql(self): + sql = "CREATE TABLE mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + results = DestructiveChangeRule().check(sql) + assert results == [] + + def test_multiple_drops(self): + sql = textwrap.dedent("""\ + DROP TABLE mydb.old_events; + ALTER TABLE mydb.users DROP COLUMN phone; + """) + results = DestructiveChangeRule().check(sql) + assert len(results) == 2 + + def test_respects_severity_off(self): + config = LintConfig(rules={"destructive_changes": Severity.OFF}) + results = DestructiveChangeRule().check("DROP TABLE foo", config=config) + assert results == [] + + def test_respects_severity_error(self): + config = LintConfig(rules={"destructive_changes": Severity.ERROR}) + results = DestructiveChangeRule().check("DROP TABLE foo", config=config) + assert len(results) == 1 + assert results[0].severity == Severity.ERROR + + def test_reports_line_number(self): + sql = "SELECT 1;\nSELECT 2;\nDROP TABLE foo;" + results = DestructiveChangeRule().check(sql) + assert results[0].line == 3 + + +# --------------------------------------------------------------------------- +# IdempotencyRule tests +# --------------------------------------------------------------------------- + + +class TestIdempotencyRule: + def test_flags_create_without_if_not_exists(self): + sql = "CREATE TABLE mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + results = IdempotencyRule().check(sql) + assert len(results) == 1 + assert "IF NOT EXISTS" in results[0].message + + def test_passes_with_if_not_exists(self): + sql = "CREATE TABLE IF NOT EXISTS mydb.users (id UInt64) ENGINE = MergeTree ORDER BY id" + results = IdempotencyRule().check(sql) + assert results == [] + + def test_passes_with_or_replace(self): + sql = "CREATE OR REPLACE DICTIONARY mydb.dict_foo (key String) PRIMARY KEY key" + results = IdempotencyRule().check(sql) + assert results == [] + + def test_flags_drop_without_if_exists(self): + sql = "DROP TABLE mydb.users" + results = IdempotencyRule().check(sql) + assert len(results) == 1 + assert "IF EXISTS" in results[0].message + + def test_passes_drop_with_if_exists(self): + sql = "DROP TABLE IF EXISTS mydb.users" + results = IdempotencyRule().check(sql) + assert results == [] + + def test_flags_create_view_without_if_not_exists(self): + sql = "CREATE VIEW mydb.v AS SELECT 1" + results = IdempotencyRule().check(sql) + assert len(results) == 1 + + def test_flags_create_materialized_view(self): + sql = "CREATE MATERIALIZED VIEW mydb.mv TO mydb.dest AS SELECT 1 FROM mydb.src" + results = IdempotencyRule().check(sql) + assert len(results) == 1 + + def test_passes_create_mv_if_not_exists(self): + sql = "CREATE MATERIALIZED VIEW IF NOT EXISTS mydb.mv TO mydb.dest AS SELECT 1" + results = IdempotencyRule().check(sql) + assert results == [] + + +# --------------------------------------------------------------------------- +# ReservedWordRule tests +# --------------------------------------------------------------------------- + + +class TestReservedWordRule: + def test_flags_reserved_column_name(self): + sql = textwrap.dedent("""\ + CREATE TABLE mydb.t ( + `id` UInt64, + `key` String, + `select` String + ) + """) + results = ReservedWordRule().check(sql) + reserved_names = {r.message.split("'")[1] for r in results} + assert "key" in reserved_names + assert "select" in reserved_names + + def test_passes_non_reserved_names(self): + sql = textwrap.dedent("""\ + CREATE TABLE mydb.t ( + `user_id` UInt64, + `event_name` String + ) + """) + results = ReservedWordRule().check(sql) + assert results == [] + + def test_case_insensitive(self): + sql = " `KEY` String" + results = ReservedWordRule().check(sql) + assert len(results) == 1 + + +# --------------------------------------------------------------------------- +# MissingOnClusterRule tests +# --------------------------------------------------------------------------- + + +class TestMissingOnClusterRule: + def test_off_by_default(self): + sql = "CREATE TABLE mydb.t (id UInt64) ENGINE = MergeTree ORDER BY id" + results = MissingOnClusterRule().check(sql) + assert results == [] + + def test_flags_when_enabled(self): + config = LintConfig(rules={"missing_on_cluster": Severity.WARN}) + sql = "CREATE TABLE mydb.t (id UInt64) ENGINE = MergeTree ORDER BY id" + results = MissingOnClusterRule().check(sql, config=config) + assert len(results) == 1 + assert "ON CLUSTER" in results[0].message + + def test_passes_with_on_cluster(self): + config = LintConfig(rules={"missing_on_cluster": Severity.WARN}) + sql = "CREATE TABLE mydb.t ON CLUSTER default (id UInt64) ENGINE = MergeTree ORDER BY id" + results = MissingOnClusterRule().check(sql, config=config) + assert results == [] + + def test_passes_with_placeholder(self): + config = LintConfig(rules={"missing_on_cluster": Severity.WARN}) + sql = "CREATE TABLE mydb.t {on_cluster} (id UInt64) ENGINE = MergeTree ORDER BY id" + results = MissingOnClusterRule().check(sql, config=config) + assert results == [] + + def test_flags_alter_and_drop(self): + config = LintConfig(rules={"missing_on_cluster": Severity.WARN}) + sql = textwrap.dedent("""\ + ALTER TABLE mydb.t ADD COLUMN foo String; + DROP TABLE mydb.t; + """) + results = MissingOnClusterRule().check(sql, config=config) + assert len(results) == 2 + + +# --------------------------------------------------------------------------- +# LargeTableMutationRule tests +# --------------------------------------------------------------------------- + + +class TestLargeTableMutationRule: + def _make_client(self, row_count: int) -> MagicMock: + client = MagicMock() + result = MagicMock() + result.result_rows = [[row_count]] + client.query.return_value = result + return client + + def test_flags_large_table(self): + client = self._make_client(200_000_000) + sql = "ALTER TABLE mydb.users ADD COLUMN phone String" + results = LargeTableMutationRule().check( + sql, client=client, database="mydb" + ) + assert len(results) == 1 + assert "200,000,000" in results[0].message + + def test_passes_small_table(self): + client = self._make_client(1000) + sql = "ALTER TABLE mydb.users ADD COLUMN phone String" + results = LargeTableMutationRule().check( + sql, client=client, database="mydb" + ) + assert results == [] + + def test_respects_custom_threshold(self): + client = self._make_client(5000) + config = LintConfig(large_table_threshold=1000) + sql = "ALTER TABLE mydb.users ADD COLUMN phone String" + results = LargeTableMutationRule().check( + sql, client=client, database="mydb", config=config + ) + assert len(results) == 1 + + def test_skips_without_client(self): + sql = "ALTER TABLE mydb.users ADD COLUMN phone String" + results = LargeTableMutationRule().check(sql) + assert results == [] + + +# --------------------------------------------------------------------------- +# MVDependencyRule tests +# --------------------------------------------------------------------------- + + +class TestMVDependencyRule: + def _make_client_with_deps(self) -> MagicMock: + """Create a mock client that returns a dependency graph with MV on events.""" + from clickhouse_alembic.introspect import ( + DepType, + DependencyEdge, + DependencyGraph, + ObjectNode, + ) + + graph = DependencyGraph() + graph.nodes = { + "events": ObjectNode(name="events", obj_type="table"), + "hourly_mv": ObjectNode(name="hourly_mv", obj_type="materialized_view"), + } + graph.edges = [ + DependencyEdge(source="events", target="hourly_mv", dep_type=DepType.SCHEMA), + ] + + client = MagicMock() + # Patch get_dependencies to return our mock graph + return client, graph + + def test_flags_drop_on_mv_source(self): + from unittest.mock import patch + + client, graph = self._make_client_with_deps() + sql = "DROP TABLE IF EXISTS events" + + with patch("clickhouse_alembic.introspect.get_dependencies", return_value=graph): + results = MVDependencyRule().check( + sql, client=client, database="mydb" + ) + + assert len(results) == 1 + assert "hourly_mv" in results[0].message + + def test_passes_on_unrelated_table(self): + from unittest.mock import patch + + client, graph = self._make_client_with_deps() + sql = "DROP TABLE IF EXISTS unrelated_table" + + with patch("clickhouse_alembic.introspect.get_dependencies", return_value=graph): + results = MVDependencyRule().check( + sql, client=client, database="mydb" + ) + + assert results == [] + + def test_skips_without_client(self): + sql = "DROP TABLE IF EXISTS events" + results = MVDependencyRule().check(sql) + assert results == [] + + +# --------------------------------------------------------------------------- +# Rule registry tests +# --------------------------------------------------------------------------- + + +class TestRuleRegistry: + def test_static_rules_dont_require_db(self): + for rule in STATIC_RULES: + assert not rule.requires_db, f"{rule.name} should not require DB" + + def test_runtime_rules_require_db(self): + for rule in RUNTIME_RULES: + assert rule.requires_db, f"{rule.name} should require DB" + + def test_all_rules_have_names(self): + for rule in ALL_RULES: + assert rule.name, f"Rule {type(rule).__name__} missing name" + + def test_all_rule_names_unique(self): + names = [r.name for r in ALL_RULES] + assert len(names) == len(set(names)) + + +# --------------------------------------------------------------------------- +# Integration: lint_migrations with temp files +# --------------------------------------------------------------------------- + + +class TestLintMigrations: + def _create_migration(self, tmp_path: Path, name: str, sql_content: str) -> Path: + """Create a migration file with embedded SQL in op.execute().""" + versions_dir = tmp_path / "versions" + versions_dir.mkdir(exist_ok=True) + + rev_id = name[:12].ljust(12, "0") + content = textwrap.dedent(f"""\ + \"\"\"Migration {name} + + Revision ID: {rev_id} + Revises: + Create Date: 2024-01-01 + + \"\"\" + from alembic import op + + revision = '{rev_id}' + down_revision = None + + def upgrade(): + op.execute(\"\"\"{sql_content}\"\"\") + + def downgrade(): + pass + """) + + file_path = versions_dir / f"{rev_id}_{name}.py" + file_path.write_text(content) + return versions_dir + + def test_static_lint_finds_issues(self, tmp_path: Path): + versions_dir = self._create_migration( + tmp_path, "drop_users", "DROP TABLE mydb.users" + ) + report = lint_migrations(versions_dir) + assert report.warning_count > 0 + + def test_static_lint_clean(self, tmp_path: Path): + versions_dir = self._create_migration( + tmp_path, + "safe_migration", + "CREATE TABLE IF NOT EXISTS mydb.t (id UInt64) ENGINE = MergeTree ORDER BY id", + ) + report = lint_migrations(versions_dir) + assert report.error_count == 0 + assert report.warning_count == 0 + + def test_lint_with_error_severity(self, tmp_path: Path): + versions_dir = self._create_migration( + tmp_path, "drop_bad", "DROP TABLE mydb.users" + ) + config = LintConfig(rules={"destructive_changes": Severity.ERROR}) + report = lint_migrations(versions_dir, config=config) + assert report.has_errors + assert report.error_count >= 1 + + def test_lint_empty_versions_dir(self, tmp_path: Path): + versions_dir = tmp_path / "versions" + versions_dir.mkdir() + report = lint_migrations(versions_dir) + assert not report.has_errors + assert report.results == [] From c5b59704f1b23809eb5c5f544dc55792bc95e7ba Mon Sep 17 00:00:00 2001 From: Daniel Young Date: Fri, 6 Mar 2026 10:07:29 -0500 Subject: [PATCH 2/2] [lint] Add dependency, diff, and hook analysis modules Adds three analysis modules (deps, diff, hooks) supporting the migration linting framework, along with CLI commands and display integration. Includes tests for all three modules plus snapshot testing. Solves: Core linting infrastructure for static and runtime analysis Tests: test_deps, test_diff, test_hooks, test_snapshot --- clickhouse_alembic/cli.py | 310 +++++++++++++++++ clickhouse_alembic/config.py | 5 + clickhouse_alembic/deps.py | 124 +++++++ clickhouse_alembic/diff.py | 225 ++++++++++++ clickhouse_alembic/display.py | 201 +++++++++++ clickhouse_alembic/env.py | 37 +- clickhouse_alembic/hooks.py | 71 ++++ .../templates/project/config.yaml.template | 9 + tests/test_deps.py | 224 ++++++++++++ tests/test_diff.py | 323 ++++++++++++++++++ tests/test_hooks.py | 242 +++++++++++++ tests/test_snapshot.py | 239 +++++++++++++ 12 files changed, 2009 insertions(+), 1 deletion(-) create mode 100644 clickhouse_alembic/deps.py create mode 100644 clickhouse_alembic/diff.py create mode 100644 clickhouse_alembic/hooks.py create mode 100644 tests/test_deps.py create mode 100644 tests/test_diff.py create mode 100644 tests/test_hooks.py create mode 100644 tests/test_snapshot.py diff --git a/clickhouse_alembic/cli.py b/clickhouse_alembic/cli.py index d890068..0bc7464 100644 --- a/clickhouse_alembic/cli.py +++ b/clickhouse_alembic/cli.py @@ -630,5 +630,315 @@ def lint(environment: str | None) -> None: sys.exit(1 if report.has_errors else 0) +@main.command() +@click.argument("environment") +@click.option("--validate", "-v", "validate_sql", type=click.Path(exists=True), + help="Validate a SQL file against the dependency graph") +def deps(environment: str, validate_sql: str | None) -> None: + """Show materialized view and dictionary dependency graph. + + Queries the live database to build a dependency graph of all tables, + views, materialized views, and dictionaries, then renders it as a tree. + + Use --validate to check if a SQL file would break any dependencies. + + \b + Examples: + ch-migrate deps dev + ch-migrate deps dev --validate migrations/sql/history/tables/users/drop.sql + """ + from clickhouse_alembic.connection import get_client + from clickhouse_alembic.deps import build_dependency_graph, validate_migration + from clickhouse_alembic.display import render_dependency_tree + + config_path = Path.cwd() / "config.yaml" + try: + env_config = get_env_config(environment, config_path) + except Exception as e: + click.echo(f"Error loading config: {e}", err=True) + sys.exit(1) + + database = env_config["database"] + + try: + client = get_client(env_config) + except Exception as e: + click.echo(f"Error connecting to {environment}: {e}", err=True) + sys.exit(1) + + click.echo(f"Building dependency graph for {environment} ({database})...") + + try: + graph = build_dependency_graph(client, database) + except Exception as e: + click.echo(f"Error building dependency graph: {e}", err=True) + sys.exit(1) + + render_dependency_tree(graph) + + if validate_sql: + sql_content = Path(validate_sql).read_text() + warnings = validate_migration(sql_content, graph) + if warnings: + click.echo() + for w in warnings: + style = "red" if w.severity == "error" else "yellow" + click.echo(click.style(f" [{w.severity.upper()}] {w.message}", fg=style)) + sys.exit(1) + else: + click.echo(click.style("\n Migration validation passed.", fg="green")) + + +@main.command(name="diff") +@click.argument("environment") +@click.option( + "--snapshot-dir", + "-s", + type=click.Path(exists=True), + help="Path to a snapshot directory to compare against. Defaults to latest snapshot.", +) +def diff_cmd(environment: str, snapshot_dir: str | None) -> None: + """Detect schema drift between local snapshot and live database. + + Compares the most recent snapshot (or a specified one) against the live + database schema. Exit code 0 if in sync, 1 if drift detected. + + \b + Examples: + ch-migrate diff dev + ch-migrate diff dev --snapshot-dir migrations/sql/snapshots/20260305_120000 + """ + from clickhouse_alembic.connection import get_client + from clickhouse_alembic.diff import DiffStatus, compare_schemas + from clickhouse_alembic.display import render_diff_report + from clickhouse_alembic.introspect import Schema, get_live_schema, parse_create_statement + + config_path = Path.cwd() / "config.yaml" + try: + env_config = get_env_config(environment, config_path) + except Exception as e: + click.echo(f"Error loading config: {e}", err=True) + sys.exit(1) + + database = env_config["database"] + + # Resolve snapshot directory + if snapshot_dir: + snap_path = Path(snapshot_dir) + else: + snapshots_base = Path.cwd() / "migrations" / "sql" / "snapshots" + if not snapshots_base.exists(): + click.echo("Error: No snapshots found. Run 'ch-migrate snapshot' first.", err=True) + sys.exit(1) + dirs = sorted(snapshots_base.iterdir()) + if not dirs: + click.echo("Error: No snapshots found. Run 'ch-migrate snapshot' first.", err=True) + sys.exit(1) + snap_path = dirs[-1] + + click.echo(f"Comparing snapshot {snap_path.name} against live {environment} ({database})...") + + # Load local schema from snapshot files + local_schema = Schema(database=database) + type_dirs = { + "tables": "table", + "views": "view", + "materialized_views": "materialized_view", + "dictionaries": "dictionary", + } + schema_attrs = { + "table": local_schema.tables, + "view": local_schema.views, + "materialized_view": local_schema.materialized_views, + "dictionary": local_schema.dictionaries, + } + + for dir_name, obj_type in type_dirs.items(): + type_path = snap_path / dir_name + if not type_path.exists(): + continue + for sql_file in sorted(type_path.glob("*.sql")): + ddl = sql_file.read_text() + name = sql_file.stem + parsed = parse_create_statement(ddl) + if parsed: + schema_attrs[obj_type][name] = parsed + else: + # Store minimal object with raw DDL + from clickhouse_alembic.introspect import ( + DictDefinition, + MVDefinition, + TableDefinition, + ViewDefinition, + ) + fallback_types = { + "table": lambda: TableDefinition(name=name, engine="", raw_ddl=ddl), + "view": lambda: ViewDefinition(name=name, select_query="", raw_ddl=ddl), + "materialized_view": lambda: MVDefinition(name=name, raw_ddl=ddl), + "dictionary": lambda: DictDefinition(name=name, raw_ddl=ddl), + } + schema_attrs[obj_type][name] = fallback_types[obj_type]() + + # Get live schema + try: + client = get_client(env_config) + live_schema = get_live_schema(client, database) + except Exception as e: + click.echo(f"Error connecting to {environment}: {e}", err=True) + sys.exit(1) + + # Compare + diffs = compare_schemas(local_schema, live_schema) + render_diff_report(diffs) + + has_drift = any(d.status != DiffStatus.IN_SYNC for d in diffs) + sys.exit(1 if has_drift else 0) + + +@main.command(name="upgrade-env") +def upgrade_env() -> None: + """Regenerate migrations/env.py from the latest ch-migrate version. + + Updates the Alembic environment file to the latest version shipped with + ch-migrate. This is needed when upgrading ch-migrate to pick up new + features like execution hooks. + + The previous env.py is backed up as env.py.bak. + """ + env_py_src = Path(__file__).parent / "env.py" + env_py_dst = Path.cwd() / "migrations" / "env.py" + + if not env_py_dst.parent.exists(): + click.echo("Error: migrations/ directory not found. Run 'ch-migrate init' first.", err=True) + sys.exit(1) + + if not env_py_src.exists(): + click.echo("Error: package env.py not found.", err=True) + sys.exit(1) + + # Back up existing env.py if present + if env_py_dst.exists(): + backup = env_py_dst.with_suffix(".py.bak") + shutil.copy(env_py_dst, backup) + click.echo(f" Backed up existing env.py to {backup.name}") + + shutil.copy(env_py_src, env_py_dst) + click.echo(f" Updated migrations/env.py") + click.echo("") + click.echo("env.py has been upgraded. If you have custom modifications,") + click.echo("compare with env.py.bak and reapply them.") + + +@main.command() +@click.argument("environment") +@click.option( + "--exclude", + "-e", + multiple=True, + help="Glob patterns to exclude (e.g., --exclude 'system_*' --exclude 'peerdb_*')", +) +@click.option( + "--filter", + "-f", + "include_filter", + multiple=True, + help="Glob patterns to include (only matching objects are captured)", +) +def snapshot(environment: str, exclude: tuple[str, ...], include_filter: tuple[str, ...]) -> None: + """Capture a schema snapshot from a live database. + + Connects to the environment and writes CREATE statements for all tables, + views, materialized views, and dictionaries to a timestamped directory. + + \b + Examples: + ch-migrate snapshot dev + ch-migrate snapshot dev --exclude 'system_*' --exclude 'peerdb_*' + ch-migrate snapshot dev --filter 'geo_*' + """ + import fnmatch + + from clickhouse_alembic.connection import get_client + from clickhouse_alembic.display import render_snapshot_progress + from clickhouse_alembic.introspect import Schema, get_live_schema + + config_path = Path.cwd() / "config.yaml" + try: + env_config = get_env_config(environment, config_path) + except Exception as e: + click.echo(f"Error loading config: {e}", err=True) + sys.exit(1) + + database = env_config["database"] + + try: + client = get_client(env_config) + except Exception as e: + click.echo(f"Error connecting to {environment}: {e}", err=True) + sys.exit(1) + + click.echo(f"Capturing schema from {environment} ({database})...") + + try: + schema = get_live_schema(client, database) + except Exception as e: + click.echo(f"Error introspecting database: {e}", err=True) + sys.exit(1) + + # Build output directory + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + snapshot_dir = Path.cwd() / "migrations" / "sql" / "snapshots" / timestamp + + # Flatten all exclude patterns (support comma-separated within a single --exclude) + exclude_patterns = [] + for pat in exclude: + exclude_patterns.extend(p.strip() for p in pat.split(",") if p.strip()) + + include_patterns = [] + for pat in include_filter: + include_patterns.extend(p.strip() for p in pat.split(",") if p.strip()) + + def should_include(name: str) -> bool: + if include_patterns and not any(fnmatch.fnmatch(name, p) for p in include_patterns): + return False + if any(fnmatch.fnmatch(name, p) for p in exclude_patterns): + return False + return True + + # Write DDL files organized by type + type_map = { + "tables": schema.tables, + "views": schema.views, + "materialized_views": schema.materialized_views, + "dictionaries": schema.dictionaries, + } + + counts: dict[str, int] = {} + excluded_count = 0 + + for type_name, objects in type_map.items(): + count = 0 + for name, obj in objects.items(): + if not should_include(name): + excluded_count += 1 + continue + type_dir = snapshot_dir / type_name + type_dir.mkdir(parents=True, exist_ok=True) + ddl = obj.raw_ddl if obj.raw_ddl else f"-- No DDL captured for {name}\n" + (type_dir / f"{name}.sql").write_text(ddl) + count += 1 + counts[type_name] = count + + if sum(counts.values()) == 0: + click.echo("No objects matched the filter criteria.", err=True) + sys.exit(1) + + render_snapshot_progress( + str(snapshot_dir.relative_to(Path.cwd())), + counts, + excluded=excluded_count, + ) + + if __name__ == "__main__": main() diff --git a/clickhouse_alembic/config.py b/clickhouse_alembic/config.py index cf2e874..3e7b5d0 100644 --- a/clickhouse_alembic/config.py +++ b/clickhouse_alembic/config.py @@ -108,4 +108,9 @@ def get_env_config(env_name: str, config_path: Path) -> dict[str, Any]: required=False, ) + # Pass through top-level hooks section (used by env.py for pre/post migrate) + hooks_config = config.get("hooks") + if hooks_config: + env_config["hooks"] = hooks_config + return env_config diff --git a/clickhouse_alembic/deps.py b/clickhouse_alembic/deps.py new file mode 100644 index 0000000..7c8bb7b --- /dev/null +++ b/clickhouse_alembic/deps.py @@ -0,0 +1,124 @@ +"""Dependency graph analysis and migration validation.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any + +from clickhouse_alembic.introspect import ( + DependencyGraph, + DepType, + ObjectNode, + get_dependencies, +) + + +@dataclass +class MigrationWarning: + severity: str # "error" or "warning" + message: str + affected_objects: list[str] + + +# Patterns that detect destructive operations +_RE_DROP_TABLE = re.compile( + r"DROP\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:`?(\w+)`?\.)?`?(\w+)`?", + re.IGNORECASE, +) + +_RE_DROP_VIEW = re.compile( + r"DROP\s+(?:MATERIALIZED\s+)?VIEW\s+(?:IF\s+EXISTS\s+)?(?:`?(\w+)`?\.)?`?(\w+)`?", + re.IGNORECASE, +) + +_RE_DROP_DICT = re.compile( + r"DROP\s+DICTIONARY\s+(?:IF\s+EXISTS\s+)?(?:`?(\w+)`?\.)?`?(\w+)`?", + re.IGNORECASE, +) + + +def validate_migration(sql: str, graph: DependencyGraph) -> list[MigrationWarning]: + """Check if migration SQL would break dependencies in the graph. + + Args: + sql: The migration SQL to validate. + graph: A live DependencyGraph from introspect.get_dependencies(). + + Returns: + List of warnings about potential dependency breakage. + """ + warnings: list[MigrationWarning] = [] + + # Check DROP TABLE + for m in _RE_DROP_TABLE.finditer(sql): + table_name = m.group(2) + if table_name in graph.nodes: + affected = graph.affected_by_drop(table_name) + if affected: + affected_names = [n.name for n in affected] + # Distinguish schema vs data_flow impact + schema_deps = [] + data_flow_deps = [] + for edge in graph.edges: + if edge.source == table_name: + if edge.dep_type == DepType.SCHEMA: + schema_deps.append(edge.target) + elif edge.dep_type == DepType.DATA_FLOW: + data_flow_deps.append(edge.target) + + if schema_deps: + warnings.append(MigrationWarning( + severity="error", + message=f"DROP TABLE {table_name} breaks schema dependencies: {', '.join(schema_deps)}", + affected_objects=schema_deps, + )) + if data_flow_deps: + warnings.append(MigrationWarning( + severity="warning", + message=f"DROP TABLE {table_name} breaks data flow to: {', '.join(data_flow_deps)}", + affected_objects=data_flow_deps, + )) + + # Check DROP VIEW / DROP MATERIALIZED VIEW + for m in _RE_DROP_VIEW.finditer(sql): + view_name = m.group(2) + if view_name in graph.nodes: + affected = graph.affected_by_drop(view_name) + if affected: + affected_names = [n.name for n in affected] + warnings.append(MigrationWarning( + severity="warning", + message=f"DROP VIEW {view_name} affects: {', '.join(affected_names)}", + affected_objects=affected_names, + )) + + # Check DROP DICTIONARY + for m in _RE_DROP_DICT.finditer(sql): + dict_name = m.group(2) + if dict_name in graph.nodes: + affected = graph.affected_by_drop(dict_name) + if affected: + affected_names = [n.name for n in affected] + warnings.append(MigrationWarning( + severity="warning", + message=f"DROP DICTIONARY {dict_name} affects: {', '.join(affected_names)}", + affected_objects=affected_names, + )) + + return warnings + + +def build_dependency_graph(client: Any, database: str) -> DependencyGraph: + """Build a dependency graph from the live database. + + Convenience wrapper around introspect.get_dependencies(). + + Args: + client: clickhouse-connect client. + database: Database name. + + Returns: + A DependencyGraph with nodes and typed edges. + """ + return get_dependencies(client, database) diff --git a/clickhouse_alembic/diff.py b/clickhouse_alembic/diff.py new file mode 100644 index 0000000..76fda11 --- /dev/null +++ b/clickhouse_alembic/diff.py @@ -0,0 +1,225 @@ +"""Schema comparison: field-by-field structured diff between two Schema objects.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Literal + +from clickhouse_alembic.introspect import ( + ColumnDefinition, + Schema, + TableDefinition, +) + + +class DiffStatus(str, Enum): + IN_SYNC = "in_sync" + MODIFIED = "modified" + LOCAL_ONLY = "local_only" + REMOTE_ONLY = "remote_only" + + +@dataclass +class FieldDiff: + field_name: str + local_value: str | None + remote_value: str | None + message: str + + +@dataclass +class SchemaDiff: + name: str + obj_type: str # "table", "view", "materialized_view", "dictionary" + status: DiffStatus + field_diffs: list[FieldDiff] = field(default_factory=list) + + +def _compare_columns( + local_cols: list[ColumnDefinition], + remote_cols: list[ColumnDefinition], +) -> list[FieldDiff]: + """Compare column lists field-by-field.""" + diffs: list[FieldDiff] = [] + local_map = {c.name: c for c in local_cols} + remote_map = {c.name: c for c in remote_cols} + + all_names = dict.fromkeys([c.name for c in local_cols] + [c.name for c in remote_cols]) + + for name in all_names: + local_col = local_map.get(name) + remote_col = remote_map.get(name) + + if local_col and not remote_col: + diffs.append(FieldDiff( + field_name=f"column '{name}'", + local_value=local_col.type, + remote_value=None, + message=f"column '{name}' exists locally but not in DB", + )) + elif remote_col and not local_col: + diffs.append(FieldDiff( + field_name=f"column '{name}'", + local_value=None, + remote_value=remote_col.type, + message=f"column '{name}' exists in DB but not locally", + )) + elif local_col and remote_col: + if local_col.type != remote_col.type: + diffs.append(FieldDiff( + field_name=f"column '{name}' type", + local_value=local_col.type, + remote_value=remote_col.type, + message=f"column '{name}' type differs: {local_col.type} vs {remote_col.type}", + )) + if local_col.default_kind != remote_col.default_kind: + diffs.append(FieldDiff( + field_name=f"column '{name}' default_kind", + local_value=local_col.default_kind, + remote_value=remote_col.default_kind, + message=f"column '{name}' default kind differs", + )) + if local_col.default_expr != remote_col.default_expr: + diffs.append(FieldDiff( + field_name=f"column '{name}' default_expr", + local_value=local_col.default_expr, + remote_value=remote_col.default_expr, + message=f"column '{name}' default expression differs", + )) + + return diffs + + +def _compare_tables(local: TableDefinition, remote: TableDefinition) -> list[FieldDiff]: + """Field-by-field comparison of two TableDefinitions.""" + diffs: list[FieldDiff] = [] + + # Engine + if local.engine != remote.engine: + diffs.append(FieldDiff( + field_name="engine", + local_value=local.engine, + remote_value=remote.engine, + message=f"engine differs: {local.engine} vs {remote.engine}", + )) + + # Columns + diffs.extend(_compare_columns(local.columns, remote.columns)) + + # ORDER BY + if local.order_by != remote.order_by: + diffs.append(FieldDiff( + field_name="order_by", + local_value=", ".join(local.order_by), + remote_value=", ".join(remote.order_by), + message=f"ORDER BY differs", + )) + + # PARTITION BY + if local.partition_by != remote.partition_by: + diffs.append(FieldDiff( + field_name="partition_by", + local_value=local.partition_by, + remote_value=remote.partition_by, + message=f"PARTITION BY differs", + )) + + # TTL + if local.ttl != remote.ttl: + diffs.append(FieldDiff( + field_name="ttl", + local_value=local.ttl, + remote_value=remote.ttl, + message=f"TTL differs", + )) + + # Settings + if local.settings != remote.settings: + diffs.append(FieldDiff( + field_name="settings", + local_value=str(local.settings), + remote_value=str(remote.settings), + message=f"SETTINGS differ", + )) + + return diffs + + +def _normalize_ddl(raw: str) -> str: + """Normalize raw DDL for fallback string comparison.""" + import re + s = raw.strip() + s = re.sub(r"\s+", " ", s) + return s + + +def _compare_raw_ddl(local_ddl: str, remote_ddl: str) -> list[FieldDiff]: + """Fallback: normalized string comparison for unparseable objects.""" + if _normalize_ddl(local_ddl) != _normalize_ddl(remote_ddl): + return [FieldDiff( + field_name="raw_ddl", + local_value=local_ddl[:200] if local_ddl else None, + remote_value=remote_ddl[:200] if remote_ddl else None, + message="DDL definition differs (raw comparison)", + )] + return [] + + +def compare_schemas(local: Schema, live: Schema) -> list[SchemaDiff]: + """Compare two Schema objects and return a list of differences. + + Args: + local: Schema from local snapshot files. + live: Schema from the live database. + + Returns: + List of SchemaDiff objects. Empty list means schemas are in sync. + """ + results: list[SchemaDiff] = [] + + type_map: list[tuple[str, dict, dict]] = [ + ("table", local.tables, live.tables), + ("view", local.views, live.views), + ("materialized_view", local.materialized_views, live.materialized_views), + ("dictionary", local.dictionaries, live.dictionaries), + ] + + for obj_type, local_objs, live_objs in type_map: + all_names = dict.fromkeys(list(local_objs.keys()) + list(live_objs.keys())) + + for name in all_names: + local_obj = local_objs.get(name) + live_obj = live_objs.get(name) + + if local_obj and not live_obj: + results.append(SchemaDiff( + name=name, obj_type=obj_type, status=DiffStatus.LOCAL_ONLY, + )) + elif live_obj and not local_obj: + results.append(SchemaDiff( + name=name, obj_type=obj_type, status=DiffStatus.REMOTE_ONLY, + )) + else: + # Both exist — compare + field_diffs: list[FieldDiff] = [] + + if obj_type == "table" and isinstance(local_obj, TableDefinition) and isinstance(live_obj, TableDefinition): + field_diffs = _compare_tables(local_obj, live_obj) + else: + # Fallback to raw DDL comparison for views, MVs, dicts + local_ddl = getattr(local_obj, "raw_ddl", "") + live_ddl = getattr(live_obj, "raw_ddl", "") + field_diffs = _compare_raw_ddl(local_ddl, live_ddl) + + if field_diffs: + results.append(SchemaDiff( + name=name, obj_type=obj_type, + status=DiffStatus.MODIFIED, field_diffs=field_diffs, + )) + else: + results.append(SchemaDiff( + name=name, obj_type=obj_type, status=DiffStatus.IN_SYNC, + )) + + return results diff --git a/clickhouse_alembic/display.py b/clickhouse_alembic/display.py index 0dccca3..7d1113b 100644 --- a/clickhouse_alembic/display.py +++ b/clickhouse_alembic/display.py @@ -8,6 +8,7 @@ from rich.panel import Panel from rich.table import Table from rich.text import Text +from rich.tree import Tree from clickhouse_alembic.rebase import RevisionGraph @@ -256,3 +257,203 @@ def render_lint_report( summary_parts.append(f"[yellow]{report.warning_count} warning(s)[/yellow]") console.print(f" {', '.join(summary_parts)}") + + +def render_snapshot_progress( + output_dir: str, + counts: dict[str, int], + excluded: int = 0, + console: Console | None = None, +) -> None: + """Render snapshot completion summary. + + Args: + output_dir: Path to the snapshot directory. + counts: Dict mapping object type to count (e.g. {"tables": 3, "views": 1}). + excluded: Number of objects excluded by filter. + console: Optional Console for testability. + """ + console = console or Console() + + table = Table(show_header=False, box=None, padding=(0, 2)) + table.add_column(style="bold") + table.add_column(justify="right") + + total = 0 + for obj_type, count in counts.items(): + if count > 0: + table.add_row(obj_type.replace("_", " ").title(), str(count)) + total += count + + status_text = Text() + status_text.append(f"\nSnapshot saved to ", style="dim") + status_text.append(output_dir, style="bold") + status_text.append(f"\n{total} objects captured", style="green") + if excluded: + status_text.append(f", {excluded} excluded", style="dim") + + panel = Panel( + Group(table, status_text), + title="[bold]Schema Snapshot[/bold]", + border_style="blue", + ) + console.print(panel) + + +def render_diff_report( + diffs: list[Any], + *, + console: Console | None = None, +) -> None: + """Render schema diff results as a Rich table. + + Args: + diffs: List of SchemaDiff objects from compare_schemas(). + console: Optional Console for testability. + """ + from clickhouse_alembic.diff import DiffStatus + + console = console or Console() + + in_sync = [d for d in diffs if d.status == DiffStatus.IN_SYNC] + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + local_only = [d for d in diffs if d.status == DiffStatus.LOCAL_ONLY] + remote_only = [d for d in diffs if d.status == DiffStatus.REMOTE_ONLY] + + has_drift = bool(modified or local_only or remote_only) + + if not has_drift: + console.print(f"[green]All {len(in_sync)} objects in sync.[/green]") + return + + table = Table(title="Schema Diff", show_lines=False) + table.add_column("Status", width=12) + table.add_column("Type", width=20) + table.add_column("Name", width=30) + table.add_column("Details") + + for d in local_only: + table.add_row( + Text("LOCAL ONLY", style="yellow"), + d.obj_type, + d.name, + "Exists in snapshot but not in DB", + ) + + for d in remote_only: + table.add_row( + Text("REMOTE ONLY", style="cyan"), + d.obj_type, + d.name, + "Exists in DB but not in snapshot", + ) + + for d in modified: + details = "; ".join(fd.message for fd in d.field_diffs) + table.add_row( + Text("MODIFIED", style="bold red"), + d.obj_type, + d.name, + details, + ) + + console.print(table) + console.print() + + parts = [] + if modified: + parts.append(f"[bold red]{len(modified)} modified[/bold red]") + if local_only: + parts.append(f"[yellow]{len(local_only)} local only[/yellow]") + if remote_only: + parts.append(f"[cyan]{len(remote_only)} remote only[/cyan]") + if in_sync: + parts.append(f"[green]{len(in_sync)} in sync[/green]") + + console.print(f" {', '.join(parts)}") + + +_OBJ_TYPE_STYLES = { + "table": ("bold", "T"), + "view": ("cyan", "V"), + "materialized_view": ("magenta", "MV"), + "dictionary": ("yellow", "D"), +} + +_DEP_TYPE_LABELS = { + "schema": "[dim]schema[/dim]", + "data_flow": "[bold blue]data_flow[/bold blue]", +} + + +def render_dependency_tree( + graph: Any, + *, + console: Console | None = None, +) -> None: + """Render a dependency graph as a Rich Tree. + + Each root node (no incoming edges) gets a tree branch. Dependent objects + are shown as children with edge type annotations. + + Args: + graph: A DependencyGraph from introspect. + console: Optional Console for testability. + """ + console = console or Console() + + if not graph.nodes: + console.print("[dim]No objects found in database.[/dim]") + return + + tree = Tree("[bold]Dependency Graph[/bold]") + + # Build adjacency: source -> [(target, dep_type)] + children_map: dict[str, list[tuple[str, str]]] = {name: [] for name in graph.nodes} + has_parent: set[str] = set() + for edge in graph.edges: + if edge.source in children_map and edge.target in graph.nodes: + children_map[edge.source].append((edge.target, edge.dep_type.value)) + has_parent.add(edge.target) + + # Roots: nodes with no incoming edges + roots = [name for name in graph.nodes if name not in has_parent] + if not roots: + # All nodes have parents (cycles) — just show all + roots = sorted(graph.nodes.keys()) + + def _add_node(parent_tree: Tree, name: str, dep_label: str | None, visited: set[str]) -> None: + node = graph.nodes[name] + style, prefix = _OBJ_TYPE_STYLES.get(node.obj_type, ("", "?")) + label = f"[{style}][{prefix}][/{style}] {name}" + if dep_label: + label += f" {dep_label}" + + if name in visited: + parent_tree.add(f"{label} [dim](circular)[/dim]") + return + + branch = parent_tree.add(label) + visited.add(name) + + for child_name, child_dep_type in children_map.get(name, []): + dep_str = _DEP_TYPE_LABELS.get(child_dep_type, child_dep_type) + _add_node(branch, child_name, dep_str, visited) + + for root_name in sorted(roots): + _add_node(tree, root_name, None, set()) + + console.print(tree) + console.print() + + # Summary + type_counts: dict[str, int] = {} + for node in graph.nodes.values(): + type_counts[node.obj_type] = type_counts.get(node.obj_type, 0) + 1 + + parts = [] + for obj_type, count in sorted(type_counts.items()): + _, prefix = _OBJ_TYPE_STYLES.get(obj_type, ("", "?")) + parts.append(f"{count} {obj_type.replace('_', ' ')}s [{prefix}]") + + console.print(f" {', '.join(parts)} — {len(graph.edges)} edges") diff --git a/clickhouse_alembic/env.py b/clickhouse_alembic/env.py index c89d95c..c3466da 100644 --- a/clickhouse_alembic/env.py +++ b/clickhouse_alembic/env.py @@ -18,6 +18,7 @@ Alembic uses to connect to ClickHouse. """ +import logging import os from logging.config import fileConfig from pathlib import Path @@ -29,6 +30,9 @@ from sqlalchemy import Connection, create_engine, pool, text from clickhouse_alembic.config import get_env_config +from clickhouse_alembic.hooks import HookRegistry, run_hooks + +logger = logging.getLogger(__name__) # Alembic Config object config = context.config @@ -55,6 +59,9 @@ if config.config_file_name is not None: fileConfig(config.config_file_name) +# Load hook registry from config +hook_registry = HookRegistry.from_config(env_config.get("hooks")) + class ClickhouseImpl(impl.DefaultImpl): """Alembic implementation for ClickHouse dialect.""" @@ -145,18 +152,46 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: """Run migrations in 'online' mode - executes against the database.""" connectable = create_engine(get_sqlalchemy_url(), poolclass=pool.NullPool) + db = DATABASE_NAME + + def _on_version_apply(ctx, step, heads, run_args): + """Fire post-migrate hooks after each migration step. + + on_version_apply fires AFTER each migration step completes, + so we use it for post-migrate hooks only. + """ + revision = step.up_revision if step.is_upgrade else ( + step.down_revisions[0] if step.down_revisions else "unknown" + ) + if hook_registry.post_migrate: + run_hooks( + ctx.connection, hook_registry.post_migrate, + db=db, phase="post_migrate", revision=revision, + ) with connectable.connect() as connection: bootstrap_version_table(connection) - context.configure( + configure_kwargs = dict( connection=connection, target_metadata=None, version_table="alembic_version", version_table_schema=DATABASE_NAME, ) + if hook_registry.has_hooks: + configure_kwargs["on_version_apply"] = _on_version_apply + + context.configure(**configure_kwargs) + with context.begin_transaction(): + # Fire pre-migrate hooks before the migration run + if hook_registry.pre_migrate: + run_hooks( + connection, hook_registry.pre_migrate, + db=db, phase="pre_migrate", revision="all", + ) + context.run_migrations() diff --git a/clickhouse_alembic/hooks.py b/clickhouse_alembic/hooks.py new file mode 100644 index 0000000..2a6f9b6 --- /dev/null +++ b/clickhouse_alembic/hooks.py @@ -0,0 +1,71 @@ +"""Pre/post migration hook support for clickhouse-alembic.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Any + +from sqlalchemy import Connection, text + +logger = logging.getLogger(__name__) + + +@dataclass +class HookRegistry: + """Registry of pre/post migration hooks loaded from config.yaml. + + Config format: + hooks: + pre_migrate: + - "SELECT 1" + post_migrate: + - "SYSTEM RELOAD DICTIONARY {db}.dict_regions ON CLUSTER default" + """ + + pre_migrate: list[str] = field(default_factory=list) + post_migrate: list[str] = field(default_factory=list) + + @classmethod + def from_config(cls, hooks_config: dict[str, Any] | None) -> HookRegistry: + """Build a HookRegistry from the hooks section of config.yaml.""" + if not hooks_config: + return cls() + + pre = hooks_config.get("pre_migrate", []) + post = hooks_config.get("post_migrate", []) + + if not isinstance(pre, list): + pre = [pre] if pre else [] + if not isinstance(post, list): + post = [post] if post else [] + + return cls(pre_migrate=pre, post_migrate=post) + + @property + def has_hooks(self) -> bool: + return bool(self.pre_migrate or self.post_migrate) + + +def run_hooks( + connection: Connection, + hooks: list[str], + *, + db: str, + phase: str, + revision: str, +) -> None: + """Execute a list of hook SQL statements. + + Args: + connection: SQLAlchemy connection to ClickHouse + hooks: List of SQL strings (may contain {db} placeholder) + db: Database name for placeholder resolution + phase: "pre_migrate" or "post_migrate" (for logging) + revision: Migration revision being processed (for logging) + """ + for i, hook_sql in enumerate(hooks, 1): + resolved = hook_sql.format(db=db) + logger.info("[%s] hook %d/%d for %s: %s", phase, i, len(hooks), revision, resolved) + connection.execute(text(resolved)) + connection.commit() diff --git a/clickhouse_alembic/templates/project/config.yaml.template b/clickhouse_alembic/templates/project/config.yaml.template index 8450a47..8e1c923 100644 --- a/clickhouse_alembic/templates/project/config.yaml.template +++ b/clickhouse_alembic/templates/project/config.yaml.template @@ -42,3 +42,12 @@ environments: host: your-prod-instance.clickhouse.cloud database: {project_name} migration_user: migration_prod + +# Optional: Pre/post migration hooks +# SQL statements executed before/after each migration run. +# Use {db} placeholder for the database name. +# hooks: +# pre_migrate: +# - "SELECT 1" # validation query +# post_migrate: +# - "SYSTEM RELOAD DICTIONARY {db}.dict_regions ON CLUSTER default" diff --git a/tests/test_deps.py b/tests/test_deps.py new file mode 100644 index 0000000..2c8eed8 --- /dev/null +++ b/tests/test_deps.py @@ -0,0 +1,224 @@ +"""Tests for dependency graph command and migration validation.""" + +from __future__ import annotations + +from io import StringIO +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from click.testing import CliRunner +from rich.console import Console + +from clickhouse_alembic.deps import MigrationWarning, validate_migration +from clickhouse_alembic.display import render_dependency_tree +from clickhouse_alembic.introspect import ( + DependencyEdge, + DependencyGraph, + DepType, + ObjectNode, +) + + +def _make_graph() -> DependencyGraph: + """Build a synthetic dependency graph for testing.""" + graph = DependencyGraph() + graph.nodes = { + "users": ObjectNode(name="users", obj_type="table"), + "events": ObjectNode(name="events", obj_type="table"), + "hourly_events": ObjectNode(name="hourly_events", obj_type="materialized_view"), + "hourly_events_dest": ObjectNode(name="hourly_events_dest", obj_type="table"), + "dict_users": ObjectNode(name="dict_users", obj_type="dictionary"), + "active_users": ObjectNode(name="active_users", obj_type="view"), + } + graph.edges = [ + DependencyEdge(source="events", target="hourly_events", dep_type=DepType.DATA_FLOW), + DependencyEdge(source="events", target="hourly_events", dep_type=DepType.SCHEMA), + DependencyEdge(source="hourly_events", target="hourly_events_dest", dep_type=DepType.DATA_FLOW), + DependencyEdge(source="users", target="dict_users", dep_type=DepType.SCHEMA), + ] + return graph + + +# --------------------------------------------------------------------------- +# validate_migration tests +# --------------------------------------------------------------------------- + + +class TestValidateMigration: + def test_drop_table_with_mv_dependency(self): + graph = _make_graph() + sql = "DROP TABLE events" + warnings = validate_migration(sql, graph) + assert len(warnings) >= 1 + # Should catch both schema and data_flow + error_msgs = [w.message for w in warnings if w.severity == "error"] + warn_msgs = [w.message for w in warnings if w.severity == "warning"] + assert any("hourly_events" in m for m in error_msgs) # schema dep + assert any("hourly_events" in m for m in warn_msgs) # data flow dep + + def test_drop_table_with_dict_dependency(self): + graph = _make_graph() + sql = "DROP TABLE IF EXISTS users" + warnings = validate_migration(sql, graph) + assert len(warnings) >= 1 + assert any("dict_users" in w.message for w in warnings) + + def test_drop_table_no_deps(self): + graph = _make_graph() + sql = "DROP TABLE hourly_events_dest" + warnings = validate_migration(sql, graph) + assert len(warnings) == 0 + + def test_drop_table_not_in_graph(self): + graph = _make_graph() + sql = "DROP TABLE nonexistent_table" + warnings = validate_migration(sql, graph) + assert len(warnings) == 0 + + def test_drop_materialized_view(self): + graph = _make_graph() + sql = "DROP MATERIALIZED VIEW hourly_events" + warnings = validate_migration(sql, graph) + # hourly_events has a child (hourly_events_dest via data_flow) + assert len(warnings) >= 1 + assert any("hourly_events_dest" in w.message for w in warnings) + + def test_drop_dictionary_no_deps(self): + graph = _make_graph() + sql = "DROP DICTIONARY dict_users" + warnings = validate_migration(sql, graph) + # dict_users has no children + assert len(warnings) == 0 + + def test_safe_migration(self): + graph = _make_graph() + sql = "ALTER TABLE users ADD COLUMN email String" + warnings = validate_migration(sql, graph) + assert len(warnings) == 0 + + def test_multiple_drops(self): + graph = _make_graph() + sql = "DROP TABLE events;\nDROP TABLE users;" + warnings = validate_migration(sql, graph) + assert len(warnings) >= 2 + + +# --------------------------------------------------------------------------- +# Tree rendering tests +# --------------------------------------------------------------------------- + + +class TestRenderDependencyTree: + def test_renders_tree(self): + output = StringIO() + console = Console(file=output, force_terminal=True, width=100) + graph = _make_graph() + + render_dependency_tree(graph, console=console) + + text = output.getvalue() + assert "Dependency Graph" in text + assert "users" in text + assert "events" in text + assert "hourly_events" in text + assert "dict_users" in text + + def test_renders_empty_graph(self): + output = StringIO() + console = Console(file=output, force_terminal=True, width=100) + graph = DependencyGraph() + + render_dependency_tree(graph, console=console) + + text = output.getvalue() + assert "No objects found" in text + + def test_shows_edge_count(self): + import re as _re + output = StringIO() + console = Console(file=output, force_terminal=True, width=100) + graph = _make_graph() + + render_dependency_tree(graph, console=console) + + # Strip ANSI escape codes for assertion + text = _re.sub(r"\x1b\[[0-9;]*m", "", output.getvalue()) + assert "4 edges" in text + + def test_handles_circular_dependency(self): + graph = DependencyGraph() + graph.nodes = { + "a": ObjectNode(name="a", obj_type="table"), + "b": ObjectNode(name="b", obj_type="materialized_view"), + } + graph.edges = [ + DependencyEdge(source="a", target="b", dep_type=DepType.SCHEMA), + DependencyEdge(source="b", target="a", dep_type=DepType.DATA_FLOW), + ] + + output = StringIO() + console = Console(file=output, force_terminal=True, width=100) + render_dependency_tree(graph, console=console) + + text = output.getvalue() + assert "circular" in text + + +# --------------------------------------------------------------------------- +# CLI command tests +# --------------------------------------------------------------------------- + + +@pytest.fixture +def deps_runner(tmp_path: Path, monkeypatch): + monkeypatch.chdir(tmp_path) + (tmp_path / "config.yaml").write_text(""" +environments: + dev: + host: localhost + database: testdb + user: test +""") + monkeypatch.setenv("CH_DEV_PASSWORD", "pass") + return CliRunner() + + +class TestDepsCommand: + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.deps.get_dependencies") + def test_shows_graph(self, mock_deps, mock_client, deps_runner): + mock_client.return_value = MagicMock() + mock_deps.return_value = _make_graph() + + from clickhouse_alembic.cli import main + result = deps_runner.invoke(main, ["deps", "dev"]) + assert result.exit_code == 0 + assert "Dependency Graph" in result.output + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.deps.get_dependencies") + def test_validate_safe_migration(self, mock_deps, mock_client, deps_runner, tmp_path): + mock_client.return_value = MagicMock() + mock_deps.return_value = _make_graph() + + sql_file = tmp_path / "safe.sql" + sql_file.write_text("ALTER TABLE users ADD COLUMN email String") + + from clickhouse_alembic.cli import main + result = deps_runner.invoke(main, ["deps", "dev", "--validate", str(sql_file)]) + assert result.exit_code == 0 + assert "passed" in result.output + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.deps.get_dependencies") + def test_validate_breaking_migration(self, mock_deps, mock_client, deps_runner, tmp_path): + mock_client.return_value = MagicMock() + mock_deps.return_value = _make_graph() + + sql_file = tmp_path / "break.sql" + sql_file.write_text("DROP TABLE events") + + from clickhouse_alembic.cli import main + result = deps_runner.invoke(main, ["deps", "dev", "--validate", str(sql_file)]) + assert result.exit_code == 1 diff --git a/tests/test_diff.py b/tests/test_diff.py new file mode 100644 index 0000000..531259d --- /dev/null +++ b/tests/test_diff.py @@ -0,0 +1,323 @@ +"""Tests for schema diff module and CLI command.""" + +from __future__ import annotations + +from io import StringIO +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from click.testing import CliRunner +from rich.console import Console + +from clickhouse_alembic.diff import ( + DiffStatus, + FieldDiff, + SchemaDiff, + compare_schemas, +) +from clickhouse_alembic.display import render_diff_report +from clickhouse_alembic.introspect import ( + ColumnDefinition, + DictDefinition, + MVDefinition, + Schema, + TableDefinition, + ViewDefinition, +) + + +# --------------------------------------------------------------------------- +# compare_schemas unit tests +# --------------------------------------------------------------------------- + + +class TestCompareSchemas: + def _base_table(self, name: str = "users", **overrides) -> TableDefinition: + defaults = dict( + name=name, + engine="MergeTree", + columns=[ + ColumnDefinition(name="id", type="UInt64"), + ColumnDefinition(name="name", type="String"), + ], + order_by=["id"], + partition_by=None, + ttl=None, + settings={}, + raw_ddl=f"CREATE TABLE {name} (id UInt64, name String) ENGINE = MergeTree ORDER BY id", + ) + defaults.update(overrides) + return TableDefinition(**defaults) + + def test_identical_schemas_in_sync(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table() + live.tables["users"] = self._base_table() + + diffs = compare_schemas(local, live) + assert len(diffs) == 1 + assert diffs[0].status == DiffStatus.IN_SYNC + + def test_local_only_table(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table() + + diffs = compare_schemas(local, live) + assert len(diffs) == 1 + assert diffs[0].status == DiffStatus.LOCAL_ONLY + assert diffs[0].name == "users" + + def test_remote_only_table(self): + local = Schema(database="db") + live = Schema(database="db") + live.tables["users"] = self._base_table() + + diffs = compare_schemas(local, live) + assert len(diffs) == 1 + assert diffs[0].status == DiffStatus.REMOTE_ONLY + assert diffs[0].name == "users" + + def test_detects_added_column(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table() + live.tables["users"] = self._base_table(columns=[ + ColumnDefinition(name="id", type="UInt64"), + ColumnDefinition(name="name", type="String"), + ColumnDefinition(name="email", type="String"), + ]) + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("email" in fd.message and "DB but not locally" in fd.message for fd in modified[0].field_diffs) + + def test_detects_removed_column(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table(columns=[ + ColumnDefinition(name="id", type="UInt64"), + ColumnDefinition(name="name", type="String"), + ColumnDefinition(name="phone", type="String"), + ]) + live.tables["users"] = self._base_table() + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("phone" in fd.message and "locally but not in DB" in fd.message for fd in modified[0].field_diffs) + + def test_detects_column_type_change(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table() + live.tables["users"] = self._base_table(columns=[ + ColumnDefinition(name="id", type="UInt64"), + ColumnDefinition(name="name", type="Nullable(String)"), + ]) + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("type differs" in fd.message for fd in modified[0].field_diffs) + + def test_detects_engine_change(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table(engine="MergeTree") + live.tables["users"] = self._base_table(engine="ReplacingMergeTree(version)") + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("engine differs" in fd.message for fd in modified[0].field_diffs) + + def test_detects_order_by_change(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table(order_by=["id"]) + live.tables["users"] = self._base_table(order_by=["id", "name"]) + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("ORDER BY" in fd.message for fd in modified[0].field_diffs) + + def test_detects_partition_by_change(self): + local = Schema(database="db") + live = Schema(database="db") + local.tables["users"] = self._base_table(partition_by=None) + live.tables["users"] = self._base_table(partition_by="toYYYYMM(created_at)") + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("PARTITION BY" in fd.message for fd in modified[0].field_diffs) + + def test_views_fallback_to_raw_ddl(self): + local = Schema(database="db") + live = Schema(database="db") + local.views["v"] = ViewDefinition(name="v", select_query="SELECT 1", raw_ddl="CREATE VIEW v AS SELECT 1") + live.views["v"] = ViewDefinition(name="v", select_query="SELECT 2", raw_ddl="CREATE VIEW v AS SELECT 2") + + diffs = compare_schemas(local, live) + modified = [d for d in diffs if d.status == DiffStatus.MODIFIED] + assert len(modified) == 1 + assert any("raw comparison" in fd.message for fd in modified[0].field_diffs) + + def test_views_in_sync(self): + local = Schema(database="db") + live = Schema(database="db") + ddl = "CREATE VIEW v AS SELECT 1" + local.views["v"] = ViewDefinition(name="v", select_query="SELECT 1", raw_ddl=ddl) + live.views["v"] = ViewDefinition(name="v", select_query="SELECT 1", raw_ddl=ddl) + + diffs = compare_schemas(local, live) + assert all(d.status == DiffStatus.IN_SYNC for d in diffs) + + def test_mixed_object_types(self): + local = Schema(database="db") + live = Schema(database="db") + + local.tables["t1"] = self._base_table("t1") + live.tables["t1"] = self._base_table("t1") + + local.views["v1"] = ViewDefinition(name="v1", select_query="SELECT 1", raw_ddl="CREATE VIEW v1 AS SELECT 1") + # v1 is local_only (not in live) + + live.dictionaries["d1"] = DictDefinition(name="d1", raw_ddl="CREATE DICTIONARY d1 (...)") + # d1 is remote_only + + diffs = compare_schemas(local, live) + statuses = {d.name: d.status for d in diffs} + assert statuses["t1"] == DiffStatus.IN_SYNC + assert statuses["v1"] == DiffStatus.LOCAL_ONLY + assert statuses["d1"] == DiffStatus.REMOTE_ONLY + + +# --------------------------------------------------------------------------- +# CLI diff command tests +# --------------------------------------------------------------------------- + + +@pytest.fixture +def diff_runner(tmp_path: Path, monkeypatch): + monkeypatch.chdir(tmp_path) + (tmp_path / "config.yaml").write_text(""" +environments: + dev: + host: localhost + database: testdb + user: test +""") + monkeypatch.setenv("CH_DEV_PASSWORD", "pass") + + # Create a snapshot + snap_dir = tmp_path / "migrations" / "sql" / "snapshots" / "20260305_120000" + tables_dir = snap_dir / "tables" + tables_dir.mkdir(parents=True) + (tables_dir / "users.sql").write_text( + "CREATE TABLE testdb.users (id UInt64, name String) ENGINE = MergeTree ORDER BY id" + ) + return CliRunner() + + +class TestDiffCommand: + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_in_sync_exits_0(self, mock_schema, mock_client, diff_runner, tmp_path): + mock_client.return_value = MagicMock() + live = Schema(database="testdb") + live.tables["users"] = TableDefinition( + name="users", + engine="MergeTree", + columns=[ + ColumnDefinition(name="id", type="UInt64"), + ColumnDefinition(name="name", type="String"), + ], + order_by=["id"], + raw_ddl="CREATE TABLE testdb.users (id UInt64, name String) ENGINE = MergeTree ORDER BY id", + ) + mock_schema.return_value = live + + from clickhouse_alembic.cli import main + result = diff_runner.invoke(main, ["diff", "dev"]) + assert result.exit_code == 0 + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_drift_exits_1(self, mock_schema, mock_client, diff_runner, tmp_path): + mock_client.return_value = MagicMock() + live = Schema(database="testdb") + live.tables["users"] = TableDefinition( + name="users", + engine="ReplacingMergeTree", + columns=[ + ColumnDefinition(name="id", type="UInt64"), + ColumnDefinition(name="name", type="String"), + ], + order_by=["id"], + raw_ddl="CREATE TABLE testdb.users (id UInt64, name String) ENGINE = ReplacingMergeTree ORDER BY id", + ) + mock_schema.return_value = live + + from clickhouse_alembic.cli import main + result = diff_runner.invoke(main, ["diff", "dev"]) + assert result.exit_code == 1 + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_no_snapshot_exits_1(self, mock_schema, mock_client, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + (tmp_path / "config.yaml").write_text(""" +environments: + dev: + host: localhost + database: testdb + user: test +""") + monkeypatch.setenv("CH_DEV_PASSWORD", "pass") + + from clickhouse_alembic.cli import main + runner = CliRunner() + result = runner.invoke(main, ["diff", "dev"]) + assert result.exit_code == 1 + assert "No snapshots found" in result.output + + +# --------------------------------------------------------------------------- +# Display rendering tests +# --------------------------------------------------------------------------- + + +class TestDiffDisplay: + def test_render_all_in_sync(self): + output = StringIO() + console = Console(file=output, force_terminal=True, width=80) + + diffs = [SchemaDiff(name="users", obj_type="table", status=DiffStatus.IN_SYNC)] + render_diff_report(diffs, console=console) + + text = output.getvalue() + assert "in sync" in text + + def test_render_with_drift(self): + output = StringIO() + console = Console(file=output, force_terminal=True, width=80) + + diffs = [ + SchemaDiff(name="users", obj_type="table", status=DiffStatus.MODIFIED, field_diffs=[ + FieldDiff("engine", "MergeTree", "ReplacingMergeTree", "engine differs"), + ]), + SchemaDiff(name="new_table", obj_type="table", status=DiffStatus.REMOTE_ONLY), + SchemaDiff(name="old_view", obj_type="view", status=DiffStatus.LOCAL_ONLY), + ] + render_diff_report(diffs, console=console) + + text = output.getvalue() + assert "MODIFIED" in text + assert "REMOTE ONLY" in text + assert "LOCAL ONLY" in text diff --git a/tests/test_hooks.py b/tests/test_hooks.py new file mode 100644 index 0000000..e872a59 --- /dev/null +++ b/tests/test_hooks.py @@ -0,0 +1,242 @@ +"""Tests for pre/post migration hook support.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, call + +import pytest + +from clickhouse_alembic.hooks import HookRegistry, run_hooks + + +class TestHookRegistry: + def test_from_empty_config(self): + registry = HookRegistry.from_config(None) + assert registry.pre_migrate == [] + assert registry.post_migrate == [] + assert not registry.has_hooks + + def test_from_empty_dict(self): + registry = HookRegistry.from_config({}) + assert not registry.has_hooks + + def test_from_post_migrate_only(self): + registry = HookRegistry.from_config({ + "post_migrate": [ + "SYSTEM RELOAD DICTIONARY {db}.dict_regions ON CLUSTER default", + ] + }) + assert registry.pre_migrate == [] + assert len(registry.post_migrate) == 1 + assert registry.has_hooks + + def test_from_pre_and_post(self): + registry = HookRegistry.from_config({ + "pre_migrate": ["SELECT 1"], + "post_migrate": [ + "SYSTEM RELOAD DICTIONARY {db}.dict_a ON CLUSTER default", + "SELECT count() FROM {db}.users", + ], + }) + assert len(registry.pre_migrate) == 1 + assert len(registry.post_migrate) == 2 + assert registry.has_hooks + + def test_coerces_single_string_to_list(self): + registry = HookRegistry.from_config({ + "post_migrate": "SYSTEM RELOAD DICTIONARY {db}.dict_a ON CLUSTER default", + }) + assert len(registry.post_migrate) == 1 + + def test_handles_none_values(self): + registry = HookRegistry.from_config({ + "pre_migrate": None, + "post_migrate": None, + }) + assert registry.pre_migrate == [] + assert registry.post_migrate == [] + + +class TestRunHooks: + def test_executes_hooks_in_order(self): + connection = MagicMock() + hooks = [ + "SYSTEM RELOAD DICTIONARY {db}.dict_a ON CLUSTER default", + "SELECT count() FROM {db}.users", + ] + + run_hooks(connection, hooks, db="mydb", phase="post_migrate", revision="abc123") + + assert connection.execute.call_count == 2 + assert connection.commit.call_count == 2 + + # Verify the SQL was resolved + executed_sql = [ + str(c.args[0]) for c in connection.execute.call_args_list + ] + assert "SYSTEM RELOAD DICTIONARY mydb.dict_a ON CLUSTER default" in executed_sql[0] + assert "SELECT count() FROM mydb.users" in executed_sql[1] + + def test_resolves_db_placeholder(self): + connection = MagicMock() + hooks = ["SELECT 1 FROM {db}.test"] + + run_hooks(connection, hooks, db="production_db", phase="pre_migrate", revision="xyz") + + executed_sql = str(connection.execute.call_args_list[0].args[0]) + assert "production_db" in executed_sql + + def test_empty_hooks_does_nothing(self): + connection = MagicMock() + run_hooks(connection, [], db="mydb", phase="post_migrate", revision="abc123") + connection.execute.assert_not_called() + + +class TestConfigIntegration: + def test_hooks_parsed_from_config_yaml(self, tmp_path: Path, monkeypatch): + """Hooks section in config.yaml is passed through to env_config.""" + from clickhouse_alembic.config import get_env_config + + config_file = tmp_path / "config.yaml" + config_file.write_text(""" +environments: + dev: + host: localhost + database: testdb + user: test + +hooks: + post_migrate: + - "SYSTEM RELOAD DICTIONARY {db}.dict_regions ON CLUSTER default" +""") + monkeypatch.setenv("CH_DEV_PASSWORD", "pass") + + env_config = get_env_config("dev", config_file) + + assert "hooks" in env_config + assert len(env_config["hooks"]["post_migrate"]) == 1 + + registry = HookRegistry.from_config(env_config.get("hooks")) + assert registry.has_hooks + assert registry.post_migrate[0] == "SYSTEM RELOAD DICTIONARY {db}.dict_regions ON CLUSTER default" + + def test_no_hooks_section_works(self, tmp_path: Path, monkeypatch): + """Existing configs without hooks section continue to work.""" + from clickhouse_alembic.config import get_env_config + + config_file = tmp_path / "config.yaml" + config_file.write_text(""" +environments: + dev: + host: localhost + database: testdb + user: test +""") + monkeypatch.setenv("CH_DEV_PASSWORD", "pass") + + env_config = get_env_config("dev", config_file) + + assert "hooks" not in env_config + registry = HookRegistry.from_config(env_config.get("hooks")) + assert not registry.has_hooks + + +class TestHookExecutionOrder: + """Verify hooks fire in the correct order during migration simulation.""" + + def test_pre_hooks_fire_before_post_hooks(self): + """Pre-migrate hooks should fire before post-migrate hooks.""" + connection = MagicMock() + registry = HookRegistry.from_config({ + "pre_migrate": ["SELECT 'pre'"], + "post_migrate": ["SELECT 'post'"], + }) + + # Simulate the env.py execution order: pre-hooks, then post-hooks + run_hooks(connection, registry.pre_migrate, db="mydb", phase="pre_migrate", revision="all") + run_hooks(connection, registry.post_migrate, db="mydb", phase="post_migrate", revision="abc123") + + assert connection.execute.call_count == 2 + calls = [str(c.args[0]) for c in connection.execute.call_args_list] + assert "pre" in calls[0] + assert "post" in calls[1] + + def test_multiple_post_hooks_fire_in_order(self): + """Multiple post-migrate hooks fire in list order.""" + connection = MagicMock() + hooks = [ + "SYSTEM RELOAD DICTIONARY {db}.dict_a ON CLUSTER default", + "SYSTEM RELOAD DICTIONARY {db}.dict_b ON CLUSTER default", + "SELECT count() FROM {db}.users", + ] + + run_hooks(connection, hooks, db="testdb", phase="post_migrate", revision="xyz") + + assert connection.execute.call_count == 3 + calls = [str(c.args[0]) for c in connection.execute.call_args_list] + assert "dict_a" in calls[0] + assert "dict_b" in calls[1] + assert "users" in calls[2] + + +class TestUpgradeEnvCommand: + """Tests for the ch-migrate upgrade-env CLI command.""" + + def test_upgrade_env_creates_new_env_py(self, tmp_path: Path): + """upgrade-env copies the package env.py when no existing env.py.""" + from click.testing import CliRunner + from clickhouse_alembic.cli import main + + migrations_dir = tmp_path / "migrations" + migrations_dir.mkdir() + + runner = CliRunner() + result = runner.invoke(main, ["upgrade-env"], catch_exceptions=False) + + # Can't run from tmp_path with CliRunner easily, but we can verify the + # command is registered and validates missing migrations dir + # Test with no migrations dir in cwd + assert result.exit_code != 0 or "Updated" in result.output or "not found" in result.output + + def test_upgrade_env_backs_up_existing(self, tmp_path: Path, monkeypatch): + """upgrade-env creates a .bak backup of existing env.py.""" + from click.testing import CliRunner + from clickhouse_alembic.cli import main + + migrations_dir = tmp_path / "migrations" + migrations_dir.mkdir() + existing_env = migrations_dir / "env.py" + existing_env.write_text("# old env.py content\n") + + monkeypatch.chdir(tmp_path) + + runner = CliRunner() + result = runner.invoke(main, ["upgrade-env"], catch_exceptions=False) + + assert result.exit_code == 0 + assert "Backed up" in result.output + assert "Updated" in result.output + + # Verify backup was created + backup = migrations_dir / "env.py.bak" + assert backup.exists() + assert backup.read_text() == "# old env.py content\n" + + # Verify new env.py was copied + new_content = existing_env.read_text() + assert "HookRegistry" in new_content + assert "run_hooks" in new_content + + def test_upgrade_env_no_migrations_dir(self, tmp_path: Path, monkeypatch): + """upgrade-env fails gracefully when no migrations dir exists.""" + from click.testing import CliRunner + from clickhouse_alembic.cli import main + + monkeypatch.chdir(tmp_path) + + runner = CliRunner() + result = runner.invoke(main, ["upgrade-env"], catch_exceptions=False) + + assert result.exit_code != 0 + assert "not found" in result.output diff --git a/tests/test_snapshot.py b/tests/test_snapshot.py new file mode 100644 index 0000000..8934171 --- /dev/null +++ b/tests/test_snapshot.py @@ -0,0 +1,239 @@ +"""Tests for ch-migrate snapshot command.""" + +from __future__ import annotations + +import re +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from click.testing import CliRunner + +from clickhouse_alembic.cli import main +from clickhouse_alembic.introspect import ( + DictDefinition, + MVDefinition, + Schema, + TableDefinition, + ViewDefinition, +) + + +def _make_schema() -> Schema: + """Build a synthetic Schema for testing.""" + schema = Schema(database="testdb", ch_version="24.3.1") + schema.tables["users"] = TableDefinition( + name="users", + engine="MergeTree", + raw_ddl="CREATE TABLE testdb.users (id UInt64) ENGINE = MergeTree ORDER BY id", + ) + schema.tables["events"] = TableDefinition( + name="events", + engine="MergeTree", + raw_ddl="CREATE TABLE testdb.events (id UInt64) ENGINE = MergeTree ORDER BY id", + ) + schema.tables["peerdb_staging"] = TableDefinition( + name="peerdb_staging", + engine="MergeTree", + raw_ddl="CREATE TABLE testdb.peerdb_staging (id UInt64) ENGINE = MergeTree ORDER BY id", + ) + schema.views["active_users"] = ViewDefinition( + name="active_users", + select_query="SELECT * FROM users WHERE active = 1", + raw_ddl="CREATE VIEW testdb.active_users AS SELECT * FROM users WHERE active = 1", + ) + schema.materialized_views["hourly_events"] = MVDefinition( + name="hourly_events", + raw_ddl="CREATE MATERIALIZED VIEW testdb.hourly_events TO testdb.hourly_dest AS SELECT count() FROM events", + ) + schema.dictionaries["dict_topic"] = DictDefinition( + name="dict_topic", + raw_ddl="CREATE DICTIONARY testdb.dict_topic (key String) PRIMARY KEY key SOURCE(CLICKHOUSE(TABLE 'topics' DB 'testdb')) LAYOUT(HASHED()) LIFETIME(300)", + ) + return schema + + +@pytest.fixture +def runner(tmp_path: Path, monkeypatch): + """Set up a CLI runner in a tmp project directory.""" + monkeypatch.chdir(tmp_path) + (tmp_path / "migrations" / "sql" / "snapshots").mkdir(parents=True) + (tmp_path / "config.yaml").write_text(""" +environments: + dev: + host: localhost + database: testdb + user: test +""") + monkeypatch.setenv("CH_DEV_PASSWORD", "pass") + return CliRunner() + + +class TestSnapshotCommand: + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_captures_all_objects(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result = runner.invoke(main, ["snapshot", "dev"]) + assert result.exit_code == 0 + + # Find the timestamped snapshot dir + snapshots_dir = tmp_path / "migrations" / "sql" / "snapshots" + snapshot_dirs = list(snapshots_dir.iterdir()) + assert len(snapshot_dirs) == 1 + + snap = snapshot_dirs[0] + assert re.match(r"\d{8}_\d{6}", snap.name) + + # Verify file structure + assert (snap / "tables" / "users.sql").exists() + assert (snap / "tables" / "events.sql").exists() + assert (snap / "tables" / "peerdb_staging.sql").exists() + assert (snap / "views" / "active_users.sql").exists() + assert (snap / "materialized_views" / "hourly_events.sql").exists() + assert (snap / "dictionaries" / "dict_topic.sql").exists() + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_ddl_content_written(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result = runner.invoke(main, ["snapshot", "dev"]) + assert result.exit_code == 0 + + snapshots_dir = tmp_path / "migrations" / "sql" / "snapshots" + snap = list(snapshots_dir.iterdir())[0] + + content = (snap / "tables" / "users.sql").read_text() + assert "CREATE TABLE testdb.users" in content + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_exclude_filter(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result = runner.invoke(main, ["snapshot", "dev", "--exclude", "peerdb_*"]) + assert result.exit_code == 0 + + snapshots_dir = tmp_path / "migrations" / "sql" / "snapshots" + snap = list(snapshots_dir.iterdir())[0] + + assert (snap / "tables" / "users.sql").exists() + assert not (snap / "tables" / "peerdb_staging.sql").exists() + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_exclude_comma_separated(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result = runner.invoke(main, ["snapshot", "dev", "--exclude", "peerdb_*,events"]) + assert result.exit_code == 0 + + snapshots_dir = tmp_path / "migrations" / "sql" / "snapshots" + snap = list(snapshots_dir.iterdir())[0] + + assert (snap / "tables" / "users.sql").exists() + assert not (snap / "tables" / "peerdb_staging.sql").exists() + assert not (snap / "tables" / "events.sql").exists() + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_include_filter(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result = runner.invoke(main, ["snapshot", "dev", "--filter", "user*,dict_*,*_users"]) + assert result.exit_code == 0 + + snapshots_dir = tmp_path / "migrations" / "sql" / "snapshots" + snap = list(snapshots_dir.iterdir())[0] + + assert (snap / "tables" / "users.sql").exists() + assert (snap / "views" / "active_users.sql").exists() + assert (snap / "dictionaries" / "dict_topic.sql").exists() + assert not (snap / "tables" / "events.sql").exists() + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_filter_no_matches_exits_1(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result = runner.invoke(main, ["snapshot", "dev", "--filter", "nonexistent_*"]) + assert result.exit_code == 1 + assert "No objects matched" in result.output + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_timestamped_dirs_no_overwrite(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = _make_schema() + + result1 = runner.invoke(main, ["snapshot", "dev"]) + assert result1.exit_code == 0 + + # Ensure a second snapshot creates a new directory (different timestamp) + import time + time.sleep(1.1) + + result2 = runner.invoke(main, ["snapshot", "dev"]) + assert result2.exit_code == 0 + + snapshots_dir = tmp_path / "migrations" / "sql" / "snapshots" + snapshot_dirs = list(snapshots_dir.iterdir()) + assert len(snapshot_dirs) == 2 + + @patch("clickhouse_alembic.connection.get_client") + @patch("clickhouse_alembic.introspect.get_live_schema") + def test_empty_schema(self, mock_schema, mock_client, runner, tmp_path): + mock_client.return_value = MagicMock() + mock_schema.return_value = Schema(database="testdb") + + result = runner.invoke(main, ["snapshot", "dev"]) + assert result.exit_code == 1 + assert "No objects matched" in result.output + + +class TestSnapshotDisplay: + def test_render_snapshot_progress(self): + from io import StringIO + from rich.console import Console + from clickhouse_alembic.display import render_snapshot_progress + + output = StringIO() + console = Console(file=output, force_terminal=True, width=80) + + render_snapshot_progress( + "migrations/sql/snapshots/20260305_120000", + {"tables": 3, "views": 1, "materialized_views": 0, "dictionaries": 2}, + excluded=5, + console=console, + ) + + text = output.getvalue() + assert "Schema Snapshot" in text + assert "6 objects captured" in text + assert "5 excluded" in text + + def test_render_snapshot_progress_no_excluded(self): + from io import StringIO + from rich.console import Console + from clickhouse_alembic.display import render_snapshot_progress + + output = StringIO() + console = Console(file=output, force_terminal=True, width=80) + + render_snapshot_progress( + "migrations/sql/snapshots/20260305_120000", + {"tables": 2, "views": 0, "materialized_views": 0, "dictionaries": 0}, + console=console, + ) + + text = output.getvalue() + assert "2 objects captured" in text + assert "excluded" not in text