Added support for custom payload transforms

rollbar/__init__.py

@@ -323,6 +323,7 @@ def _get_fastapi_request():
     'request_pool_maxsize': None,
     'request_max_retries': None,
     'batch_transforms': False,
+    'custom_transforms': [],
 }
 
 _CURRENT_LAMBDA_CONTEXT = None
@@ -426,12 +427,23 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
                                    keys=shortener_keys,
                                    **SETTINGS['locals']['sizes'])
     _transforms = [
-        shortener,
-        ScrubRedactTransform(),
-        _serialize_transform,
-        ScrubTransform(suffixes=[(field,) for field in SETTINGS['scrub_fields']], redact_char='*'),
-        ScrubUrlTransform(suffixes=[(field,) for field in SETTINGS['url_fields']], params_to_scrub=SETTINGS['scrub_fields'])
+        shortener,  # priority: 10
+        ScrubRedactTransform(),  # priority: 20
+        _serialize_transform,  # priority: 30
+        ScrubTransform(suffixes=[(field,) for field in SETTINGS['scrub_fields']], redact_char='*'),  # priority: 40
+        ScrubUrlTransform(
+            suffixes=[(field,) for field in SETTINGS['url_fields']],
+            params_to_scrub=SETTINGS['scrub_fields'],
+        )  # priority: 50
     ]
+
+    # Add custom transforms
+    if len(SETTINGS['custom_transforms']) > 0:
+        _transforms.extend(SETTINGS['custom_transforms'])
+
+    # Sort the transforms by priority
+    _transforms = sorted(_transforms, key=lambda x: x.priority)
+
     _threads = queue.Queue()
     events.reset()
     filters.add_builtin_filters(SETTINGS)

rollbar/lib/transform.py

@@ -1,5 +1,6 @@
 class Transform(object):
     depth_first = True
+    priority = 100
 
     def default(self, o, key=None):
         return o

rollbar/lib/transforms/scrub.py

@@ -6,6 +6,7 @@
 
 class ScrubTransform(Transform):
     suffix_matcher = None
+    priority = 40
     def __init__(self, suffixes=None, redact_char='*', randomize_len=True):
         super(ScrubTransform, self).__init__()
         if suffixes is not None and len(suffixes) > 0:

rollbar/lib/transforms/scrub_redact.py

@@ -9,6 +9,7 @@ class RedactRef(object):
 
 
 class ScrubRedactTransform(ScrubTransform):
+    priority = 20
     def default(self, o, key=None):
         if o is REDACT_REF:
             return self.redact(o)

rollbar/lib/transforms/scruburl.py

@@ -9,6 +9,7 @@
 
 
 class ScrubUrlTransform(ScrubTransform):
+    priority = 50
     def __init__(self,
                  suffixes=None,
                  scrub_username=False,

rollbar/lib/transforms/serializable.py

@@ -9,6 +9,7 @@
 
 
 class SerializableTransform(Transform):
+    priority = 30
     def __init__(self, safe_repr=True, safelist_types=None):
         super(SerializableTransform, self).__init__()
         self.safe_repr = safe_repr

rollbar/lib/transforms/shortener.py

@@ -108,6 +108,7 @@ def shorten_tuple(obj: tuple, max_len: int) -> tuple:
 
 class ShortenerTransform(Transform):
     depth_first = False
+    priority = 10
 
     def __init__(self, safe_repr=True, keys=None, **sizes):
         super(ShortenerTransform, self).__init__()

rollbar/test/test_transform.py

@@ -0,0 +1,93 @@
+import copy
+
+import rollbar
+from rollbar.test import BaseTest
+from rollbar.lib.transforms import Transform
+from rollbar.lib.transforms.shortener import ShortenerTransform
+from rollbar.lib.transforms.scrub_redact import ScrubRedactTransform
+from rollbar.lib.transforms.serializable import SerializableTransform
+from rollbar.lib.transforms.scrub import ScrubTransform
+from rollbar.lib.transforms.scruburl import ScrubUrlTransform
+
+_test_access_token = 'aaaabbbbccccddddeeeeffff00001111'
+_default_settings = copy.deepcopy(rollbar.SETTINGS)
+
+
+class TransformTest(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_default_transforms(self):
+        transforms = {transform.__class__ for transform in rollbar._transforms}
+
+        self.assertEqual({
+            ShortenerTransform,
+            ScrubRedactTransform,
+            SerializableTransform,
+            ScrubTransform,
+            ScrubUrlTransform,
+        }, transforms)
+
+    def test_add_custom_transform(self):
+        class CustomTransform(Transform):
+            pass
+
+        rollbar._initialized = False
+        rollbar.SETTINGS = copy.deepcopy(_default_settings)
+        rollbar.init(_test_access_token,
+                     locals={'enabled': True},
+                     handler='blocking',
+                     timeout=12345,
+                     custom_transforms=[CustomTransform()])
+
+        transforms = {transform.__class__ for transform in rollbar._transforms}
+        transforms_ordered = [transform.__class__ for transform in rollbar._transforms]
+
+        self.assertEqual({
+            ShortenerTransform,
+            ScrubRedactTransform,
+            SerializableTransform,
+            ScrubTransform,
+            ScrubUrlTransform,
+            CustomTransform,
+        }, transforms)
+
+        # CustomTransform should be last because it has the default priority of 100
+        self.assertEqual([
+            ShortenerTransform,
+            ScrubRedactTransform,
+            SerializableTransform,
+            ScrubTransform,
+            ScrubUrlTransform,
+            CustomTransform,
+        ], transforms_ordered)
+
+    def test_add_custom_transform_first(self):
+        class CustomTransform(Transform):
+            priority = 1
+
+        class CustomTransformTwo(Transform):
+            priority = 35
+
+        rollbar._initialized = False
+        rollbar.SETTINGS = copy.deepcopy(_default_settings)
+        rollbar.init(_test_access_token,
+                     locals={'enabled': True},
+                     handler='blocking',
+                     timeout=12345,
+                     custom_transforms=[CustomTransform(), CustomTransformTwo()])
+
+        transforms = [transform.__class__ for transform in rollbar._transforms]
+
+        # Custom transforms should be first and fifth.
+        self.assertEqual([
+            CustomTransform,  # priority 1
+            ShortenerTransform,  # priority 10
+            ScrubRedactTransform,  # priority 20
+            SerializableTransform,  # priority 30
+            CustomTransformTwo,  # priority 35
+            ScrubTransform,  # priority 40
+            ScrubUrlTransform,  # priority 50
+        ], transforms)