Smarter session migration

delivrance · delivrance · commit 8c96e5f46aac · 2019-06-20T03:31:37.000+02:00
diff --git a/pyrogram/client/storage/file_storage.py b/pyrogram/client/storage/file_storage.py
@@ -19,12 +19,9 @@
 import base64
 import json
 import logging
-import os
 import sqlite3
 from pathlib import Path
-from sqlite3 import DatabaseError
 from threading import Lock
-from typing import Union
 
 from .memory_storage import MemoryStorage
 
@@ -43,25 +40,18 @@ def __init__(self, name: str, workdir: Path):
         self.lock = Lock()
 
     # noinspection PyAttributeOutsideInit
-    def migrate_from_json(self, path: Union[str, Path]):
-        log.warning("JSON session storage detected! Pyrogram will now convert it into an SQLite session storage...")
-
-        with open(path, encoding="utf-8") as f:
-            json_session = json.load(f)
-
-        os.remove(path)
-
+    def migrate_from_json(self, session_json: dict):
         self.open()
 
-        self.dc_id = json_session["dc_id"]
-        self.test_mode = json_session["test_mode"]
-        self.auth_key = base64.b64decode("".join(json_session["auth_key"]))
-        self.user_id = json_session["user_id"]
-        self.date = json_session.get("date", 0)
-        self.is_bot = json_session.get("is_bot", False)
+        self.dc_id = session_json["dc_id"]
+        self.test_mode = session_json["test_mode"]
+        self.auth_key = base64.b64decode("".join(session_json["auth_key"]))
+        self.user_id = session_json["user_id"]
+        self.date = session_json.get("date", 0)
+        self.is_bot = session_json.get("is_bot", False)
 
-        peers_by_id = json_session.get("peers_by_id", {})
-        peers_by_phone = json_session.get("peers_by_phone", {})
+        peers_by_id = session_json.get("peers_by_id", {})
+        peers_by_phone = session_json.get("peers_by_phone", {})
 
         peers = {}
 
@@ -81,22 +71,40 @@ def migrate_from_json(self, path: Union[str, Path]):
         # noinspection PyTypeChecker
         self.update_peers(peers.values())
 
-        log.warning("Done! The session has been successfully converted from JSON to SQLite storage")
-
     def open(self):
-        database_exists = os.path.isfile(self.database)
+        path = self.database
+        file_exists = path.is_file()
+
+        if file_exists:
+            try:
+                with open(path, encoding="utf-8") as f:
+                    session_json = json.load(f)
+            except ValueError:
+                pass
+            else:
+                log.warning("JSON session storage detected! Converting it into an SQLite session storage...")
+
+                path.rename(path.name + ".OLD")
+
+                log.warning('The old session file has been renamed to "{}.OLD"'.format(path.name))
+
+                self.migrate_from_json(session_json)
+
+                log.warning("Done! The session has been successfully converted from JSON to SQLite storage")
+
+                return
+
+        if Path(path.name + ".OLD").is_file():
+            log.warning('Old session file detected: "{}.OLD". You can remove this file now'.format(path.name))
 
         self.conn = sqlite3.connect(
-            str(self.database),
+            path,
             timeout=1,
             check_same_thread=False
         )
 
-        try:
-            if not database_exists:
-                self.create()
+        if not file_exists:
+            self.create()
 
-            with self.conn:
-                self.conn.execute("VACUUM")
-        except DatabaseError:
-            self.migrate_from_json(self.database)
+        with self.conn:
+            self.conn.execute("VACUUM")
