start

api/db/subscription.py

@@ -0,0 +1,121 @@
+from api.db.connection import engine, SessionLocal
+from api.models.subscription import Subscription
+from api.models.base import Base
+from api.config import settings
+from datetime import datetime, timedelta
+from sqlalchemy import select
+import asyncpg
+
+
+async def init_db():
+    try:
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+        return
+    except Exception as exc:  # try to create database if it doesn't exist
+        msg = str(exc).lower()
+        if "does not exist" in msg or "invalidcatalogname" in msg or "database" in msg:
+            try:
+                admin_conn = await asyncpg.connect(
+                    user=settings.DB_USER,
+                    password=settings.DB_PASS,
+                    database="postgres",
+                    host=settings.DB_HOST,
+                    port=int(settings.DB_PORT),
+                )
+                try:
+                    await admin_conn.execute(f'CREATE DATABASE "{settings.DB_NAME}"')
+                except Exception:
+                    pass
+                await admin_conn.close()
+            except Exception:
+                raise
+            async with engine.begin() as conn:
+                await conn.run_sync(Base.metadata.create_all)
+            return
+        raise
+
+
+async def create_subscription(user_id: str, uuid: str, until: datetime, subscription_link: str | None = None, connect_link: str | None = None, active: bool = True):
+    async with SessionLocal() as session:
+        # if user already has a subscription, remove it (replace)
+        q = select(Subscription).where(Subscription.user_id == user_id)
+        res = await session.execute(q)
+        existing = res.scalars().first()
+        if existing:
+            await session.delete(existing)
+        # avoid NULL in DB: replace None links with empty strings
+        safe_subscription_link = subscription_link if subscription_link is not None else ""
+        safe_connect_link = connect_link if connect_link is not None else ""
+        sub = Subscription(user_id=user_id, uuid=uuid, until=until, active=active, subscription_link=safe_subscription_link, connect_link=safe_connect_link)
+        session.add(sub)
+        await session.commit()
+        await session.refresh(sub)
+        return sub
+
+
+async def get_subscription_by_user(user_id: str):
+    async with SessionLocal() as session:
+        q = select(Subscription).where(Subscription.user_id == user_id)
+        res = await session.execute(q)
+        sub = res.scalars().first()
+        return sub
+
+
+async def get_subscription_by_uuid(sub_uuid: str):
+    async with SessionLocal() as session:
+        q = select(Subscription).where(Subscription.uuid == sub_uuid)
+        res = await session.execute(q)
+        return res.scalars().first()
+
+
+async def update_subscription(sub_uuid: str, until: datetime = None, user_id: str | None = None, active: bool | None = None, subscription_link: str | None = None, connect_link: str | None = None):
+    async with SessionLocal() as session:
+        q = select(Subscription).where(Subscription.uuid == sub_uuid)
+        res = await session.execute(q)
+        sub = res.scalars().first()
+        if not sub:
+            return None
+        if until is not None:
+            sub.until = until
+        if user_id is not None:
+            sub.user_id = user_id
+        if active is not None:
+            sub.active = active
+        if subscription_link is not None:
+            sub.subscription_link = subscription_link
+        if connect_link is not None:
+            sub.connect_link = connect_link
+        await session.commit()
+        return sub
+
+
+async def delete_subscription(sub_uuid: str):
+    async with SessionLocal() as session:
+        q = select(Subscription).where(Subscription.uuid == sub_uuid)
+        res = await session.execute(q)
+        sub = res.scalars().first()
+        if not sub:
+            return False
+        await session.delete(sub)
+        await session.commit()
+        return True
+
+
+async def list_subscriptions():
+    async with SessionLocal() as session:
+        q = select(Subscription)
+        res = await session.execute(q)
+        return res.scalars().all()
+
+
+async def extend_subscription_by_user(user_id: str, months: int):
+    sub = await get_subscription_by_user(user_id)
+    if sub:
+        new_until = max(sub.until, datetime.now()) + timedelta(days=30 * months)
+        return await update_subscription(sub.uuid, until=new_until, active=True)
+    return None
+
+
+async def deactivate_subscription(sub_uuid: str):
+    return await update_subscription(sub_uuid, active=False)