From e226bdbf8931cab7b88916c0af03415f644b0304 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Sat, 14 Sep 2019 19:39:47 +0200 Subject: [PATCH] Create SQLiteStorage for better abstraction --- pyrogram/client/storage/sqlite_storage.py | 184 ++++++++++++++++++++++ 1 file changed, 184 insertions(+) create mode 100644 pyrogram/client/storage/sqlite_storage.py diff --git a/pyrogram/client/storage/sqlite_storage.py b/pyrogram/client/storage/sqlite_storage.py new file mode 100644 index 00000000..1c36c427 --- /dev/null +++ b/pyrogram/client/storage/sqlite_storage.py @@ -0,0 +1,184 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-2019 Dan Tès +# +# This file is part of Pyrogram. +# +# Pyrogram 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, +# 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 . + +import inspect +import sqlite3 +import time +from pathlib import Path +from threading import Lock +from typing import List, Tuple, Any + +from pyrogram.api import types +from pyrogram.client.ext import utils +from .storage import Storage + + +def get_input_peer(peer_id: int, access_hash: int, peer_type: str): + if peer_type in ["user", "bot"]: + return types.InputPeerUser( + user_id=peer_id, + access_hash=access_hash + ) + + if peer_type == "group": + return types.InputPeerChat( + chat_id=-peer_id + ) + + if peer_type in ["channel", "supergroup"]: + return types.InputPeerChannel( + channel_id=utils.get_channel_id(peer_id), + access_hash=access_hash + ) + + raise ValueError("Invalid peer type: {}".format(peer_type)) + + +class SQLiteStorage(Storage): + VERSION = 2 + USERNAME_TTL = 8 * 60 * 60 + + def __init__(self, name: str): + super().__init__(name) + + self.conn = None # type: sqlite3.Connection + self.lock = Lock() + + def create(self): + with self.lock, self.conn: + with open(str(Path(__file__).parent / "schema.sql"), "r") as schema: + self.conn.executescript(schema.read()) + + self.conn.execute( + "INSERT INTO version VALUES (?)", + (self.VERSION,) + ) + + self.conn.execute( + "INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?)", + (2, None, None, 0, None, None) + ) + + def open(self): + raise NotImplementedError + + def save(self): + self.date(int(time.time())) + + with self.lock: + self.conn.commit() + + def close(self): + with self.lock: + self.conn.close() + + def delete(self): + raise NotImplementedError + + def update_peers(self, peers: List[Tuple[int, int, str, str, str]]): + with self.lock: + self.conn.executemany( + "REPLACE INTO peers (id, access_hash, type, username, phone_number)" + "VALUES (?, ?, ?, ?, ?)", + peers + ) + + def get_peer_by_id(self, peer_id: int): + r = self.conn.execute( + "SELECT id, access_hash, type FROM peers WHERE id = ?", + (peer_id,) + ).fetchone() + + if r is None: + raise KeyError("ID not found: {}".format(peer_id)) + + return get_input_peer(*r) + + def get_peer_by_username(self, username: str): + r = self.conn.execute( + "SELECT id, access_hash, type, last_update_on FROM peers WHERE username = ?", + (username,) + ).fetchone() + + if r is None: + raise KeyError("Username not found: {}".format(username)) + + if abs(time.time() - r[3]) > self.USERNAME_TTL: + raise KeyError("Username expired: {}".format(username)) + + return get_input_peer(*r[:3]) + + def get_peer_by_phone_number(self, phone_number: str): + r = self.conn.execute( + "SELECT id, access_hash, type FROM peers WHERE phone_number = ?", + (phone_number,) + ).fetchone() + + if r is None: + raise KeyError("Phone number not found: {}".format(phone_number)) + + return get_input_peer(*r) + + def _get(self): + attr = inspect.stack()[2].function + + return self.conn.execute( + "SELECT {} FROM sessions".format(attr) + ).fetchone()[0] + + def _set(self, value: Any): + attr = inspect.stack()[2].function + + with self.lock, self.conn: + self.conn.execute( + "UPDATE sessions SET {} = ?".format(attr), + (value,) + ) + + def _accessor(self, value: Any = object): + return self._get() if value == object else self._set(value) + + def dc_id(self, value: int = object): + return self._accessor(value) + + def test_mode(self, value: bool = object): + return self._accessor(value) + + def auth_key(self, value: bytes = object): + return self._accessor(value) + + def date(self, value: int = object): + return self._accessor(value) + + def user_id(self, value: int = object): + return self._accessor(value) + + def is_bot(self, value: bool = object): + return self._accessor(value) + + def version(self, value: int = object): + if value == object: + return self.conn.execute( + "SELECT number FROM version" + ).fetchone()[0] + else: + with self.lock, self.conn: + self.conn.execute( + "UPDATE version SET number = ?", + (value,) + )