Working on sqlite integration tests

Author: k2bd

src/flux/migration/read_migration.py

@@ -14,8 +14,8 @@ def read_migrations(*, config: FluxConfig) -> list[Migration]:
     """
     migrations = []
     for migration_file in os.listdir(config.migration_directory):
-        if migration_file.endswith("_up.sql"):
-            migration_id = migration_file[:-7]
+        if migration_file.endswith(".sql") and not migration_file.endswith(".undo.sql"):
+            migration_id = migration_file[:-4]
             migrations.append(
                 read_sql_migration(config=config, migration_id=migration_id)
             )
@@ -88,8 +88,8 @@ def read_sql_migration(*, config: FluxConfig, migration_id: str) -> Migration:
     """
     Read a pair of SQL migration files and return a Migration object
     """
-    up_file = os.path.join(config.migration_directory, f"{migration_id}_up.sql")
-    down_file = os.path.join(config.migration_directory, f"{migration_id}_down.sql")
+    up_file = os.path.join(config.migration_directory, f"{migration_id}.sql")
+    down_file = os.path.join(config.migration_directory, f"{migration_id}.undo.sql")
 
     try:
         with open(up_file) as f:
@@ -128,6 +128,13 @@ def read_repeatable_sql_migration(
     except Exception as e:
         raise MigrationLoadingError("Error reading up migration") from e
 
+    if os.path.exists(
+        os.path.join(
+            config.migration_directory, migration_subdir, f"{migration_id}.undo.sql"
+        )
+    ):
+        raise MigrationLoadingError("Repeatable migrations cannot have a down")
+
     return Migration(id=migration_id, up=up, down=None)
 
 
@@ -139,19 +146,19 @@ def read_python_migration(*, config: FluxConfig, migration_id: str) -> Migration
 
     with temporary_module(migration_file) as module:
         try:
-            up_migration = module.up()
+            up_migration = module.apply()
         except Exception as e:
             raise MigrationLoadingError("Error reading up migration") from e
         if not isinstance(up_migration, str):
             raise MigrationLoadingError("Up migration must return a string")
-        if hasattr(module, "down"):
+        if hasattr(module, "undo"):
             try:
-                down_migration = module.down()
+                down_migration = module.undo()
             except Exception as e:
                 raise MigrationLoadingError("Error reading down migration") from e
-            if not isinstance(down_migration, str) and down_migration is not None:
+            if not isinstance(down_migration, str):
                 raise MigrationLoadingError(
-                    "Down migration must return a string or None"
+                    "Down migration must return a string"
                 )
         else:
             down_migration = None
@@ -174,12 +181,12 @@ def read_repeatable_python_migration(
 
     with temporary_module(migration_file) as module:
         try:
-            up_migration = module.up()
+            up_migration = module.apply()
         except Exception as e:
             raise MigrationLoadingError("Error reading up migration") from e
         if not isinstance(up_migration, str):
             raise MigrationLoadingError("Up migration must return a string")
-        if hasattr(module, "down"):
+        if hasattr(module, "undo"):
             raise MigrationLoadingError("Repeatable migrations cannot have a down")
 
     return Migration(id=migration_id, up=up_migration, down=None)

tests/integration/sqlite/backend.py

@@ -2,8 +2,8 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import Any
-
 import aiosqlite
+from aiosqlite import Connection
 
 from flux.backend.applied_migration import AppliedMigration
 from flux.backend.base import MigrationBackend
@@ -18,7 +18,7 @@ class SQLiteBackend(MigrationBackend):
     db_path: str
     migrations_table: str
 
-    _conn: Any = field(init=False, repr=False)
+    _conn: Connection = field(init=False, repr=False)
 
     @classmethod
     def from_config(cls, config: FluxConfig) -> "SQLiteBackend":
@@ -64,10 +64,11 @@ async def transaction(self):
         """
         try:
             yield
-            await self._conn.commit()
         except Exception:
             await self._conn.rollback()
             raise
+        else:
+            await self._conn.commit()
 
     @asynccontextmanager
     async def migration_lock(self):
@@ -156,3 +157,19 @@ async def get_applied_migrations(self) -> set[AppliedMigration]:
             AppliedMigration(id=row[0], hash=row[1], applied_at=row[2])
             for row in await cursor.fetchall()
         }
+
+    # -- Testing methods
+
+    async def table_info(self, table_name: str):
+        """
+        Get information about a table in the database
+        """
+        cursor = await self._conn.execute(f"pragma table_info({table_name})")
+        return await cursor.fetchall()
+
+    async def get_all_rows(self, table_name: str):
+        """
+        Get all rows from a table
+        """
+        cursor = await self._conn.execute(f"select * from {table_name}")
+        return await cursor.fetchall()

tests/integration/sqlite/conftest.py

@@ -5,6 +5,8 @@
 import pytest
 
 from tests.integration.sqlite.backend import SQLiteBackend
+from tests.integration.sqlite.constants import MIGRATIONS_1_DIR
+import shutil
 
 
 @pytest.fixture
@@ -17,3 +19,15 @@ def temp_db_path() -> Generator[str, None, None]:
 @pytest.fixture
 def sqlite_backend(temp_db_path: str):
     yield SQLiteBackend(db_path=temp_db_path, migrations_table="_flux_migrations")
+
+
+@pytest.fixture
+def example_migrations_dir() -> Generator[str, None, None]:
+    """
+    Create a temporary copy of the migrations directory that can be freely
+    modified by tests
+    """
+    with TemporaryDirectory() as tempdir:
+        migrations_dir = os.path.join(tempdir, "migrations")
+        shutil.copytree(MIGRATIONS_1_DIR, migrations_dir)
+        yield migrations_dir

tests/integration/sqlite/data/migrations-1/20200101_001_add_description_to_simple_table_up.sql tests/integration/sqlite/data/migrations-1/20200101_001_add_description_to_simple_table.sql

@@ -0,0 +1 @@
+alter table simple_table drop column description

tests/integration/sqlite/data/migrations-1/20200101_001_add_description_to_simple_table.undo.sql

@@ -1,23 +1,10 @@
-def up():
+def apply():
     return """
     alter table another_table add column timestamp text;
     """
 
 
-def down():
-    """
-    pragma foreign_keys=off;
-
-    create table another_table_new (
-        id integer primary key autoincrement,
-        value integer
-    );
-
-    insert into another_table_new (id, value)
-    select id, value from another_table;
-
-    drop table another_table;
-    alter table another_table_new rename to another_table;
-
-    pragma foreign_keys=on;
+def undo():
+    return """
+    alter table another_table drop column timestamp;
     """

tests/integration/sqlite/data/migrations-1/20200101_001_add_description_to_simple_table_down.sql

@@ -1,4 +1,4 @@
-def up():
+def apply():
     return """
     create table if not exists new_table (
         id integer primary key autoincrement,

tests/integration/sqlite/data/migrations-1/20200101_002_add_timestamp_to_another_table.py

@@ -1,4 +1,4 @@
-def up():
+def apply():
     return """
     create view if not exists view2 as
     select id, value from another_table where value > 10;

tests/integration/sqlite/data/migrations-1/20200102_001_create_new_table.py

@@ -1,4 +1,4 @@
-def up():
+def apply():
     return """
     create table if not exists simple_table (
         id integer primary key autoincrement,