Add jsonpath_ng.ext.parse support for SqsSensor

airflow/providers/amazon/aws/sensors/sqs.py

@@ -22,13 +22,12 @@
 from typing import TYPE_CHECKING, Any, Collection, Sequence
 
 from deprecated import deprecated
-from typing_extensions import Literal
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
-from airflow.providers.amazon.aws.utils.sqs import process_response
+from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType, process_response
 from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
@@ -60,9 +59,9 @@ class SqsSensor(BaseSensorOperator):
     :param visibility_timeout: Visibility timeout, a period of time during which
         Amazon SQS prevents other consumers from receiving and processing the message.
     :param message_filtering: Specified how received messages should be filtered. Supported options are:
-        `None` (no filtering, default), `'literal'` (message Body literal match) or `'jsonpath'`
-        (message Body filtered using a JSONPath expression).
-        You may add further methods by overriding the relevant class methods.
+        `None` (no filtering, default), `'literal'` (message Body literal match), `'jsonpath'`
+        (message Body filtered using a JSONPath expression), or `'jsonpath-ext'` (like `'jsonpath'`, but with
+        an expanded query grammar). You may add further methods by overriding the relevant class methods.
     :param message_filtering_match_values: Optional value/s for the message filter to match on.
         For example, with literal matching, if a message body matches any of the specified values
         then it is included. For JSONPath matching, the result of the JSONPath expression is used
