From a032ac7df63cdd5479f4a383cbd7a70ed51ad8ee Mon Sep 17 00:00:00 2001 From: Aleh Strakachuk Date: Mon, 1 Jul 2024 22:59:22 +0300 Subject: [PATCH] Allow default db name #24 (#26) * Allow default db name (from db URL) --------- Co-authored-by: Aleh Strakachuk --- src/clickhouse_migrations/__init__.py | 2 +- .../clickhouse_cluster.py | 34 ++++++++++++++--- src/clickhouse_migrations/command_line.py | 19 +++++++++- src/clickhouse_migrations/migration.py | 18 +++++++-- src/clickhouse_migrations/migrator.py | 15 ++++++-- src/tests/test_clickhouse_migration.py | 38 +++++++++++++++++++ src/tests/test_init_clickhouse_cluster.py | 7 ++-- 7 files changed, 112 insertions(+), 21 deletions(-) diff --git a/src/clickhouse_migrations/__init__.py b/src/clickhouse_migrations/__init__.py index 8276820..aebd650 100644 --- a/src/clickhouse_migrations/__init__.py +++ b/src/clickhouse_migrations/__init__.py @@ -1,4 +1,4 @@ """ Simple file-based migrations for clickhouse """ -__version__ = "0.7.0" +__version__ = "0.7.1" diff --git a/src/clickhouse_migrations/clickhouse_cluster.py b/src/clickhouse_migrations/clickhouse_cluster.py index 25414c4..15df263 100644 --- a/src/clickhouse_migrations/clickhouse_cluster.py +++ b/src/clickhouse_migrations/clickhouse_cluster.py @@ -16,12 +16,16 @@ def __init__( db_password: str = DB_PASSWORD, db_port: str = DB_PORT, db_url: Optional[str] = None, + db_name: Optional[str] = None, **kwargs, ): self.db_url: Optional[str] = db_url + self.default_db_name: Optional[str] = db_name + if db_url: parts = self.db_url.split("/") if len(parts) == 4: + self.default_db_name = parts[-1] parts = parts[0:-1] self.db_url = "/".join(parts) @@ -32,7 +36,9 @@ def __init__( self.db_password = db_password self.connection_kwargs = kwargs - def connection(self, db_name: str) -> Client: + def connection(self, db_name: Optional[str] = None) -> Client: + db_name = db_name if db_name is not None else self.default_db_name + if self.db_url: db_url = self.db_url if db_name: @@ -49,7 +55,11 @@ def connection(self, db_name: str) -> Client: ) return ch_client - def create_db(self, db_name, cluster_name=None): + def create_db( + self, db_name: Optional[str] = None, cluster_name: Optional[str] = None + ): + db_name = db_name if db_name is not None else self.default_db_name + with self.connection("") as conn: if cluster_name is None: conn.execute(f'CREATE DATABASE IF NOT EXISTS "{db_name}"') @@ -58,27 +68,37 @@ def create_db(self, db_name, cluster_name=None): f'CREATE DATABASE IF NOT EXISTS "{db_name}" ON CLUSTER "{cluster_name}"' ) - def init_schema(self, db_name, cluster_name=None): + def init_schema( + self, db_name: Optional[str] = None, cluster_name: Optional[str] = None + ): + db_name = db_name if db_name is not None else self.default_db_name + with self.connection(db_name) as conn: migrator = Migrator(conn) migrator.init_schema(cluster_name) def show_tables(self, db_name): + db_name = db_name if db_name is not None else self.default_db_name + with self.connection(db_name) as conn: result = conn.execute("show tables") return [t[0] for t in result] def migrate( self, - db_name: str, + db_name: Optional[str], migration_path: Path, + explicit_migrations: Optional[List[str]] = None, cluster_name: Optional[str] = None, create_db_if_no_exists: bool = True, multi_statement: bool = True, dryrun: bool = False, + fake: bool = False, ): + db_name = db_name if db_name is not None else self.default_db_name + storage = MigrationStorage(migration_path) - migrations = storage.migrations() + migrations = storage.migrations(explicit_migrations) return self.apply_migrations( db_name, @@ -87,6 +107,7 @@ def migrate( create_db_if_no_exists=create_db_if_no_exists, multi_statement=multi_statement, dryrun=dryrun, + fake=fake, ) def apply_migrations( @@ -97,6 +118,7 @@ def apply_migrations( cluster_name: Optional[str] = None, create_db_if_no_exists: bool = True, multi_statement: bool = True, + fake: bool = False, ) -> List[Migration]: if create_db_if_no_exists: if cluster_name is None: @@ -107,4 +129,4 @@ def apply_migrations( with self.connection(db_name) as conn: migrator = Migrator(conn, dryrun) migrator.init_schema(cluster_name) - return migrator.apply_migration(migrations, multi_statement) + return migrator.apply_migration(migrations, multi_statement, fake) diff --git a/src/clickhouse_migrations/command_line.py b/src/clickhouse_migrations/command_line.py index 8664a6a..e921101 100644 --- a/src/clickhouse_migrations/command_line.py +++ b/src/clickhouse_migrations/command_line.py @@ -7,7 +7,6 @@ from clickhouse_migrations.clickhouse_cluster import ClickhouseCluster from clickhouse_migrations.defaults import ( DB_HOST, - DB_NAME, DB_PASSWORD, DB_PORT, DB_USER, @@ -64,7 +63,7 @@ def get_context(args): ) parser.add_argument( "--db-name", - default=os.environ.get("DB_NAME", DB_NAME), + default=os.environ.get("DB_NAME", None), help="Clickhouse database name", ) parser.add_argument( @@ -96,6 +95,20 @@ def get_context(args): type=bool, help="Dry run mode", ) + parser.add_argument( + "--fake", + default=os.environ.get("FAKE", "0"), + type=bool, + help="Marks the migrations up to the target one (following the rules above) as applied, " + "but without actually running the SQL to change your database schema.", + ) + parser.add_argument( + "--migrations", + default=os.environ.get("MIGRATIONS", "").split(","), + type=str, + nargs="+", + help="Dry run mode", + ) parser.add_argument( "--secure", default=os.environ.get("SECURE", "0"), @@ -120,9 +133,11 @@ def migrate(ctx) -> int: cluster.migrate( db_name=ctx.db_name, migration_path=ctx.migrations_dir, + explicit_migrations=ctx.migrations, cluster_name=ctx.cluster_name, multi_statement=ctx.multi_statement, dryrun=ctx.dry_run, + fake=ctx.fake, ) return 0 diff --git a/src/clickhouse_migrations/migration.py b/src/clickhouse_migrations/migration.py index a3b39be..03018ec 100644 --- a/src/clickhouse_migrations/migration.py +++ b/src/clickhouse_migrations/migration.py @@ -2,7 +2,7 @@ import os from collections import namedtuple from pathlib import Path -from typing import List +from typing import List, Optional Migration = namedtuple("Migration", ["version", "md5", "script"]) @@ -19,17 +19,27 @@ def filenames(self) -> List[Path]: return l - def migrations(self) -> List[Migration]: + def migrations( + self, explicit_migrations: Optional[List[str]] = None + ) -> List[Migration]: migrations: List[Migration] = [] for full_path in self.filenames(): + version_string = full_path.name.split("_")[0] + version_number = int(version_string) migration = Migration( - version=int(full_path.name.split("_")[0]), + version=version_number, script=str(full_path.read_text(encoding="utf8")), md5=hashlib.md5(full_path.read_bytes()).hexdigest(), ) - migrations.append(migration) + if ( + explicit_migrations is None + or full_path.stem in explicit_migrations + or version_string in explicit_migrations + or str(version_number) in explicit_migrations + ): + migrations.append(migration) migrations.sort(key=lambda m: m.version) diff --git a/src/clickhouse_migrations/migrator.py b/src/clickhouse_migrations/migrator.py index d764c0f..ed318de 100644 --- a/src/clickhouse_migrations/migrator.py +++ b/src/clickhouse_migrations/migrator.py @@ -90,7 +90,10 @@ def migrations_to_apply(self, incoming: List[Migration]) -> List[Migration]: return sorted(to_apply, key=lambda x: x.version) def apply_migration( - self, migrations: List[Migration], multi_statement: bool + self, + migrations: List[Migration], + multi_statement: bool, + fake: bool = False, ) -> List[Migration]: new_migrations = self.migrations_to_apply(migrations) @@ -106,10 +109,14 @@ def apply_migration( logging.info("Migration contains %s statements to apply", len(statements)) for statement in statements: - if not self._dryrun: - self._execute(statement) - else: + if fake: + logging.warning( + "Fake mode, statement will be skipped: %s", statement + ) + elif self._dryrun: logging.info("Dry run mode, would have executed: %s", statement) + else: + self._execute(statement) logging.info("Migration applied, need to update schema version table.") if not self._dryrun: diff --git a/src/tests/test_clickhouse_migration.py b/src/tests/test_clickhouse_migration.py index 8207c58..8fc64dd 100644 --- a/src/tests/test_clickhouse_migration.py +++ b/src/tests/test_clickhouse_migration.py @@ -159,6 +159,19 @@ def test_main_pass_db_name_ok(): ) +def test_main_pass_db_url_ok(): + migrate( + get_context( + [ + "--db-url", + "clickhouse://default:@localhost:9000/pytest", + "--migrations-dir", + str(TESTS_DIR / "migrations"), + ] + ) + ) + + def test_check_multistatement_arg(): context = get_context(["--multi-statement", "false"]) assert context.multi_statement is False @@ -168,3 +181,28 @@ def test_check_multistatement_arg(): context = get_context(["--multi-statement", "0"]) assert context.multi_statement is False + + +def test_check_explicit_migrations_ok(): + migrate( + get_context( + [ + "--migrations", + "001_init", + "002", + "3", + "--migrations-dir", + str(TESTS_DIR / "complex_migrations"), + ] + ) + ) + + +def test_check_explicit_migrations_args_ok(): + context = get_context(["--migrations", "001_init", "002_test2"]) + assert context.migrations == ["001_init", "002_test2"] + + +def test_check_fake_ok(): + context = get_context(["--fake"]) + assert context.fake is True diff --git a/src/tests/test_init_clickhouse_cluster.py b/src/tests/test_init_clickhouse_cluster.py index c8c5414..7bfa1d9 100644 --- a/src/tests/test_init_clickhouse_cluster.py +++ b/src/tests/test_init_clickhouse_cluster.py @@ -3,7 +3,6 @@ import pytest from clickhouse_migrations.clickhouse_cluster import ClickhouseCluster -from clickhouse_migrations.defaults import DB_URL from clickhouse_migrations.migration import Migration, MigrationStorage TESTS_DIR = Path(__file__).parent @@ -12,13 +11,13 @@ @pytest.fixture def cluster(): - return ClickhouseCluster(db_url=DB_URL) + return ClickhouseCluster(db_url="clickhouse://default:@localhost:9000/pytest") def test_apply_new_migration_ok(cluster): - cluster.init_schema("pytest") + cluster.init_schema() - with cluster.connection("pytest") as conn: + with cluster.connection() as conn: conn.execute( "INSERT INTO schema_versions(version, script, md5) VALUES", [{"version": 1, "script": "SHOW TABLES", "md5": "12345"}],