Nightly, bank, award & dinks

database/crud/currency.py

@@ -0,0 +1,43 @@
+from datetime import datetime
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from database.crud import users
+from database.exceptions import currency as exceptions
+from database.models import Bank
+
+
+NIGHTLY_AMOUNT = 420
+
+
+async def get_bank(session: AsyncSession, user_id: int) -> Bank:
+    """Get a user's bank info"""
+    user = await users.get_or_add(session, user_id)
+    return user.bank
+
+
+async def add_dinks(session: AsyncSession, user_id: int, amount: int):
+    """Increase the Dinks counter for a user"""
+    bank = await get_bank(session, user_id)
+    bank.dinks += amount
+    session.add(bank)
+    await session.commit()
+
+
+async def claim_nightly(session: AsyncSession, user_id: int):
+    """Claim daily Dinks"""
+    user = await users.get_or_add(session, user_id)
+    nightly_data = user.nightly_data
+
+    now = datetime.now()
+
+    if nightly_data.last_nightly is not None and nightly_data.last_nightly.date() == now.date():
+        raise exceptions.DoubleNightly
+
+    bank = user.bank
+    bank.dinks += NIGHTLY_AMOUNT
+    nightly_data.last_nightly = now
+
+    session.add(bank)
+    session.add(nightly_data)
+    await session.commit()

database/crud/users.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from database.models import User, Bank, NightlyData
+
+
+async def get_or_add(session: AsyncSession, user_id: int) -> User:
+    """Get a user's profile
+    If it doesn't exist yet, create it (along with all linked datastructures).
+    """
+    statement = select(User).where(User.user_id == user_id)
+    user: Optional[User] = (await session.execute(statement)).scalar_one_or_none()
+
+    # User exists
+    if user is not None:
+        return user
+
+    # Create new user
+    user = User(user_id=user_id)
+    session.add(user)
+    await session.commit()
+
+    # Add bank & nightly info
+    bank = Bank(user_id=user_id)
+    nightly_data = NightlyData(user_id=user_id)
+    user.bank = bank
+    user.nightly_data = nightly_data
+
+    session.add(bank)
+    session.add(nightly_data)
+    session.add(user)
+
+    await session.commit()
+
+    return user