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

rollbar/__init__.py

@@ -536,6 +536,49 @@ def wait(f=None):
         return f()
 
 
+def feature_flag(flag_key, variation=None, user=None):
+    """
+    A context manager interface that sets the tags used to model a feature flag.
+
+    key: the key of the feature flag.
+    variation: (optional) the evaluated feature flag variation.
+    user: (optional) the user being evaluated.
+
+    Example usage:
+
+        with rollbar.featureflag('flag1', variation=True, user='[email protected]'):
+            code()
+
+    Tags generated from the above example:
+
+    [
+        {'key': feature_flag.key', 'value': 'flag1'},
+        {'key': feature_flag.data.flag1.order', 'value': 0},
+        {'key': feature_flag.data.flag1.variation', 'value': True},
+        {'key': feature_flag.data.flag1.user, 'value': '[email protected]'}
+    ]
+    """
+    flag_key_tag = _create_tag('feature_flag.key', flag_key)
+
+    # create the feature flag order tag
+    order_key = _feature_flag_data_key(flag_key, 'order')
+    order_tag = _create_tag(order_key, len(_tags.value))
+
+    tags = [flag_key_tag, order_tag]
+
+    if variation is not None:
+        variation_key = _feature_flag_data_key(flag_key, 'variation')
+        variation_tag = _create_tag(variation_key, variation)
+        tags.append(variation_tag)
+
+    if user is not None:
+        user_key = _feature_flag_data_key(flag_key, 'user')
+        user_tag = _create_tag(user_key, user)
+        tags.append(user_tag)
+
+    return _TagManager(tags)
+
+
 class ApiException(Exception):
     """
     This exception will be raised if there was a problem decoding the
@@ -702,6 +745,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'] = _flatten_nested_lists(tags)
+
     request = _get_actual_request(request)
     _add_request_data(data, request)
     _add_person_data(data, request)
@@ -788,6 +837,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'] = _flatten_nested_lists(_tags.value)
+
     if payload_data:
         data = dict_merge(data, payload_data, silence_errors=True)
 
@@ -1606,3 +1658,62 @@ def _wsgi_extract_user_ip(environ):
     if real_ip:
         return real_ip
     return environ['REMOTE_ADDR']
+
+
+def _create_tag(key, value):
+    return {'key': key, 'value': value}
+
+
+def _feature_flag_data_key(flag_key, attribute):
+    return 'feature_flag.data.%s.%s' % (flag_key, attribute)
+
+
+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()