Debug toolbar extended with current mongoengine commands

flask_mongoengine/operation_tracker.py

@@ -25,8 +25,14 @@
 
 _original_methods = {
     "insert": pymongo.collection.Collection.insert,
+    "insert_one": pymongo.collection.Collection.insert_one,
+    "insert_many": pymongo.collection.Collection.insert_many,
     "update": pymongo.collection.Collection.update,
+    "update_one": pymongo.collection.Collection.update_one,
+    "update_many": pymongo.collection.Collection.update_many,
     "remove": pymongo.collection.Collection.remove,
+    "delete_one": pymongo.collection.Collection.delete_one,
+    "delete_many": pymongo.collection.Collection.delete_many,
     "refresh": pymongo.cursor.Cursor._refresh,
     "_unpack_response": pymongo.command_cursor.CommandCursor._unpack_response,
 }
@@ -52,18 +58,29 @@ def _unpack_response(response, *args, **kwargs):
 # Wrap Cursor.insert for getting queries
 @functools.wraps(_original_methods["insert"])
 def _insert(
-    collection_self, doc_or_docs, manipulate=True, safe=None, check_keys=True, **kwargs
+    collection_self,
+    doc_or_docs,
+    manipulate=True,
+    check_keys=True,
+    continue_on_error=False,
+    **kwargs,
 ):
     start_time = time.time()
     result = _original_methods["insert"](
-        collection_self, doc_or_docs, check_keys=check_keys, **kwargs
+        collection_self,
+        doc_or_docs,
+        manipulate=manipulate,
+        check_keys=check_keys,
+        continue_on_error=continue_on_error,
+        **kwargs,
     )
     total_time = (time.time() - start_time) * 1000
 
     __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     inserts.append(
         {
+            "operation": "insert(deprecated)",
             "document": doc_or_docs,
             "time": total_time,
             "stack_trace": stack_trace,
@@ -74,32 +91,199 @@ def _insert(
     return result
 
 
+# Wrap Cursor.insert_one for getting queries
+@functools.wraps(_original_methods["insert_one"])
+def _insert_one(
+    collection_self, document, bypass_document_validation=False, session=None, **kwargs
+):
+    start_time = time.time()
+    result = _original_methods["insert_one"](
+        collection_self,
+        document,
+        bypass_document_validation=bypass_document_validation,
+        session=session,
+        **kwargs,
+    )
+    total_time = (time.time() - start_time) * 1000
+
+    __traceback_hide__ = True  # noqa
+    stack_trace, internal = _tidy_stacktrace()
+    inserts.append(
+        {
+            "operation": "insert one",
+            "document": document,
+            "time": total_time,
+            "stack_trace": stack_trace,
+            "size": response_sizes[-1] if response_sizes else 0,
+            "internal": internal,
+        }
+    )
+    return result
+
+
+# Wrap Cursor.insert_many for getting queries
+@functools.wraps(_original_methods["insert_many"])
+def _insert_many(
+    collection_self,
+    documents,
+    ordered=True,
+    bypass_document_validation=False,
+    session=None,
+    **kwargs,
+):
+    start_time = time.time()
+    result = _original_methods["insert_many"](
+        collection_self,
+        documents,
+        ordered=ordered,
+        bypass_document_validation=bypass_document_validation,
+        session=session,
+        **kwargs,
+    )
+    total_time = (time.time() - start_time) * 1000
+
+    __traceback_hide__ = True  # noqa
+    stack_trace, internal = _tidy_stacktrace()
+    inserts.append(
+        {
+            "operation": "insert many",
+            "document": documents,
+            "time": total_time,
+            "stack_trace": stack_trace,
+            "size": response_sizes[-1] if response_sizes else 0,
+            "internal": internal,
+        }
+    )
+    return result
+
+
 # Wrap Cursor.update for getting queries
 @functools.wraps(_original_methods["update"])
 def _update(
     collection_self,
     spec,
     document,
     upsert=False,
-    maniuplate=False,
-    safe=None,
+    manipulate=False,
     multi=False,
+    check_keys=True,
     **kwargs,
 ):
     start_time = time.time()
     result = _original_methods["update"](
-        collection_self, spec, document, upsert=upsert, multi=multi, **kwargs
+        collection_self,
+        spec,
+        document,
+        upsert=upsert,
+        manipulate=manipulate,
+        multi=multi,
+        check_keys=check_keys,
+        **kwargs,
     )
     total_time = (time.time() - start_time) * 1000
 
     __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     updates.append(
         {
-            "document": document,
+            "operation": "update(deprecated)",
+            "flt": spec,
+            "upd": document,
             "upsert": upsert,
             "multi": multi,
-            "spec": spec,
+            "time": total_time,
+            "stack_trace": stack_trace,
+            "size": response_sizes[-1] if response_sizes else 0,
+            "internal": internal,
+        }
+    )
+    return result
+
+
+# Wrap Cursor.update_one for getting queries
+@functools.wraps(_original_methods["update_one"])
+def _update_one(
+    collection_self,
+    filter,
+    update,
+    upsert=False,
+    bypass_document_validation=False,
+    collation=None,
+    array_filters=None,
+    hint=None,
+    session=None,
+    **kwargs,
+):
+    start_time = time.time()
+    result = _original_methods["update_one"](
+        collection_self,
+        filter,
+        update,
+        upsert=upsert,
+        bypass_document_validation=bypass_document_validation,
+        collation=collation,
+        array_filters=array_filters,
+        hint=hint,
+        session=session,
+        **kwargs,
+    )
+    total_time = (time.time() - start_time) * 1000
+
+    __traceback_hide__ = True  # noqa
+    stack_trace, internal = _tidy_stacktrace()
+    updates.append(
+        {
+            "operation": "update one",
+            "flt": filter,
+            "upd": update,
+            "upsert": upsert,
+            "multi": False,
+            "time": total_time,
+            "stack_trace": stack_trace,
+            "size": response_sizes[-1] if response_sizes else 0,
+            "internal": internal,
+        }
+    )
+    return result
+
+
+# Wrap Cursor.update_many for getting queries
+@functools.wraps(_original_methods["update_many"])
+def _update_many(
+    collection_self,
+    filter,
+    update,
+    upsert=False,
+    array_filters=None,
+    bypass_document_validation=False,
+    collation=None,
+    hint=None,
+    session=None,
+    **kwargs,
+):
+    start_time = time.time()
+    result = _original_methods["update_many"](
+        collection_self,
+        filter,
+        update,
+        upsert=upsert,
+        array_filters=array_filters,
+        bypass_document_validation=bypass_document_validation,
+        collation=collation,
+        hint=hint,
+        session=session,
+        **kwargs,
+    )
+    total_time = (time.time() - start_time) * 1000
+
+    __traceback_hide__ = True  # noqa
+    stack_trace, internal = _tidy_stacktrace()
+    updates.append(
+        {
+            "operation": "update many",
+            "flt": filter,
+            "upd": update,
+            "multi": True,
             "time": total_time,
             "stack_trace": stack_trace,
             "size": response_sizes[-1] if response_sizes else 0,
@@ -111,18 +295,86 @@ def _update(
 
 # Wrap Cursor.remove for getting queries
 @functools.wraps(_original_methods["remove"])
-def _remove(collection_self, spec_or_id, safe=None, **kwargs):
+def _remove(collection_self, spec_or_id, multi=True, **kwargs):
     start_time = time.time()
-    result = _original_methods["remove"](collection_self, spec_or_id, **kwargs)
+    result = _original_methods["remove"](
+        collection_self, spec_or_id=spec_or_id, multi=multi, **kwargs
+    )
     total_time = (time.time() - start_time) * 1000
 
     __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     removes.append(
         {
+            "operation": "remove(deprecated)",
             "spec_or_id": spec_or_id,
+            "multi": multi,
+            "time": total_time,
+            "stack_trace": stack_trace,
+            "size": response_sizes[-1] if response_sizes else 0,
+            "internal": internal,
+        }
+    )
+    return result
+
+
+# Wrap Cursor.delete_one for getting queries
+@functools.wraps(_original_methods["delete_one"])
+def _delete_one(
+    collection_self, filter, collation=None, hint=None, session=None, **kwargs
+):
+    start_time = time.time()
+    result = _original_methods["delete_one"](
+        collection_self,
+        filter,
+        collation=collation,
+        hint=hint,
+        session=session,
+        **kwargs,
+    )
+    total_time = (time.time() - start_time) * 1000
+
+    __traceback_hide__ = True  # noqa
+    stack_trace, internal = _tidy_stacktrace()
+    removes.append(
+        {
+            "operation": "delete one",
+            "spec_or_id": filter,
+            "multi": False,
             "time": total_time,
-            "   ": stack_trace,
+            "stack_trace": stack_trace,
+            "size": response_sizes[-1] if response_sizes else 0,
+            "internal": internal,
+        }
+    )
+    return result
+
+
+# Wrap Cursor.delete_many for getting queries
+@functools.wraps(_original_methods["delete_many"])
+def _delete_many(
+    collection_self, filter, collation=None, hint=None, session=None, **kwargs
+):
+    start_time = time.time()
+    result = _original_methods["delete_many"](
+        collection_self,
+        filter,
+        collation=collation,
+        hint=hint,
+        session=session,
+        **kwargs,
+    )
+    total_time = (time.time() - start_time) * 1000
+
+    __traceback_hide__ = True  # noqa
+    stack_trace, internal = _tidy_stacktrace()
+    removes.append(
+        {
+            "operation": "delete many",
+            "spec_or_id": filter,
+            "multi": True,
+            "time": total_time,
+            "stack_trace": stack_trace,
             "size": response_sizes[-1] if response_sizes else 0,
             "internal": internal,
         }
@@ -215,10 +467,22 @@ def privar(name):
 def install_tracker():
     if pymongo.collection.Collection.insert != _insert:
         pymongo.collection.Collection.insert = _insert
+    if pymongo.collection.Collection.insert_one != _insert_one:
+        pymongo.collection.Collection.insert_one = _insert_one
+    if pymongo.collection.Collection.insert_many != _insert_many:
+        pymongo.collection.Collection.insert_many = _insert_many
     if pymongo.collection.Collection.update != _update:
         pymongo.collection.Collection.update = _update
+    if pymongo.collection.Collection.update_one != _update_one:
+        pymongo.collection.Collection.update_one = _update_one
+    if pymongo.collection.Collection.update_many != _update_many:
+        pymongo.collection.Collection.update_many = _update_many
     if pymongo.collection.Collection.remove != _remove:
         pymongo.collection.Collection.remove = _remove
+    if pymongo.collection.Collection.delete_one != _delete_one:
+        pymongo.collection.Collection.delete_one = _delete_one
+    if pymongo.collection.Collection.delete_many != _delete_many:
+        pymongo.collection.Collection.delete_many = _delete_many
     if pymongo.cursor.Cursor._refresh != _cursor_refresh:
         pymongo.cursor.Cursor._refresh = _cursor_refresh
     if pymongo.command_cursor.CommandCursor._unpack_response != _unpack_response:
@@ -228,10 +492,22 @@ def install_tracker():
 def uninstall_tracker():
     if pymongo.collection.Collection.insert == _insert:
         pymongo.collection.Collection.insert = _original_methods["insert"]
+    if pymongo.collection.Collection.insert_one == _insert_one:
+        pymongo.collection.Collection.insert_one = _original_methods["insert_one"]
+    if pymongo.collection.Collection.insert_many == _insert_many:
+        pymongo.collection.Collection.insert_many = _original_methods["insert_many"]
     if pymongo.collection.Collection.update == _update:
         pymongo.collection.Collection.update = _original_methods["update"]
+    if pymongo.collection.Collection.update_one == _update_one:
+        pymongo.collection.Collection.update_one = _original_methods["update_one"]
+    if pymongo.collection.Collection.update_many == _update_many:
+        pymongo.collection.Collection.update_many = _original_methods["update_many"]
     if pymongo.collection.Collection.remove == _remove:
         pymongo.collection.Collection.remove = _original_methods["remove"]
+    if pymongo.collection.Collection.delete_one == _delete_one:
+        pymongo.collection.Collection.delete_one = _original_methods["delete_one"]
+    if pymongo.collection.Collection.delete_many == _delete_many:
+        pymongo.collection.Collection.delete_many = _original_methods["delete_many"]
     if pymongo.cursor.Cursor._refresh == _cursor_refresh:
         pymongo.cursor.Cursor._refresh = _original_methods["cursor_refresh"]
     if pymongo.command_cursor.CommandCursor._unpack_response == _unpack_response: