6️⃣ CHAPTER_06 unit of work

README.md

@@ -17,3 +17,4 @@ Reference https://github.com/cosmicpython/code
 - [✅ CHAPTER_01 domain model](https://github.com/sawaca96/architecture-patterns-with-python/commit/fb604a0bc25b70a98e16dc4185eb8c9eb96ceb3d)
 - [✅ CHAPTER_02 repository](https://github.com/sawaca96/architecture-patterns-with-python/commit/fb604a0bc25b70a98e16dc4185eb8c9eb96ceb3d)
 - [4️⃣ CHAPTER_04 usecase](https://github.com/sawaca96/architecture-patterns-with-python/pull/1)
+- [6️⃣ CHAPTER_06 unit of work](https://github.com/sawaca96/architecture-patterns-with-python/pull/2#pullrequestreview-1265028411)

app/db.py → app/allocation/adapters/db.py

@@ -12,7 +12,6 @@ def __init__(self, url: str) -> None:
         self._session_factory = async_scoped_session(
             sessionmaker(
                 autocommit=False,
-                autoflush=False,
                 class_=AsyncSession,
                 bind=self._engine,
             ),

app/orm.py → app/allocation/adapters/orm.py

@@ -2,7 +2,7 @@
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import registry, relationship
 
-from app import models
+from app.allocation.domain import models
 
 mapper_registry = registry()
 metadata = mapper_registry.metadata
@@ -20,7 +20,7 @@
     metadata,
     sa.Column("id", UUID(as_uuid=True), primary_key=True),
     sa.Column("sku", sa.String),
-    sa.Column("purchased_quantity", sa.Integer),
+    sa.Column("qty", sa.Integer),
     sa.Column("eta", sa.Date, nullable=True),
 )
 

app/repository.py → app/allocation/adapters/repository.py

@@ -3,12 +3,12 @@
 
 import sqlalchemy as sa
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import subqueryload
+from sqlalchemy.orm import selectinload, subqueryload
 
-from app import models
+from app.allocation.domain import models
 
 
-class BatchAbstractRepository(abc.ABC):
+class AbstractBatchRepository(abc.ABC):
     @abc.abstractmethod
     async def add(self, batch: models.Batch) -> None:
         raise NotImplementedError
@@ -22,7 +22,7 @@ async def list(self) -> list[models.Batch]:
         raise NotImplementedError
 
 
-class PGBatchRepository(BatchAbstractRepository):
+class PGBatchRepository(AbstractBatchRepository):
     def __init__(self, session: AsyncSession) -> None:
         self._session = session
 
@@ -31,7 +31,14 @@ async def add(self, batch: models.Batch) -> None:
         await self._session.flush()
 
     async def get(self, id: UUID) -> models.Batch:
-        return await self._session.get(models.Batch, id)
+        # async sqlalchemy doesn't support relationship
+        # It raise 'greenlet_spawn has not been called; can't call await_() here. Was IO attempted in an unexpected place?' # noqa E501
+        result = await self._session.execute(
+            sa.select(models.Batch)
+            .where(models.Batch.id == id)
+            .options(selectinload(models.Batch.allocations))
+        )
+        return result.scalar_one()
 
     async def list(self) -> list[models.Batch]:
         result = await self._session.execute(

app/models.py → app/allocation/domain/models.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import date
-from uuid import UUID
+from uuid import UUID, uuid4
 
 
 class OutOfStock(Exception):
@@ -18,21 +18,20 @@ def allocate(line: OrderLine, batches: list[Batch]) -> UUID:
         raise OutOfStock(f"Out of stock for sku {line.sku}")
 
 
-@dataclass(unsafe_hash=True, kw_only=True)  # TODO: kw_only를 언제 써야 할까?
+@dataclass(unsafe_hash=True, kw_only=True)
 class OrderLine:
-    id: UUID
+    id: UUID = field(default_factory=uuid4)
     sku: str
     qty: int
 
 
+@dataclass(kw_only=True)
 class Batch:
-    def __init__(self, id: UUID, sku: str, qty: int, eta: date | None) -> None:
-        # TODO: id 값 업으면 기본 값 채우기
-        self.id = id
-        self.sku = sku
-        self.eta = eta
-        self.purchased_quantity = qty
-        self.allocations: set[OrderLine] = set()
+    id: UUID = field(default_factory=uuid4)
+    sku: str
+    eta: date = None
+    qty: int
+    allocations: set[OrderLine] = field(default_factory=lambda: set())
 
     def __repr__(self) -> str:
         return f"<Batch {self.id}>"
@@ -66,7 +65,7 @@ def allocated_quantity(self) -> int:
 
     @property
     def available_quantity(self) -> int:
-        return self.purchased_quantity - self.allocated_quantity
+        return self.qty - self.allocated_quantity
 
     def can_allocate(self, line: OrderLine) -> bool:
         return (

app/dependencies.py → app/allocation/routers/dependencies.py

@@ -4,9 +4,10 @@
 from fastapi import Depends
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from app.allocation.adapters.db import DB
+from app.allocation.adapters.repository import AbstractBatchRepository, PGBatchRepository
+from app.allocation.service_layer.unit_of_work import AbstractUnitOfWork, BatchUnitOfWork
 from app.config import get_config
-from app.db import DB
-from app.repository import BatchAbstractRepository, PGBatchRepository
 
 config = get_config()
 
@@ -23,5 +24,9 @@ async def session(db: DB = Depends(db)) -> AsyncGenerator[AsyncSession, None]:
 
 def repository(
     session: AsyncSession = Depends(session),
-) -> BatchAbstractRepository:
+) -> AbstractBatchRepository:
     return PGBatchRepository(session)
+
+
+def batch_uow() -> AbstractUnitOfWork[AbstractBatchRepository]:
+    return BatchUnitOfWork()

app/allocation/routers/main.py

@@ -0,0 +1,44 @@
+from datetime import date
+from uuid import UUID
+
+from fastapi import Body, Depends, FastAPI, HTTPException
+
+from app.allocation.adapters.repository import AbstractBatchRepository
+from app.allocation.domain import models
+from app.allocation.routers.dependencies import batch_uow
+from app.allocation.service_layer import services
+from app.allocation.service_layer.unit_of_work import AbstractUnitOfWork
+
+app = FastAPI()
+# start_mappers() # TODO: 운영환경에서는 실행되어야 함
+
+
+@app.get("/")
+async def root() -> dict[str, str]:
+    return {"message": "Hello World"}
+
+
+@app.post("/batches", status_code=201)
+async def add_batch(
+    batch_id: UUID = Body(),
+    sku: str = Body(),
+    quantity: int = Body(),
+    eta: date = Body(default=None),
+    uow: AbstractUnitOfWork[AbstractBatchRepository] = Depends(batch_uow),
+) -> dict[str, str]:
+    await services.add_batch(batch_id, sku, quantity, eta, uow)
+    return {"message": "success"}
+
+
+@app.post("/allocate", response_model=dict[str, str], status_code=201)
+async def allocate(
+    line_id: UUID = Body(),
+    sku: str = Body(),
+    quantity: int = Body(),
+    uow: AbstractUnitOfWork[AbstractBatchRepository] = Depends(batch_uow),
+) -> dict[str, str]:
+    try:
+        batch_id = await services.allocate(line_id, sku, quantity, uow)
+    except (models.OutOfStock, services.InvalidSku) as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    return {"batch_id": str(batch_id)}

app/allocation/service_layer/services.py

@@ -0,0 +1,39 @@
+from datetime import date
+from uuid import UUID
+
+from app.allocation.adapters.repository import AbstractBatchRepository
+from app.allocation.domain import models
+from app.allocation.service_layer import unit_of_work
+
+
+class InvalidSku(Exception):
+    pass
+
+
+async def add_batch(
+    batch_id: UUID,
+    sku: str,
+    qty: int,
+    eta: date | None,
+    uow: unit_of_work.AbstractUnitOfWork[AbstractBatchRepository],
+) -> None:
+    async with uow:
+        await uow.repo.add(models.Batch(id=batch_id, sku=sku, qty=qty, eta=eta))
+        await uow.commit()
+
+
+async def allocate(
+    line_id: UUID, sku: str, qty: int, uow: unit_of_work.AbstractUnitOfWork[AbstractBatchRepository]
+) -> UUID:
+    line = models.OrderLine(id=line_id, sku=sku, qty=qty)
+    async with uow:
+        batches = await uow.repo.list()
+        if not _is_valid_sku(line.sku, batches):
+            raise InvalidSku(f"Invalid sku {line.sku}")
+        batch_id = models.allocate(line, batches)
+        await uow.commit()
+    return batch_id
+
+
+def _is_valid_sku(sku: str, batches: list[models.Batch]) -> bool:
+    return sku in {b.sku for b in batches}

app/allocation/service_layer/unit_of_work.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import abc
+from asyncio import current_task
+from typing import Any, Generic, TypeVar
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_scoped_session, create_async_engine
+from sqlalchemy.orm import sessionmaker
+
+from app.allocation.adapters.repository import AbstractBatchRepository, PGBatchRepository
+from app.config import get_config
+
+config = get_config()
+
+Repo = TypeVar("Repo")
+
+
+class AbstractUnitOfWork(abc.ABC, Generic[Repo]):
+    async def __aenter__(self) -> AbstractUnitOfWork[Repo]:
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self.rollback()
+
+    @abc.abstractproperty
+    def repo(self) -> Repo:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def commit(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def rollback(self) -> None:
+        raise NotImplementedError
+
+
+class BatchUnitOfWork(AbstractUnitOfWork[AbstractBatchRepository]):
+    def __init__(self) -> None:
+        self._engine = create_async_engine(config.PG_DSN, echo=False)
+        self._session_factory = async_scoped_session(
+            sessionmaker(
+                autocommit=False,
+                autoflush=False,
+                class_=AsyncSession,
+                bind=self._engine,
+            ),
+            scopefunc=current_task,
+        )
+
+    @property
+    def repo(self) -> AbstractBatchRepository:
+        return self._batches
+
+    async def __aenter__(self) -> AbstractUnitOfWork[AbstractBatchRepository]:
+        self._session: AsyncSession = self._session_factory()
+        self._batches = PGBatchRepository(self._session)
+        return await super().__aenter__()
+
+    async def __aexit__(self, *args: Any) -> None:
+        await super().__aexit__(*args)
+        await self._session.close()
+
+    async def commit(self) -> None:
+        await self._session.commit()
+
+    async def rollback(self) -> None:
+        await self._session.rollback()

tests/conftest.py

@@ -6,9 +6,9 @@
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from sqlalchemy.orm import clear_mappers
 
+from app.allocation.adapters.db import DB
+from app.allocation.adapters.orm import metadata, start_mappers
 from app.config import get_config
-from app.db import DB
-from app.orm import metadata, start_mappers
 
 config = get_config()
 db = DB(config.PG_DSN)