Improve performance by caching MongoClient

Author: timgraham

django_mongodb_backend/base.py

@@ -89,6 +89,7 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         "endswith": "LIKE '%%' || {}",
         "iendswith": "LIKE '%%' || UPPER({})",
     }
+    _connection_pools = {}
 
     def _isnull_operator(a, b):
         is_null = {
@@ -176,7 +177,12 @@ def get_connection_params(self):
 
     @async_unsafe
     def get_new_connection(self, conn_params):
-        return MongoClient(**conn_params, driver=self._driver_info())
+        if self.alias not in self._connection_pools:
+            conn = MongoClient(**conn_params, driver=self._driver_info())
+            # setdefault() ensures that multiple threads don't set this in
+            # parallel.
+            self._connection_pools.setdefault(self.alias, conn)
+        return self._connection_pools[self.alias]
 
     def _driver_info(self):
         if not os.environ.get("RUNNING_DJANGOS_TEST_SUITE"):
@@ -189,14 +195,23 @@ def _commit(self):
     def _rollback(self):
         pass
 
-    def set_autocommit(self, autocommit, force_begin_transaction_with_broken_autocommit=False):
-        self.autocommit = autocommit
+    def _close(self):
+        # MongoClient is a connection pool and, unlike database drivers that
+        # implement PEP 249, shouldn't be closed by connection.close().
+        pass
 
-    @async_unsafe
-    def close(self):
-        super().close()
+    def close_pool(self):
+        connection = self.connection
+        # Remove any reference to the connection.
+        self.connection = None
         with contextlib.suppress(AttributeError):
             del self.database
+        del self._connection_pools[self.alias]
+        # Then close it.
+        connection.close()
+
+    def set_autocommit(self, autocommit, force_begin_transaction_with_broken_autocommit=False):
+        self.autocommit = autocommit
 
     @async_unsafe
     def cursor(self):

django_mongodb_backend/features.py

@@ -75,6 +75,13 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # Connection creation doesn't follow the usual Django API.
         "backends.tests.ThreadTests.test_pass_connection_between_threads",
         "backends.tests.ThreadTests.test_default_connection_thread_local",
+        # Disallowed query protection doesn't work because this backend doesn't
+        # use cursor(), chunked_cursor(), etc.
+        # https://github.com/django/django/blob/045110ff3089aefd9c3e65c707df465bacfed986/django/test/testcases.py#L195-L206
+        "test_utils.test_testcase.TestTestCase.test_disallowed_database_queries",
+        "test_utils.test_transactiontestcase.DisallowedDatabaseQueriesTests.test_disallowed_database_queries",
+        "test_utils.tests.DisallowedDatabaseQueriesTests.test_disallowed_database_chunked_cursor_queries",
+        "test_utils.tests.DisallowedDatabaseQueriesTests.test_disallowed_database_queries",
         "test_utils.tests.DisallowedDatabaseQueriesTests.test_disallowed_thread_database_connection",
         # Object of type ObjectId is not JSON serializable.
         "auth_tests.test_views.LoginTest.test_login_session_without_hash_session_key",

tests/backend_/test_base.py

@@ -33,11 +33,11 @@ def receiver(sender, connection, **kwargs):  # noqa: ARG001
             data["database"] = connection.database
 
         connection_created.connect(receiver)
-        connection.close()
+        connection.close_pool()
         # Accessing database implicitly connects.
         connection.database  # noqa: B018
         self.assertIs(data["database"], connection.database)
-        connection.close()
+        connection.close_pool()
         connection_created.disconnect(receiver)
         data.clear()
         connection.connect()