feat: add custom jmespath functions support

Author: heitorlessa

aws_lambda_powertools/utilities/validation/base.py

@@ -1,3 +1,4 @@
+import json
 from typing import Dict
 
 import fastjsonschema
@@ -6,6 +7,13 @@
 
 try:
     import jmespath
+
+    class PowertoolsJson(jmespath.functions.Functions):
+        @jmespath.functions.signature({"types": ["string"]})
+        def _func_powertools_json(self, value):
+            return json.loads(value)
+
+
 except ModuleNotFoundError:
     jmespath = None
 
@@ -20,11 +28,14 @@ def validate_data_against_schema(data: Dict, schema: Dict):
         raise InvalidSchemaError(e)
 
 
-def unwrap_event_from_envelope(data: Dict, envelope: str):
+def unwrap_event_from_envelope(data: Dict, envelope: str, jmespath_options: Dict):
     if not jmespath:
         raise ModuleNotFoundError("This feature require aws-lambda-powertools[jmespath] extra package")
 
+    if not jmespath_options:
+        jmespath_options = {"custom_functions": PowertoolsJson()}
+
     try:
-        return jmespath.search(envelope, data)
+        return jmespath.search(envelope, data, options=jmespath.Options(**jmespath_options))
     except jmespath.exceptions.LexerError as e:
         raise InvalidEnvelopeExpressionError(e)

aws_lambda_powertools/utilities/validation/validator.py

@@ -16,12 +16,13 @@ def validator(
     inbound_schema: Dict = None,
     outbound_schema: Dict = None,
     envelope: str = None,
+    jmespath_options: Dict = None,
 ):
     if isinstance(event, str):
         event = json.loads(event)
 
     if envelope:
-        event = unwrap_event_from_envelope(data=event, envelope=envelope)
+        event = unwrap_event_from_envelope(data=event, envelope=envelope, jmespath_options=jmespath_options)
 
     if inbound_schema:
         logger.debug("Validating inbound event")
@@ -36,9 +37,9 @@ def validator(
     return response
 
 
-def validate(event: Dict, schema: Dict = None, envelope: str = None):
+def validate(event: Dict, schema: Dict = None, envelope: str = None, jmespath_options: Dict = None):
     if envelope:
-        event = unwrap_event_from_envelope(data=event, envelope=envelope)
+        event = unwrap_event_from_envelope(data=event, envelope=envelope, jmespath_options=jmespath_options)
 
     if isinstance(event, str):
         event = json.loads(event)