Step 9：药品管理（IMP）

backend/app/crud/imp_inventory.py

@@ -0,0 +1,46 @@
+import uuid
+from typing import Sequence
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.models.imp_inventory import ImpInventory
+from app.models.imp_batch import ImpBatch
+
+
+async def get_inventory_for_update(db: AsyncSession, study_id: uuid.UUID, site_id: uuid.UUID, batch_id: uuid.UUID) -> ImpInventory | None:
+    result = await db.execute(
+        select(ImpInventory)
+        .where(
+            ImpInventory.study_id == study_id,
+            ImpInventory.site_id == site_id,
+            ImpInventory.batch_id == batch_id,
+        )
+        .with_for_update()
+    )
+    return result.scalar_one_or_none()
+
+
+async def get_or_create_inventory(db: AsyncSession, study_id: uuid.UUID, site_id: uuid.UUID, batch_id: uuid.UUID) -> ImpInventory:
+    inv = await get_inventory_for_update(db, study_id, site_id, batch_id)
+    if inv:
+        return inv
+    inv = ImpInventory(study_id=study_id, site_id=site_id, batch_id=batch_id, quantity_on_hand=0)
+    db.add(inv)
+    await db.flush()
+    return inv
+
+
+async def list_inventory(
+    db: AsyncSession,
+    study_id: uuid.UUID,
+    site_id: uuid.UUID | None = None,
+    product_id: uuid.UUID | None = None,
+) -> Sequence[ImpInventory]:
+    stmt = select(ImpInventory).where(ImpInventory.study_id == study_id)
+    if site_id:
+        stmt = stmt.where(ImpInventory.site_id == site_id)
+    if product_id:
+        stmt = stmt.join(ImpBatch, ImpBatch.id == ImpInventory.batch_id).where(ImpBatch.product_id == product_id)
+    result = await db.execute(stmt)
+    return result.scalars().all()