refactor connection creation to use Django's APIs

Author: timgraham

django_mongodb/base.py

@@ -1,5 +1,6 @@
+import contextlib
+
 from django.db.backends.base.base import BaseDatabaseWrapper
-from django.db.backends.signals import connection_created
 from pymongo.collection import Collection
 from pymongo.mongo_client import MongoClient
 
@@ -128,11 +129,6 @@ def _isnull_operator(a, b):
     introspection_class = DatabaseIntrospection
     ops_class = DatabaseOperations
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.connected = False
-        del self.connection
-
     def get_collection(self, name, **kwargs):
         collection = Collection(self.database, name, **kwargs)
         if self.queries_logged:
@@ -145,44 +141,45 @@ def get_database(self):
         return self.database
 
     def __getattr__(self, attr):
-        """
-        Connect to the database the first time `connection` or `database` are
-        accessed.
-        """
-        if attr in ["connection", "database"]:
-            assert not self.connected
-            self._connect()
+        """Connect to the database the first time `database` is accessed."""
+        if attr == "database":
+            if self.connection is None:
+                self.connect()
             return getattr(self, attr)
         raise AttributeError(attr)
 
-    def _connect(self):
-        settings_dict = self.settings_dict
-        self.connection = MongoClient(
-            host=settings_dict["HOST"] or None,
-            port=int(settings_dict["PORT"] or 27017),
-            username=settings_dict.get("USER"),
-            password=settings_dict.get("PASSWORD"),
-            **settings_dict["OPTIONS"],
-        )
-        db_name = settings_dict["NAME"]
+    def init_connection_state(self):
+        db_name = self.settings_dict["NAME"]
         if db_name:
             self.database = self.connection[db_name]
+        super().init_connection_state()
+
+    def get_connection_params(self):
+        settings_dict = self.settings_dict
+        return {
+            "host": settings_dict["HOST"] or None,
+            "port": int(settings_dict["PORT"] or 27017),
+            "username": settings_dict.get("USER"),
+            "password": settings_dict.get("PASSWORD"),
+            **settings_dict["OPTIONS"],
+        }
 
-        self.connected = True
-        connection_created.send(sender=self.__class__, connection=self)
+    def get_new_connection(self, conn_params):
+        return MongoClient(**conn_params)
 
     def _commit(self):
         pass
 
     def _rollback(self):
         pass
 
+    def set_autocommit(self, autocommit, force_begin_transaction_with_broken_autocommit=False):
+        pass
+
     def close(self):
-        if self.connected:
-            self.connection.close()
-            del self.connection
+        super().close()
+        with contextlib.suppress(AttributeError):
             del self.database
-            self.connected = False
 
     def cursor(self):
         return Cursor()

django_mongodb/features.py

@@ -71,7 +71,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "backends.tests.BackendTestCase.test_is_usable_after_database_disconnects",
         # Connection creation doesn't follow the usual Django API.
         "backends.tests.ThreadTests.test_pass_connection_between_threads",
-        "backends.tests.ThreadTests.test_closing_non_shared_connections",
         "backends.tests.ThreadTests.test_default_connection_thread_local",
         # Union as subquery is not mapping the parent parameter and collections:
         # https://github.com/mongodb-labs/django-mongodb/issues/156

django_mongodb/query.py

@@ -9,6 +9,7 @@
 from django.db.models.sql.constants import INNER
 from django.db.models.sql.datastructures import Join
 from django.db.models.sql.where import AND, OR, XOR, ExtraWhere, NothingNode, WhereNode
+from django.utils.functional import cached_property
 from pymongo.errors import BulkWriteError, DuplicateKeyError, PyMongoError
 
 
@@ -47,7 +48,6 @@ def __init__(self, compiler):
         self.query = compiler.query
         self._negated = False
         self.ordering = []
-        self.collection = self.compiler.collection
         self.collection_name = self.compiler.collection_name
         self.mongo_query = getattr(compiler.query, "raw_query", {})
         self.subqueries = None
@@ -60,6 +60,12 @@ def __init__(self, compiler):
         # subquery.
         self.subquery_lookup = None
 
+    @cached_property
+    def collection(self):
+        # Initialize this lazily since `compiler.collection` connects to the
+        # database.
+        return self.compiler.collection
+
     def __repr__(self):
         return f"<MongoQuery: {self.mongo_query!r} ORDER {self.ordering!r}>"