PyrogramBotBase/classes/pyrouser_sqlite.py

import logging
from dataclasses import dataclass
from typing import Union

from modules.database_sqlite import cursor

logger = logging.getLogger(__name__)


@dataclass
class PyroUser:
    """Dataclass of DB entry of a user"""

    __slots__ = ("id", "locale")

    id: int
    locale: Union[str, None]

    @classmethod
    async def find(cls, id: int, locale: Union[str, None] = None):
        """Find user in database and create new record if user does not exist.

        ### Args:
            * id (`int`): User's Telegram ID
            * locale (`Union[str, None]`, *optional*): User's locale. Defaults to `None`.

        ### Raises:
            * `RuntimeError`: Raised when user entry after insertion could not be found.

        ### Returns:
            * `PyroUser`: User with its database data.
        """
        db_entry = cursor.execute(
            "SELECT id, locale FROM users WHERE id = ?", (id,)
        ).fetchone()

        if db_entry is None:
            cursor.execute("INSERT INTO users VALUES (?, ?)", (id, locale))
            cursor.connection.commit()
            db_entry = cursor.execute(
                "SELECT id, locale FROM users WHERE id = ?", (id,)
            ).fetchone()

        if db_entry is None:
            raise RuntimeError("Could not find inserted user entry.")

        return cls(*db_entry)

    async def update_locale(self, locale: Union[str, None]) -> None:
        """Change user's locale stored in the database.

        ### Args:
            * locale (`Union[str, None]`): New locale to be set.
        """
        logger.debug("%s's locale has been set to %s", self.id, locale)
        cursor.execute(
            "UPDATE users SET locale = ? WHERE id = ?",
            (locale, self.id),
        )
        cursor.connection.commit()