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

rollbar/__init__.py

@@ -536,6 +536,23 @@ def wait(f=None):
         return f()
 
 
+def watch_feature_flag(key):
+    """
+    Sets the feature flag payload for the code wrapped by this method.
+
+    key: key of the feature flag.
+
+    Usage:
+
+        with rollbar.watch_feature_flag('feature_flag_a'):
+            do_something_risky()
+    """
+    feature_flag_obj = {'key': key}
+
+    return _FeatureFlagManager(feature_flag_obj)
+
+
+
 class ApiException(Exception):
     """
     This exception will be raised if there was a problem decoding the
@@ -702,6 +719,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 feature flags attached to the exception, append that to the feature flags on
+    # singleton to create the full feature flags stack.
+    feature_flags = _feature_flags + getattr(exc_info[1], '_rollbar_feature_flags', [])[::-1]
+    if feature_flags:
+        data['feature_flags'] = feature_flags
+
     request = _get_actual_request(request)
     _add_request_data(data, request)
     _add_person_data(data, request)
@@ -788,6 +811,9 @@ def _report_message(message, level, request, extra_data, payload_data):
     _add_lambda_context_data(data)
     data['server'] = _build_server_data()
 
+    if _feature_flags:
+        data['feature_flags'] = _feature_flags
+
     if payload_data:
         data = dict_merge(data, payload_data, silence_errors=True)
 
@@ -1606,3 +1632,32 @@ def _wsgi_extract_user_ip(environ):
     if real_ip:
         return real_ip
     return environ['REMOTE_ADDR']
+
+
+_feature_flags = []
+
+
+class _FeatureFlagManager(object):
+    """
+    Context manager object that interfaces with the `_feature_flags` stack:
+
+        On enter, puts the feature flag object at top of the stack.
+        On exit, pops off the top element of the stack.
+          - If there is an exception, attach the feature flag object to the exception
+            for rebuilding of the `_feature_flags` stack before reporting.
+    """
+    def __init__(self, feature_flag_obj):
+        self.feature_flag = feature_flag_obj
+
+    def __enter__(self):
+        _feature_flags.append(self.feature_flag)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+
+        if exc_value:
+            if not hasattr(exc_value, '_rollbar_feature_flags'):
+                exc_value._rollbar_feature_flags = []
+
+            exc_value._rollbar_feature_flags.append(self.feature_flag)
+
+        _feature_flags.pop()