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

rollbar/__init__.py

@@ -536,6 +536,21 @@ def wait(f=None):
         return f()
 
 
+def watch(key, value):
+    """
+    Sets a tag for the code wrapped by this method.
+
+    key: key to look up the tag.
+    value: value associated with tag.
+
+    Usage:
+
+        with rollbar.watch('feature_flag', 'feature_flag_a'):
+            do_something_risky()
+    """
+    return _TagManager({'key': key, 'value': value})
+
+
 class ApiException(Exception):
     """
     This exception will be raised if there was a problem decoding the
@@ -702,6 +717,12 @@ 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
 
+    # if there are tags attached to the exception, reverse the order, and append to `_tags`
+    # stack to send the full chain of tags to Rollbar.
+    tags = _tags.value + getattr(exc_info[1], '_rollbar_tags', [])[::-1]
+    if tags:
+        data['tags'] = tags
+
     request = _get_actual_request(request)
     _add_request_data(data, request)
     _add_person_data(data, request)
@@ -788,6 +809,9 @@ def _report_message(message, level, request, extra_data, payload_data):
     _add_lambda_context_data(data)
     data['server'] = _build_server_data()
 
+    if _tags.value:
+        data['tags'] = _tags.value
+
     if payload_data:
         data = dict_merge(data, payload_data, silence_errors=True)
 
@@ -1606,3 +1630,54 @@ def _wsgi_extract_user_ip(environ):
     if real_ip:
         return real_ip
     return environ['REMOTE_ADDR']
+
+
+class _LocalTags(object):
+    """
+    An object to ensure thread safety.
+    """
+    def __init__(self):
+        self._registry = threading.local()
+        self._registry.tags = []
+
+    def append(self, value):
+        self._registry.tags.append(value)
+
+    def pop(self):
+        self._registry.tags.pop()
+
+    @property
+    def value(self):
+        if hasattr(self._registry, 'tags'):
+            return self._registry.tags
+
+        return []
+
+
+_tags = _LocalTags()
+
+
+class _TagManager(object):
+    """
+    Context manager object that interfaces with the `_tags` stack:
+
+        On enter, puts the tag object at top of the stack.
+        On exit, pops off the top element of the stack.
+          - If there is an exception, attach the tag object to the exception
+            for rebuilding of the `_tags` stack before reporting.
+    """
+    def __init__(self, tag):
+        self.tag = tag
+
+    def __enter__(self):
+        _tags.append(self.tag)
+
+    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()