Skip to content

Commit 57a48da

Browse files
insspbinsspb
committed
Add modern mongoengine operations to debug panel
1 parent 085f9f8 commit 57a48da

File tree

2 files changed

+297
-25
lines changed

2 files changed

+297
-25
lines changed

‎flask_mongoengine/operation_tracker.py‎

Lines changed: 286 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,14 @@
2525

2626
_original_methods = {
2727
"insert": pymongo.collection.Collection.insert,
28+
"insert_one": pymongo.collection.Collection.insert_one,
29+
"insert_many": pymongo.collection.Collection.insert_many,
2830
"update": pymongo.collection.Collection.update,
31+
"update_one": pymongo.collection.Collection.update_one,
32+
"update_many": pymongo.collection.Collection.update_many,
2933
"remove": pymongo.collection.Collection.remove,
34+
"delete_one": pymongo.collection.Collection.delete_one,
35+
"delete_many": pymongo.collection.Collection.delete_many,
3036
"refresh": pymongo.cursor.Cursor._refresh,
3137
"_unpack_response": pymongo.command_cursor.CommandCursor._unpack_response,
3238
}
@@ -52,18 +58,29 @@ def _unpack_response(response, *args, **kwargs):
5258
# Wrap Cursor.insert for getting queries
5359
@functools.wraps(_original_methods["insert"])
5460
def _insert(
55-
collection_self, doc_or_docs, manipulate=True, safe=None, check_keys=True, **kwargs
61+
collection_self,
62+
doc_or_docs,
63+
manipulate=True,
64+
check_keys=True,
65+
continue_on_error=False,
66+
**kwargs,
5667
):
5768
start_time = time.time()
5869
result = _original_methods["insert"](
59-
collection_self, doc_or_docs, check_keys=check_keys, **kwargs
70+
collection_self,
71+
doc_or_docs,
72+
manipulate=manipulate,
73+
check_keys=check_keys,
74+
continue_on_error=continue_on_error,
75+
**kwargs,
6076
)
6177
total_time = (time.time() - start_time) * 1000
6278

