Allow default db name #24 (#26)

* Allow default db name (from db URL)

---------

Co-authored-by: Aleh Strakachuk <[email protected]>

src/clickhouse_migrations/__init__.py

@@ -1,4 +1,4 @@
 """
 Simple file-based migrations for clickhouse
 """
-__version__ = "0.7.0"
+__version__ = "0.7.1"

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)

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
 

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)
 

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:

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

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"}],