Pyrofork: Storage: SqliteStorage: Save fragment username(s) to sessio…

…ns database

Signed-off-by: wulan17 <wulan17@nusantararom.org>

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
@@ -48,6 +49,13 @@
     last_update_on INTEGER NOT NULL DEFAULT (CAST(STRFTIME('%s', 'now') AS INTEGER))
 );
 
+CREATE TABLE usernames
+(
+    id             TEXT PRIMARY KEY,
+    peer_id        INTEGER NOT NULL,
+    last_update_on INTEGER NOT NULL DEFAULT (CAST(STRFTIME('%s', 'now') AS INTEGER))
+);
+
 CREATE TABLE version
 (
     number INTEGER PRIMARY KEY
@@ -65,6 +73,34 @@
     SET last_update_on = CAST(STRFTIME('%s', 'now') AS INTEGER)
     WHERE id = NEW.id;
 END;
+
+CREATE TRIGGER 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;
+"""
+
+
+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;
 """
 
 
@@ -133,7 +169,19 @@ async def update_peers(self, peers: List[Tuple[int, int, str, str, str]]):
         )
 
     async def update_usernames(self, usernames: List[Tuple[int, str]]):
-        pass
+        self.conn.executescript(UNAME_SCHEMA)
+        user_ids = []
+        for user in usernames:
+            user_ids.append((user[0]))
+        self.conn.executemany(
+            "DELETE FROM username WHERE peer_id=?",
+            user_ids
+        )
+        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 +202,22 @@ async def get_peer_by_username(self, username: str):
         ).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}")