From eeb1a6cf42ca692ec9eeba17d65ca3fe94921311 Mon Sep 17 00:00:00 2001 From: KurimuzonAkuma Date: Fri, 24 Nov 2023 01:31:51 +0300 Subject: [PATCH] Allow Custom Storage Engines --- pyrogram/client.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/pyrogram/client.py b/pyrogram/client.py index 72daabfe..7cec1059 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -49,7 +49,7 @@ from pyrogram.errors import ( from pyrogram.handlers.handler import Handler from pyrogram.methods import Methods from pyrogram.session import Auth, Session -from pyrogram.storage import FileStorage, MemoryStorage +from pyrogram.storage import Storage, FileStorage, MemoryStorage from pyrogram.types import User, TermsOfService from pyrogram.utils import ainput from .dispatcher import Dispatcher @@ -177,6 +177,10 @@ class Client(Methods): Set the maximum amount of concurrent transmissions (uploads & downloads). A value that is too high may result in network related issues. Defaults to 1. + + storage_engine (:obj:`~pyrogram.storage.Storage`, *optional*): + Pass an instance of your own implementation of session storage engine. + Useful when you want to store your session in databases like Mongo, Redis, etc. """ APP_VERSION = f"Pyrogram {__version__}" @@ -225,7 +229,8 @@ class Client(Methods): takeout: bool = None, sleep_threshold: int = Session.SLEEP_THRESHOLD, hide_password: bool = False, - max_concurrent_transmissions: int = MAX_CONCURRENT_TRANSMISSIONS + max_concurrent_transmissions: int = MAX_CONCURRENT_TRANSMISSIONS, + storage_engine: Storage = None ): super().__init__() @@ -261,6 +266,8 @@ class Client(Methods): self.storage = MemoryStorage(self.name, self.session_string) elif self.in_memory: self.storage = MemoryStorage(self.name) + elif isinstance(storage_engine, Storage): + self.storage = storage_engine else: self.storage = FileStorage(self.name, self.workdir)