Smarter session migration
This commit is contained in:
parent
fa1976c8a0
commit
8c96e5f46a
@ -19,12 +19,9 @@
|
|||||||
import base64
|
import base64
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
|
||||||
import sqlite3
|
import sqlite3
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from sqlite3 import DatabaseError
|
|
||||||
from threading import Lock
|
from threading import Lock
|
||||||
from typing import Union
|
|
||||||
|
|
||||||
from .memory_storage import MemoryStorage
|
from .memory_storage import MemoryStorage
|
||||||
|
|
||||||
@ -43,25 +40,18 @@ class FileStorage(MemoryStorage):
|
|||||||
self.lock = Lock()
|
self.lock = Lock()
|
||||||
|
|
||||||
# noinspection PyAttributeOutsideInit
|
# noinspection PyAttributeOutsideInit
|
||||||
def migrate_from_json(self, path: Union[str, Path]):
|
def migrate_from_json(self, session_json: dict):
|
||||||
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)
|
|
||||||
|
|
||||||
self.open()
|
self.open()
|
||||||
|
|
||||||
self.dc_id = json_session["dc_id"]
|
self.dc_id = session_json["dc_id"]
|
||||||
self.test_mode = json_session["test_mode"]
|
self.test_mode = session_json["test_mode"]
|
||||||
self.auth_key = base64.b64decode("".join(json_session["auth_key"]))
|
self.auth_key = base64.b64decode("".join(session_json["auth_key"]))
|
||||||
self.user_id = json_session["user_id"]
|
self.user_id = session_json["user_id"]
|
||||||
self.date = json_session.get("date", 0)
|
self.date = session_json.get("date", 0)
|
||||||
self.is_bot = json_session.get("is_bot", False)
|
self.is_bot = session_json.get("is_bot", False)
|
||||||
|
|
||||||
peers_by_id = json_session.get("peers_by_id", {})
|
peers_by_id = session_json.get("peers_by_id", {})
|
||||||
peers_by_phone = json_session.get("peers_by_phone", {})
|
peers_by_phone = session_json.get("peers_by_phone", {})
|
||||||
|
|
||||||
peers = {}
|
peers = {}
|
||||||
|
|
||||||
@ -81,22 +71,40 @@ class FileStorage(MemoryStorage):
|
|||||||
# noinspection PyTypeChecker
|
# noinspection PyTypeChecker
|
||||||
self.update_peers(peers.values())
|
self.update_peers(peers.values())
|
||||||
|
|
||||||
log.warning("Done! The session has been successfully converted from JSON to SQLite storage")
|
|
||||||
|
|
||||||
def open(self):
|
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(
|
self.conn = sqlite3.connect(
|
||||||
str(self.database),
|
path,
|
||||||
timeout=1,
|
timeout=1,
|
||||||
check_same_thread=False
|
check_same_thread=False
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
if not file_exists:
|
||||||
if not database_exists:
|
self.create()
|
||||||
self.create()
|
|
||||||
|
|
||||||
with self.conn:
|
with self.conn:
|
||||||
self.conn.execute("VACUUM")
|
self.conn.execute("VACUUM")
|
||||||
except DatabaseError:
|
|
||||||
self.migrate_from_json(self.database)
|
|
||||||
|
Loading…
Reference in New Issue
Block a user