first commit

bot/database/orm.py

@@ -0,0 +1,323 @@
+# sqlalchemy import
+from sqlalchemy import update, select, delete, func
+
+# Database engine
+from database.engine import create_async_engine, get_session_maker
+
+# DB Models
+from database.db_models import *
+
+# Config
+from decouple import config
+
+# Another
+from datetime import datetime
+from decimal import Decimal
+from typing import Any
+
+
+class ORM:
+
+    def __init__(self):
+        self.async_engine = create_async_engine(
+            url=f"postgresql+asyncpg://{config('POSTGRES_USER')}:{config('POSTGRES_PASSWORD')}@{config('POSTGRES_HOST')}:{config('POSTGRES_PORT')}/{config('POSTGRES_DB')}"
+        )
+        self.session_maker = get_session_maker(self.async_engine)
+
+    async def proceed_schemas(self) -> None:
+        async with self.async_engine.begin() as conn:
+            await conn.run_sync(BaseModel.metadata.create_all)
+
+    # *############################
+    # *#          USERS           #
+    # *############################
+
+    async def is_user_exists(self, user_id: int) -> bool:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.execute(
+                    select(User.user_id).where(User.user_id == user_id)
+                )
+
+                return query.one_or_none() is not None
+
+    async def create_user(
+        self, user_id: int, username: str, fullname: str, register_date: datetime
+    ) -> int:
+        async with self.session_maker() as session:
+            async with session.begin():
+                if not await self.is_user_exists(user_id):
+                    user = User(
+                        user_id=user_id,
+                        username=username,
+                        fullname=fullname,
+                        register_date=register_date,
+                    )
+
+                    session.add(user)
+                    await session.flush()
+                    return user.user_id
+                else:
+                    return
+
+    async def set_users_field(
+        self, user_id: int, field: str, value: int | str | bool
+    ) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                await session.execute(
+                    update(User)
+                    .where(User.user_id == user_id)
+                    .values({getattr(User, field): value})
+                )
+
+    async def get_user(self, user_id: int) -> User:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(
+                    select(User).where(User.user_id == user_id)
+                )
+
+                return query.one_or_none()
+
+    async def get_all_users(self) -> list[User]:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(select(User))
+
+                return query.all()
+
+    async def get_users_count(self) -> int:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(select(func.count()).select_from(User))
+
+                return query.one_or_none()
+
+    async def get_all_user_ids(self) -> list[int]:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(select(User.user_id))
+
+                return query.all()
+
+    # *############################
+    # *#          ADMINS          #
+    # *############################
+
+    async def is_admin_exists(self, user_id: int) -> bool:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.execute(
+                    select(Admin.user_id).where(Admin.user_id == user_id)
+                )
+
+                return query.one_or_none() is not None
+
+    async def create_admin(self, user_id: int, username: str, fullname: str) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                admin = Admin(user_id=user_id, username=username, fullname=fullname)
+
+                await session.merge(admin)
+
+    async def get_admin(self, user_id: int) -> Admin:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(
+                    select(Admin).where(Admin.user_id == user_id)
+                )
+
+                return query.one_or_none()
+
+    async def get_all_admins(self) -> list[Admin]:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(select(Admin))
+
+                return query.all()
+
+    async def delete_admin(self, user_id: int) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                await session.execute(delete(Admin).where(Admin.user_id == user_id))
+
+    async def set_admin_field(
+        self, user_id: int, field: str, value: int | str | bool
+    ) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                await session.execute(
+                    update(Admin)
+                    .where(Admin.user_id == user_id)
+                    .values({getattr(Admin, field): value})
+                )
+
+    # *############################
+    # *#         SETTINGS         #
+    # *############################
+
+    async def is_setting_exists(self, name: str) -> bool:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.execute(
+                    select(Setting).where(Setting.name == name)
+                )
+
+                return query.one_or_none() is not None
+
+    async def create_setting(self, name: str, value: Any) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                setting = Setting(name=name, value=value)
+
+                await session.merge(setting)
+
+    async def init_settings(self) -> None: ...
+
+    async def get_setting_value(self, name: str) -> Any:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(
+                    select(Setting.value).where(Setting.name == name)
+                )
+
+                return query.one_or_none()
+
+    async def update_setting_value(self, name: str, value: dict | list) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                await session.execute(
+                    update(Setting)
+                    .where(Setting.name == name)
+                    .values({getattr(Setting, "value"): value})
+                )
+
+    # *############################
+    # *#        BLACKLIST         #
+    # *############################
+
+    async def is_blacklisted(self, user_id: int) -> bool:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.execute(
+                    select(Blacklist).where(Blacklist.user_id == user_id)
+                )
+
+                return query.one_or_none() is not None
+
+    async def create_blacklist(self, user_id: int) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                blacklist = Blacklist(user_id=user_id)
+
+                await session.merge(blacklist)
+
+    async def get_all_blacklist(self) -> list[int]:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(
+                    select(Blacklist.user_id).order_by(Blacklist.user_id)
+                )
+
+                return query.all()
+
+    async def delete_blacklist(self, user_id: int) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                await session.execute(
+                    delete(Blacklist).where(Blacklist.user_id == user_id)
+                )
+
+    # *############################
+    # *#         PAYMENTS         #
+    # *############################
+
+    async def create_payment(
+        self,
+        user_id: int,
+        order_id: str,
+        amount: Decimal,
+        currency: str,
+        description: str,
+        created_at: datetime,
+    ) -> int:
+        async with self.session_maker() as session:
+            async with session.begin():
+                payment = Payment(
+                    user_id=user_id,
+                    order_id=order_id,
+                    amount=amount,
+                    currency=currency,
+                    description=description,
+                    status="created",
+                    created_at=created_at,
+                    updated_at=created_at,
+                )
+
+                session.add(payment)
+                await session.flush()
+
+                return payment.id
+
+    async def get_payment_by_order_id(self, order_id: str) -> Payment:
+        async with self.session_maker() as session:
+            async with session.begin():
+                query = await session.scalars(
+                    select(Payment).where(Payment.order_id == order_id)
+                )
+
+                return query.one_or_none()
+
+    async def update_payment_link(
+        self,
+        order_id: str,
+        payment_link_id: str,
+        payment_url: str,
+        payment_link_status: str,
+        updated_at: datetime,
+    ) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                await session.execute(
+                    update(Payment)
+                    .where(Payment.order_id == order_id)
+                    .values(
+                        payment_link_id=payment_link_id,
+                        payment_url=payment_url,
+                        payment_link_status=payment_link_status,
+                        status="link_created",
+                        updated_at=updated_at,
+                    )
+                )
+
+    async def update_payment_status(
+        self,
+        order_id: str,
+        status: str,
+        transaction_status: str | None,
+        updated_at: datetime,
+        transaction_id: str | None = None,
+        error_code: str | None = None,
+        error_description: str | None = None,
+        paid_at: datetime | None = None,
+    ) -> None:
+        async with self.session_maker() as session:
+            async with session.begin():
+                values = {
+                    "status": status,
+                    "transaction_status": transaction_status,
+                    "updated_at": updated_at,
+                    "error_code": error_code,
+                    "error_description": error_description,
+                }
+
+                if transaction_id:
+                    values["transaction_id"] = transaction_id
+
+                if paid_at:
+                    values["paid_at"] = paid_at
+
+                await session.execute(
+                    update(Payment).where(Payment.order_id == order_id).values(values)
+                )