Added additional SQLite option

classes/pyroclient.py

@@ -3,7 +3,10 @@ from typing import Union
 from libbot.pyrogram.classes import PyroClient as LibPyroClient
 from pyrogram.types import User
 
-from classes.pyrouser import PyroUser
+# PyroClient uses MongoDB implementation of PyroUser
+# but you can also select SQLite one below
+# from classes.pyrouser_sqlite import PyroUser
+from classes.pyrouser_mongo import PyroUser
 
 
 class PyroClient(LibPyroClient):

classes/pyrouser_mongo.py

@@ -4,7 +4,7 @@ from typing import Union
 
 from bson import ObjectId
 
-from modules.database import col_users
+from modules.database_mongo import col_users
 
 logger = logging.getLogger(__name__)
 

classes/pyrouser_sqlite.py

@@ -0,0 +1,60 @@
+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()