1 | # Licensed to the StackStorm, Inc ('StackStorm') under one or more |
||
2 | # contributor license agreements. See the NOTICE file distributed with |
||
3 | # this work for additional information regarding copyright ownership. |
||
4 | # The ASF licenses this file to You under the Apache License, Version 2.0 |
||
5 | # (the "License"); you may not use this file except in compliance with |
||
6 | # the License. You may obtain a copy of the License at |
||
7 | # |
||
8 | # http://www.apache.org/licenses/LICENSE-2.0 |
||
9 | # |
||
10 | # Unless required by applicable law or agreed to in writing, software |
||
11 | # distributed under the License is distributed on an "AS IS" BASIS, |
||
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||
13 | # See the License for the specific language governing permissions and |
||
14 | # limitations under the License. |
||
15 | |||
16 | """ |
||
17 | Module containing MongoDB profiling related functionality. |
||
18 | """ |
||
19 | |||
20 | from __future__ import absolute_import |
||
21 | from mongoengine.queryset import QuerySet |
||
22 | |||
23 | from st2common import log as logging |
||
24 | |||
25 | __all__ = [ |
||
26 | 'enable_profiling', |
||
27 | 'disable_profiling', |
||
28 | 'is_enabled', |
||
29 | 'log_query_and_profile_data_for_queryset' |
||
30 | ] |
||
31 | |||
32 | LOG = logging.getLogger(__name__) |
||
33 | |||
34 | ENABLE_PROFILING = False |
||
35 | |||
36 | |||
37 | def enable_profiling(): |
||
38 | global ENABLE_PROFILING |
||
0 ignored issues
–
show
|
|||
39 | ENABLE_PROFILING = True |
||
40 | return ENABLE_PROFILING |
||
41 | |||
42 | |||
43 | def disable_profiling(): |
||
44 | global ENABLE_PROFILING |
||
0 ignored issues
–
show
|
|||
45 | ENABLE_PROFILING = False |
||
46 | return ENABLE_PROFILING |
||
47 | |||
48 | |||
49 | def is_enabled(): |
||
50 | global ENABLE_PROFILING |
||
0 ignored issues
–
show
|
|||
51 | return ENABLE_PROFILING |
||
52 | |||
53 | |||
54 | def log_query_and_profile_data_for_queryset(queryset): |
||
55 | """ |
||
56 | Function which logs MongoDB query and profile data for the provided mongoengine queryset object. |
||
57 | |||
58 | Keep in mind that this method needs to be called at the very end after all the mongoengine |
||
59 | methods are chained. |
||
60 | |||
61 | For example: |
||
62 | |||
63 | result = model.object(...) |
||
64 | result = model.limit(...) |
||
65 | result = model.order_by(...) |
||
66 | """ |
||
67 | if not ENABLE_PROFILING: |
||
68 | # Profiling is disabled |
||
69 | return queryset |
||
70 | |||
71 | if not isinstance(queryset, QuerySet): |
||
72 | # Note: Some mongoengine methods don't return queryset (e.g. count) |
||
73 | return queryset |
||
74 | |||
75 | query = getattr(queryset, '_query', None) |
||
76 | mongo_query = getattr(queryset, '_mongo_query', query) |
||
77 | ordering = getattr(queryset, '_ordering', None) |
||
78 | limit = getattr(queryset, '_limit', None) |
||
79 | collection = getattr(queryset, '_collection', None) |
||
80 | collection_name = getattr(collection, 'name', None) |
||
81 | only_fields = getattr(queryset, 'only_fields', None) |
||
82 | |||
83 | # Note: We need to clone the queryset when using explain because explain advances the cursor |
||
84 | # internally which changes the function result |
||
85 | cloned_queryset = queryset.clone() |
||
86 | explain_info = cloned_queryset.explain(format=True) |
||
87 | |||
88 | if mongo_query is not None and collection_name is not None: |
||
89 | mongo_shell_query = construct_mongo_shell_query(mongo_query=mongo_query, |
||
90 | collection_name=collection_name, |
||
91 | ordering=ordering, |
||
92 | limit=limit, |
||
93 | only_fields=only_fields) |
||
94 | extra = {'mongo_query': mongo_query, 'mongo_shell_query': mongo_shell_query} |
||
95 | LOG.debug('MongoDB query: %s' % (mongo_shell_query), extra=extra) |
||
0 ignored issues
–
show
|
|||
96 | LOG.debug('MongoDB explain data: %s' % (explain_info)) |
||
0 ignored issues
–
show
|
|||
97 | |||
98 | return queryset |
||
99 | |||
100 | |||
101 | def construct_mongo_shell_query(mongo_query, collection_name, ordering, limit, |
||
102 | only_fields=None): |
||
103 | result = [] |
||
104 | |||
105 | # Select collection |
||
106 | part = 'db.{collection}'.format(collection=collection_name) |
||
107 | result.append(part) |
||
108 | |||
109 | # Include filters (if any) |
||
110 | if mongo_query: |
||
111 | filter_predicate = mongo_query |
||
112 | else: |
||
113 | filter_predicate = '' |
||
114 | |||
115 | part = 'find({filter_predicate})'.format(filter_predicate=filter_predicate) |
||
116 | |||
117 | # Include only fields (projection) |
||
118 | if only_fields: |
||
119 | projection_items = ['\'%s\': 1' % (field) for field in only_fields] |
||
120 | projection = ', '.join(projection_items) |
||
121 | part = 'find({filter_predicate}, {{{projection}}})'.format( |
||
122 | filter_predicate=filter_predicate, projection=projection) |
||
123 | else: |
||
124 | part = 'find({filter_predicate})'.format(filter_predicate=filter_predicate) |
||
125 | |||
126 | result.append(part) |
||
127 | |||
128 | # Include ordering info (if any) |
||
129 | if ordering: |
||
130 | sort_predicate = [] |
||
131 | for field_name, direction in ordering: |
||
132 | sort_predicate.append('{name}: {direction}'.format(name=field_name, |
||
133 | direction=direction)) |
||
134 | |||
135 | sort_predicate = ', '.join(sort_predicate) |
||
136 | part = 'sort({{{sort_predicate}}})'.format(sort_predicate=sort_predicate) |
||
137 | result.append(part) |
||
138 | |||
139 | # Include limit info (if any) |
||
140 | if limit is not None: |
||
141 | part = 'limit({limit})'.format(limit=limit) |
||
142 | result.append(part) |
||
143 | |||
144 | result = '.'.join(result) + ';' |
||
145 | return result |
||
146 |
Usage of
global
can make code hard to read and test, its usage is generally not recommended unless you are dealing with legacy code.