@@ -90,7 +89,7 @@ def __init__(
         num_batches: int = 1,
         wait_time_seconds: int = 1,
         visibility_timeout: int | None = None,
-        message_filtering: Literal["literal", "jsonpath"] | None = None,
+        message_filtering: MessageFilteringType | None = None,
         message_filtering_match_values: Any = None,
         message_filtering_config: Any = None,
         delete_message_on_reception: bool = True,

airflow/providers/amazon/aws/triggers/sqs.py

@@ -19,11 +19,9 @@
 import asyncio
 from typing import TYPE_CHECKING, Any, AsyncIterator, Collection
 
-from typing_extensions import Literal
-
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
-from airflow.providers.amazon.aws.utils.sqs import process_response
+from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType, process_response
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -66,7 +64,7 @@ def __init__(
         num_batches: int = 1,
         wait_time_seconds: int = 1,
         visibility_timeout: int | None = None,
-        message_filtering: Literal["literal", "jsonpath"] | None = None,
+        message_filtering: MessageFilteringType | None = None,
         message_filtering_match_values: Any = None,
         message_filtering_config: Any = None,
         delete_message_on_reception: bool = True,

airflow/providers/amazon/aws/utils/sqs.py

@@ -20,15 +20,19 @@
 import logging
 from typing import Any
 
-from jsonpath_ng import parse
+import jsonpath_ng
+import jsonpath_ng.ext
 from typing_extensions import Literal
 
 log = logging.getLogger(__name__)
 
 
+MessageFilteringType = Literal["literal", "jsonpath", "jsonpath-ext"]
+
+
 def process_response(
     response: Any,
-    message_filtering: Literal["literal", "jsonpath"] | None = None,
+    message_filtering: MessageFilteringType | None = None,
     message_filtering_match_values: Any = None,
     message_filtering_config: Any = None,
 ) -> Any:
@@ -60,7 +64,13 @@ def filter_messages(
     if message_filtering == "literal":
         return filter_messages_literal(messages, message_filtering_match_values)
     if message_filtering == "jsonpath":
-        return filter_messages_jsonpath(messages, message_filtering_match_values, message_filtering_config)
+        return filter_messages_jsonpath(
+            messages, message_filtering_match_values, message_filtering_config, jsonpath_ng.parse
+        )
+    if message_filtering == "jsonpath-ext":
+        return filter_messages_jsonpath(
+            messages, message_filtering_match_values, message_filtering_config, jsonpath_ng.ext.parse
+        )
     else:
         raise NotImplementedError("Override this method to define custom filters")
 
@@ -69,7 +79,9 @@ def filter_messages_literal(messages, message_filtering_match_values) -> list[An
     return [message for message in messages if message["Body"] in message_filtering_match_values]
 
 
-def filter_messages_jsonpath(messages, message_filtering_match_values, message_filtering_config) -> list[Any]:
+def filter_messages_jsonpath(
+    messages, message_filtering_match_values, message_filtering_config, parse
+) -> list[Any]:
     jsonpath_expr = parse(message_filtering_config)
     filtered_messages = []
     for message in messages:

tests/providers/amazon/aws/sensors/test_sqs.py

@@ -288,6 +288,95 @@ def mock_delete_message_batch(**kwargs):
         ]
         mock_conn.assert_has_calls(calls_delete_message_batch)
 
+    @patch("airflow.providers.amazon.aws.hooks.sqs.SqsHook.conn", new_callable=mock.PropertyMock)
+    def test_poke_message_filtering_jsonpath_ext(self, mock_conn):
+        self.sqs_hook.create_queue(QUEUE_NAME)
+        matching = [
+            {"id": 11, "key": "a", "value": "b"},
+        ]
+        non_matching = [
+            {"id": 14, "key": "a"},
+            {"id": 14, "value": "b"},
+        ]
+        all = matching + non_matching
+
+        def mock_receive_message(**kwargs):
+            messages = []
+            for message in all:
+                messages.append(
+                    {
+                        "MessageId": message["id"],
+                        "ReceiptHandle": 100 + message["id"],
+                        "Body": json.dumps(message),
+                    }
+                )
+            return {"Messages": messages}
+
+        mock_conn.return_value.receive_message.side_effect = mock_receive_message
+
+        def mock_delete_message_batch(**kwargs):
+            return {"Successful"}
+
+        mock_conn.return_value.delete_message_batch.side_effect = mock_delete_message_batch
+
+        # Test that messages are filtered
+        self.sensor.message_filtering = "jsonpath-ext"
+        self.sensor.message_filtering_config = "$.key + $.value"
+        result = self.sensor.poke(self.mock_context)
+        assert result
+
+        # Test that only filtered messages are deleted
+        delete_entries = [{"Id": x["id"], "ReceiptHandle": 100 + x["id"]} for x in matching]
+        calls_delete_message_batch = [
+            mock.call().delete_message_batch(QueueUrl=QUEUE_URL, Entries=delete_entries)
+        ]
+        mock_conn.assert_has_calls(calls_delete_message_batch)
+
+    @patch("airflow.providers.amazon.aws.hooks.sqs.SqsHook.conn", new_callable=mock.PropertyMock)
+    def test_poke_message_filtering_jsonpath_ext_values(self, mock_conn):
+        self.sqs_hook.create_queue(QUEUE_NAME)
+        matching = [
+            {"id": 11, "key": "a1", "value": "b1"},
+        ]
+        non_matching = [
+            {"id": 22, "key": "a2", "value": "b1"},
+            {"id": 33, "key": "a1", "value": "b2"},
+        ]
+        all = matching + non_matching
+
+        def mock_receive_message(**kwargs):
+            messages = []
+            for message in all:
+                messages.append(
+                    {
+                        "MessageId": message["id"],
+                        "ReceiptHandle": 100 + message["id"],
+                        "Body": json.dumps(message),
+                    }
+                )
+            return {"Messages": messages}
+
+        mock_conn.return_value.receive_message.side_effect = mock_receive_message
+
+        def mock_delete_message_batch(**kwargs):
+            return {"Successful"}
+
+        mock_conn.return_value.delete_message_batch.side_effect = mock_delete_message_batch
+
+        # Test that messages are filtered
+        self.sensor.message_filtering = "jsonpath-ext"
+        self.sensor.message_filtering_config = "$.key + ' ' + $.value"
+        self.sensor.message_filtering_match_values = ["a1 b1"]
+        result = self.sensor.poke(self.mock_context)
+        assert result
+
+        # Test that only filtered messages are deleted
+        delete_entries = [{"Id": x["id"], "ReceiptHandle": 100 + x["id"]} for x in matching]
+        calls_delete_message_batch = [
+            mock.call().delete_message_batch(QueueUrl="https://test-queue", Entries=delete_entries)
+        ]
+        mock_conn.assert_has_calls(calls_delete_message_batch)
+
     @patch("airflow.providers.amazon.aws.hooks.sqs.SqsHook.conn", new_callable=mock.PropertyMock)
     def test_poke_do_not_delete_message_on_received(self, mock_conn):
         self.sqs_hook.create_queue(QUEUE_NAME)