Add support for custom config for example using local mongoDB server.

Changes to be committed:
	modified:   pyrogram/client.py
	modified:   pyrogram/storage/mongo_storage.py
This commit is contained in:
Juan Simon .D 2023-05-17 14:58:19 +00:00
parent 5bc419ce2c
commit 8c462af6a2
2 changed files with 20 additions and 15 deletions

View file

@ -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)

View file

@ -11,28 +11,33 @@ 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.
- remove_peers (bool): 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.
- **kwargs: Additional keyword arguments to pass to the AsyncIOMotorClient.
Note:
The `kwargs` parameter is used to pass additional arguments to the underlying AsyncIOMotorClient
instance. Refer to the AsyncIOMotorClient documentation for the list of available arguments.
Example:
session = MongoStorage("my_session", remove_peers=True, uri="mongodb://...")
"""
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, remove_peers: bool = None, **kwargs):
super().__init__(name=name)
database = AsyncIOMotorClient(**kwargs)[name]
self.lock = asyncio.Lock()
self.database = database
self._peer = database['peers']
self._session = database['session']
self._remove_peers = remove_peers
self._remove_peers = remove_peers or False
async def open(self):
"""