From 9e3103b6f6c2d618fd955487d9b28e208f1d919f Mon Sep 17 00:00:00 2001
From: wulan17 <wulan17@nusantararom.org>
Date: Fri, 27 Oct 2023 13:38:49 +0700
Subject: [PATCH] Pyrofork: Storage: SqliteStorage: Save fragment username(s)
 to sessions database

Signed-off-by: wulan17 <wulan17@nusantararom.org>
---
 pyrogram/storage/sqlite_storage.py | 60 ++++++++++++++++++++++++++----
 1 file changed, 53 insertions(+), 7 deletions(-)

diff --git a/pyrogram/storage/sqlite_storage.py b/pyrogram/storage/sqlite_storage.py
index 7daa88bd..74d785aa 100644
--- a/pyrogram/storage/sqlite_storage.py
+++ b/pyrogram/storage/sqlite_storage.py
@@ -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}")