gigacook

1
from os import environ
2

3
from sqlalchemy import delete, insert, inspect, select, update
4
from sqlalchemy.ext.asyncio import AsyncAttrs, AsyncSession, async_sessionmaker, create_async_engine
5
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
6

7
async_session = AsyncSession
8

9

10
class Base(AsyncAttrs, DeclarativeBase):
11
    @classmethod
12
    def __get_pk(cls):
13
        return getattr(cls, inspect(cls).primary_key[0].name)
14

15
    @classmethod
16
    async def create(cls, **kwargs):
17
        async with async_session() as session:
18
            await session.execute(insert(cls).values(**kwargs))
19
            await session.commit()
20

21
    @classmethod
22
    async def read(cls, user):
23
        async with async_session() as session:
24
            return list(await session.scalars(select(cls).where(cls.user == user)))
25

26
    @classmethod
27
    async def update(cls, pk, **kwargs):
28
        async with async_session() as session:
29
            await session.execute(update(cls).where(cls.__get_pk() == pk).values(**kwargs))
30
            await session.commit()
31

32
    @classmethod
33
    async def delete(cls, pk):
34
        async with async_session() as session:
35
            await session.execute(delete(cls).where(cls.__get_pk() == pk))
36
            await session.commit()
37

38

39
class Stage(Base):
40
    __tablename__ = "stage"
41

42
    user: Mapped[int] = mapped_column(primary_key=True)
43
    name: Mapped[str]
44

45
    @classmethod
46
    async def set(cls, user: int, name: str):
47
        if await cls.read(user):
48
            await cls.update(user, name=name)
49
        else:
50
            await cls.create(user=user, name=name)
51

52

53
class Product(Base):
54
    __tablename__ = "product"
55

56
    id: Mapped[int] = mapped_column(primary_key=True)
57
    user: Mapped[int]
58
    name: Mapped[str]
59

60

61
class Bookmark(Base):
62
    __tablename__ = "bookmark"
63

64
    id: Mapped[int] = mapped_column(primary_key=True)
65
    user: Mapped[int]
66
    text: Mapped[str]
67

68

69
def run_async_session() -> None:
70
    global async_session
71

72
    engine = create_async_engine(environ["DB_URL"], echo=True)
73
    async_session = async_sessionmaker(engine, expire_on_commit=False)
74