add debug logging for SchemaEditor's create_collection() calls

Author: timgraham

django_mongodb/base.py

@@ -12,7 +12,7 @@
 from .operations import DatabaseOperations
 from .query_utils import regex_match
 from .schema import DatabaseSchemaEditor
-from .utils import CollectionDebugWrapper
+from .utils import OperationDebugWrapper
 
 
 class Cursor:
@@ -137,9 +137,14 @@ def __init__(self, *args, **kwargs):
     def get_collection(self, name, **kwargs):
         collection = Collection(self.database, name, **kwargs)
         if self.queries_logged:
-            collection = CollectionDebugWrapper(collection, self)
+            collection = OperationDebugWrapper(self, collection)
         return collection
 
+    def get_database(self):
+        if self.queries_logged:
+            return OperationDebugWrapper(self)
+        return self.database
+
     def __getattr__(self, attr):
         """
         Connect to the database the first time `connection` or `database` are

django_mongodb/schema.py

@@ -7,9 +7,12 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
     def get_collection(self, name):
         return self.connection.get_collection(name)
 
+    def get_database(self):
+        return self.connection.get_database()
+
     @wrap_database_errors
     def create_model(self, model):
-        self.connection.database.create_collection(model._meta.db_table)
+        self.get_database().create_collection(model._meta.db_table)
         # Make implicit M2M tables.
         for field in model._meta.local_many_to_many:
             if field.remote_field.through._meta.auto_created:

django_mongodb/utils.py

@@ -25,13 +25,15 @@ def check_django_compatability():
         )
 
 
-class CollectionDebugWrapper:
-    def __init__(self, collection, db):
+class OperationDebugWrapper:
+    def __init__(self, db, collection=None):
         self.collection = collection
         self.db = db
 
     def __getattr__(self, attr):
-        return getattr(self.collection, attr)
+        if self.collection is not None:
+            return getattr(self.collection, attr)
+        return getattr(self.db, attr)
 
     def profile_call(self, func, args=(), kwargs=None):
         start = time.monotonic()
@@ -44,7 +46,8 @@ def log(self, op, duration, args, kwargs=None):
         # added to this logging.
         msg = "(%.3f) %s"
         args = ", ".join(str(arg) for arg in args)
-        operation = f"{self.collection.name}.{op}({args})"
+        collection_name = f"{self.collection.name}." if self.collection is not None else ""
+        operation = f"db.{collection_name}{op}({args})"
         if len(settings.DATABASES) > 1:
             msg += f"; alias={self.db.alias}"
         self.db.queries_log.append(
@@ -66,7 +69,10 @@ def log(self, op, duration, args, kwargs=None):
 
     def logging_wrapper(method):
         def wrapper(self, *args, **kwargs):
-            func = getattr(self.collection, method)
+            if self.collection is not None:
+                func = getattr(self.collection, method)
+            else:
+                func = getattr(self.db.database, method)
             # Collection.insert_many() mutates args (the documents) by adding
             #  _id. deepcopy() to avoid logging that version.
             original_args = copy.deepcopy(args)
@@ -78,6 +84,7 @@ def wrapper(self, *args, **kwargs):
 
     # These are the operations that this backend uses.
     aggregate = logging_wrapper("aggregate")
+    create_collection = logging_wrapper("create_collection")
     drop = logging_wrapper("drop")
     insert_many = logging_wrapper("insert_many")
     delete_many = logging_wrapper("delete_many")