Jail timer

database/crud/jail.py

@@ -1,12 +1,19 @@
 from datetime import datetime
 from typing import Optional
 
-from sqlalchemy import select
+from sqlalchemy import delete, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.schemas import Jail
 
-__all__ = ["get_jail", "get_user_jail", "imprison"]
+__all__ = [
+    "get_jail",
+    "get_jail_entry_by_id",
+    "get_next_jail_release",
+    "get_user_jail",
+    "imprison",
+    "delete_prisoner_by_id",
+]
 
 
 async def get_jail(session: AsyncSession) -> list[Jail]:
@@ -15,14 +22,36 @@ async def get_jail(session: AsyncSession) -> list[Jail]:
     return list((await session.execute(statement)).scalars().all())
 
 
+async def get_jail_entry_by_id(session: AsyncSession, jail_id: int) -> Optional[Jail]:
+    """Get a jail entry by its id"""
+    statement = select(Jail).where(Jail.jail_entry_id == jail_id)
+    return (await session.execute(statement)).scalar_one_or_none()
+
+
+async def get_next_jail_release(session: AsyncSession) -> Optional[Jail]:
+    """Get the next person being released from jail"""
+    statement = select(Jail).order_by(Jail.until)
+    return (await session.execute(statement)).scalars().first()
+
+
 async def get_user_jail(session: AsyncSession, user_id: int) -> Optional[Jail]:
     """Check how long a given user is still in jail for"""
     statement = select(Jail).where(Jail.user_id == user_id)
     return (await session.execute(statement)).scalar_one_or_none()
 
 
-async def imprison(session: AsyncSession, user_id: int, until: datetime):
+async def imprison(session: AsyncSession, user_id: int, until: datetime) -> Jail:
     """Put a user in Didier Jail"""
     jail = Jail(user_id=user_id, until=until)
     session.add(jail)
     await session.commit()
+    await session.refresh(jail)
+
+    return jail
+
+
+async def delete_prisoner_by_id(session: AsyncSession, jail_id: int):
+    """Release a user from jail using their jail entry id"""
+    statement = delete(Jail).where(Jail.jail_entry_id == jail_id)
+    await session.execute(statement)
+    await session.commit()

database/exceptions/currency.py

@@ -1,4 +1,4 @@
-__all__ = ["DoubleNightly", "NotEnoughDinks"]
+__all__ = ["DoubleNightly", "NotEnoughDinks", "SavingsCapExceeded"]
 
 
 class DoubleNightly(Exception):

database/schemas.py

@@ -235,7 +235,7 @@ class Jail(Base):
 
     __tablename__ = "jail"
 
-    jail_entry_i: Mapped[int] = mapped_column(primary_key=True)
+    jail_entry_id: Mapped[int] = mapped_column(primary_key=True)
     user_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("users.user_id"))
     until: Mapped[datetime] = mapped_column(nullable=False)