ch-migrate-cli 0.5.0__py3-none-any.whl

ch_migrate/env.py

@@ -0,0 +1,226 @@
+"""
+ClickHouse Alembic Environment.
+
+This module is copied to user projects and configures Alembic for ClickHouse.
+It loads configuration from config.yaml and secrets from environment variables.
+
+Why SQLAlchemy?
+--------------
+Alembic is built on SQLAlchemy as its core database abstraction layer. While we use
+clickhouse-connect for direct ClickHouse operations (bootstrap, helpers), Alembic
+requires SQLAlchemy for:
+- Connection management and pooling
+- Transaction handling (even though ClickHouse DDL isn't transactional)
+- Database dialect registration (clickhouse-sqlalchemy provides the ClickHouse dialect)
+- The `op.execute()` interface in migrations
+
+The clickhouse-sqlalchemy package provides the `clickhouse+http://` URL dialect that
+Alembic uses to connect to ClickHouse.
+"""
+
+import logging
+import os
+from logging.config import fileConfig
+from pathlib import Path
+from urllib.parse import quote_plus
+
+from alembic import context
+from alembic.ddl import impl
+from clickhouse_sqlalchemy import Table, engines
+from clickhouse_sqlalchemy.drivers.compilers.sqlcompiler import ClickHouseSQLCompiler
+from dotenv import load_dotenv
+from sqlalchemy import Column, Connection, DateTime, MetaData, String, create_engine, pool, text
+
+from ch_migrate.config import get_env_config
+from ch_migrate.hooks import HookRegistry, run_hooks
+
+logger = logging.getLogger(__name__)
+
+# Alembic Config object
+config = context.config
+
+# Get environment name from CH_ENVIRONMENT (set by ch-migrate CLI)
+env_name = os.environ.get("CH_ENVIRONMENT", "dev")
+
+# Load .env.local for secrets (if it exists)
+env_local = Path.cwd() / ".env.local"
+if env_local.exists():
+    load_dotenv(env_local)
+
+# Load configuration
+config_path = Path.cwd() / "config.yaml"
+try:
+    env_config = get_env_config(env_name, config_path)
+except FileNotFoundError:
+    raise FileNotFoundError(
+        f"config.yaml not found in {Path.cwd()}\n"
+        f"Run 'ch-migrate init' to create project structure."
+    )
+
+# Setup logging
+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."""
+
+    __dialect__ = "clickhouse"
+    transactional_ddl = False
+
+    def version_table_impl(self, *, version_table, version_table_schema, **kw):
+        """Match bootstrap's version table when Alembic renders offline DDL."""
+        return Table(
+            version_table,
+            MetaData(),
+            Column("updated", DateTime, server_default=text("now()")),
+            Column("version_num", String, nullable=False),
+            engines.ReplacingMergeTree(order_by="updated"),
+            schema=version_table_schema,
+        )
+
+
+def get_sqlalchemy_url() -> str:
+    """Build SQLAlchemy-compatible ClickHouse URL from config."""
+    host = env_config["host"]
+    # Support both new 'migration_user' and legacy 'user' field
+    user = env_config.get("migration_user") or env_config.get("user")
+    password = env_config["password"]
+    database = env_config["database"]
+    port = env_config.get("port", 8443)
+    secure = env_config.get("secure", True)
+
+    # URL-encode password to handle special characters
+    encoded_password = quote_plus(password)
+
+    # Build SQLAlchemy URL
+    protocol_param = "?protocol=https" if secure else ""
+    return f"clickhouse+http://{user}:{encoded_password}@{host}:{port}/{database}{protocol_param}"
+
+
+# Export database name for use in migrations
+DATABASE_NAME = env_config["database"]
+os.environ["CH_DATABASE"] = DATABASE_NAME
+
+# Export cluster name if configured (for ON CLUSTER support)
+if env_config.get("cluster"):
+    os.environ["CH_CLUSTER"] = env_config["cluster"]
+
+
+def bootstrap_version_table(connection: Connection) -> None:
+    """
+    Create alembic_version table with ClickHouse-compatible engine.
+
+    Must be called before Alembic tries to use its version table.
+    The database must already exist (created by 'ch-migrate bootstrap').
+    """
+    db = DATABASE_NAME
+
+    # Check if version table exists
+    result = connection.execute(
+        text(
+            f"""
+        SELECT count() FROM system.tables
+        WHERE database = '{db}' AND name = 'alembic_version'
+    """
+        )
+    )
+    count = result.scalar()
+    exists = count is not None and count > 0
+
+    if not exists:
+        # Create with ClickHouse Cloud compatible engine
+        connection.execute(
+            text(
+                f"""
+            CREATE TABLE {db}.alembic_version
+            (
+                `updated` DateTime DEFAULT now(),
+                `version_num` String
+            )
+            ENGINE = ReplacingMergeTree()
+            ORDER BY updated
+            SETTINGS index_granularity = 8192
+        """
+            )
+        )
+        connection.commit()
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode - generates SQL without executing."""
+    url = get_sqlalchemy_url()
+    context.configure(
+        url=url,
+        target_metadata=None,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        version_table="alembic_version",
+        version_table_schema=DATABASE_NAME,
+    )
+    # No DBAPI interpolates offline output; its percent escaping would corrupt literals.
+    context.get_context().dialect.statement_compiler = _OfflineSQLCompiler
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+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)
+
+        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()
+
+
+class _OfflineSQLCompiler(ClickHouseSQLCompiler):
+    def post_process_text(self, text):
+        """Do not double percent signs in SQL intended for direct execution."""
+        return text
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

ch_migrate/helpers.py

@@ -0,0 +1,162 @@
+"""Helper functions for ClickHouse migrations."""
+
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+from typing import Any
+
+
+def _get_sql_dir() -> Path:
+    """Get the SQL directory relative to current working directory."""
+    return Path.cwd() / "migrations" / "sql"
+
+
+def read_sql(path: str, **kwargs: Any) -> str:
+    """
+    Read a SQL file and substitute placeholders.
+
+    Args:
+        path: Relative path from migrations/sql/ (e.g., "history/tables/users/001_abc.sql")
+        **kwargs: Values to substitute (e.g., db="mydb", password="secret")
+
+    Returns:
+        SQL string with placeholders replaced
+
+    Raises:
+        FileNotFoundError: If the SQL file doesn't exist
+
+    Example:
+        >>> read_sql("history/tables/users/001_abc.sql", db="mydb")
+        'CREATE TABLE mydb.users ...'
+    """
+    sql_path = _get_sql_dir() / path
+    if not sql_path.exists():
+        raise FileNotFoundError(f"SQL file not found: {sql_path}")
+
+    sql = sql_path.read_text()
+    return sql.format(**kwargs)
+
+
+def get_db() -> str:
+    """
+    Get database name from environment.
+
+    Returns:
+        Database name from CH_DATABASE env var, or "default" if not set
+    """
+    return os.environ.get("CH_DATABASE", "default")
+
+
+def get_cluster() -> str | None:
+    """
+    Get the cluster name from environment.
+
+    Returns:
+        Cluster name from CH_CLUSTER env var, or None if not set
+    """
+    return os.environ.get("CH_CLUSTER") or None
+
+
+def on_cluster() -> str:
+    """
+    Get the ON CLUSTER clause for use in DDL statements.
+
+    This is an opt-in template variable. Not all DDL supports ON CLUSTER
+    equally — dictionary creation, some ALTER operations, and system queries
+    have version-dependent ON CLUSTER support. Use this explicitly in
+    statements where ON CLUSTER is appropriate.
+
+    Returns:
+        "ON CLUSTER cluster_name" if cluster is configured, empty string otherwise
+
+    Example:
+        >>> read_sql("tables/users.sql", db=get_db(), on_cluster=on_cluster())
+        # In SQL: CREATE TABLE {db}.users {on_cluster} (...)
+    """
+    cluster = get_cluster()
+    if cluster:
+        return f"ON CLUSTER {cluster}"
+    return ""
+
+
+def get_config_value(key: str) -> str | None:
+    """
+    Get a configuration value from environment.
+
+    Args:
+        key: Environment variable name
+
+    Returns:
+        Value or None if not set
+    """
+    return os.environ.get(key)
+
+
+def create_dictionary(path: str, *, db: str | None = None, password: str | None = None) -> None:
+    """
+    Create a dictionary with automatic SELECT grant for dict_reader.
+
+    This helper:
+    1. Reads the dictionary SQL file
+    2. Parses the source table from the SQL
+    3. Grants SELECT on that table to dict_reader
+    4. Creates the dictionary
+
+    Args:
+        path: Relative path to dictionary SQL file
+        db: Database name (defaults to get_db())
+        password: dict_reader password (defaults to DICT_READER_PASSWORD env var)
+
+    Example:
+        create_dictionary("history/dictionaries/dict_users/001_abc.sql")
+    """
+    # Import here to avoid circular imports and allow usage without alembic context
+    from alembic import op
+
+    if db is None:
+        db = get_db()
+    if password is None:
+        password = os.environ.get("DICT_READER_PASSWORD")
+        if not password:
+            raise ValueError(
+                "DICT_READER_PASSWORD environment variable is required for dictionaries."
+            )
+
+    # Read and format the dictionary SQL
+    dict_sql = read_sql(path, db=db, password=password)
+
+    # Parse the source table from the SQL
+    source_table = _parse_source_table(dict_sql)
+
+    if source_table:
+        # Grant SELECT to dict_reader before creating dictionary
+        op.execute(f"GRANT SELECT ON {db}.{source_table} TO dict_reader")
+
+    # Create the dictionary
+    op.execute(dict_sql)
+
+
+def _parse_source_table(dict_sql: str) -> str | None:
+    """
+    Parse the source table name from dictionary SQL.
+
+    Supports two patterns:
+    1. TABLE 'table_name' (simple table source)
+    2. QUERY '...FROM db.table_name...' (query source)
+
+    Returns:
+        Table name or None if not found
+    """
+    # Try TABLE 'table_name' pattern first
+    table_match = re.search(r"TABLE\s+'(\w+)'", dict_sql, re.IGNORECASE)
+    if table_match:
+        return table_match.group(1)
+
+    # Try QUERY pattern: FROM db.table_name
+    query_match = re.search(r"FROM\s+\w+\.(\w+)", dict_sql, re.IGNORECASE)
+    if query_match:
+        return query_match.group(1)
+
+    return None

ch_migrate/hooks.py

@@ -0,0 +1,74 @@
+"""Pre/post migration hook support for ch-migrate-cli."""
+
+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):
+        # Use explicit replace instead of str.format() to avoid KeyError on
+        # unknown placeholders (e.g. {cluster}) and ClickHouse parameterized
+        # query syntax like {param:String}.
+        resolved = hook_sql.replace("{db}", db)
+        logger.info("[%s] hook %d/%d for %s: %s", phase, i, len(hooks), revision, resolved)
+        connection.execute(text(resolved))
+        connection.commit()