Add explicit migrations and fake option to update applied migrations

dev/docker-compose.yaml

@@ -17,8 +17,8 @@ services:
         ipv4_address: 172.23.0.10
 
   clickhouse01:
-    image: yandex/clickhouse-server:20.4
-    # image: altinity/clickhouse-server:23.8.11.29.altinitystable
+    # image: yandex/clickhouse-server:20.4
+    image: altinity/clickhouse-server:23.8.11.29.altinitystable
     container_name: clickhouse01
     hostname: clickhouse01
     networks:
@@ -33,8 +33,8 @@ services:
       - zookeeper
 
   clickhouse02:
-    image: yandex/clickhouse-server:20.4
-    # image: altinity/clickhouse-server:23.8.11.29.altinitystable
+    # image: yandex/clickhouse-server:20.4
+    image: altinity/clickhouse-server:23.8.11.29.altinitystable
     container_name: clickhouse02
     hostname: clickhouse02
     networks:
@@ -46,8 +46,8 @@ services:
       - zookeeper
 
   clickhouse03:
-    image: yandex/clickhouse-server:20.4
-    # image: altinity/clickhouse-server:23.8.11.29.altinitystable
+    # image: yandex/clickhouse-server:20.4
+    image: altinity/clickhouse-server:23.8.11.29.altinitystable
     container_name: clickhouse03
     hostname: clickhouse03
     networks:
@@ -59,8 +59,8 @@ services:
       - zookeeper
 
   clickhouse04:
-    image: yandex/clickhouse-server:20.4
-    # image: altinity/clickhouse-server:23.8.11.29.altinitystable
+    # image: yandex/clickhouse-server:20.4
+    image: altinity/clickhouse-server:23.8.11.29.altinitystable
     container_name: clickhouse04
     hostname: clickhouse04
     networks:

src/clickhouse_migrations/clickhouse_cluster.py

@@ -88,15 +88,17 @@ def migrate(
         self,
         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,
@@ -105,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(
@@ -115,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:
@@ -125,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

@@ -95,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"),
@@ -119,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

@@ -181,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