diff --git a/airflow/providers/amazon/aws/sensors/sqs.py b/airflow/providers/amazon/aws/sensors/sqs.py index cc026ec5e8306..c1b6c4c74769f 100644 --- a/airflow/providers/amazon/aws/sensors/sqs.py +++ b/airflow/providers/amazon/aws/sensors/sqs.py @@ -18,12 +18,13 @@ """Reads and then deletes the message from SQS queue""" import json import warnings -from typing import TYPE_CHECKING, Any, Optional, Sequence +from typing import TYPE_CHECKING, Any, Collection, List, Optional, Sequence from jsonpath_ng import parse from typing_extensions import Literal from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection from airflow.providers.amazon.aws.hooks.sqs import SqsHook from airflow.sensors.base import BaseSensorOperator @@ -34,9 +35,13 @@ class SqsSensor(BaseSensorOperator): """ Get messages from an Amazon SQS queue and then delete the messages from the queue. - If deletion of messages fails an AirflowException is thrown. Otherwise, the messages + If deletion of messages fails, an AirflowException is thrown. Otherwise, the messages are pushed through XCom with the key ``messages``. + By default,the sensor performs one and only one SQS call per poke, which limits the result to + a maximum of 10 messages. However, the total number of SQS API calls per poke can be controlled + by num_batches param. + .. seealso:: For more information on how to use this sensor, take a look at the guide: :ref:`howto/sensor:SqsSensor` @@ -44,6 +49,7 @@ class SqsSensor(BaseSensorOperator): :param aws_conn_id: AWS connection id :param sqs_queue: The SQS queue url (templated) :param max_messages: The maximum number of messages to retrieve for each poke (templated) + :param num_batches: The number of times the sensor will call the SQS API to receive messages (default: 1) :param wait_time_seconds: The time in seconds to wait for receiving messages (default: 1 second) :param visibility_timeout: Visibility timeout, a period of time during which Amazon SQS prevents other consumers from receiving and processing the message. @@ -72,6 +78,7 @@ def __init__( sqs_queue, aws_conn_id: str = 'aws_default', max_messages: int = 5, + num_batches: int = 1, wait_time_seconds: int = 1, visibility_timeout: Optional[int] = None, message_filtering: Optional[Literal["literal", "jsonpath"]] = None, @@ -84,6 +91,7 @@ def __init__( self.sqs_queue = sqs_queue self.aws_conn_id = aws_conn_id self.max_messages = max_messages + self.num_batches = num_batches self.wait_time_seconds = wait_time_seconds self.visibility_timeout = visibility_timeout @@ -104,15 +112,13 @@ def __init__( self.hook: Optional[SqsHook] = None - def poke(self, context: 'Context'): + def poll_sqs(self, sqs_conn: BaseAwsConnection) -> Collection: """ - Check for message on subscribed queue and write to xcom the message with key ``messages`` + Poll SQS queue to retrieve messages. - :param context: the context object - :return: ``True`` if message is available or ``False`` + :param sqs_conn: SQS connection + :return: A list of messages retrieved from SQS """ - sqs_conn = self.get_hook().get_conn() - self.log.info('SqsSensor checking for message on queue: %s', self.sqs_queue) receive_message_kwargs = { @@ -126,7 +132,7 @@ def poke(self, context: 'Context'): response = sqs_conn.receive_message(**receive_message_kwargs) if "Messages" not in response: - return False + return [] messages = response['Messages'] num_messages = len(messages) @@ -136,28 +142,47 @@ def poke(self, context: 'Context'): messages = self.filter_messages(messages) num_messages = len(messages) self.log.info("There are %d messages left after filtering", num_messages) + return messages - if not num_messages: - return False + def poke(self, context: 'Context'): + """ + Check subscribed queue for messages and write them to xcom with the ``messages`` key. - if not self.delete_message_on_reception: - context['ti'].xcom_push(key='messages', value=messages) - return True + :param context: the context object + :return: ``True`` if message is available or ``False`` + """ + sqs_conn = self.get_hook().get_conn() - self.log.info("Deleting %d messages", num_messages) + message_batch: List[Any] = [] - entries = [ - {'Id': message['MessageId'], 'ReceiptHandle': message['ReceiptHandle']} for message in messages - ] - response = sqs_conn.delete_message_batch(QueueUrl=self.sqs_queue, Entries=entries) + # perform multiple SQS call to retrieve messages in series + for _ in range(self.num_batches): + messages = self.poll_sqs(sqs_conn=sqs_conn) - if 'Successful' in response: - context['ti'].xcom_push(key='messages', value=messages) - return True - else: - raise AirflowException( - 'Delete SQS Messages failed ' + str(response) + ' for messages ' + str(messages) - ) + if not len(messages): + continue + + message_batch.extend(messages) + + if self.delete_message_on_reception: + + self.log.info("Deleting %d messages", len(messages)) + + entries = [ + {'Id': message['MessageId'], 'ReceiptHandle': message['ReceiptHandle']} + for message in messages + ] + response = sqs_conn.delete_message_batch(QueueUrl=self.sqs_queue, Entries=entries) + + if 'Successful' not in response: + raise AirflowException( + 'Delete SQS Messages failed ' + str(response) + ' for messages ' + str(messages) + ) + if not len(message_batch): + return False + + context['ti'].xcom_push(key='messages', value=message_batch) + return True def get_hook(self) -> SqsHook: """Create and return an SqsHook""" diff --git a/tests/providers/amazon/aws/sensors/test_sqs.py b/tests/providers/amazon/aws/sensors/test_sqs.py index c2a6b3aa2ffba..2254779535544 100644 --- a/tests/providers/amazon/aws/sensors/test_sqs.py +++ b/tests/providers/amazon/aws/sensors/test_sqs.py @@ -305,3 +305,31 @@ def test_poke_do_not_delete_message_on_received(self, mock_conn): ) self.sensor.poke(self.mock_context) assert mock_conn.delete_message_batch.called is False + + @mock_sqs + def test_poke_batch_messages(self): + messages = ["hello", "brave", "world"] + + self.sqs_hook.create_queue(QUEUE_NAME) + # Do publish 3 messages + for message in messages: + self.sqs_hook.send_message(queue_url=QUEUE_URL, message_body=message) + + # Init batch sensor to get 1 message for each SQS poll + # and perform 3 polls + self.sensor = SqsSensor( + task_id='test_task3', + dag=self.dag, + sqs_queue=QUEUE_URL, + aws_conn_id='aws_default', + max_messages=1, + num_batches=3, + ) + result = self.sensor.poke(self.mock_context) + assert result + + # expect all messages are retrieved + for message in messages: + assert f"'Body': '{message}'" in str( + self.mock_context['ti'].method_calls + ), "context call should contain message '{message}'" diff --git a/tests/system/providers/amazon/aws/example_sqs.py b/tests/system/providers/amazon/aws/example_sqs.py index 09f697d4727af..0aabe9e118608 100644 --- a/tests/system/providers/amazon/aws/example_sqs.py +++ b/tests/system/providers/amazon/aws/example_sqs.py @@ -52,8 +52,13 @@ def delete_queue(queue_url): sqs_queue = create_queue() # [START howto_operator_sqs] - publish_to_queue = SqsPublishOperator( - task_id='publish_to_queue', + publish_to_queue_1 = SqsPublishOperator( + task_id='publish_to_queue_1', + sqs_queue=sqs_queue, + message_content='{{ task_instance }}-{{ logical_date }}', + ) + publish_to_queue_2 = SqsPublishOperator( + task_id='publish_to_queue_2', sqs_queue=sqs_queue, message_content='{{ task_instance }}-{{ logical_date }}', ) @@ -64,14 +69,26 @@ def delete_queue(queue_url): task_id='read_from_queue', sqs_queue=sqs_queue, ) + # Retrieve multiple batches of messages from SQS. + # The SQS API only returns a maximum of 10 messages per poll. + read_from_queue_in_batch = SqsSensor( + task_id='read_from_queue_in_batch', + sqs_queue=create_queue, + # Get maximum 10 messages each poll + max_messages=10, + # Combine 3 polls before returning results + num_batches=3, + ) # [END howto_sensor_sqs] chain( # TEST SETUP sqs_queue, # TEST BODY - publish_to_queue, + publish_to_queue_1, read_from_queue, + publish_to_queue_2, + read_from_queue_in_batch, # TEST TEARDOWN delete_queue(sqs_queue), )