diff --git a/pyrogram/storage/file_storage.py b/pyrogram/storage/file_storage.py index e35e01eb..986787cd 100644 --- a/pyrogram/storage/file_storage.py +++ b/pyrogram/storage/file_storage.py @@ -21,7 +21,6 @@ import os import sqlite3 from pathlib import Path -from . import sqlite from .sqlite_storage import SQLiteStorage log = logging.getLogger(__name__) @@ -35,36 +34,39 @@ class FileStorage(SQLiteStorage): self.database = workdir / (self.name + self.FILE_EXTENSION) - async def update(self): - version = await self.version() + def update(self): + version = self.version() if version == 1: - await self.conn.execute("DELETE FROM peers") + with self.lock, self.conn: + self.conn.execute("DELETE FROM peers") version += 1 if version == 2: - await self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER") + with self.lock, self.conn: + self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER") version += 1 - await self.version(version) + self.version(version) async def open(self): path = self.database file_exists = path.is_file() - self.conn = sqlite.AsyncSqlite(database=str(path), timeout=1, check_same_thread=False) + self.conn = sqlite3.connect(str(path), timeout=1, check_same_thread=False) if not file_exists: - await self.create() + self.create() else: - await self.update() + self.update() - try: # Python 3.6.0 (exactly this version) is bugged and won't successfully execute the vacuum - await self.conn.execute("VACUUM") - except sqlite3.OperationalError: - pass + with self.conn: + try: # Python 3.6.0 (exactly this version) is bugged and won't successfully execute the vacuum + self.conn.execute("VACUUM") + except sqlite3.OperationalError: + pass async def delete(self): os.remove(self.database) diff --git a/pyrogram/storage/memory_storage.py b/pyrogram/storage/memory_storage.py index 8893f31a..2c01f447 100644 --- a/pyrogram/storage/memory_storage.py +++ b/pyrogram/storage/memory_storage.py @@ -18,9 +18,9 @@ import base64 import logging +import sqlite3 import struct -from . import sqlite from .sqlite_storage import SQLiteStorage log = logging.getLogger(__name__) @@ -33,8 +33,8 @@ class MemoryStorage(SQLiteStorage): self.session_string = session_string async def open(self): - self.conn = sqlite.AsyncSqlite(database=":memory:", check_same_thread=False) - await self.create() + self.conn = sqlite3.connect(":memory:", check_same_thread=False) + self.create() if self.session_string: # Old format diff --git a/pyrogram/storage/sqlite/__init__.py b/pyrogram/storage/sqlite/__init__.py deleted file mode 100644 index c7d70200..00000000 --- a/pyrogram/storage/sqlite/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -# Pyrofork - Telegram MTProto API Client Library for Python -# Copyright (C) 2022-present Mayuri-Chan -# -# This file is part of Pyrofork. -# -# Pyrofork is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# Pyrofork is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with Pyrofork. If not, see . - -from .cursor import AsyncCursor -from .sqlite import AsyncSqlite - -__all__ = [AsyncSqlite, AsyncCursor] diff --git a/pyrogram/storage/sqlite/cursor.py b/pyrogram/storage/sqlite/cursor.py deleted file mode 100644 index 6e1d5103..00000000 --- a/pyrogram/storage/sqlite/cursor.py +++ /dev/null @@ -1,29 +0,0 @@ -# Pyrofork - Telegram MTProto API Client Library for Python -# Copyright (C) 2022-present Mayuri-Chan -# -# This file is part of Pyrofork. -# -# Pyrofork is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# Pyrofork is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with Pyrofork. If not, see . - -from pyrogram.utils import run_sync -from sqlite3 import Cursor -from threading import Thread - -class AsyncCursor(Thread): - def __init__(self, cursor: Cursor): - super().__init__() - self.cursor = cursor - - async def fetchone(self): - return await run_sync(self.cursor.fetchone) diff --git a/pyrogram/storage/sqlite/sqlite.py b/pyrogram/storage/sqlite/sqlite.py deleted file mode 100644 index 4c7a3f3c..00000000 --- a/pyrogram/storage/sqlite/sqlite.py +++ /dev/null @@ -1,46 +0,0 @@ -# Pyrofork - Telegram MTProto API Client Library for Python -# Copyright (C) 2022-present Mayuri-Chan -# -# This file is part of Pyrofork. -# -# Pyrofork is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# Pyrofork is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with Pyrofork. If not, see . - -import sqlite3 -from .cursor import AsyncCursor -from pathlib import Path -from pyrogram.utils import run_sync -from threading import Thread -from typing import Union - -class AsyncSqlite(Thread): - def __init__(self, database: Union[str, Path], *args, **kwargs): - super().__init__() - self.connection = sqlite3.connect(database, *args, **kwargs) - - async def commit(self): - return await run_sync(self.connection.commit) - - async def close(self): - return await run_sync(self.connection.close) - - async def execute(self, *args, **kwargs): - r = await run_sync(self.connection.execute, *args, **kwargs) - return AsyncCursor(r) - - async def executemany(self, *args, **kwargs): - r = await run_sync(self.connection.executemany, *args, **kwargs) - return AsyncCursor(r) - - async def executescript(self, *args, **kwargs): - r = await run_sync(self.connection.executescript, *args, **kwargs) diff --git a/pyrogram/storage/sqlite_storage.py b/pyrogram/storage/sqlite_storage.py index a3786c38..15e5ddc0 100644 --- a/pyrogram/storage/sqlite_storage.py +++ b/pyrogram/storage/sqlite_storage.py @@ -17,7 +17,9 @@ # along with Pyrogram. If not, see . import inspect +import sqlite3 import time +from threading import Lock from typing import List, Tuple, Any from pyrogram import raw @@ -66,6 +68,7 @@ BEGIN END; """ + def get_input_peer(peer_id: int, access_hash: int, peer_type: str): if peer_type in ["user", "bot"]: return raw.types.InputPeerUser( @@ -94,20 +97,22 @@ class SQLiteStorage(Storage): def __init__(self, name: str): super().__init__(name) - self.conn = None + self.conn = None # type: sqlite3.Connection + self.lock = Lock() - async def create(self): - await self.conn.executescript(SCHEMA) + def create(self): + with self.lock, self.conn: + self.conn.executescript(SCHEMA) - await self.conn.execute( - "INSERT INTO version VALUES (?)", - (self.VERSION,) - ) + self.conn.execute( + "INSERT INTO version VALUES (?)", + (self.VERSION,) + ) - await self.conn.execute( - "INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?)", - (2, None, None, None, 0, None, None) - ) + self.conn.execute( + "INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?)", + (2, None, None, None, 0, None, None) + ) async def open(self): raise NotImplementedError @@ -115,27 +120,29 @@ class SQLiteStorage(Storage): async def save(self): await self.date(int(time.time())) - await self.conn.commit() + with self.lock: + self.conn.commit() async def close(self): - await self.conn.close() + with self.lock: + self.conn.close() async def delete(self): raise NotImplementedError async def update_peers(self, peers: List[Tuple[int, int, str, str, str]]): - await self.conn.executemany( - "REPLACE INTO peers (id, access_hash, type, username, phone_number)" - "VALUES (?, ?, ?, ?, ?)", - peers - ) + with self.lock: + self.conn.executemany( + "REPLACE INTO peers (id, access_hash, type, username, phone_number)" + "VALUES (?, ?, ?, ?, ?)", + peers + ) async def get_peer_by_id(self, peer_id: int): - q = await self.conn.execute( + r = self.conn.execute( "SELECT id, access_hash, type FROM peers WHERE id = ?", (peer_id,) - ) - r = await q.fetchone() + ).fetchone() if r is None: raise KeyError(f"ID not found: {peer_id}") @@ -143,12 +150,11 @@ class SQLiteStorage(Storage): return get_input_peer(*r) async def get_peer_by_username(self, username: str): - q = await self.conn.execute( + r = self.conn.execute( "SELECT id, access_hash, type, last_update_on FROM peers WHERE username = ?" "ORDER BY last_update_on DESC", (username,) - ) - r = await q.fetchone() + ).fetchone() if r is None: raise KeyError(f"Username not found: {username}") @@ -159,65 +165,64 @@ class SQLiteStorage(Storage): return get_input_peer(*r[:3]) async def get_peer_by_phone_number(self, phone_number: str): - q = await self.conn.execute( + r = self.conn.execute( "SELECT id, access_hash, type FROM peers WHERE phone_number = ?", (phone_number,) - ) - r = await q.fetchone() + ).fetchone() if r is None: raise KeyError(f"Phone number not found: {phone_number}") return get_input_peer(*r) - async def _get(self): + def _get(self): attr = inspect.stack()[2].function - q = await self.conn.execute( + return self.conn.execute( f"SELECT {attr} FROM sessions" - ) - return (await q.fetchone())[0] + ).fetchone()[0] - async def _set(self, value: Any): + def _set(self, value: Any): attr = inspect.stack()[2].function - await self.conn.execute( - f"UPDATE sessions SET {attr} = ?", - (value,) - ) - - async def _accessor(self, value: Any = object): - return await self._get() if value == object else await self._set(value) - - async def dc_id(self, value: int = object): - return await self._accessor(value) - - async def api_id(self, value: int = object): - return await self._accessor(value) - - async def test_mode(self, value: bool = object): - return await self._accessor(value) - - async def auth_key(self, value: bytes = object): - return await self._accessor(value) - - async def date(self, value: int = object): - return await self._accessor(value) - - async def user_id(self, value: int = object): - return await self._accessor(value) - - async def is_bot(self, value: bool = object): - return await self._accessor(value) - - async def version(self, value: int = object): - if value == object: - q = await self.conn.execute( - "SELECT number FROM version" - ) - return (await q.fetchone())[0] - else: - await self.conn.execute( - "UPDATE version SET number = ?", + with self.lock, self.conn: + self.conn.execute( + f"UPDATE sessions SET {attr} = ?", (value,) ) + + def _accessor(self, value: Any = object): + return self._get() if value == object else self._set(value) + + async def dc_id(self, value: int = object): + return self._accessor(value) + + async def api_id(self, value: int = object): + return self._accessor(value) + + async def test_mode(self, value: bool = object): + return self._accessor(value) + + async def auth_key(self, value: bytes = object): + return self._accessor(value) + + async def date(self, value: int = object): + return self._accessor(value) + + async def user_id(self, value: int = object): + return self._accessor(value) + + async def is_bot(self, value: bool = object): + return self._accessor(value) + + def version(self, value: int = object): + if value == object: + return self.conn.execute( + "SELECT number FROM version" + ).fetchone()[0] + else: + with self.lock, self.conn: + self.conn.execute( + "UPDATE version SET number = ?", + (value,) + )