esotericelf/image_edit_creation

1import os2from collections.abc import AsyncGenerator3from datetime import datetime, timezone4 5from sqlalchemy import Boolean, DateTime, String, func, select6from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine7from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column8 9DATABASE_URL = os.getenv(10    "DATABASE_URL",11    "postgresql+asyncpg://nano_banana:nano_banana@db:5432/nano_banana",12)13if DATABASE_URL.startswith("postgres://"):14    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)15if DATABASE_URL.startswith("postgresql://"):16    DATABASE_URL = DATABASE_URL.replace("postgresql://", "postgresql+asyncpg://", 1)17 18 19class Base(DeclarativeBase):20    pass21 22 23class GiftToken(Base):24    __tablename__ = "gift_tokens"25 26    token: Mapped[str] = mapped_column(String(36), primary_key=True)27    is_used: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)28    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)29    created_at: Mapped[datetime] = mapped_column(30        DateTime(timezone=True),31        nullable=False,32        server_default=func.now(),33    )34 35 36engine = create_async_engine(DATABASE_URL, pool_pre_ping=True)37AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False)38 39 40async def get_db_session() -> AsyncGenerator[AsyncSession, None]:41    async with AsyncSessionLocal() as session:42        yield session43 44 45async def init_db() -> None:46    async with engine.begin() as conn:47        await conn.run_sync(Base.metadata.create_all)48 49 50async def count_active_gift_tokens(session: AsyncSession) -> int:51    now = datetime.now(timezone.utc)52    result = await session.execute(53        select(func.count())54        .select_from(GiftToken)55        .where(56            GiftToken.is_used.is_(False),57            GiftToken.expires_at > now,58        )59    )60    return int(result.scalar_one())61