This repository has been archived on 2024-10-15. You can view files and clone it, but cannot push or open issues or pull requests.
PyrogramBotBase/classes/pyrouser_sqlite.py

61 lines
1.8 KiB
Python

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()