From 8272c706a6ba3da564f3ea678ba384cb75a6a991 Mon Sep 17 00:00:00 2001
From: wulan17 <wulan17@nusantararom.org>
Date: Mon, 22 May 2023 19:18:55 +0700
Subject: [PATCH] PyroFork: storage: mongo: Use existing database connection

support both async_pymongo and motor

Signed-off-by: wulan17 <wulan17@nusantararom.org>
---
 docs/source/topics/storage-engines.rst | 20 ++++++++----
 pyrogram/client.py                     |  2 +-
 pyrogram/storage/mongo_storage.py      | 45 ++++++++++++++++++++++----
 requirements.txt                       |  1 -
 4 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/docs/source/topics/storage-engines.rst b/docs/source/topics/storage-engines.rst
index b739ed9a..b148d483 100644
--- a/docs/source/topics/storage-engines.rst
+++ b/docs/source/topics/storage-engines.rst
@@ -67,15 +67,23 @@ Mongodb Storage
 In case you want to have persistent session but you don't have persistent storage you can use mongodb storage by passing
 mongodb config as ``dict`` to the ``mongodb`` parameter of the :obj:`~pyrogram.Client` constructor:
 
+Using async_pymongo (Recommended for python3.9+):
 .. code-block:: python
-
+    from async_pymongo import AsyncClient
     from pyrogram import Client
 
-    # uri (``str``):
-    #   mongodb database uri
-    # remove_peers (``bool``, *optional*):
-    #   remove peers collection on logout, default = False
-    async with Client("my_account", mongodb=dict(uri="mongodb://...", remove_peers=False)) as app:
+    conn = AsyncClient("mongodb://...")
+    async with Client("my_account", mongodb=dict(connection=conn, remove_peers=False)) as app:
+        print(await app.get_me())
+
+
+Using motor:
+.. code-block:: python
+    from motor.motor_asyncio import AsyncIOMotorClient
+    from pyrogram import Client
+
+    conn = AsyncIOMotorClient("mongodb://...")
+    async with Client("my_account", mongodb=dict(connection=conn, remove_peers=False)) as app:
         print(await app.get_me())
 
 This storage engine is backed by MongoDB, a session will be created and saved to mongodb database. Any subsequent client
diff --git a/pyrogram/client.py b/pyrogram/client.py
index 2995ba93..2567db48 100644
--- a/pyrogram/client.py
+++ b/pyrogram/client.py
@@ -121,7 +121,7 @@ class Client(Methods):
             Defaults to False.
 
         mongodb (``dict``, *optional*):
-            Mongodb config as dict, e.g.: *dict(uri="mongodb://...", remove_peers=False)*.
+            Mongodb config as dict, e.g.: *dict(connection=async_pymongo.AsyncClient("mongodb://..."), remove_peers=False)*.
             Only applicable for new sessions.
 
         phone_number (``str``, *optional*):
diff --git a/pyrogram/storage/mongo_storage.py b/pyrogram/storage/mongo_storage.py
index 10841a27..416e6453 100644
--- a/pyrogram/storage/mongo_storage.py
+++ b/pyrogram/storage/mongo_storage.py
@@ -3,7 +3,6 @@ import inspect
 import time
 from typing import List, Tuple, Any
 
-from motor.motor_asyncio import AsyncIOMotorClient
 from pymongo import UpdateOne
 from pyrogram.storage.storage import Storage
 from pyrogram.storage.sqlite_storage import get_input_peer
@@ -17,9 +16,9 @@ class MongoStorage(Storage):
         - 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
+        - connection (`obj`):
+            Mongodb connections object.
+            ~async_pymongo.AsyncClient or ~motor.motor_asyncio.AsyncIOMotorClient object
 
         - remove_peers (`bool`, *optional*):
             Flag to remove data in the peers collection. If set to True, 
@@ -27,14 +26,46 @@ class MongoStorage(Storage):
             If set to False or None, the data will not be removed.
 
     Example:
-        session = MongoStorage("my_session", uri="mongodb://...", remove_peers=True)
+        import async_pymongo
+
+        conn = async_pymongo.AsyncClient("mongodb://...")
+        bot_db = conn["my_bot"]
+        session = MongoStorage("my_session", connection=conn, remove_peers=True)
     """
     lock: asyncio.Lock
     USERNAME_TTL = 8 * 60 * 60
 
-    def __init__(self, name: str, uri: str, remove_peers: bool = False):
+    def __init__(
+        self,
+        name: str,
+        connection: object,
+        remove_peers: bool = False
+    ):
         super().__init__(name=name)
-        database = AsyncIOMotorClient(uri)[name]
+        database = None
+        try:
+            import async_pymongo
+        except ImportError:
+            pass
+        else:
+            if isinstance(connection, async_pymongo.AsyncClient):
+                database = connection[name]
+
+        try:
+            from motor.motor_asyncio import AsyncIOMotorClient
+        except ImportError:
+            pass
+        else:
+            if database:
+                pass
+            elif isinstance(connection, AsyncIOMotorClient):
+                database = connection[name]
+            else:
+                raise Exception("Wrong connection object type! please pass valid connection object to connection parameter!")
+
+        if not database:
+            raise Exception("Please install one of following modules!: async_pymongo, motor")
+
         self.lock = asyncio.Lock()
         self.database = database
         self._peer = database['peers']
diff --git a/requirements.txt b/requirements.txt
index d0b2161a..6605d469 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,4 @@
 aiosqlite>=0.17.0,<0.19.0
-motor==3.1.2
 pyaes==1.6.1
 pymediainfo==6.0.1
 pymongo==4.3.3