refactor!: deprecate extract_data_from_envelope in jmespath_utils and…

… replace with query
aws-powertools · Aug 5, 2024 · 83e19e3 · 83e19e3
1 parent 0f610dc
commit 83e19e3
Show file tree

Hide file tree

Showing 15 changed files with 56 additions and 34 deletions.
diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py
@@ -440,7 +440,7 @@ def decorate(event, context, *args, **kwargs):
 
             if correlation_id_path:
                 self.set_correlation_id(
-                    jmespath_utils.extract_data_from_envelope(envelope=correlation_id_path, data=event),
+                    jmespath_utils.query(envelope=correlation_id_path, data=event),
                 )
 
             if log_event:

diff --git a/aws_lambda_powertools/utilities/feature_flags/appconfig.py b/aws_lambda_powertools/utilities/feature_flags/appconfig.py
@@ -103,7 +103,7 @@ def get_configuration(self) -> Dict[str, Any]:
 
         if self.envelope:
             self.logger.debug("Envelope enabled; extracting data from config", extra={"envelope": self.envelope})
-            config = jmespath_utils.extract_data_from_envelope(
+            config = jmespath_utils.query(
                 data=config,
                 envelope=self.envelope,
                 jmespath_options=self.jmespath_options,

diff --git a/aws_lambda_powertools/utilities/jmespath_utils/__init__.py b/aws_lambda_powertools/utilities/jmespath_utils/__init__.py
@@ -2,6 +2,7 @@
 import gzip
 import json
 import logging
+import warnings
 from typing import Any, Dict, Optional, Union
 
 import jmespath
@@ -33,6 +34,29 @@ def _func_powertools_base64_gzip(self, value):
 def extract_data_from_envelope(data: Union[Dict, str], envelope: str, jmespath_options: Optional[Dict] = None) -> Any:
     """Searches and extracts data using JMESPath
 
+    *Deprecated*: Use query instead
+    """
+    warnings.warn(
+        "The extract_data_from_envelope method is deprecated and will be "
+        "removed in the next major version. Use query instead.",
+        category=DeprecationWarning,
+        stacklevel=2,
+    )
+
+    if not jmespath_options:
+        jmespath_options = {"custom_functions": PowertoolsFunctions()}
+
+    try:
+        logger.debug(f"Envelope detected: {envelope}. JMESPath options: {jmespath_options}")
+        return jmespath.search(envelope, data, options=jmespath.Options(**jmespath_options))
+    except (LexerError, TypeError, UnicodeError) as e:
+        message = f"Failed to unwrap event from envelope using expression. Error: {e} Exp: {envelope}, Data: {data}"  # noqa: B306, E501
+        raise InvalidEnvelopeExpressionError(message)
+
+
+def query(data: Union[Dict, str], envelope: str, jmespath_options: Optional[Dict] = None) -> Any:
+    """Searches and extracts data using JMESPath
+
     Envelope being the JMESPath expression to extract the data you're after
 
     Built-in JMESPath functions include: powertools_json, powertools_base64, powertools_base64_gzip
@@ -42,13 +66,13 @@ def extract_data_from_envelope(data: Union[Dict, str], envelope: str, jmespath_o
 
     **Deserialize JSON string and extracts data from body key**
 
-        from aws_lambda_powertools.utilities.jmespath_utils import extract_data_from_envelope
+        from aws_lambda_powertools.utilities.jmespath_utils import query
         from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
         def handler(event: dict, context: LambdaContext):
             # event = {"body": "{\"customerId\":\"dd4649e6-2484-4993-acb8-0f9123103394\"}"}  # noqa: ERA001
-            payload = extract_data_from_envelope(data=event, envelope="powertools_json(body)")
+            payload = query(data=event, envelope="powertools_json(body)")
             customer = payload.get("customerId")  # now deserialized
             ...
 

diff --git a/aws_lambda_powertools/utilities/validation/validator.py b/aws_lambda_powertools/utilities/validation/validator.py
@@ -119,7 +119,7 @@ def handler(event, context):
         When JMESPath expression to unwrap event is invalid
     """  # noqa: E501
     if envelope:
-        event = jmespath_utils.extract_data_from_envelope(
+        event = jmespath_utils.query(
             data=event,
             envelope=envelope,
             jmespath_options=jmespath_options,
@@ -223,7 +223,7 @@ def handler(event, context):
         When JMESPath expression to unwrap event is invalid
     """  # noqa: E501
     if envelope:
-        event = jmespath_utils.extract_data_from_envelope(
+        event = jmespath_utils.query(
             data=event,
             envelope=envelope,
             jmespath_options=jmespath_options,

diff --git a/docs/utilities/jmespath_functions.md b/docs/utilities/jmespath_functions.md
@@ -30,7 +30,7 @@ Powertools for AWS Lambda (Python) also have utilities like [validation](validat
 
 ### Extracting data
 
-You can use the `extract_data_from_envelope` function with any [JMESPath expression](https://jmespath.org/tutorial.html){target="_blank" rel="nofollow"}.
+You can use the `query` function with any [JMESPath expression](https://jmespath.org/tutorial.html){target="_blank" rel="nofollow"}.
 
 ???+ tip
 	Another common use case is to fetch deeply nested data, filter, flatten, and more.
@@ -69,15 +69,15 @@ These are all built-in envelopes you can use along with their expression as a re
 | **`API_GATEWAY_HTTP`**            | `powertools_json(body)`                                                                   |
 | **`API_GATEWAY_REST`**            | `powertools_json(body)`                                                                   |
 | **`CLOUDWATCH_EVENTS_SCHEDULED`** | `detail`                                                                                  |
-| **`CLOUDWATCH_LOGS`**             | `awslogs.powertools_base64_gzip(data)                                                     | powertools_json(@).logEvents[*]` |
+| **`CLOUDWATCH_LOGS`**             | `awslogs.powertools_base64_gzip(data)                                                     |
 | **`EVENTBRIDGE`**                 | `detail`                                                                                  |
 | **`KINESIS_DATA_STREAM`**         | `Records[*].kinesis.powertools_json(powertools_base64(data))`                             |
 | **`S3_EVENTBRIDGE_SQS`**          | `Records[*].powertools_json(body).detail`                                                 |
 | **`S3_KINESIS_FIREHOSE`**         | `records[*].powertools_json(powertools_base64(data)).Records[0]`                          |
 | **`S3_SNS_KINESIS_FIREHOSE`**     | `records[*].powertools_json(powertools_base64(data)).powertools_json(Message).Records[0]` |
 | **`S3_SNS_SQS`**                  | `Records[*].powertools_json(body).powertools_json(Message).Records[0]`                    |
 | **`S3_SQS`**                      | `Records[*].powertools_json(body).Records[0]`                                             |
-| **`SNS`**                         | `Records[0].Sns.Message                                                                   | powertools_json(@)`              |
+| **`SNS`**                         | `Records[0].Sns.Message                                                                   |
 | **`SQS`**                         | `Records[*].powertools_json(body)`                                                        |
 
 ???+ tip "Using SNS?"

diff --git a/examples/jmespath_functions/src/extract_data_from_builtin_envelope.py b/examples/jmespath_functions/src/extract_data_from_builtin_envelope.py
@@ -3,15 +3,15 @@
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.utilities.jmespath_utils import (
     envelopes,
-    extract_data_from_envelope,
+    query,
 )
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 logger = Logger()
 
 
 def handler(event: dict, context: LambdaContext) -> dict:
-    records: list = extract_data_from_envelope(data=event, envelope=envelopes.SQS)
+    records: list = query(data=event, envelope=envelopes.SQS)
     for record in records:  # records is a list
         logger.info(record.get("customerId"))  # now deserialized
 

diff --git a/examples/jmespath_functions/src/extract_data_from_envelope.py b/examples/jmespath_functions/src/extract_data_from_envelope.py
@@ -1,12 +1,12 @@
-from aws_lambda_powertools.utilities.jmespath_utils import extract_data_from_envelope
+from aws_lambda_powertools.utilities.jmespath_utils import query
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
 def handler(event: dict, context: LambdaContext) -> dict:
-    payload = extract_data_from_envelope(data=event, envelope="powertools_json(body)")
+    payload = query(data=event, envelope="powertools_json(body)")
     customer_id = payload.get("customerId")  # now deserialized
 
     # also works for fetching and flattening deeply nested data
-    some_data = extract_data_from_envelope(data=event, envelope="deeply_nested[*].some_data[]")
+    some_data = query(data=event, envelope="deeply_nested[*].some_data[]")
 
     return {"customer_id": customer_id, "message": "success", "context": some_data, "statusCode": 200}
diff --git a/examples/jmespath_functions/src/powertools_base64_gzip_jmespath_function.py b/examples/jmespath_functions/src/powertools_base64_gzip_jmespath_function.py
@@ -14,7 +14,7 @@ def lambda_handler(event, context: LambdaContext) -> dict:
     try:
         validate(event=event, schema=schemas.INPUT, envelope="powertools_base64_gzip(payload) | powertools_json(@)")
 
-        # Alternatively, extract_data_from_envelope works here too
+        # Alternatively, query works here too
         encoded_payload = base64.b64decode(event["payload"])
         uncompressed_payload = gzip.decompress(encoded_payload).decode()
         log: dict = json.loads(uncompressed_payload)

diff --git a/examples/jmespath_functions/src/powertools_base64_jmespath_function.py b/examples/jmespath_functions/src/powertools_base64_jmespath_function.py
@@ -35,7 +35,7 @@ def lambda_handler(event, context: LambdaContext) -> dict:
     try:
         validate(event=event, schema=schemas.INPUT, envelope="powertools_json(powertools_base64(payload))")
 
-        # alternatively, extract_data_from_envelope works here too
+        # alternatively, query works here too
         payload_decoded = base64.b64decode(event["payload"]).decode()
 
         order_payload: dict = json.loads(payload_decoded)

diff --git a/examples/jmespath_functions/src/powertools_custom_jmespath_function.py b/examples/jmespath_functions/src/powertools_custom_jmespath_function.py
@@ -7,7 +7,7 @@
 
 from aws_lambda_powertools.utilities.jmespath_utils import (
     PowertoolsFunctions,
-    extract_data_from_envelope,
+    query,
 )
 
 
@@ -27,7 +27,7 @@ def lambda_handler(event, context) -> dict:
     try:
         logs = []
         logs.append(
-            extract_data_from_envelope(
+            query(
                 data=event,
                 # NOTE: Use the prefix `_func_` before the name of the function
                 envelope="Records[*].decode_zlib_compression(log)",

diff --git a/examples/jmespath_functions/src/powertools_json_jmespath_function.py b/examples/jmespath_functions/src/powertools_json_jmespath_function.py
@@ -34,7 +34,7 @@ def lambda_handler(event, context: LambdaContext) -> dict:
         validate(event=event, schema=schemas.INPUT, envelope="powertools_json(payload)")
 
         # Deserialize JSON string order as dict
-        # alternatively, extract_data_from_envelope works here too
+        # alternatively, query works here too
         order_payload: dict = json.loads(event.get("payload"))
 
         return {

diff --git a/examples/middleware_factory/src/combining_powertools_utilities_function.py b/examples/middleware_factory/src/combining_powertools_utilities_function.py
@@ -11,7 +11,7 @@
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
 from aws_lambda_powertools.shared.types import JSONType
 from aws_lambda_powertools.utilities.feature_flags import AppConfigStore, FeatureFlags
-from aws_lambda_powertools.utilities.jmespath_utils import extract_data_from_envelope
+from aws_lambda_powertools.utilities.jmespath_utils import query
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from aws_lambda_powertools.utilities.validation import SchemaValidationError, validate
 
@@ -42,8 +42,8 @@ def middleware_custom(
         }
 
     # extracting headers and requestContext from event
-    headers = extract_data_from_envelope(data=event, envelope="headers")
-    request_context = extract_data_from_envelope(data=event, envelope="requestContext")
+    headers = query(data=event, envelope="headers")
+    request_context = query(data=event, envelope="requestContext")
 
     logger.debug(f"X-Customer-Id => {headers.get('X-Customer-Id')}")
     tracer.put_annotation(key="CustomerId", value=headers.get("X-Customer-Id"))

diff --git a/examples/middleware_factory/src/getting_started_middleware_before_logic_function.py b/examples/middleware_factory/src/getting_started_middleware_before_logic_function.py
@@ -5,7 +5,7 @@
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
 from aws_lambda_powertools.utilities.jmespath_utils import (
     envelopes,
-    extract_data_from_envelope,
+    query,
 )
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
@@ -19,8 +19,7 @@ class Payment:
     payment_id: str = field(default_factory=lambda: f"{uuid4()}")
 
 
-class PaymentError(Exception):
-    ...
+class PaymentError(Exception): ...
 
 
 @lambda_handler_decorator
@@ -30,7 +29,7 @@ def middleware_before(
     context: LambdaContext,
 ) -> dict:
     # extract payload from a EventBridge event
-    detail: dict = extract_data_from_envelope(data=event, envelope=envelopes.EVENTBRIDGE)
+    detail: dict = query(data=event, envelope=envelopes.EVENTBRIDGE)
 
     # check if status_id exists in payload, otherwise add default state before processing payment
     if "status_id" not in detail:
@@ -44,7 +43,7 @@ def middleware_before(
 @middleware_before
 def lambda_handler(event: dict, context: LambdaContext) -> dict:
     try:
-        payment_payload: dict = extract_data_from_envelope(data=event, envelope=envelopes.EVENTBRIDGE)
+        payment_payload: dict = query(data=event, envelope=envelopes.EVENTBRIDGE)
         return {
             "order": Payment(**payment_payload).__dict__,
             "message": "payment created",

diff --git a/examples/middleware_factory/src/getting_started_middleware_with_params_function.py b/examples/middleware_factory/src/getting_started_middleware_with_params_function.py
@@ -6,7 +6,7 @@
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
 from aws_lambda_powertools.utilities.jmespath_utils import (
     envelopes,
-    extract_data_from_envelope,
+    query,
 )
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
@@ -23,8 +23,7 @@ class Booking:
     booking_id: str = field(default_factory=lambda: f"{uuid4()}")
 
 
-class BookingError(Exception):
-    ...
+class BookingError(Exception): ...
 
 
 @lambda_handler_decorator
@@ -35,7 +34,7 @@ def obfuscate_sensitive_data(
     fields: List,
 ) -> dict:
     # extracting payload from a EventBridge event
-    detail: dict = extract_data_from_envelope(data=event, envelope=envelopes.EVENTBRIDGE)
+    detail: dict = query(data=event, envelope=envelopes.EVENTBRIDGE)
     guest_data: Any = detail.get("guest")
 
     # Obfuscate fields (email, vat, passport) before calling Lambda handler
@@ -56,7 +55,7 @@ def obfuscate_data(value: str) -> bytes:
 @obfuscate_sensitive_data(fields=["email", "passport", "vat"])
 def lambda_handler(event: dict, context: LambdaContext) -> dict:
     try:
-        booking_payload: dict = extract_data_from_envelope(data=event, envelope=envelopes.EVENTBRIDGE)
+        booking_payload: dict = query(data=event, envelope=envelopes.EVENTBRIDGE)
         return {
             "book": Booking(**booking_payload).__dict__,
             "message": "booking created",

diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py
@@ -11,7 +11,7 @@
 
 from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer
 from aws_lambda_powertools.utilities.idempotency.idempotency import IdempotencyConfig
-from aws_lambda_powertools.utilities.jmespath_utils import extract_data_from_envelope
+from aws_lambda_powertools.utilities.jmespath_utils import query
 from aws_lambda_powertools.utilities.validation import envelopes
 from tests.functional.idempotency.utils import hash_idempotency_key
 from tests.functional.utils import json_serialize, load_event
@@ -195,7 +195,7 @@ def hashed_idempotency_key(request, lambda_apigw_event, default_jmespath, lambda
 
 @pytest.fixture
 def hashed_idempotency_key_with_envelope(request, lambda_apigw_event):
-    event = extract_data_from_envelope(
+    event = query(
         data=lambda_apigw_event,
         envelope=envelopes.API_GATEWAY_HTTP,
         jmespath_options={},