cherry picking commit from other fork

cavallium/faster-pyrogram@12f9207

Co-authored-by: Andrea Cavalli<[email protected]>

Author: SpEcHiDe

pyrogram/storage/file_storage.py

@@ -21,6 +21,7 @@
 #  You should have received a copy of the GNU Lesser General Public License
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+
 import logging
 import os
 import sqlite3
@@ -31,14 +32,14 @@
 log = logging.getLogger(__name__)
 
 USERNAMES_SCHEMA = """
-CREATE TABLE IF NOT EXISTS usernames
+CREATE TABLE usernames
 (
     id       INTEGER,
     username TEXT,
     FOREIGN KEY (id) REFERENCES peers(id)
 );
 
-CREATE INDEX IF NOT EXISTS idx_usernames_username ON usernames (username);
+CREATE INDEX idx_usernames_username ON usernames (username);
 """
 
 UPDATE_STATE_SCHEMA = """
@@ -61,9 +62,33 @@ def __init__(self, name: str, workdir: Path):
 
         self.database = workdir / (self.name + self.FILE_EXTENSION)
 
+    def _vacuum(self):
+        with self.conn:
+            self.conn.execute("VACUUM")
+
+    def _update_from_one_impl(self):
+        with self.conn:
+            self.conn.execute("DELETE FROM peers")
+
+    def _update_from_two_impl(self):
+        with self.conn:
+            self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER")
+
+    def _update_from_three_impl(self):
+        with self.conn:
+            self.conn.executescript(USERNAMES_SCHEMA)
+
+    def _update_from_four_impl(self):
+        with self.conn:
+            self.conn.executescript(UPDATE_STATE_SCHEMA)
+
+    def _update_from_five_impl(self):
+        with self.conn:
+            self.conn.executescript("CREATE INDEX idx_usernames_id ON usernames (id);")
+
     def _connect_impl(self, path):
-        self.conn = sqlite3.connect(path)
-        
+        self.conn = sqlite3.connect(str(path), timeout=1, check_same_thread=False)
+
         with self.conn:
             self.conn.execute("PRAGMA journal_mode=WAL").close()
             self.conn.execute("PRAGMA synchronous=NORMAL").close()
@@ -73,36 +98,23 @@ async def update(self):
         version = await self.version()
 
         if version == 1:
-            with self.conn:
-                self.conn.execute("DELETE FROM peers")
-
+            await self.loop.run_in_executor(self.executor, self._update_from_one_impl)
             version += 1
 
         if version == 2:
-            with self.conn:
-                try:
-                    self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER")
-                except Exception as e:
-                    log.exception(e)
-
+            await self.loop.run_in_executor(self.executor, self._update_from_two_impl)
             version += 1
 
         if version == 3:
-            with self.conn:
-                self.conn.executescript(USERNAMES_SCHEMA)
-
+            await self.loop.run_in_executor(self.executor, self._update_from_three_impl)
             version += 1
 
         if version == 4:
-            with self.conn:
-                self.conn.executescript(UPDATE_STATE_SCHEMA)
-
+            await self.loop.run_in_executor(self.executor, self._update_from_four_impl)
             version += 1
 
         if version == 5:
-            with self.conn:
-                self.conn.executescript("CREATE INDEX IF NOT EXISTS idx_usernames_id ON usernames (id);")
-
+            await self.loop.run_in_executor(self.executor, self._update_from_five_impl)
             version += 1
 
         await self.version(version)
@@ -111,15 +123,14 @@ async def open(self):
         path = self.database
         file_exists = path.is_file()
 
-        self.conn = sqlite3.connect(str(path), timeout=1, check_same_thread=False)
+        self.executor.submit(self._connect_impl, path).result()
 
         if not file_exists:
             await self.create()
         else:
             await self.update()
 
-        with self.conn:
-            self.conn.execute("VACUUM")
+        await self.loop.run_in_executor(self.executor, self._vacuum)
 
     async def delete(self):
         os.remove(self.database)