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

@@ -47,7 +47,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
@@ -68,15 +67,15 @@ def delete(self):
         """Execute a delete query."""
         if self.compiler.subqueries:
             raise NotSupportedError("Cannot use QuerySet.delete() when a subquery is required.")
-        return self.collection.delete_many(self.mongo_query).deleted_count
+        return self.compiler.collection.delete_many(self.mongo_query).deleted_count
 
     @wrap_database_errors
     def get_cursor(self):
         """
         Return a pymongo CommandCursor that can be iterated on to give the
         results of the query.
         """
-        return self.collection.aggregate(self.get_pipeline())
+        return self.compiler.collection.aggregate(self.get_pipeline())
 
     def get_pipeline(self):
         pipeline = []