yasirarism/pyrofork
1
0
Fork 0
mirror of https://github.com/Mayuri-Chan/pyrofork.git synced 2025-12-29 12:04:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(storage): migrated to aiosqlite

Browse source 
Signed-off-by: wulan17 <wulan17@nusantararom.org>
This commit is contained in:
Hitalo M 2023-11-06 19:24:56 -03:00 committed by wulan17
parent ce46d49ec2
commit 8193d87e2c
No known key found for this signature in database
GPG key ID: 318CD6CD3A6AC0A5
4 changed files with 86 additions and 87 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
pyproject.toml
View file

@ -3,7 +3,12 @@ name = "pyrofork"
dynamic = ["version"] dynamic = ["version"]
description = "Fork of pyrogram. Elegant, modern and asynchronous Telegram MTProto API framework in Python for users and bots" description = "Fork of pyrogram. Elegant, modern and asynchronous Telegram MTProto API framework in Python for users and bots"
authors = [{ name = "wulan17", email = "mayuri@mayuri.my.id" }] authors = [{ name = "wulan17", email = "mayuri@mayuri.my.id" }]
dependencies = ["pyaes==1.6.1", "pysocks==1.7.1", "pymediainfo-pyrofork>=6.0.1,<7.0.0"] dependencies = [
"aiosqlite>=0.19.0",
"pyaes==1.6.1",
"pysocks==1.7.1",
"pymediainfo-pyrofork>=6.0.1,<7.0.0"
]
readme = "README.md" readme = "README.md"
license = "LGPL-3.0-or-later" license = "LGPL-3.0-or-later"
requires-python = "~=3.8" requires-python = "~=3.8"

32
pyrogram/storage/file_storage.py
View file

@ -18,10 +18,10 @@
# along with Pyrofork. If not, see <http://www.gnu.org/licenses/>. # along with Pyrofork. If not, see <http://www.gnu.org/licenses/>.
import logging import logging
import os
import sqlite3
from pathlib import Path from pathlib import Path
import aiosqlite
from .sqlite_storage import SQLiteStorage from .sqlite_storage import SQLiteStorage
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -35,36 +35,38 @@ class FileStorage(SQLiteStorage):
self.database = workdir / (self.name + self.FILE_EXTENSION) self.database = workdir / (self.name + self.FILE_EXTENSION)
def update(self): async def update(self):
version = self.version() version = await self.version()
if version == 1: if version == 1:
with self.conn: await self.conn.execute("DELETE FROM peers")
self.conn.execute("DELETE FROM peers") await self.conn.commit()
version += 1 version += 1
if version == 2: if version == 2:
with self.conn: await self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER")
self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER") await self.conn.commit()
version += 1 version += 1
self.version(version) await self.version(version)
async def open(self): async def open(self):
path = self.database path = self.database
file_exists = path.is_file() file_exists = path.is_file()
self.conn = sqlite3.connect(str(path), timeout=1, check_same_thread=False) self.conn = await aiosqlite.connect(str(path), timeout=1)
await self.conn.execute("PRAGMA journal_mode=WAL")
if not file_exists: if not file_exists:
self.create() await self.create()
else: else:
self.update() await self.update()
with self.conn: await self.conn.execute("VACUUM")
self.conn.execute("VACUUM") await self.conn.commit()
async def delete(self): async def delete(self):
os.remove(self.database) Path(self.database).unlink()

7
pyrogram/storage/memory_storage.py
View file

@ -19,9 +19,10 @@
import base64 import base64
import logging import logging
import sqlite3
import struct import struct
import aiosqlite
from .sqlite_storage import SQLiteStorage from .sqlite_storage import SQLiteStorage
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -34,8 +35,8 @@ class MemoryStorage(SQLiteStorage):
self.session_string = session_string self.session_string = session_string
async def open(self): async def open(self):
self.conn = sqlite3.connect(":memory:", check_same_thread=False) self.conn = await aiosqlite.connect(":memory:")
self.create() await self.create()
if self.session_string: if self.session_string:
# Old format # Old format

127
pyrogram/storage/sqlite_storage.py
View file

