From 79898a45d5f2e68f9cf5dd11ae5494fb593a89d7 Mon Sep 17 00:00:00 2001 From: Chris Lieb Date: Mon, 11 Dec 2023 10:14:38 -0500 Subject: [PATCH 1/2] Add jsonpath-ext option for message_filtering in SqsSensor --- airflow/providers/amazon/aws/sensors/sqs.py | 11 +++++------ airflow/providers/amazon/aws/triggers/sqs.py | 6 ++---- airflow/providers/amazon/aws/utils/sqs.py | 20 ++++++++++++++++---- 3 files changed, 23 insertions(+), 14 deletions(-) diff --git a/airflow/providers/amazon/aws/sensors/sqs.py b/airflow/providers/amazon/aws/sensors/sqs.py index 95ecdd8369ac0..5394704b1acd6 100644 --- a/airflow/providers/amazon/aws/sensors/sqs.py +++ b/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, diff --git a/airflow/providers/amazon/aws/triggers/sqs.py b/airflow/providers/amazon/aws/triggers/sqs.py index 23bbb44a1828d..2a7c3c2d0dc41 100644 --- a/airflow/providers/amazon/aws/triggers/sqs.py +++ b/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, diff --git a/airflow/providers/amazon/aws/utils/sqs.py b/airflow/providers/amazon/aws/utils/sqs.py index 078baa27482d7..293aa1b898d27 100644 --- a/airflow/providers/amazon/aws/utils/sqs.py +++ b/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: From e9eebee1df6f8a99a2b8724a831fdc84a1342d11 Mon Sep 17 00:00:00 2001 From: Chris Lieb Date: Mon, 11 Dec 2023 10:48:30 -0500 Subject: [PATCH 2/2] Add tests for jsonpath-ext --- .../providers/amazon/aws/sensors/test_sqs.py | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/tests/providers/amazon/aws/sensors/test_sqs.py b/tests/providers/amazon/aws/sensors/test_sqs.py index 86727b32893b8..e0b87bf6c151a 100644 --- a/tests/providers/amazon/aws/sensors/test_sqs.py +++ b/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)