Skip to content

Commit

Permalink
Add support for custom config for example using local mongoDB server.
Browse files Browse the repository at this point in the history
Changes to be committed:
	modified:   pyrogram/client.py
	modified:   pyrogram/storage/mongo_storage.py
  • Loading branch information
jusidama18 authored and wulan17 committed May 18, 2023
1 parent a736a18 commit 6d8a810
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 14 deletions.
2 changes: 1 addition & 1 deletion pyrogram/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -255,7 +255,7 @@ def __init__(
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)

Expand Down
31 changes: 18 additions & 13 deletions pyrogram/storage/mongo_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,23 +11,28 @@

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, 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']
Expand Down

0 comments on commit 6d8a810

Please sign in to comment.