Use model to insert all fixtures in the same transaction. Initial async fixtures implementation.

Author: febus982

.idea/dataSources.xml

@@ -1,10 +1,15 @@
+import hashlib
+import importlib
 import logging
+import sys
 from asyncio import get_event_loop
-from datetime import datetime
+from datetime import datetime, timezone
+from os import listdir, path
+from os.path import isfile, join
 
-from sqlalchemy import Table, Column, String, DateTime
-from sqlalchemy.ext.asyncio import AsyncEngine
-from sqlalchemy.orm import registry
+from sqlalchemy import Table, Column, String, DateTime, UniqueConstraint, text, select
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+from sqlalchemy.orm import registry, sessionmaker, Mapped, mapped_column
 
 from alembic import context
 from common.bootstrap import application_init
@@ -35,17 +40,21 @@
 db_names = target_metadata.keys()
 config.set_main_option("databases", ",".join(db_names))
 
-def inject_fixture_tables(registry_mapper: registry):
-    Table(
-        "alembic_fixtures",
-        registry_mapper.metadata,
-        Column("filename", String(), primary_key=True),
-        Column("signature", String(), nullable=False),
-        Column("processed_at", DateTime(timezone=True), nullable=False, default=datetime.now),
-    )
+def generate_fixture_migration_model(declarative_base: type):
+    class FixtureMigration(declarative_base):
+        __tablename__ = "alembic_fixtures"
+
+        bind: Mapped[str] = mapped_column(String(), primary_key=True)
+        filename: Mapped[str] = mapped_column(String(), primary_key=True)
+        signature: Mapped[str] = mapped_column(String(), nullable=False)
+        processed_at: Mapped[datetime] = mapped_column(DateTime(), nullable=False, default=datetime.now)
+    return FixtureMigration
 
+fixture_migration_models = {}
 for name in db_names:
-    inject_fixture_tables(sa_manager.get_bind(name).registry_mapper)
+    fixture_migration_models[name] = generate_fixture_migration_model(
+        sa_manager.get_bind(name).declarative_base
+    )
 
 
 # add your model's MetaData objects here
@@ -66,6 +75,51 @@ def inject_fixture_tables(registry_mapper: registry):
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
+def calculate_md5(file_path):
+    hasher = hashlib.md5()
+    with open(file_path, 'rb') as file:
+        hasher.update(file.read())
+    return hasher.hexdigest()
+
+async def a_migrate_fixtures(bind_name: str, Session: async_sessionmaker[AsyncSession]) -> str:
+    alembic_path = path.dirname(path.realpath(__file__))
+    fixtures_path = alembic_path + "/fixtures"
+    sys.path.append(alembic_path)
+    module_names = [
+        f[:-3] for f in listdir(fixtures_path)
+        if isfile(join(fixtures_path, f))
+        and f.endswith(".py")
+        and f != "__init__.py"
+    ]
+    async with Session() as session:
+        for module_name in module_names:
+            logging.debug(f"Creating {module_name} fixtures for {bind_name}")
+            m = importlib.import_module(f"fixtures.{module_name}")
+            fixture_migration = await session.get(
+                fixture_migration_models[bind_name],
+                (bind_name, f"{module_name}.py")
+            )
+            signature = calculate_md5(f"{fixtures_path}/{module_name}.py")
+            if fixture_migration:
+                if signature != fixture_migration.signature:
+                    logging.warning(f"Signature mismatch for {fixture_migration.filename} fixture. The file has been modified after being processed.")
+                logging.debug(f"{module_name} fixtures already migrated for {bind_name}")
+                continue
+
+            session.add_all(m.fixtures().get(bind_name, []))
+            session.add(fixture_migration_models[bind_name](
+                bind=bind_name,
+                filename=f"{module_name}.py",
+                signature=signature,
+            ))
+            try:
+                await session.commit()
+                logging.info(f"{module_name} fixtures correctly created for {bind_name}")
+            except:
+                await session.rollback()
+                logging.error(f"{module_name} fixtures failed to apply to {bind_name}")
+
+
 
 def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode.
@@ -149,13 +203,17 @@ async def run_migrations_online() -> None:
         for name, rec in engines.items():
             logger.info(f"Migrating database {name}")
             if isinstance(rec["engine"], AsyncEngine):
-
                 def migration_callable(*args, **kwargs):
                     return do_run_migration(*args, name=name, **kwargs)
 
                 await rec["connection"].run_sync(migration_callable)
+                Session = async_sessionmaker(bind=rec["connection"])
+                await a_migrate_fixtures(bind_name=name, Session=Session)
+
             else:
                 do_run_migration(rec["connection"], name)
+                # Session = sessionmaker(bind=rec["connection"])
+                # session = Session()
 
         if USE_TWOPHASE:
             for rec in engines.values():

src/alembic/env.py

@@ -0,0 +1,25 @@
+"""
+`fixtures` is a dictionary following the format:
+
+"BIND_NAME": "LIST_OF_FACTORIES"
+"""
+from typing import Dict, List
+
+from factory import Factory
+
+from domains.books._models import BookModel
+
+
+def fixtures() -> Dict[str, List]:
+    class BookFactory(Factory):
+        class Meta:
+            model = BookModel
+
+    return {
+        "default": [
+            BookFactory(
+                title="The Shining",
+                author_name="Stephen King",
+            ),
+        ],
+    }

src/alembic/fixtures/__init__.py

@@ -27,10 +27,11 @@ def downgrade(engine_name: str) -> None:
 
 def upgrade_default() -> None:
     op.create_table('alembic_fixtures',
+        sa.Column('bind', sa.String(), nullable=False),
         sa.Column('filename', sa.String(), nullable=False),
         sa.Column('signature', sa.String(), nullable=False),
         sa.Column('processed_at', sa.DateTime(timezone=True), nullable=False),
-        sa.PrimaryKeyConstraint('filename')
+        sa.PrimaryKeyConstraint('bind', 'filename'),
     )