From 2bf362ca0a59767d88767872df430c67b9d48536 Mon Sep 17 00:00:00 2001 From: "Juan Simon .D" <72490350+jusidama18@users.noreply.github.com> Date: Wed, 17 May 2023 14:58:19 +0000 Subject: [PATCH] Pyrofork: Use session name as database name, add some parameters informations and some cleanup (#3) Changes to be committed: modified: pyrogram/client.py modified: pyrogram/storage/mongo_storage.py Signed-off-by: wulan17 --- pyrogram/client.py | 2 +- pyrogram/storage/mongo_storage.py | 33 +++++++++++++++++++------------ 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/pyrogram/client.py b/pyrogram/client.py index b9477f56..3d5021a7 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -255,7 +255,7 @@ class Client(Methods): elif self.in_memory: self.storage = MemoryStorage(self.name) elif self.mongodb: - self.storage = MongoStorage(self.mongodb) + self.storage = MongoStorage(self.name, **self.mongodb) else: self.storage = FileStorage(self.name, self.workdir) diff --git a/pyrogram/storage/mongo_storage.py b/pyrogram/storage/mongo_storage.py index 9b27df89..10841a27 100644 --- a/pyrogram/storage/mongo_storage.py +++ b/pyrogram/storage/mongo_storage.py @@ -11,23 +11,30 @@ from pyrogram.storage.sqlite_storage import get_input_peer class MongoStorage(Storage): """ - config (``dict``) - Mongodb config as dict, e.g.: *dict(uri="mongodb://...", db_name="pyrofork-session", remove_peers=False)*. - Only applicable for new sessions. + Initializes a new session. + + Parameters: + - name (`str`): + The session name used for database name. + + - uri (`str`): + MongoDB Connection String URI. + For more information refer to https://www.mongodb.com/docs/manual/reference/connection-string + + - remove_peers (`bool`, *optional*): + Flag to remove data in the peers collection. If set to True, + the data related to peers will be removed everytime client log out. + If set to False or None, the data will not be removed. + + Example: + session = MongoStorage("my_session", uri="mongodb://...", remove_peers=True) """ lock: asyncio.Lock USERNAME_TTL = 8 * 60 * 60 - def __init__(self, config: dict): - super().__init__('') - db_name = "pyrofork-session" - db_uri = config["uri"] - remove_peers = False - if "db_name" in config: - db_name = config["db_name"] - if "remove_peers" in config: - remove_peers = config["remove_peers"] - database = AsyncIOMotorClient(db_uri)[db_name] + def __init__(self, name: str, uri: str, remove_peers: bool = False): + super().__init__(name=name) + database = AsyncIOMotorClient(uri)[name] self.lock = asyncio.Lock() self.database = database self._peer = database['peers']