From 866ab2429f7232a2b9cc68b70ec159df2657b34f Mon Sep 17 00:00:00 2001 From: profitroll Date: Wed, 23 Aug 2023 15:45:15 +0200 Subject: [PATCH] Added additional SQLite option --- classes/pyroclient.py | 5 +- classes/{pyrouser.py => pyrouser_mongo.py} | 2 +- classes/pyrouser_sqlite.py | 60 ++++++++++++++++++++++ commands.json | 8 +++ main.py | 6 ++- modules/{database.py => database_mongo.py} | 0 modules/database_sqlite.py | 11 ++++ modules/{migrator.py => migrator_mongo.py} | 0 modules/migrator_sqlite.py | 26 ++++++++++ plugins/commands/shutdown.py | 16 ++++++ requirements.txt | 3 +- validation/users.json | 2 +- 12 files changed, 134 insertions(+), 5 deletions(-) rename classes/{pyrouser.py => pyrouser_mongo.py} (97%) create mode 100644 classes/pyrouser_sqlite.py rename modules/{database.py => database_mongo.py} (100%) create mode 100644 modules/database_sqlite.py rename modules/{migrator.py => migrator_mongo.py} (100%) create mode 100644 modules/migrator_sqlite.py create mode 100644 plugins/commands/shutdown.py diff --git a/classes/pyroclient.py b/classes/pyroclient.py index e7800ad..e00765f 100644 --- a/classes/pyroclient.py +++ b/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): diff --git a/classes/pyrouser.py b/classes/pyrouser_mongo.py similarity index 97% rename from classes/pyrouser.py rename to classes/pyrouser_mongo.py index 2fb9b73..ad6df40 100644 --- a/classes/pyrouser.py +++ b/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__) diff --git a/classes/pyrouser_sqlite.py b/classes/pyrouser_sqlite.py new file mode 100644 index 0000000..8c3b27b --- /dev/null +++ b/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() diff --git a/commands.json b/commands.json index 7a79e6a..db9aa29 100644 --- a/commands.json +++ b/commands.json @@ -10,6 +10,14 @@ } ] }, + "shutdown": { + "scopes": [ + { + "name": "BotCommandScopeChat", + "chat_id": "owner" + } + ] + }, "remove_commands": { "scopes": [ { diff --git a/main.py b/main.py index 0f502b5..5f18be7 100644 --- a/main.py +++ b/main.py @@ -7,7 +7,11 @@ from pathlib import Path from libbot import sync from classes.pyroclient import PyroClient -from modules.migrator import migrate_database + +# Main uses MongoDB implementation of DB +# but you can also select SQLite one below +# from modules.migrator_sqlite import migrate_database +from modules.migrator_mongo import migrate_database from modules.scheduler import scheduler # Uncomment this and the line below client declaration diff --git a/modules/database.py b/modules/database_mongo.py similarity index 100% rename from modules/database.py rename to modules/database_mongo.py diff --git a/modules/database_sqlite.py b/modules/database_sqlite.py new file mode 100644 index 0000000..48325bc --- /dev/null +++ b/modules/database_sqlite.py @@ -0,0 +1,11 @@ +"""Module that provides database access""" + +import sqlite3 +from pathlib import Path + +from libbot.sync import config_get + +db: sqlite3.Connection = sqlite3.connect(Path(config_get("database"))) +cursor: sqlite3.Cursor = db.cursor() + +cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER, card TEXT, locale TEXT)") diff --git a/modules/migrator.py b/modules/migrator_mongo.py similarity index 100% rename from modules/migrator.py rename to modules/migrator_mongo.py diff --git a/modules/migrator_sqlite.py b/modules/migrator_sqlite.py new file mode 100644 index 0000000..db252c6 --- /dev/null +++ b/modules/migrator_sqlite.py @@ -0,0 +1,26 @@ +from os import rename +from pathlib import Path +from typing import Mapping + +from libbot.sync import json_read + +from modules.database_sqlite import cursor + + +def migrate_database() -> None: + """Apply migrations from old JSON database to SQLite""" + if not Path("data/database.json").exists(): + return + + db_old: Mapping[str, Mapping[str, str]] = json_read(Path("data/database.json")) + + for user, keys in db_old.items(): + user_locale = None if "locale" not in keys else keys["locale"] + user_card = None if "card" not in keys else keys["card"] + + cursor.execute( + "INSERT INTO users VALUES (?, ?)", (int(user), user_card, user_locale) + ) + + cursor.connection.commit() + rename(Path("data/database.json"), Path("data/database.migrated.json")) diff --git a/plugins/commands/shutdown.py b/plugins/commands/shutdown.py new file mode 100644 index 0000000..aa971c5 --- /dev/null +++ b/plugins/commands/shutdown.py @@ -0,0 +1,16 @@ +import asyncio + +from pyrogram import filters +from pyrogram.types import Message + +from classes.pyroclient import PyroClient + + +@PyroClient.on_message( + ~filters.scheduled + & filters.private + & filters.command(["shutdown", "reboot", "restart"], prefixes=["/"]) # type: ignore +) +async def command_shutdown(app: PyroClient, msg: Message): + if msg.from_user.id == app.owner: + asyncio.get_event_loop().create_task(app.stop()) diff --git a/requirements.txt b/requirements.txt index b53b25b..c1b6be8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,10 @@ apscheduler~=3.10.3 convopyro==0.5 -mongodb-migrations==1.3.0 pykeyboard==0.1.5 tgcrypto==1.2.5 uvloop==0.17.0 +# If uses MongoDB: +mongodb-migrations==1.3.0 --extra-index-url https://git.end-play.xyz/api/packages/profitroll/pypi/simple async_pymongo==0.1.4 libbot[speed,pyrogram]==2.0.1 \ No newline at end of file diff --git a/validation/users.json b/validation/users.json index 6560dab..6f27166 100644 --- a/validation/users.json +++ b/validation/users.json @@ -17,7 +17,7 @@ "string", "null" ], - "description": "Preferred messages language according to " + "description": "Preferred messages language according to user's preference" } } }