Completed
Pull Request — master (#2258)
by Manas
09:20 queued 03:35
created

st2common.services.get_trace_component_for_rule()   A

Complexity

Conditions 2

Size

Total Lines 21

Duplication

Lines 0
Ratio 0 %
Metric Value
dl 0
loc 21
rs 9.3143
cc 2
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
from mongoengine import ValidationError
17
18
from st2common import log as logging
19
from st2common.constants.triggers import ACTION_SENSOR_TRIGGER, NOTIFY_TRIGGER
20
from st2common.constants.trace import TRACE_CONTEXT
21
from st2common.exceptions.db import StackStormDBObjectNotFoundError
22
from st2common.exceptions.trace import UniqueTraceNotFoundException
23
from st2common.models.api.trace import TraceContext
24
from st2common.models.db.trace import TraceDB, TraceComponentDB
25
from st2common.models.system.common import ResourceReference
26
from st2common.persistence.execution import ActionExecution
27
from st2common.persistence.trace import Trace
28
29
LOG = logging.getLogger(__name__)
30
31
__all__ = [
32
    'get_trace_db_by_action_execution',
33
    'get_trace_db_by_rule',
34
    'get_trace_db_by_trigger_instance',
35
    'get_trace',
36
    'add_or_update_given_trace_context',
37
    'add_or_update_given_trace_db',
38
    'get_trace_component_for_action_execution',
39
    'get_trace_component_for_rule',
40
    'get_trace_component_for_trigger_instance'
41
]
42
43
44
ACTION_SENSOR_TRIGGER_REF = ResourceReference.to_string_reference(
45
    pack=ACTION_SENSOR_TRIGGER['pack'], name=ACTION_SENSOR_TRIGGER['name'])
46
NOTIFY_TRIGGER_REF = ResourceReference.to_string_reference(
47
    pack=NOTIFY_TRIGGER['pack'], name=NOTIFY_TRIGGER['name'])
48
49
50
def _get_valid_trace_context(trace_context):
51
    """
52
    Check if tarce_context is a valid type and returns a TraceContext object.
53
    """
54
    assert isinstance(trace_context, (TraceContext, dict))
55
56
    # Pretty much abuse the dynamic nature of python to make it possible to support
57
    # both dict and TraceContext types.
58
    if isinstance(trace_context, dict):
59
        trace_context = TraceContext(**trace_context)
60
61
    return trace_context
62
63
64
def _get_single_trace_by_component(**component_filter):
65
    """
66
    Tries to return a single Trace mathing component_filter. Raises an exception
67
    when a filter matches multiple.
68
    """
69
    traces = Trace.query(**component_filter)
70
    if len(traces) == 0:
71
        return None
72
    elif len(traces) > 1:
73
        raise UniqueTraceNotFoundException(
74
            'More than 1 trace matching %s found.' % component_filter)
75
    return traces[0]
76
77
78
def get_trace_db_by_action_execution(action_execution=None, action_execution_id=None):
79
    if action_execution:
80
        action_execution_id = str(action_execution.id)
81
    return _get_single_trace_by_component(action_executions__object_id=action_execution_id)
82
83
84
def get_trace_db_by_rule(rule=None, rule_id=None):
85
    if rule:
86
        rule_id = str(rule.id)
87
    # by rule could return multiple traces
88
    return Trace.query(rules__object_id=rule_id)
89
90
91
def get_trace_db_by_trigger_instance(trigger_instance=None, trigger_instance_id=None):
92
    if trigger_instance:
93
        trigger_instance_id = str(trigger_instance.id)
94
    return _get_single_trace_by_component(trigger_instances__object_id=trigger_instance_id)
95
96
97
def get_trace(trace_context, ignore_trace_tag=False):
98
    """
99
    :param trace_context: context object using which a trace can be found.
100
    :type trace_context: ``dict`` or ``TraceContext``
101
102
    :param ignore_trace_tag: Even if a trace_tag is provided will be ignored.
103
    :type ignore_trace_tag: ``str``
104
105
    :rtype: ``TraceDB``
106
    """
107
108
    trace_context = _get_valid_trace_context(trace_context)
109
110
    if not trace_context.id_ and not trace_context.trace_tag:
111
        raise ValueError('Atleast one of id_ or trace_tag should be specified.')
112
113
    if trace_context.id_:
114
        try:
115
            return Trace.get_by_id(trace_context.id_)
116
        except (ValidationError, ValueError):
117
            LOG.warning('Database lookup for Trace with id="%s" failed.',
118
                        trace_context.id_, exc_info=True)
119
            raise StackStormDBObjectNotFoundError(
120
                'Unable to find Trace with id="%s"' % trace_context.id_)
121
122
    if ignore_trace_tag:
123
        return None
124
125
    traces = Trace.query(trace_tag=trace_context.trace_tag)
126
127
    # Assume this method only handles 1 trace.
128
    if len(traces) > 1:
129
        raise UniqueTraceNotFoundException(
130
            'More than 1 Trace matching %s found.' % trace_context.trace_tag)
131
132
    return traces[0]
133
134
135
def get_trace_db_by_live_action(liveaction):
136
    """
137
    Given a liveaction does the best attempt to return a TraceDB.
138
    1. From trace_context in liveaction.context
139
    2. From parent in liveaction.context
140
    3. From action_execution associated with provided liveaction
141
    4. Creates a new TraceDB (which calling method is on the hook to persist).
142
143
    :param liveaction: liveaction from which to figure out a TraceDB.
144
    :type liveaction: ``LiveActionDB``
145
146
    :returns: (boolean, TraceDB) if the TraceDB was created(but not saved to DB) or
147
               retrieved from the DB and the TraceDB itself.
148
    :rtype: ``tuple``
149
    """
150
    trace_db = None
151
    created = False
152
    # 1. Try to get trace_db from liveaction context.
153
    #    via trigger_instance + rule or via user specified trace_context
154
    trace_context = liveaction.context.get(TRACE_CONTEXT, None)
155
    if trace_context:
156
        trace_context = _get_valid_trace_context(trace_context)
157
        trace_db = get_trace(trace_context=trace_context, ignore_trace_tag=True)
158
        # found a trace_context but no trace_db. This implies a user supplied
159
        # trace_tag so create a new trace_db
160
        if not trace_db:
161
            trace_db = TraceDB(trace_tag=trace_context.trace_tag)
162
            created = True
163
        return (created, trace_db)
164
    # 2. If not found then check if parent context contains an execution_id.
165
    #    This cover case for child execution of a workflow.
166
    if not trace_context and 'parent' in liveaction.context:
167
        parent_execution_id = liveaction.context['parent'].get('execution_id', None)
168
        if parent_execution_id:
169
            # go straight to a trace_db. If there is a parent execution then that must
170
            # be associated with a Trace.
171
            trace_db = get_trace_db_by_action_execution(action_execution_id=parent_execution_id)
172
            if not trace_db:
173
                raise StackStormDBObjectNotFoundError('No trace found for execution %s' %
174
                                                      parent_execution_id)
175
            return (created, trace_db)
176
    # 3. Check if the action_execution associated with liveaction leads to a trace_db
177
    execution = ActionExecution.get(liveaction__id=str(liveaction.id))
178
    if execution:
179
        trace_db = get_trace_db_by_action_execution(action_execution=execution)
180
    # 4. No trace_db found, therefore create one. This typically happens
181
    #    when execution is run by hand.
182
    if not trace_db:
183
        trace_db = TraceDB(trace_tag='execution-%s' % str(liveaction.id))
184
        created = True
185
    return (created, trace_db)
186
187
188
def add_or_update_given_trace_context(trace_context, action_executions=None, rules=None,
189
                                      trigger_instances=None):
190
    """
191
    Will update an existing Trace or add a new Trace. This method will only look for exact
192
    Trace as identified by the trace_context. Even if the trace_context contain a trace_tag
193
    it shall not be used to lookup a Trace.
194
195
    * If an exact matching Trace is not found a new Trace is created
196
    * Whenever only a trace_tag is supplied a new Trace is created.
197
198
    :param trace_context: context object using which a trace can be found. If not found
199
                          trace_context.trace_tag is used to start new trace.
200
    :type trace_context: ``dict`` or ``TraceContext``
201
202
    :param action_executions: The action_execution to be added to the Trace. Should a list
203
                              of object_ids or a dict containing object_ids and caused_by.
204
    :type action_executions: ``list``
205
206
    :param rules: The rules to be added to the Trace.  Should a list of object_ids or a dict
207
                  containing object_ids and caused_by.
208
    :type rules: ``list``
209
210
    :param trigger_instances: The trigger_instances to be added to the Trace. Should a list
211
                              of object_ids or a dict containing object_ids and caused_by.
212
    :type trigger_instances: ``list``
213
214
    :rtype: ``TraceDB``
215
    """
216
    trace_db = get_trace(trace_context=trace_context, ignore_trace_tag=True)
217
    if not trace_db:
218
        # since trace_db is None need to end up with a valid trace_context
219
        trace_context = _get_valid_trace_context(trace_context)
220
        trace_db = TraceDB(trace_tag=trace_context.trace_tag)
221
    return add_or_update_given_trace_db(trace_db=trace_db,
222
                                        action_executions=action_executions,
223
                                        rules=rules,
224
                                        trigger_instances=trigger_instances)
225
226
227
def add_or_update_given_trace_db(trace_db, action_executions=None, rules=None,
228
                                 trigger_instances=None):
229
    """
230
    Will update an existing Trace.
231
232
    :param trace_db: The TraceDB to update.
233
    :type trace_db: ``TraceDB``
234
235
    :param action_executions: The action_execution to be added to the Trace. Should a list
236
                              of object_ids or a dict containing object_ids and caused_by.
237
    :type action_executions: ``list``
238
239
    :param rules: The rules to be added to the Trace. Should a list of object_ids or a dict
240
                  containing object_ids and caused_by.
241
    :type rules: ``list``
242
243
    :param trigger_instances: The trigger_instances to be added to the Trace. Should a list
244
                              of object_ids or a dict containing object_ids and caused_by.
245
    :type trigger_instances: ``list``
246
247
    :rtype: ``TraceDB``
248
    """
249
    if trace_db is None:
250
        raise ValueError('trace_db should be non-None.')
251
252
    if not action_executions:
253
        action_executions = []
254
    action_executions = [_to_trace_component_db(component=action_execution)
255
                         for action_execution in action_executions]
256
257
    if not rules:
258
        rules = []
259
    rules = [_to_trace_component_db(component=rule) for rule in rules]
260
261
    if not trigger_instances:
262
        trigger_instances = []
263
    trigger_instances = [_to_trace_component_db(component=trigger_instance)
264
                         for trigger_instance in trigger_instances]
265
266
    # If an id exists then this is an update and we do not want to perform
267
    # an upsert so use push_components which will use the push operator.
268
    if trace_db.id:
269
        return Trace.push_components(trace_db,
270
                                     action_executions=action_executions,
271
                                     rules=rules,
272
                                     trigger_instances=trigger_instances)
273
274
    trace_db.action_executions = action_executions
275
    trace_db.rules = rules
276
    trace_db.trigger_instances = trigger_instances
277
278
    return Trace.add_or_update(trace_db)
279
280
281
def get_trace_component_for_action_execution(action_execution_db):
282
    """
283
    Returns the trace_component compatible dict representation of an actionexecution.
284
285
    :param action_execution_db: ActionExecution to translate
286
    :type action_execution_db: ActionExecutionDB
287
288
    :rtype: ``dict``
289
    """
290
    if not action_execution_db:
291
        raise ValueError('action_execution_db expected.')
292
    trace_component = {'id': str(action_execution_db.id)}
293
    caused_by = {}
294
    if action_execution_db.rule and action_execution_db.trigger_instance:
295
        # Once RuleEnforcement is available that can be used instead.
296
        caused_by['type'] = 'rule'
297
        caused_by['id'] = '%s:%s' % (action_execution_db.rule['id'],
298
                                     action_execution_db.trigger_instance['id'])
299
    trace_component['caused_by'] = caused_by
300
    return trace_component
301
302
303
def get_trace_component_for_rule(rule_db, trigger_instance_db):
304
    """
305
    Returns the trace_component compatible dict representation of a rule.
306
307
    :param rule_db: The rule to translate
308
    :type rule_db: RuleDB
309
310
    :param trigger_instance_db: The TriggerInstance with causal relation to rule_db
311
    :type trigger_instance_db: TriggerInstanceDB
312
313
    :rtype: ``dict``
314
    """
315
    trace_component = {}
316
    trace_component = {'id': str(rule_db.id)}
317
    caused_by = {}
318
    if trigger_instance_db:
319
        # Once RuleEnforcement is available that can be used instead.
320
        caused_by['type'] = 'trigger_instance'
321
        caused_by['id'] = str(trigger_instance_db.id)
322
    trace_component['caused_by'] = caused_by
323
    return trace_component
324
325
326
def get_trace_component_for_trigger_instance(trigger_instance_db):
327
    """
328
    Returns the trace_component compatible dict representation of a triggerinstance.
329
330
    :param trigger_instance_db: The TriggerInstance to translate
331
    :type trigger_instance_db: TriggerInstanceDB
332
333
    :rtype: ``dict``
334
    """
335
    trace_component = {}
336
    trace_component = {'id': str(trigger_instance_db.id)}
337
    caused_by = {}
338
    # Special handling for ACTION_SENSOR_TRIGGER and NOTIFY_TRIGGER where we
339
    # know how to maintain the links.
340
    if trigger_instance_db.trigger == ACTION_SENSOR_TRIGGER_REF or \
341
       trigger_instance_db.trigger == NOTIFY_TRIGGER_REF:
342
        # Once RuleEnforcement is available that can be used instead.
343
        caused_by['type'] = 'action_execution'
344
        # For both action trigger and notidy trigger execution_id is stored in the payload.
345
        caused_by['id'] = trigger_instance_db.payload['execution_id']
346
    trace_component['caused_by'] = caused_by
347
    return trace_component
348
349
350
def _to_trace_component_db(component):
351
    """
352
    Take the component as string or a dict and will construct a TraceComponentDB.
353
354
    :param component: Should identify the component. If a string should be id of the
355
                      component. If a dict should contain id and the caused_by.
356
    :type component: ``bson.ObjectId`` or ``dict``
357
358
    :rtype: ``TraceComponentDB``
359
    """
360
    if not isinstance(component, (basestring, dict)):
361
        print type(component)
362
        raise ValueError('Expected component to be str or dict')
363
364
    object_id = component if isinstance(component, basestring) else component['id']
365
    caused_by = component.get('caused_by', {}) if isinstance(component, dict) else {}
366
367
    return TraceComponentDB(object_id=object_id, caused_by=caused_by)
368