@ -18,13 +18,13 @@
# along with Pyrofork. If not, see <http://www.gnu.org/licenses/>. # along with Pyrofork. If not, see <http://www.gnu.org/licenses/>.
import inspect import inspect
import sqlite3
import time import time
from typing import List, Tuple, Any from typing import Any, List, Tuple
from pyrogram import raw import aiosqlite
from pyrogram import raw, utils
from .storage import Storage from .storage import Storage
from .. import utils
# language=SQLite # language=SQLite
SCHEMA = """ SCHEMA = """
@ -116,61 +116,55 @@ class SQLiteStorage(Storage):
def __init__(self, name: str): def __init__(self, name: str):
super().__init__(name) super().__init__(name)
self.conn = None # type: sqlite3.Connection self.conn: aiosqlite.Connection = None
def create(self): async def create(self):
with self.conn: await self.conn.executescript(SCHEMA)
self.conn.executescript(SCHEMA) await self.conn.execute("INSERT INTO version VALUES (?)", (self.VERSION,))
self.conn.executescript(UNAME_SCHEMA) await self.conn.execute(
"INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?)",
self.conn.execute( (2, None, None, None, 0, None, None),
"INSERT INTO version VALUES (?)", )
(self.VERSION,) await self.conn.commit()
)
self.conn.execute(
"INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?)",
(2, None, None, None, 0, None, None)
)
async def open(self): async def open(self):
raise NotImplementedError raise NotImplementedError
async def save(self): async def save(self):
await self.date(int(time.time())) await self.date(int(time.time()))
self.conn.commit() await self.conn.commit()
async def close(self): async def close(self):
self.conn.close() await self.conn.close()
async def delete(self): async def delete(self):
raise NotImplementedError raise NotImplementedError
async def update_peers(self, peers: List[Tuple[int, int, str, str, str]]): async def update_peers(self, peers: List[Tuple[int, int, str, str, str]]):
self.conn.executemany( await self.conn.executemany(
"REPLACE INTO peers (id, access_hash, type, username, phone_number)" "REPLACE INTO peers (id, access_hash, type, username, phone_number)"
"VALUES (?, ?, ?, ?, ?)", "VALUES (?, ?, ?, ?, ?)",
peers peers
) )
async def update_usernames(self, usernames: List[Tuple[int, str]]): async def update_usernames(self, usernames: List[Tuple[int, str]]):
self.conn.executescript(UNAME_SCHEMA) await self.conn.executescript(UNAME_SCHEMA)
for user in usernames: for user in usernames:
self.conn.execute( await self.conn.execute(
"DELETE FROM usernames WHERE peer_id=?", "DELETE FROM usernames WHERE peer_id=?",
(user[0],) (user[0],)
) )
self.conn.executemany( await self.conn.executemany(
"REPLACE INTO usernames (peer_id, id)" "REPLACE INTO usernames (peer_id, id)"
"VALUES (?, ?)", "VALUES (?, ?)",
usernames usernames
) )
async def get_peer_by_id(self, peer_id: int): async def get_peer_by_id(self, peer_id: int):
r = self.conn.execute( q = await self.conn.execute(
"SELECT id, access_hash, type FROM peers WHERE id = ?", "SELECT id, access_hash, type FROM peers WHERE id = ?", (peer_id,)
(peer_id,) )
).fetchone() r = await q.fetchone()
if r is None: if r is None:
raise KeyError(f"ID not found: {peer_id}") raise KeyError(f"ID not found: {peer_id}")
@ -178,27 +172,30 @@ class SQLiteStorage(Storage):
return get_input_peer(*r) return get_input_peer(*r)
async def get_peer_by_username(self, username: str): async def get_peer_by_username(self, username: str):
r = self.conn.execute( q = await self.conn.execute(
"SELECT id, access_hash, type, last_update_on FROM peers WHERE username = ?" "SELECT id, access_hash, type, last_update_on FROM peers WHERE username = ?"
"ORDER BY last_update_on DESC", "ORDER BY last_update_on DESC",
(username,) (username,),
).fetchone() )
r = await q.fetchone()
if r is None: if r is None:
r2 = self.conn.execute( r2 = self.conn.execute(
"SELECT peer_id, last_update_on FROM usernames WHERE id = ?" "SELECT peer_id, last_update_on FROM usernames WHERE id = ?"
"ORDER BY last_update_on DESC", "ORDER BY last_update_on DESC",
(username,) (username,)
).fetchone() )
r2 = await r2.fetchone()
if r2 is None: if r2 is None:
raise KeyError(f"Username not found: {username}") raise KeyError(f"Username not found: {username}")
if abs(time.time() - r2[1]) > self.USERNAME_TTL: if abs(time.time() - r2[1]) > self.USERNAME_TTL:
raise KeyError(f"Username expired: {username}") raise KeyError(f"Username expired: {username}")
r = r = self.conn.execute( r = await self.conn.execute(
"SELECT id, access_hash, type, last_update_on FROM peers WHERE id = ?" "SELECT id, access_hash, type, last_update_on FROM peers WHERE id = ?"
"ORDER BY last_update_on DESC", "ORDER BY last_update_on DESC",
(r2[0],) (r2[0],)
).fetchone() )
r = await r.fetchone()
if r is None: if r is None:
raise KeyError(f"Username not found: {username}") raise KeyError(f"Username not found: {username}")
@ -208,64 +205,58 @@ class SQLiteStorage(Storage):
return get_input_peer(*r[:3]) return get_input_peer(*r[:3])
async def get_peer_by_phone_number(self, phone_number: str): async def get_peer_by_phone_number(self, phone_number: str):
r = self.conn.execute( q = await self.conn.execute(
"SELECT id, access_hash, type FROM peers WHERE phone_number = ?", "SELECT id, access_hash, type FROM peers WHERE phone_number = ?",
(phone_number,) (phone_number,),
).fetchone() )
r = await q.fetchone()
if r is None: if r is None:
raise KeyError(f"Phone number not found: {phone_number}") raise KeyError(f"Phone number not found: {phone_number}")
return get_input_peer(*r) return get_input_peer(*r)
def _get(self): async def _get(self):
attr = inspect.stack()[2].function attr = inspect.stack()[2].function
return self.conn.execute( q = await self.conn.execute(f"SELECT {attr} FROM sessions")
f"SELECT {attr} FROM sessions" row = await q.fetchone()
).fetchone()[0] return row[0] if row else None
def _set(self, value: Any): async def _set(self, value: Any):
attr = inspect.stack()[2].function attr = inspect.stack()[2].function
await self.conn.execute(f"UPDATE sessions SET {attr} = ?", (value,))
await self.conn.commit()
with self.conn: async def _accessor(self, value: Any = object):
self.conn.execute( return await self._get() if value == object else await self._set(value)
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): async def dc_id(self, value: int = object):
return self._accessor(value) return await self._accessor(value)
async def api_id(self, value: int = object): async def api_id(self, value: int = object):
return self._accessor(value) return await self._accessor(value)
async def test_mode(self, value: bool = object): async def test_mode(self, value: bool = object):
return self._accessor(value) return await self._accessor(value)
async def auth_key(self, value: bytes = object): async def auth_key(self, value: bytes = object):
return self._accessor(value) return await self._accessor(value)
async def date(self, value: int = object): async def date(self, value: int = object):
return self._accessor(value) return await self._accessor(value)
async def user_id(self, value: int = object): async def user_id(self, value: int = object):
return self._accessor(value) return await self._accessor(value)
async def is_bot(self, value: bool = object): async def is_bot(self, value: bool = object):
return self._accessor(value) return await self._accessor(value)
def version(self, value: int = object): async def version(self, value: int = object):
if value == object: if value == object:
return self.conn.execute( q = await self.conn.execute("SELECT number FROM version")
"SELECT number FROM version" row = await q.fetchone()
).fetchone()[0] return row[0] if row else None
else: await self.conn.execute("UPDATE version SET number = ?", (value,))
with self.conn: await self.conn.commit()
self.conn.execute( return None
"UPDATE version SET number = ?",
(value,)
)