mir/pyrogram
2
0
Fork 0
mirror of https://github.com/pyrogram/pyrogram synced 2025-08-28 12:57:52 +00:00
Code Issues Releases Wiki Activity

Add session migrating from json; add some indexes to sqlite sessions

Browse Source
This commit is contained in:
bakatrouble 2019-02-26 21:43:23 +03:00
parent 03b92b3302
commit 10fc340eff
3 changed files with 28 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pyrogram/client/client.py
View File

@ -214,7 +214,7 @@ class Client(Methods, BaseClient):
elif session_name.startswith(':'): elif session_name.startswith(':'):
session_storage = StringSessionStorage(self, session_name) session_storage = StringSessionStorage(self, session_name)
else: else:
session_storage = JsonSessionStorage(self, session_name) session_storage = SQLiteSessionStorage(self, session_name)
elif isinstance(session_name, SessionStorage): elif isinstance(session_name, SessionStorage):
session_storage = session_name session_storage = session_name
else: else:

3
pyrogram/client/session_storage/sqlite/0001.sql
View File

@ -18,4 +18,7 @@ create table migrations (
name text primary key name text primary key
); );
create index username_idx on peers_cache(username);
create index phone_idx on peers_cache(phone);
insert into migrations (name) values ('0001'); insert into migrations (name) values ('0001');

29
pyrogram/client/session_storage/sqlite/__init__.py
View File

@ -18,17 +18,18 @@
import logging import logging
import os import os
import shutil
import sqlite3 import sqlite3
import pyrogram import pyrogram
from ....api import types from ....api import types
from ...ext import utils from ...ext import utils
from .. import MemorySessionStorage, SessionDoesNotExist from .. import MemorySessionStorage, SessionDoesNotExist, JsonSessionStorage
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
EXTENSION = '.session.sqlite3' EXTENSION = '.session'
MIGRATIONS = ['0001'] MIGRATIONS = ['0001']
@ -54,13 +55,32 @@ class SQLiteSessionStorage(MemorySessionStorage):
with open(os.path.join(os.path.dirname(__file__), '{}.sql'.format(name))) as script: with open(os.path.join(os.path.dirname(__file__), '{}.sql'.format(name))) as script:
self._conn.executescript(script.read()) self._conn.executescript(script.read())
def _migrate_from_json(self):
jss = JsonSessionStorage(self._client, self._session_name)
jss.load()
file_path = self._get_file_name(self._session_name)
self._conn = sqlite3.connect(file_path + '.tmp')
self._apply_migrations(new_db=True)
self._dc_id, self._test_mode, self._auth_key, self._user_id, self._date, self._is_bot = \
jss.dc_id, jss.test_mode, jss.auth_key, jss.user_id, jss.date, jss.is_bot
self.save()
self._conn.close()
shutil.move(file_path + '.tmp', file_path)
log.warning('Session was migrated from JSON, loading...')
self.load()
def load(self): def load(self):
file_path = self._get_file_name(self._session_name) file_path = self._get_file_name(self._session_name)
log.info('Loading SQLite session from {}'.format(file_path)) log.info('Loading SQLite session from {}'.format(file_path))
if os.path.isfile(file_path): if os.path.isfile(file_path):
self._conn = sqlite3.connect(file_path) try:
self._apply_migrations() self._conn = sqlite3.connect(file_path)
self._apply_migrations()
except sqlite3.DatabaseError:
log.warning('Trying to migrate session from JSON...')
self._migrate_from_json()
return
else: else:
self._conn = sqlite3.connect(file_path) self._conn = sqlite3.connect(file_path)
self._apply_migrations(new_db=True) self._apply_migrations(new_db=True)
@ -88,7 +108,6 @@ class SQLiteSessionStorage(MemorySessionStorage):
access_hash = entity.access_hash access_hash = entity.access_hash
elif isinstance(entity, (types.Chat, types.ChatForbidden)): elif isinstance(entity, (types.Chat, types.ChatForbidden)):
peer_id = -entity.id peer_id = -entity.id
# input_peer = types.InputPeerChat(chat_id=entity.id)
elif isinstance(entity, (types.Channel, types.ChannelForbidden)): elif isinstance(entity, (types.Channel, types.ChannelForbidden)):
peer_id = int('-100' + str(entity.id)) peer_id = int('-100' + str(entity.id))
username = entity.username.lower() if hasattr(entity, 'username') and entity.username else None username = entity.username.lower() if hasattr(entity, 'username') and entity.username else None