diff --git a/rollbar/__init__.py b/rollbar/__init__.py index c40f4b78..c821f99b 100644 --- a/rollbar/__init__.py +++ b/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='foobar@rollbar.com'): + 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': 'foobar@rollbar.com'} + ] + """ + 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() diff --git a/rollbar/test/test_feature_flag_context_manager.py b/rollbar/test/test_feature_flag_context_manager.py new file mode 100644 index 00000000..b549b29d --- /dev/null +++ b/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='atran@rollbar.com') + + 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'], 'atran@rollbar.com') + + @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)