35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
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
|