🏭 Factories

.gitignore

@@ -1,4 +1,4 @@
 __pycache__
 
 # Development sqlite db
-chowda.sqlite
+*.sqlite

chowda/app.py

@@ -1,16 +1,13 @@
 from fastapi import FastAPI
-from sqlalchemy import create_engine
 from starlette.responses import HTMLResponse
 from starlette.routing import Route
 from starlette_admin.contrib.sqlmodel import Admin, ModelView
 
 from sqlmodel import SQLModel
 
-from .config import ENGINE_URI
 from .models import User, MediaFile, Collection, ClamsApp, Pipeline, Batch, ClamsEvent
 from ._version import __version__
-
-engine = create_engine(ENGINE_URI, connect_args={'check_same_thread': False}, echo=True)
+from .db import engine
 
 
 def init_database() -> None:

chowda/db.py

@@ -0,0 +1,12 @@
+from sqlalchemy import create_engine
+from .config import ENGINE_URI
+
+engine = create_engine(ENGINE_URI, connect_args={'check_same_thread': False}, echo=True)
+
+
+def create_async_engine():
+    from sqlmodel.ext.asyncio.session import AsyncEngine
+
+    return AsyncEngine(
+        create_engine(ENGINE_URI, connect_args={'check_same_thread': False}, echo=True)
+    )

chowda/factories.py

@@ -0,0 +1,86 @@
+from chowda.db import engine
+from chowda.models import (
+    MediaFile,
+    Collection,
+    ClamsApp,
+    User,
+    MediaFileCollectionLink,
+    MediaFileBatchLink,
+    Pipeline,
+    ClamsEvent,
+)
+from pydantic_factories import (
+    ModelFactory,
+    SyncPersistenceProtocol,
+    AsyncPersistenceProtocol,
+)
+from typing import TypeVar, List
+from pydantic import BaseModel
+from sqlmodel import Session
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class SyncPersistenceHandler(SyncPersistenceProtocol[T]):
+    def save(self, data: T) -> T:
+        with Session(engine) as session:
+            session.add(data)
+            session.commit()
+
+    def save_many(self, data: List[T]) -> List[T]:
+        with Session(engine) as session:
+            for d in data:
+                session.add(d)
+            session.commit()
+
+
+class AsyncPersistenceHandler(AsyncPersistenceProtocol[T]):
+    async def save(self, data: T) -> T:
+        async with AsyncSession(engine) as session:
+            session.add(data)
+            await session.commit()
+
+    async def save_many(self, data: List[T]) -> List[T]:
+        async with AsyncSession(engine) as session:
+            for d in data:
+                session.add(d)
+            await session.commit()
+
+
+class ChowdaFactory(ModelFactory):
+    __allow_none_optionals__ = False
+    __sync_persistence__ = SyncPersistenceHandler
+    __async_persistence__ = AsyncPersistenceHandler
+
+
+class MediaFileFactory(ChowdaFactory):
+    __model__ = MediaFile
+
+
+class CollectionFactory(ChowdaFactory):
+    __model__ = Collection
+
+
+class ClamsAppFactory(ChowdaFactory):
+    __model__ = ClamsApp
+
+
+class UserFactory(ChowdaFactory):
+    __model__ = User
+
+
+class MediaFileCollectionLinkFactory(ChowdaFactory):
+    __model__ = MediaFileCollectionLink
+
+
+class MediaFileBatchLinkFactory(ChowdaFactory):
+    __model__ = MediaFileBatchLink
+
+
+class PipelineFactory(ChowdaFactory):
+    __model__ = Pipeline
+
+
+class ClamsEventFactory(ChowdaFactory):
+    __model__ = ClamsEvent

pyproject.toml

@@ -17,6 +17,7 @@ pydantic = {extras = ["email"], version = "^1.10.5"}
 uvicorn = "^0.20.0"
 black = "^23.1.0"
 flake8 = "^6.0.0"
+pydantic-factories = "^1.17.2"
 
 [tool.poetry.group.test]
 optional = true

requirements.txt

@@ -2,3 +2,4 @@ starlette-admin
 uvicorn
 sqlmodel
 pydantic[email]
+pydantic-factories