From b04b4a49ff14e13989aaea23a32fdfbb476fa5ec Mon Sep 17 00:00:00 2001 From: Daniel Morell Date: Fri, 17 Jan 2025 16:01:16 -0600 Subject: [PATCH] Issue #245 Added support for custom payload transforms. --- rollbar/__init__.py | 22 ++++-- rollbar/lib/transform.py | 1 + rollbar/lib/transforms/scrub.py | 1 + rollbar/lib/transforms/scrub_redact.py | 1 + rollbar/lib/transforms/scruburl.py | 1 + rollbar/lib/transforms/serializable.py | 1 + rollbar/lib/transforms/shortener.py | 1 + rollbar/test/test_transform.py | 93 ++++++++++++++++++++++++++ 8 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 rollbar/test/test_transform.py diff --git a/rollbar/__init__.py b/rollbar/__init__.py index 2a3a62d7..d26f6fd4 100644 --- a/rollbar/__init__.py +++ b/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) diff --git a/rollbar/lib/transform.py b/rollbar/lib/transform.py index 86b4e39d..70b957b4 100644 --- a/rollbar/lib/transform.py +++ b/rollbar/lib/transform.py @@ -1,5 +1,6 @@ class Transform(object): depth_first = True + priority = 100 def default(self, o, key=None): return o diff --git a/rollbar/lib/transforms/scrub.py b/rollbar/lib/transforms/scrub.py index ff0af810..e3b4264d 100644 --- a/rollbar/lib/transforms/scrub.py +++ b/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: diff --git a/rollbar/lib/transforms/scrub_redact.py b/rollbar/lib/transforms/scrub_redact.py index a81234b9..09b9a76d 100644 --- a/rollbar/lib/transforms/scrub_redact.py +++ b/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) diff --git a/rollbar/lib/transforms/scruburl.py b/rollbar/lib/transforms/scruburl.py index 0ab922cc..a8366143 100644 --- a/rollbar/lib/transforms/scruburl.py +++ b/rollbar/lib/transforms/scruburl.py @@ -9,6 +9,7 @@ class ScrubUrlTransform(ScrubTransform): + priority = 50 def __init__(self, suffixes=None, scrub_username=False, diff --git a/rollbar/lib/transforms/serializable.py b/rollbar/lib/transforms/serializable.py index 49f95d29..f426f2d4 100644 --- a/rollbar/lib/transforms/serializable.py +++ b/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 diff --git a/rollbar/lib/transforms/shortener.py b/rollbar/lib/transforms/shortener.py index 383b1636..17176737 100644 --- a/rollbar/lib/transforms/shortener.py +++ b/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__() diff --git a/rollbar/test/test_transform.py b/rollbar/test/test_transform.py new file mode 100644 index 00000000..3851997f --- /dev/null +++ b/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)