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

rollbar/__init__.py

@@ -536,6 +536,44 @@ def wait(f=None):
         return f()
 
 
+def feature_flag(flag_key, variation=None, user=None):
+    """
+    A context manager interface that creates a list of tags to represent 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.feature_flag('flag1', variation=True, user='[email protected]'):
+            code()
+
+    Tags generated from the above example:
+
+    [
+        {'key': feature_flag.key', 'value': 'flag1'},
+        {'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)
+
+    tags = [flag_key_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 +740,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
 
+    tags = _in_context_tags.to_list()
+
+    # if there are tags attached to the exception, reverse the order, and append to `tags`
+    # to send the full chain of tags to Rollbar.
+    if hasattr(exc_info[1], '_rollbar_tags'):
+        tags += getattr(exc_info[1], '_rollbar_tags').to_list(reverse=True)
+
+    if tags:
+        data['tags'] = tags
+
     request = _get_actual_request(request)
     _add_request_data(data, request)
     _add_person_data(data, request)
@@ -788,6 +836,9 @@ def _report_message(message, level, request, extra_data, payload_data):
     _add_lambda_context_data(data)
     data['server'] = _build_server_data()
 
+    if _in_context_tags.to_list():
+        data['tags'] = _in_context_tags.to_list()
+
     if payload_data:
         data = dict_merge(data, payload_data, silence_errors=True)
 
@@ -1606,3 +1657,66 @@ 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()
+
+    def append(self, value):
+        if not hasattr(self._registry, 'tags'):
+            self._registry.tags = []
+
+        self._registry.tags.append(value)
+
+    def pop(self):
+        self._registry.tags.pop()
+
+    def to_list(self, reverse=False):
+        if not hasattr(self._registry, 'tags'):
+            self._registry.tags = []
+
+        if reverse:
+            return _flatten_nested_lists(self._registry.tags[::-1])
+
+        return _flatten_nested_lists(self._registry.tags)
+
+
+_in_context_tags = _LocalTags()
+
+
+class _TagManager(object):
+    """
+    Context manager object that interfaces with the `_in_context_tags` stack:
+
+        On enter, puts the tags at top of the stack.
+        On exit, pops off the top element of the stack.
+          - If there is an exception, attach the tags to the exception
+            for rebuilding the tags in the context before reporting.
+    """
+    def __init__(self, tags):
+        self.tags = tags
+
+    def __enter__(self):
+        _in_context_tags.append(self.tags)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+
+        if exc_value:
+            if not hasattr(exc_value, '_rollbar_tags'):
+                exc_value._rollbar_tags = _LocalTags()
+
+            exc_value._rollbar_tags.append(self.tags)
+
+        _in_context_tags.pop()

rollbar/test/test_feature_flag_context_manager.py

@@ -0,0 +1,129 @@
+import copy
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+import rollbar
+
+from rollbar.test import BaseTest
+
+
+_test_access_token = 'aaaabbbbccccddddeeeeffff00001111'
+_default_settings = copy.deepcopy(rollbar.SETTINGS)
+
+
+class FeatureFlagContextManagerTest(BaseTest):
+    def setUp(self):
+        rollbar._initialized = False
+        rollbar.SETTINGS = copy.deepcopy(_default_settings)
+        rollbar.init(_test_access_token, locals={'enabled': True}, handler='blocking', timeout=12345)
+
+    def test_feature_flag_generates_correct_tag_payload(self):
+        cm = rollbar.feature_flag('feature-foo', variation=True, user='[email protected]')
+
+        tags = cm.tags
+        self.assertEqual(len(tags), 3)
+
+        key, variation, user = tags
+
+        self.assertEqual(key['key'], 'feature_flag.key')
+        self.assertEqual(key['value'], 'feature-foo')
+
+        self.assertEqual(variation['key'], 'feature_flag.data.feature-foo.variation')
+        self.assertEqual(variation['value'], True)
+
+        self.assertEqual(user['key'], 'feature_flag.data.feature-foo.user')
+        self.assertEqual(user['value'], '[email protected]')
+
+    @mock.patch('rollbar.send_payload')
+    def test_report_message_inside_feature_flag_context_manager(self, send_payload):
+        with rollbar.feature_flag('feature-foo'):
+            rollbar.report_message('hello world')
+
+        self.assertEqual(send_payload.called, True)
+
+        # [0][0] is used here to index into the mocked objects `call_args` and get the
+        # right payload for comparison.
+        payload_data = send_payload.call_args[0][0]['data']
+        self.assertIn('tags', payload_data)
+
+        tags = payload_data['tags']
+        self.assertEquals(len(tags), 1)
+
+        self._assert_tag_equals(tags[0], 'feature_flag.key', 'feature-foo')
+
+        self._report_message_and_assert_no_tags(send_payload)
+
+    @mock.patch('rollbar.send_payload')
+    def test_report_exc_info_inside_feature_flag_context_manager(self, send_payload):
+        with rollbar.feature_flag('feature-foo'):
+            try:
+                raise Exception('foo')
+            except:
+                rollbar.report_exc_info()
+
+        self.assertEqual(send_payload.called, True)
+
+        payload_data = send_payload.call_args[0][0]['data']
+        self.assertIn('tags', payload_data)
+
+        tags = payload_data['tags']
+        self.assertEquals(len(tags), 1)
+
+        self._assert_tag_equals(tags[0], 'feature_flag.key', 'feature-foo')
+
+        self._report_message_and_assert_no_tags(send_payload)
+
+    @mock.patch('rollbar.send_payload')
+    def test_report_exc_info_outside_feature_flag_context_manager(self, send_payload):
+        try:
+            with rollbar.feature_flag('feature-foo'):
+                raise Exception('foo')
+        except:
+            rollbar.report_exc_info()
+
+        self.assertEqual(send_payload.called, True)
+
+        payload_data = send_payload.call_args[0][0]['data']
+        self.assertIn('tags', payload_data)
+
+        tags = payload_data['tags']
+        self.assertEquals(len(tags), 1)
+
+        self._assert_tag_equals(tags[0], 'feature_flag.key', 'feature-foo')
+
+        self._report_message_and_assert_no_tags(send_payload)
+
+    @mock.patch('rollbar.send_payload')
+    def test_report_exc_info_inside_nested_feature_flag_context_manager(self, send_payload):
+        try:
+            with rollbar.feature_flag('feature-foo'):
+                with rollbar.feature_flag('feature-bar'):
+                    raise Exception('foo')
+        except:
+            rollbar.report_exc_info()
+
+        self.assertEqual(send_payload.called, True)
+
+        payload_data = send_payload.call_args[0][0]['data']
+        self.assertIn('tags', payload_data)
+
+        tags = payload_data['tags']
+        self.assertEquals(len(tags), 2)
+
+        self._assert_tag_equals(tags[0], 'feature_flag.key', 'feature-foo')
+        self._assert_tag_equals(tags[1], 'feature_flag.key', 'feature-bar')
+
+        self._report_message_and_assert_no_tags(send_payload)
+
+    def _assert_tag_equals(self, tag, key, value):
+        self.assertEqual(tag, {'key': key, 'value': value})
+
+    def _report_message_and_assert_no_tags(self, mocked_send):
+        rollbar.report_message('this report message is to check that there are no tags')
+        self.assertEqual(mocked_send.called, True)
+
+        payload_data = mocked_send.call_args[0][0]['data']
+        self.assertNotIn('tags', payload_data)