6379
__traceback_hide__ = True # noqa
6480
stack_trace, internal = _tidy_stacktrace()
6581
inserts.append(
6682
{
83+
"operation": "insert(deprecated)",
6784
"document": doc_or_docs,
6885
"time": total_time,
6986
"stack_trace": stack_trace,
@@ -74,32 +91,199 @@ def _insert(
7491
return result
7592

7693

94+
# Wrap Cursor.insert_one for getting queries
95+
@functools.wraps(_original_methods["insert_one"])
96+
def _insert_one(
97+
collection_self, document, bypass_document_validation=False, session=None, **kwargs
98+
):
99+
start_time = time.time()
100+
result = _original_methods["insert_one"](
101+
collection_self,
102+
document,
103+
bypass_document_validation=bypass_document_validation,
104+
session=session,
105+
**kwargs,
106+
)
107+
total_time = (time.time() - start_time) * 1000
108+
109+
__traceback_hide__ = True # noqa
110+
stack_trace, internal = _tidy_stacktrace()
111+
inserts.append(
112+
{
113+
"operation": "insert one",
114+
"document": document,
115+
"time": total_time,
116+
"stack_trace": stack_trace,
117+
"size": response_sizes[-1] if response_sizes else 0,
118+
"internal": internal,
119+
}
120+
)
121+
return result
122+
123+
124+
# Wrap Cursor.insert_many for getting queries
125+
@functools.wraps(_original_methods["insert_many"])
126+
def _insert_many(
127+
collection_self,
128+
documents,
129+
ordered=True,
130+
bypass_document_validation=False,
131+
session=None,
132+
**kwargs,
133+
):
134+
start_time = time.time()
135+
result = _original_methods["insert_many"](
136+
collection_self,
137+
documents,
138+
ordered=ordered,
139+
bypass_document_validation=bypass_document_validation,
140+
session=session,
141+
**kwargs,
142+
)
143+
total_time = (time.time() - start_time) * 1000
144+
145+
__traceback_hide__ = True # noqa
146+
stack_trace, internal = _tidy_stacktrace()
147+
inserts.append(
148+
{
149+
"operation": "insert many",
150+
"document": documents,
151+
"time": total_time,
152+
"stack_trace": stack_trace,
153+
"size": response_sizes[-1] if response_sizes else 0,
154+
"internal": internal,
155+
}
156+
)
157+
return result
158+
159+
77160
# Wrap Cursor.update for getting queries
78161
@functools.wraps(_original_methods["update"])
79162
def _update(
80163
collection_self,
81164
spec,
82165
document,
83166
upsert=False,
84-
maniuplate=False,
85-
safe=None,
167+
manipulate=False,
86168
multi=False,
169+
check_keys=True,
87170
**kwargs,
88171
):
89172
start_time = time.time()
90173
result = _original_methods["update"](
91-
collection_self, spec, document, upsert=upsert, multi=multi, **kwargs
174+
collection_self,
175+
spec,
176+
document,
177+
upsert=upsert,
178+
manipulate=manipulate,
179+
multi=multi,
180+
check_keys=check_keys,
181+
**kwargs,
92182
)
93183
total_time = (time.time() - start_time) * 1000
94184

95185
__traceback_hide__ = True # noqa
96186
stack_trace, internal = _tidy_stacktrace()
97187
updates.append(
98188
{
99-
"document": document,
189+
"operation": "update(deprecated)",
190+
"flt": spec,
191+
"upd": document,
100192
"upsert": upsert,
101193
"multi": multi,
102-
"spec": spec,
194+
"time": total_time,
195+
"stack_trace": stack_trace,
196+
"size": response_sizes[-1] if response_sizes else 0,
197+
"internal": internal,
198+
}
199+
)
200+
return result
201+
202+
203+
# Wrap Cursor.update_one for getting queries
204+
@functools.wraps(_original_methods["update_one"])
205+
def _update_one(
206+
collection_self,
207+
filter,
208+
update,
209+
upsert=False,
210+
bypass_document_validation=False,
211+
collation=None,
212+
array_filters=None,
213+
hint=None,
214+
session=None,
215+
**kwargs,
216+
):
217+
start_time = time.time()
218+
result = _original_methods["update_one"](
219+
collection_self,
220+
filter,
221+
update,
222+
upsert=upsert,
223+
bypass_document_validation=bypass_document_validation,
224+
collation=collation,
225+
array_filters=array_filters,
226+
hint=hint,
227+
session=session,
228+
**kwargs,
229+
)
230+
total_time = (time.time() - start_time) * 1000
231+
232+
__traceback_hide__ = True # noqa
233+
stack_trace, internal = _tidy_stacktrace()
234+
updates.append(
235+
{
236+
"operation": "update one",
237+
"flt": filter,
238+
"upd": update,
239+
"upsert": upsert,
240+
"multi": False,
241+
"time": total_time,
242+
"stack_trace": stack_trace,
243+
"size": response_sizes[-1] if response_sizes else 0,
244+
"internal": internal,
245+
}
246+
)
247+
return result
248+
249+
250+
# Wrap Cursor.update_many for getting queries
251+
@functools.wraps(_original_methods["update_many"])
252+
def _update_many(
253+
collection_self,
254+
filter,
255+
update,
256+
upsert=False,
257+
array_filters=None,
258+
bypass_document_validation=False,
259+
collation=None,
260+
hint=None,
261+
session=None,
262+
**kwargs,
263+
):
264+
start_time = time.time()
265+
result = _original_methods["update_many"](
266+
collection_self,
267+
filter,
268+
update,
269+
upsert=upsert,
270+
array_filters=array_filters,
271+
bypass_document_validation=bypass_document_validation,
272+
collation=collation,
273+
hint=hint,
274+
session=session,
275+
**kwargs,
276+
)
277+
total_time = (time.time() - start_time) * 1000
278+
279+
__traceback_hide__ = True # noqa
280+
stack_trace, internal = _tidy_stacktrace()
281+
updates.append(
282+
{
283+
"operation": "update many",
284+
"flt": filter,
285+
"upd": update,
286+
"multi": True,
103287
"time": total_time,
104288
"stack_trace": stack_trace,
105289
"size": response_sizes[-1] if response_sizes else 0,
@@ -111,18 +295,86 @@ def _update(
111295

112296
# Wrap Cursor.remove for getting queries
113297
@functools.wraps(_original_methods["remove"])
114-
def _remove(collection_self, spec_or_id, safe=None, **kwargs):
298+
def _remove(collection_self, spec_or_id, multi=True, **kwargs):
115299
start_time = time.time()
116-
result = _original_methods["remove"](collection_self, spec_or_id, **kwargs)
300+
result = _original_methods["remove"](
301+
collection_self, spec_or_id=spec_or_id, multi=multi, **kwargs
302+
)
117303
total_time = (time.time() - start_time) * 1000
118304

119305
__traceback_hide__ = True # noqa
120306
stack_trace, internal = _tidy_stacktrace()
121307
removes.append(
122308
{
309+
"operation": "remove(deprecated)",
123310
"spec_or_id": spec_or_id,
311+
"multi": multi,
312+
"time": total_time,
313+
"stack_trace": stack_trace,
314+
"size": response_sizes[-1] if response_sizes else 0,
315+
"internal": internal,
316+
}
317+
)
318+
return result
319+
320+
321+
# Wrap Cursor.delete_one for getting queries
322+
@functools.wraps(_original_methods["delete_one"])
323+
def _delete_one(
324+
collection_self, filter, collation=None, hint=None, session=None, **kwargs
325+
):
326+
start_time = time.time()
327+
result = _original_methods["delete_one"](
328+
collection_self,
329+
filter,
330+
collation=collation,
331+
hint=hint,
332+
session=session,
333+
**kwargs,
334+
)
335+
total_time = (time.time() - start_time) * 1000
336+
337+
__traceback_hide__ = True # noqa
338+
stack_trace, internal = _tidy_stacktrace()
339+
removes.append(
340+
{
341+
"operation": "delete one",
342+
"spec_or_id": filter,
343+
"multi": False,
124344
"time": total_time,
125-
" ": stack_trace,
345+
"stack_trace": stack_trace,
346+
"size": response_sizes[-1] if response_sizes else 0,
347+
"internal": internal,
348+
}
349+
)
350+
return result
351+
352+
353+
# Wrap Cursor.delete_many for getting queries
354+
@functools.wraps(_original_methods["delete_many"])
355+
def _delete_many(
356+
collection_self, filter, collation=None, hint=None, session=None, **kwargs
357+
):
358+
start_time = time.time()
359+
result = _original_methods["delete_many"](
360+
collection_self,
361+
filter,
362+
collation=collation,
363+
hint=hint,
364+
session=session,
365+
**kwargs,
366+
)
367+
total_time = (time.time() - start_time) * 1000
368+
369+
__traceback_hide__ = True # noqa
370+
stack_trace, internal = _tidy_stacktrace()
371+
removes.append(
372+
{
373+
"operation": "delete many",
374+
"spec_or_id": filter,
375+
"multi": True,
376+
"time": total_time,
377+
"stack_trace": stack_trace,
126378
"size": response_sizes[-1] if response_sizes else 0,
127379
"internal": internal,
128380
}
@@ -215,10 +467,22 @@ def privar(name):
215467
def install_tracker():
216468
if pymongo.collection.Collection.insert != _insert:
217469
pymongo.collection.Collection.insert = _insert
470+
if pymongo.collection.Collection.insert_one != _insert_one:
471+
pymongo.collection.Collection.insert_one = _insert_one
472+
if pymongo.collection.Collection.insert_many != _insert_many:
473+
pymongo.collection.Collection.insert_many = _insert_many
218474
if pymongo.collection.Collection.update != _update:
219475
pymongo.collection.Collection.update = _update
476+
if pymongo.collection.Collection.update_one != _update_one:
477+
pymongo.collection.Collection.update_one = _update_one
478+
if pymongo.collection.Collection.update_many != _update_many:
479+
pymongo.collection.Collection.update_many = _update_many
220480
if pymongo.collection.Collection.remove != _remove:
221481
pymongo.collection.Collection.remove = _remove
482+
if pymongo.collection.Collection.delete_one != _delete_one:
483+
pymongo.collection.Collection.delete_one = _delete_one
484+
if pymongo.collection.Collection.delete_many != _delete_many:
485+
pymongo.collection.Collection.delete_many = _delete_many
222486
if pymongo.cursor.Cursor._refresh != _cursor_refresh:
223487
pymongo.cursor.Cursor._refresh = _cursor_refresh
224488
if pymongo.command_cursor.CommandCursor._unpack_response != _unpack_response:
@@ -228,10 +492,22 @@ def install_tracker():
228492
def uninstall_tracker():
229493
if pymongo.collection.Collection.insert == _insert:
230494
pymongo.collection.Collection.insert = _original_methods["insert"]
495+
if pymongo.collection.Collection.insert_one == _insert_one:
496+
pymongo.collection.Collection.insert_one = _original_methods["insert_one"]
497+
if pymongo.collection.Collection.insert_many == _insert_many:
498+
pymongo.collection.Collection.insert_many = _original_methods["insert_many"]
231499
if pymongo.collection.Collection.update == _update:
232500
pymongo.collection.Collection.update = _original_methods["update"]
501+
if pymongo.collection.Collection.update_one == _update_one:
502+
pymongo.collection.Collection.update_one = _original_methods["update_one"]
503+
if pymongo.collection.Collection.update_many == _update_many:
504+
pymongo.collection.Collection.update_many = _original_methods["update_many"]
233505
if pymongo.collection.Collection.remove == _remove:
234506
pymongo.collection.Collection.remove = _original_methods["remove"]
507+
if pymongo.collection.Collection.delete_one == _delete_one:
508+
pymongo.collection.Collection.delete_one = _original_methods["delete_one"]
509+
if pymongo.collection.Collection.delete_many == _delete_many:
510+
pymongo.collection.Collection.delete_many = _original_methods["delete_many"]
235511
if pymongo.cursor.Cursor._refresh == _cursor_refresh:
236512
pymongo.cursor.Cursor._refresh = _original_methods["cursor_refresh"]
237513
if pymongo.command_cursor.CommandCursor._unpack_response == _unpack_response:

0 commit comments

Comments
 (0)