Add a context manager object and a method to call said object

rollbar/__init__.py

@@ -536,6 +536,22 @@ def wait(f=None):
         return f()
 
 
+def watch(tag_name, extra_data=None):
+    """
+    Sets the tag for the code wrapped by this method.
+
+    tag_name: name of the tag.
+    extra_data: optional, dictionary of params will be included in the 'tag' object.
+                'name' is reserved for the tag_name.
+
+    Usage:
+
+        with Rollbar.watch('foobar'):
+            do_something_risky()
+    """
+    return _TagManager(tag_name, extra_data)
+
+
 class ApiException(Exception):
     """
     This exception will be raised if there was a problem decoding the
@@ -702,6 +718,16 @@ def _report_exc_info(exc_info, request, extra_data, payload_data, level=None):
     if extra_trace_data and not extra_data:
         data['custom'] = extra_trace_data
 
+    attached_tags = getattr(exc_info[1], '_rollbar_tags', None)
+    if attached_tags:
+        # if there are tags attached to the exception, take the first one as that's
+        # the deepest tag where the exception occurred
+        data['tag'] = attached_tags[0]
+    elif tags:
+        # if there are no tags attached to the exception and there are `tags`, that
+        # means we haven't exited yet - in this case, take the top tag of the stack.
+        data['tag'] = tags[-1]
+
     request = _get_actual_request(request)
     _add_request_data(data, request)
     _add_person_data(data, request)
@@ -787,6 +813,9 @@ def _report_message(message, level, request, extra_data, payload_data):
     _add_person_data(data, request)
     _add_lambda_context_data(data)
     data['server'] = _build_server_data()
+
+    if tags:
+        data['tag'] = tags[-1]
 
     if payload_data:
         data = dict_merge(data, payload_data, silence_errors=True)
@@ -1606,3 +1635,39 @@ def _wsgi_extract_user_ip(environ):
     if real_ip:
         return real_ip
     return environ['REMOTE_ADDR']
+
+
+tags = []
+
+
+class _TagManager(object):
+    """
+    Context manager object that interfaces with the `tags` stack:
+
+        On enter, puts current tag object at top of the stack.
+        On exit, pops off top element of the stack.
+          - If there is an exception, append the tag to exception._rollbar_tags.
+    """
+    def __init__(self, name, extra_data):
+        self.tag = {}
+
+        if extra_data:
+            if not isinstance(extra_data, dict):
+                raise TypeError('expected \'extra_data\' to be a dictonary')
+
+            self.tag.update(extra_data)
+
+        self.tag['name'] = name
+
+    def __enter__(self):
+        tags.append(self.tag)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_value:
+            if not hasattr(exc_value, '_rollbar_tags'):
+                exc_value._rollbar_tags = []
+
+            exc_value._rollbar_tags.append(self.tag)
+
+        tags.pop()