test api

api/db/connection.py

@@ -0,0 +1,7 @@
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import sessionmaker
+from api.config import settings
+
+DATABASE_URL = settings.DATABASE_URL
+engine = create_async_engine(DATABASE_URL, echo=False, future=True)
+SessionLocal = sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)

api/db/subscription.py

@@ -0,0 +1,34 @@
+from api.db.connection import engine, SessionLocal
+from api.models.subscription import Subscription
+from api.models.base import Base
+from datetime import datetime, timedelta
+
+async def init_db():
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+async def set_subscription(user_id: int, uuid: str, until: datetime):
+    async with SessionLocal() as session:
+        sub = await session.get(Subscription, user_id)
+        if sub:
+            sub.uuid = uuid
+            sub.until = until
+        else:
+            sub = Subscription(user_id=user_id, uuid=uuid, until=until)
+            session.add(sub)
+        await session.commit()
+
+async def get_subscription(user_id: int):
+    async with SessionLocal() as session:
+        sub = await session.get(Subscription, user_id)
+        if sub:
+            return {"uuid": sub.uuid, "until": sub.until}
+        return None
+
+async def extend_subscription(user_id: int, months: int):
+    sub = await get_subscription(user_id)
+    if sub:
+        new_until = max(sub["until"], datetime.now()) + timedelta(days=30*months)
+        await set_subscription(user_id, sub["uuid"], new_until)
+        return new_until
+    return None