support custom migration table

Author: 3nids

pum/schema_migrations.py

@@ -77,8 +77,17 @@ def create(self, conn: Connection, commit: bool = True):
         if self.exists(conn):
             logger.debug(f"{self.config.schema_migrations_table} table already exists")
             return
+        
+        # Create the schema if it doesn't exist
+        create_schema_query = None
+        schema = "public"
+        table_identifiers = self.config.schema_migrations_table.split(".")
+        if len(table_identifiers) == 2:
+            schema = table_identifiers[0]
+        if schema != "public":
+            create_schema_query = sql.SQL("CREATE SCHEMA IF NOT EXISTS {schema};").format(schema=sql.Identifier(schema))
 
-        create_query = sql.SQL(
+        create_table_query = sql.SQL(
             """CREATE TABLE IF NOT EXISTS {schema_migrations_table}
             (
             id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
@@ -102,7 +111,9 @@ def create(self, conn: Connection, commit: bool = True):
             schema_migrations_table=sql.Identifier(*self.config.schema_migrations_table.split("."))
         )
 
-        execute_sql(conn, create_query)
+        if create_schema_query:
+            execute_sql(conn, create_schema_query)
+        execute_sql(conn, create_table_query)
         execute_sql(conn, comment_query)
 
         logger.info(f"Created {self.config.schema_migrations_table} table")

test/data/custom_migration_table/.pum-config.yaml

@@ -0,0 +1,3 @@
+
+schema_migrations_table:
+  pum_test.custom_migration

test/data/custom_migration_table/changelogs/1.2.3/create_northwind.sql

@@ -0,0 +1,9 @@
+CREATE SCHEMA IF NOT EXISTS pum_test_data;
+
+CREATE TABLE pum_test_data.some_table (
+    id INT PRIMARY KEY,
+    name VARCHAR(100) NOT NULL,
+    created_date DATE DEFAULT CURRENT_DATE,
+    is_active BOOLEAN DEFAULT TRUE,
+    amount NUMERIC(10,2)
+);

test/test_upgrader.py

@@ -19,6 +19,7 @@ class TestUpgrader(unittest.TestCase):
     def tearDown(self):
         self.cur.execute("DROP SCHEMA IF EXISTS pum_test_data CASCADE;")
         self.cur.execute("DROP TABLE IF EXISTS public.pum_migrations;")
+        self.cur.execute("DROP TABLE IF EXISTS pum_test.custom_migration;")
         self.conn.commit()
         self.conn.close()
 
@@ -35,6 +36,7 @@ def setUp(self):
         self.cur = self.conn.cursor()
         self.cur.execute("DROP SCHEMA IF EXISTS pum_test_data CASCADE;")
         self.cur.execute("DROP TABLE IF EXISTS public.pum_migrations;")
+        self.cur.execute("DROP TABLE IF EXISTS pum_test.custom_migration;")
         self.conn.commit()
 
         self.tmpdir = tempfile.TemporaryDirectory()
@@ -92,6 +94,16 @@ def test_install_custom_directory(self):
         upgrader.install()
         self.assertTrue(sm.exists(self.conn))
 
+    def test_install_custom_migration_table(self):
+        cfg = PumConfig.from_yaml(str(Path("test") / "data" / "custom_migration_table" / ".pum-config.yaml"))
+        sm = SchemaMigrations(cfg)
+        self.assertFalse(sm.exists(self.conn))
+        upgrader = Upgrader(
+            pg_service=self.pg_service, config=cfg, dir=str(Path("test") / "data" / "custom_migration_table")
+        )
+        upgrader.install()
+        self.assertTrue(sm.exists(self.conn))
+
     def test_install_multiple_changelogs(self):
         cfg = PumConfig()
         sm = SchemaMigrations(cfg)