Pyrofork: Storage: SqliteStorage: Save fragment username(s) to sessions database

Signed-off-by: wulan17 <wulan17@nusantararom.org>
This commit is contained in:
wulan17 2023-10-27 13:38:49 +07:00
parent 2900a03dc5
commit 9e3103b6f6
No known key found for this signature in database
GPG key ID: 318CD6CD3A6AC0A5

View file

@ -1,20 +1,21 @@
# Pyrogram - Telegram MTProto API Client Library for Python
# Pyrofork - Telegram MTProto API Client Library for Python
# Copyright (C) 2017-present Dan <https://github.com/delivrance>
# Copyright (C) 2022-present Mayuri-Chan <https://github.com/Mayuri-Chan>
#
# This file is part of Pyrogram.
# This file is part of Pyrofork.
#
# Pyrogram is free software: you can redistribute it and/or modify
# 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.
#
# Pyrogram is distributed in the hope that it will be useful,
# 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 Pyrogram. If not, see <http://www.gnu.org/licenses/>.
# along with Pyrofork. If not, see <http://www.gnu.org/licenses/>.
import inspect
import sqlite3
@ -68,6 +69,25 @@ END;
"""
UNAME_SCHEMA = """
CREATE TABLE IF NOT EXISTS usernames
(
id TEXT PRIMARY KEY,
peer_id INTEGER NOT NULL,
last_update_on INTEGER NOT NULL DEFAULT (CAST(STRFTIME('%s', 'now') AS INTEGER))
);
CREATE TRIGGER IF NOT EXISTS trg_usernames_last_update_on
AFTER UPDATE
ON usernames
BEGIN
UPDATE usernames
SET last_update_on = CAST(STRFTIME('%s', 'now') AS INTEGER)
WHERE id = NEW.id;
END;
"""
def get_input_peer(peer_id: int, access_hash: int, peer_type: str):
if peer_type in ["user", "bot"]:
return raw.types.InputPeerUser(
@ -101,6 +121,7 @@ class SQLiteStorage(Storage):
def create(self):
with self.conn:
self.conn.executescript(SCHEMA)
self.conn.executescript(UNAME_SCHEMA)
self.conn.execute(
"INSERT INTO version VALUES (?)",
@ -133,7 +154,17 @@ class SQLiteStorage(Storage):
)
async def update_usernames(self, usernames: List[Tuple[int, str]]):
pass
self.conn.executescript(UNAME_SCHEMA)
for user in usernames:
self.conn.execute(
"DELETE FROM usernames WHERE peer_id=?",
(user[0],)
)
self.conn.executemany(
"REPLACE INTO usernames (peer_id, id)"
"VALUES (?, ?)",
usernames
)
async def get_peer_by_id(self, peer_id: int):
r = self.conn.execute(
@ -154,7 +185,22 @@ class SQLiteStorage(Storage):
).fetchone()
if r is None:
raise KeyError(f"Username not found: {username}")
r2 = self.conn.execute(
"SELECT peer_id, last_update_on FROM usernames WHERE id = ?"
"ORDER BY last_update_on DESC",
(username,)
).fetchone()
if r2 is None:
raise KeyError(f"Username not found: {username}")
if abs(time.time() - r2[1]) > self.USERNAME_TTL:
raise KeyError(f"Username expired: {username}")
r = r = self.conn.execute(
"SELECT id, access_hash, type, last_update_on FROM peers WHERE id = ?"
"ORDER BY last_update_on DESC",
(r2[0],)
).fetchone()
if r is None:
raise KeyError(f"Username not found: {username}")
if abs(time.time() - r[3]) > self.USERNAME_TTL:
raise KeyError(f"Username expired: {username}")