Add support for connection pooling

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"):
@@ -192,11 +198,27 @@ def _rollback(self):
     def set_autocommit(self, autocommit, force_begin_transaction_with_broken_autocommit=False):
         self.autocommit = autocommit
 
+    def _close(self):
+        pass
+
     @async_unsafe
     def close(self):
-        super().close()
+        self.validate_thread_sharing()
+        # MongoClient is a connection pool and, unlike database drivers that
+        # implement PEP 249, shouldn't be closed by connection.close().
+
+    def close_pool(self):
+        """Close the MongoClient."""
+        connection = self.connection
+        if connection is None:
+            return
+        # Remove all references to the connection.
+        self.connection = None
         with contextlib.suppress(AttributeError):
             del self.database
+        del self._connection_pools[self.alias]
+        # Then close it.
+        connection.close()
 
     @async_unsafe
     def cursor(self):

django_mongodb_backend/creation.py

@@ -8,6 +8,9 @@
 
 class DatabaseCreation(BaseDatabaseCreation):
     def _execute_create_test_db(self, cursor, parameters, keepdb=False):
+        # Close the connection (which may point to the non-test database) so
+        # that a new connection to the test database can be established later.
+        self.connection.close_pool()
         if not keepdb:
             self._destroy_test_db(parameters["dbname"], verbosity=0)
 
@@ -29,3 +32,8 @@ def create_test_db(self, *args, **kwargs):
             database=self.connection.alias, verbosity=kwargs["verbosity"]
         )
         return test_database_name
+
+    def destroy_test_db(self, old_database_name=None, verbosity=1, keepdb=False, suffix=None):
+        super().destroy_test_db(old_database_name, verbosity, keepdb, suffix)
+        # Close the connection to the test database.
+        self.connection.close_pool()

django_mongodb_backend/features.py

@@ -72,7 +72,6 @@ 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",
-        "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",
         # GenericRelation.value_to_string() assumes integer pk.
@@ -544,6 +543,19 @@ def django_test_expected_failures(self):
             "custom_lookups.tests.LookupTests.test_div3_extract",
             "custom_lookups.tests.SubqueryTransformTests.test_subquery_usage",
         },
+        "connection.close() does not close the connection.": {
+            "servers.test_liveserverthread.LiveServerThreadTest.test_closes_connections",
+            "servers.tests.LiveServerTestCloseConnectionTest.test_closes_connections",
+        },
+        "Disallowed query protection doesn't work on MongoDB.": {
+            # 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",
+        },
     }
 
     @cached_property

docs/source/index.rst

@@ -39,6 +39,7 @@ Models
 - :doc:`ref/models/querysets`
 - :doc:`ref/models/models`
 - :doc:`ref/models/indexes`
+- :doc:`ref/database`
 
 **Topic guides:**
 

docs/source/ref/database.rst

@@ -0,0 +1,43 @@
+==================
+Database reference
+==================
+
+This document supplements :doc:`Django's documentation on databases
+<django:ref/databases>`.
+
+Persistent connections
+======================
+
+Persistent connections avoid the overhead of reestablishing a connection to
+the database in each HTTP request. They're normally controlled by the
+:setting:`CONN_MAX_AGE` parameter which defines the maximum lifetime of a
+connection. However, this parameter is unnecessary and has no effect with
+Django MongoDB Backend because Django's API for connection-closing
+(``django.db.connection.close()``) has no effect. In other words, persistent
+connections are enabled by default.
+
+.. versionadded:: 5.2.0b0
+
+    Support for connection pooling was added.  In older versions, use
+    :setting:`CONN_MAX_AGE` to enable persistent connections.
+
+.. _connection-management:
+
+Connection management
+=====================
+
+Django uses this backend to open a connection pool to the database when it
+first makes a database query. It keeps this pool open and reuses it in
+subsequent requests.
+
+The underlying :class:`~pymongo.mongo_client.MongoClient` takes care connection
+management, so the :setting:`CONN_HEALTH_CHECKS` setting is unnecessary and has
+no effect.
+
+Django's API for connection-closing (``django.db.connection.close()``) has no
+effect. Rather, if you need to close the connection pool, use
+``django.db.connection.close_pool()``.
+
+.. versionadded:: 5.2.0b0
+
+    Support for connection pooling and ``connection.close_pool()`` were added.

docs/source/ref/index.rst

@@ -7,5 +7,6 @@ API reference
 
    models/index
    forms
+   database
    django-admin
    utils

docs/source/releases/5.2.x.rst

@@ -20,6 +20,8 @@ New features
 
 - Added :class:`.SearchIndex` and :class:`.VectorSearchIndex` for use on
   a model's :attr:`Meta.indexes <django.db.models.Options.indexes>`.
+- PyMongo's connection pooling is now used by default. See
+  :ref:`connection-management`.
 
 Backwards incompatible changes
 ------------------------------

tests/backend_/test_base.py

@@ -23,6 +23,19 @@ def test_set_autocommit(self):
         connection.set_autocommit(True)
         self.assertIs(connection.get_autocommit(), True)
 
+    def test_close(self):
+        """connection.close() doesn't close the connection."""
+        conn = connection.connection
+        self.assertIsNotNone(conn)
+        connection.close()
+        self.assertEqual(connection.connection, conn)
+
+    def test_close_pool(self):
+        """connection.close_pool() closes the connection."""
+        self.assertIsNotNone(connection.connection)
+        connection.close_pool()
+        self.assertIsNone(connection.connection)
+
     def test_connection_created_database_attr(self):
         """
         connection.database is available in the connection_created signal.
@@ -33,11 +46,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()