add test to check changelog does not contain BEGIN; or COMMIT;

Author: 3nids

pum/exceptions.py

@@ -25,6 +25,11 @@ class PumSqlException(Exception):
 
     pass
 
+class PumInvalidChangelog(Exception):
+    """Exception raised for invalid changelog."""
+
+    pass
+
 
 class PumConfigError(PumException):
     """Exception raised for errors in the PUM configuration."""

pum/utils/execute_sql.py

@@ -4,7 +4,7 @@
 from psycopg import Connection, Cursor
 from psycopg.errors import SyntaxError
 
-from ..exceptions import PumSqlException
+from ..exceptions import PumSqlException, PumInvalidChangelog
 import re
 
 logger = logging.getLogger(__name__)
@@ -42,6 +42,13 @@ def remove_sql_comments(sql):
                 sql = re.sub(r'(?m)(^|;)\s*--.*?(\r\n|\r|\n)', r'\1', sql)
                 return sql
             sql_content = remove_sql_comments(sql_content)
+
+            # Check for forbidden transaction statements
+            forbidden_statements = ['BEGIN;', 'COMMIT;']
+            for forbidden in forbidden_statements:
+                if re.search(rf'\b{forbidden[:-1]}\b\s*;', sql_content, re.IGNORECASE):
+                    raise PumInvalidChangelog(f"SQL contains forbidden transaction statement: {forbidden}")
+
             if parameters:
                 for key, value in parameters.items():
                     sql_content = sql_content.replace(f"{{{{ {key} }}}}", str(value))

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

@@ -0,0 +1,13 @@
+BEGIN;
+
+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)
+);
+
+COMMIT;

test/test_upgrader.py

@@ -114,6 +114,7 @@ def test_install_complex_files_content(self):
         )
         upgrader.install()
         self.assertTrue(sm.exists(self.conn))
+
     def test_install_multiple_changelogs(self):
         cfg = PumConfig()
         sm = SchemaMigrations(cfg)
@@ -139,6 +140,18 @@ def test_install_multiple_changelogs(self):
             ],
         )
 
+    def test_invalid_changelog(self):
+        cfg = PumConfig()
+        sm = SchemaMigrations(cfg)
+        self.assertFalse(sm.exists(self.conn))
+        upgrader = Upgrader(
+            pg_service=self.pg_service,
+            config=cfg,
+            dir=str(Path("test") / "data" / "invalid_changelog"),
+        )
+        with self.assertRaises(Exception) as context:
+            upgrader.install()
+        self.assertTrue("SQL contains forbidden transaction statement: BEGIN;" in str(context.exception))
 
 if __name__ == "__main__":
     unittest.main()