diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index a6db72c55ea7d..afbdb188dc395 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -790,6 +790,7 @@ hiveserver hmsclient Hoc hoc +Hohpe homebrew honoured hookable diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py index d5cb666411e05..d6f6d38c71d86 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/asb.py @@ -17,7 +17,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, Callable -from uuid import uuid4 +from uuid import UUID, uuid4 from azure.core.exceptions import ResourceNotFoundError from azure.servicebus import ( @@ -468,7 +468,15 @@ def get_conn(self) -> ServiceBusClient: self.log.info("Create and returns ServiceBusClient") return client - def send_message(self, queue_name: str, messages: str | list[str], batch_message_flag: bool = False): + def send_message( + self, + queue_name: str, + messages: str | list[str], + batch_message_flag: bool = False, + message_id: str | None = None, + reply_to: str | None = None, + message_headers: dict[str | bytes, int | float | bytes | bool | str | UUID] | None = None, + ): """ Use ServiceBusClient Send to send message(s) to a Service Bus Queue. @@ -478,38 +486,49 @@ def send_message(self, queue_name: str, messages: str | list[str], batch_message :param messages: Message which needs to be sent to the queue. It can be string or list of string. :param batch_message_flag: bool flag, can be set to True if message needs to be sent as batch message. + :param message_id: Message ID to set on message being sent to the queue. Please note, message_id may only be + set when a single message is sent. + :param reply_to: Reply to which needs to be sent to the queue. + :param message_headers: Headers to add to the message's application_properties field for Azure Service Bus. """ if queue_name is None: raise TypeError("Queue name cannot be None.") if not messages: raise ValueError("Messages list cannot be empty.") + if message_id and not isinstance(messages, str): + raise TypeError("Message ID can only be set if a single message is sent.") with ( self.get_conn() as service_bus_client, service_bus_client.get_queue_sender(queue_name=queue_name) as sender, sender, ): - if isinstance(messages, str): - if not batch_message_flag: - msg = ServiceBusMessage(messages) - sender.send_messages(msg) - else: - self.send_batch_message(sender, [messages]) + message_creator = lambda msg_body: ServiceBusMessage( + msg_body, message_id=message_id, reply_to=reply_to, application_properties=message_headers + ) + message_list = [messages] if isinstance(messages, str) else messages + if not batch_message_flag: + self.send_list_messages(sender, message_list, message_creator) else: - if not batch_message_flag: - self.send_list_messages(sender, messages) - else: - self.send_batch_message(sender, messages) + self.send_batch_message(sender, message_list, message_creator) @staticmethod - def send_list_messages(sender: ServiceBusSender, messages: list[str]): - list_messages = [ServiceBusMessage(message) for message in messages] + def send_list_messages( + sender: ServiceBusSender, + messages: list[str], + message_creator: Callable[[str], ServiceBusMessage], + ): + list_messages = [message_creator(body) for body in messages] sender.send_messages(list_messages) # type: ignore[arg-type] @staticmethod - def send_batch_message(sender: ServiceBusSender, messages: list[str]): + def send_batch_message( + sender: ServiceBusSender, + messages: list[str], + message_creator: Callable[[str], ServiceBusMessage], + ): batch_message = sender.create_message_batch() for message in messages: - batch_message.add_message(ServiceBusMessage(message)) + batch_message.add_message(message_creator(message)) sender.send_messages(batch_message) def receive_message( diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/asb.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/asb.py index aa8eecb3f2442..9707e351f36c5 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/asb.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/asb.py @@ -18,15 +18,21 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Callable +from uuid import UUID, uuid4 from airflow.models import BaseOperator from airflow.providers.microsoft.azure.hooks.asb import AdminClientHook, MessageHook +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError +from azure.servicebus.management import ( + AuthorizationRule, + CorrelationRuleFilter, + SqlRuleFilter, +) if TYPE_CHECKING: import datetime from azure.servicebus import ServiceBusMessage - from azure.servicebus.management import AuthorizationRule, CorrelationRuleFilter, SqlRuleFilter from airflow.utils.context import Context @@ -100,6 +106,11 @@ class AzureServiceBusSendMessageOperator(BaseOperator): as batch message it can be set to True. :param azure_service_bus_conn_id: Reference to the :ref: `Azure Service Bus connection`. + :param message_id: Message ID to set on message being sent to the queue. Please note, message_id may only be + set when a single message is sent. + :param reply_to: Name of queue or topic the receiver should reply to. Determination of if the reply will be sent to + a queue or a topic should be made out-of-band. + :param message_headers: Headers to add to the message's application_properties field for Azure Service Bus. """ template_fields: Sequence[str] = ("queue_name",) @@ -112,6 +123,9 @@ def __init__( message: str | list[str], batch: bool = False, azure_service_bus_conn_id: str = "azure_service_bus_default", + message_id: str | None = None, + reply_to: str | None = None, + message_headers: dict[str | bytes, int | float | bytes | bool | str | UUID] | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -119,6 +133,9 @@ def __init__( self.batch = batch self.message = message self.azure_service_bus_conn_id = azure_service_bus_conn_id + self.message_id = message_id + self.reply_to = reply_to + self.message_headers = message_headers def execute(self, context: Context) -> None: """Send Message to the specific queue in Service Bus namespace.""" @@ -126,7 +143,9 @@ def execute(self, context: Context) -> None: hook = MessageHook(azure_service_bus_conn_id=self.azure_service_bus_conn_id) # send message - hook.send_message(self.queue_name, self.message, self.batch) + hook.send_message( + self.queue_name, self.message, self.batch, self.message_id, self.reply_to, self.message_headers + ) class AzureServiceBusReceiveMessageOperator(BaseOperator): @@ -370,6 +389,7 @@ class AzureServiceBusSubscriptionCreateOperator(BaseOperator): :param filter_rule_name: Optional rule name to use applying the rule filter to the subscription :param azure_service_bus_conn_id: Reference to the :ref:`Azure Service Bus connection`. + :param filter: The filter expression used to match messages. For SQLFilter, the expression """ template_fields: Sequence[str] = ("topic_name", "subscription_name") @@ -639,3 +659,162 @@ def execute(self, context: Context) -> None: self.log.info("Topic %s deleted.", self.topic_name) else: self.log.info("Topic %s does not exist.", self.topic_name) + + +class AzureServiceBusRequestReplyOperator(BaseOperator): + """ + Implement request-reply pattern using Azure Service Bus. + + Send a message to an Azure Service Bus Queue and receive a reply by correlation id from an Azure Service + Bus Topic. This implements the Request-Reply pattern from Enterprise Integration Patterns, Hohpe, Woolf, + Addison-Wesley, 2003: https://www.enterpriseintegrationpatterns.com/patterns/messaging/RequestReply.html + + Steps are: + 1. Generate a unique ID for the message. The subscription needs to exist before the request message is + sent or there will be a race condition where the reply message could be sent before the + subscription is created. By default, a UUID is used for the unique ID and this should be + reasonably unique over the life of the universe. + 2. Create a subscription to the reply topic for messages where the correlation ID equals the unique ID + created in #1. + 3. Send the message to the request queue with (a) the reply-to property set to the reply topic, + (b) the reply type property set to topic, and (c) the message ID set to the unique ID. + 4. Wait for a reply message on the reply topic with the correlation ID set to the unique ID. + 5. Remove the subscription on the reply topic. + + The caller must pass in a generator function to create the request message body (request_body_generator) + from the context and an optional callback can process the reply message (reply_message_callback). This + callback could either detect errors and abort processing by throwing an exception or could process the + message body and add information into XComs for downstream tasks to use. The remote service can send back + the message in any format it chooses. The supplied callback should be able to handle the message format. + + The remote service should reply to the topic or queue specified in the reply_to property of the request + message. The remote service can tell if the reply should go to a topic or a queue based on the reply_type + although the current implementation expects all replies to be sent through a topic. The reply message + should have the correlation ID set to the message ID of the request message. + + :param request_queue_name: Name of the queue to send the request to. This queue must be reachable from + a connection created from the connection name specified in the param `azure_service_bus_conn_id`. + :param request_body_generator: A method to generate the request message body from the context. + :param reply_topic_name: Name of the topic to send the reply to. This topic must be reachable from + a connection created from the connection name specified in the param "azure_service_bus_conn_id". + :param reply_correlation_id: a string to use to correlate the request and reply. This will also be the + message ID of the request message. If not specified, a UUID will be generated and used. Generally, + the default behavior should be used. + :param max_wait_time: maximum wait for a reply in seconds. This should be set to some small multiple of + the expected processing time. Perhaps 3x the expected processing time. Default is 60 seconds. + :param reply_message_callback: An optional callback to handle the response message. This takes the service + bus message and the context as parameters and can raise an exception to abort processing or insert + values into the XCOM for downstream tasks to use. + :param azure_service_bus_conn_id: ID of the airflow connection to the Azure Service Bus. It defaults to + `azure_service_bus_default`, + """ + + REPLY_SUBSCRIPTION_PREFIX = "reply-" + REPLY_RULE_SUFFIX = "-rule" + + def __init__( + self, + *, + request_queue_name: str, + request_body_generator: Callable[[Context], str], + reply_topic_name: str, + reply_correlation_id: str | None = None, + max_wait_time: float = 60, + reply_message_callback: MessageCallback | None = None, + azure_service_bus_conn_id: str = "azure_service_bus_default", + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.request_queue_name = request_queue_name + self.request_body_generator = request_body_generator + self.reply_topic_name = reply_topic_name + self.reply_correlation_id = reply_correlation_id + self.max_wait_time = max_wait_time + self.reply_message_callback = reply_message_callback + self.azure_service_bus_conn_id = azure_service_bus_conn_id + + if not self.reply_correlation_id: + self.reply_correlation_id = str(uuid4()) + self.subscription_name = self.REPLY_SUBSCRIPTION_PREFIX + self.reply_correlation_id + + def execute(self, context: Context) -> None: + """Implement the request-reply pattern using existing hooks.""" + self._validate_params() + admin_hook = AdminClientHook(azure_service_bus_conn_id=self.azure_service_bus_conn_id) + self._create_reply_subscription_for_correlation_id(admin_hook, context) + + message_hook = MessageHook(azure_service_bus_conn_id=self.azure_service_bus_conn_id) + try: + # send the request message + message_hook.send_message( + self.request_queue_name, + self.request_body_generator(context), + message_id=self.reply_correlation_id, + reply_to=self.reply_topic_name, + message_headers={"reply_type": "topic"}, + ) + + # Wait for and receive the reply message + message_hook.receive_subscription_message( + self.reply_topic_name, + self.subscription_name, + context, + max_message_count=1, + max_wait_time=self.max_wait_time, + message_callback=self.reply_message_callback, + ) + finally: + # Remove the subscription on the reply topic + self._remove_reply_subscription(admin_hook) + + def _create_reply_subscription_for_correlation_id( + self, admin_hook: AdminClientHook, context: Context + ) -> None: + """Create subscription on the reply topic for the correlation ID.""" + self.log.info("Creating subscription %s on topic %s", self.subscription_name, self.reply_topic_name) + rule_name = self.subscription_name + self.REPLY_RULE_SUFFIX + filter = CorrelationRuleFilter(correlation_id=self.reply_correlation_id) + try: + admin_hook.create_subscription( + topic_name=self.reply_topic_name, + subscription_name=self.subscription_name, + default_message_time_to_live="PT1H", # 1 hour + dead_lettering_on_message_expiration=True, + dead_lettering_on_filter_evaluation_exceptions=True, + enable_batched_operations=False, + user_metadata=f"Subscription for reply to {self.reply_correlation_id} for task ID {context['task'].task_id}", + auto_delete_on_idle="PT6H", # 6 hours + filter_rule_name=rule_name, + filter_rule=filter, + ) + except ResourceExistsError: + # subscription already created, so return + self.log.info( + "Subscription %s on topic %s already existed.", + self.subscription_name, + self.reply_topic_name, + ) + return + + def _remove_reply_subscription(self, admin_hook: AdminClientHook) -> None: + try: + admin_hook.delete_subscription(self.subscription_name, self.reply_topic_name) + self.log.debug("Subscription removed!") + except ResourceNotFoundError: + # already deleted, ignore. + self.log.debug("Subscription already removed!") + self.log.info("Removed subscription %s", self.subscription_name) + + def _validate_params(self): + error_message: str = "" + + if not self.request_queue_name: + error_message += "Request queue name is required. " + if not self.request_body_generator: + error_message += "Request body creator is required. " + if not self.reply_topic_name: + error_message += "Reply topic name is required. " + + if error_message: + self.log.error(error_message) + raise TypeError(error_message) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_asb.py b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_asb.py index 667f40971987e..3d352e03e1721 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_asb.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_asb.py @@ -19,6 +19,7 @@ from unittest import mock import pytest +from azure.core.exceptions import ResourceExistsError from azure.servicebus import ServiceBusMessage try: @@ -26,11 +27,14 @@ except ImportError: pytest.skip("Azure Service Bus not available", allow_module_level=True) +from azure.core.exceptions import ResourceNotFoundError + from airflow.providers.microsoft.azure.operators.asb import ( ASBReceiveSubscriptionMessageOperator, AzureServiceBusCreateQueueOperator, AzureServiceBusDeleteQueueOperator, AzureServiceBusReceiveMessageOperator, + AzureServiceBusRequestReplyOperator, AzureServiceBusSendMessageOperator, AzureServiceBusSubscriptionCreateOperator, AzureServiceBusSubscriptionDeleteOperator, @@ -131,52 +135,80 @@ def test_delete_queue(self, mock_get_conn): class TestAzureServiceBusSendMessageOperator: @pytest.mark.parametrize( - "mock_message, mock_batch_flag", + "mock_message, mock_batch_flag, mock_message_id, mock_reply_to, mock_headers", [ - (MESSAGE, True), - (MESSAGE, False), - (MESSAGE_LIST, True), - (MESSAGE_LIST, False), + (MESSAGE, True, None, None, None), + (MESSAGE, False, "test_message_id", "test_reply_to", {"test_header": "test_value"}), + (MESSAGE_LIST, True, None, None, None), + (MESSAGE_LIST, False, None, None, None), ], ) - def test_init(self, mock_message, mock_batch_flag): + def test_init(self, mock_message, mock_batch_flag, mock_message_id, mock_reply_to, mock_headers): """ Test init by creating AzureServiceBusSendMessageOperator with task id, queue_name, message, - batch and asserting with values + batch, message_id, reply_to, and message headers and asserting with values """ asb_send_message_queue_operator = AzureServiceBusSendMessageOperator( task_id="asb_send_message_queue_without_batch", queue_name=QUEUE_NAME, message=mock_message, batch=mock_batch_flag, + message_id=mock_message_id, + reply_to=mock_reply_to, + message_headers=mock_headers, ) assert asb_send_message_queue_operator.task_id == "asb_send_message_queue_without_batch" assert asb_send_message_queue_operator.queue_name == QUEUE_NAME assert asb_send_message_queue_operator.message == mock_message assert asb_send_message_queue_operator.batch is mock_batch_flag + assert asb_send_message_queue_operator.message_id == mock_message_id + assert asb_send_message_queue_operator.reply_to == mock_reply_to + assert asb_send_message_queue_operator.message_headers == mock_headers - @mock.patch("airflow.providers.microsoft.azure.hooks.asb.MessageHook.get_conn") - def test_send_message_queue(self, mock_get_conn): + @mock.patch("airflow.providers.microsoft.azure.hooks.asb.MessageHook.send_message") + def test_send_message_queue(self, mock_send_message): """ Test AzureServiceBusSendMessageOperator with queue name, batch boolean flag, mock the send_messages of azure service bus function """ + TASK_ID = "task-id" + MSG_BODY = "test message body" + MSG_ID = None + REPLY_TO = None + HDRS = None asb_send_message_queue_operator = AzureServiceBusSendMessageOperator( - task_id="asb_send_message_queue", + task_id=TASK_ID, queue_name=QUEUE_NAME, - message="Test message", + message=MSG_BODY, batch=False, ) asb_send_message_queue_operator.execute(None) - expected_calls = [ - mock.call() - .__enter__() - .get_queue_sender(QUEUE_NAME) - .__enter__() - .send_messages(ServiceBusMessage("Test message")) - .__exit__() - ] - mock_get_conn.assert_has_calls(expected_calls, any_order=False) + expected_calls = [mock.call(QUEUE_NAME, MSG_BODY, False, MSG_ID, REPLY_TO, HDRS)] + mock_send_message.assert_has_calls(expected_calls, any_order=False) + + @mock.patch("airflow.providers.microsoft.azure.hooks.asb.MessageHook.send_message") + def test_send_message_queue_with_id_hdrs_and_reply_to(self, mock_send_message): + """ + Test AzureServiceBusSendMessageOperator with queue name, batch boolean flag, mock + the send_messages of azure service bus function + """ + TASK_ID = "task-id" + MSG_ID = "test_message_id" + MSG_BODY = "test message body" + REPLY_TO = "test_reply_to" + HDRS = {"test_header": "test_value"} + asb_send_message_queue_operator = AzureServiceBusSendMessageOperator( + task_id=TASK_ID, + queue_name=QUEUE_NAME, + message=MSG_BODY, + batch=False, + message_id=MSG_ID, + reply_to=REPLY_TO, + message_headers=HDRS, + ) + asb_send_message_queue_operator.execute(None) + expected_calls = [mock.call(QUEUE_NAME, MSG_BODY, False, MSG_ID, REPLY_TO, HDRS)] + mock_send_message.assert_has_calls(expected_calls, any_order=False) class TestAzureServiceBusReceiveMessageOperator: @@ -593,3 +625,269 @@ def test_delete_topic_exception(self, mock_sb_admin_client): ) with pytest.raises(TypeError): asb_delete_topic_exception.execute(None) + + +class TestAzureServiceBusRequestReplyOperator: + # tests for AzureServiceBusRequestReplyOperator._remove_reply_subscription + # use mock for the admin_hook passed into _remove_reply_subscription to + # ensure delete_subscription is called with correct parameters + def test_remove_reply_subscription(self): + with mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") as mock_admin_hook: + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name="test_queue", + request_body_generator=lambda: "test_body", + reply_topic_name="reply-topic-name", + ) + + # Set the subscription_name attribute for the operator + operator.subscription_name = "test_subscription" + + operator._remove_reply_subscription(mock_admin_hook) + + mock_admin_hook.delete_subscription.assert_called_once_with( + operator.subscription_name, operator.reply_topic_name + ) + + def test_remove_reply_subscription_ignores_resource_not_found_error(self): + with mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") as mock_admin_hook: + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name="test_queue", + request_body_generator=lambda: "test_body", + reply_topic_name="reply-topic-name", + ) + + # Set the subscription_name attribute for the operator + operator.subscription_name = "test_subscription" + + # Mock the delete_subscription method to raise ResourceNotFoundError + mock_admin_hook.delete_subscription.side_effect = ResourceNotFoundError + + operator._remove_reply_subscription(mock_admin_hook) + + mock_admin_hook.delete_subscription.assert_called_once_with( + operator.subscription_name, operator.reply_topic_name + ) + + # tests for AzureServiceBusRequestReplyOperator._validate_params with different combinations of parameters + # and expected results + @pytest.mark.parametrize( + "request_queue_name, request_body_generator, reply_topic_name, expected_exception, expected_message", + [ + (None, lambda: "test_body", "test_topic", TypeError, "Request queue name is required. "), + ("test_queue", None, "test_topic", TypeError, "Request body creator is required. "), + ("test_queue", lambda: "test_body", None, TypeError, "Reply topic name is required. "), + ( + None, + None, + None, + TypeError, + "Request queue name is required. Request body creator is required. Reply topic name is required. ", + ), + ("test_queue", lambda: "test_body", "test_topic", None, None), + ], + ) + def test_validate_params( + self, + request_queue_name, + request_body_generator, + reply_topic_name, + expected_exception, + expected_message, + ): + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name=request_queue_name, + request_body_generator=request_body_generator, + reply_topic_name=reply_topic_name, + ) + + if expected_exception: + with pytest.raises(expected_exception) as exc_info: + operator._validate_params() + assert str(exc_info.value) == expected_message + else: + operator._validate_params() # Should not raise any exception + + @mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") + def test_create_reply_subscription_for_correlation_id(self, mock_admin_hook): + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name="test_queue", + request_body_generator=lambda: "test_body", + reply_topic_name="reply-topic-name", + ) + + context = mock.MagicMock() + context["task"] = mock.MagicMock() + context["task"].task_id = 345 + + operator._create_reply_subscription_for_correlation_id(mock_admin_hook, context) + + mock_admin_hook.create_subscription.assert_called_once_with( + topic_name="reply-topic-name", + subscription_name=operator.subscription_name, + default_message_time_to_live="PT1H", # 1 hour + dead_lettering_on_message_expiration=True, + dead_lettering_on_filter_evaluation_exceptions=True, + enable_batched_operations=False, + user_metadata=f"Subscription for reply to {operator.reply_correlation_id} for task ID {context['task'].task_id}", + auto_delete_on_idle="PT6H", # 6 hours + filter_rule_name=operator.subscription_name + operator.REPLY_RULE_SUFFIX, + filter_rule=mock.ANY, + ) + + @mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") + def test_create_reply_subscription_for_correlation_id_subscription_exists(self, mock_admin_hook): + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name="test_queue", + request_body_generator=lambda: "test_body", + reply_topic_name="reply-topic-name", + ) + + context = mock.MagicMock() + context["task"] = mock.MagicMock() + context["task"].task_id = 987 + + mock_admin_hook.create_subscription.side_effect = ResourceExistsError + + operator._create_reply_subscription_for_correlation_id(mock_admin_hook, context) + + mock_admin_hook.create_subscription.assert_called_once_with( + topic_name="reply-topic-name", + subscription_name=operator.subscription_name, + default_message_time_to_live="PT1H", # 1 hour + dead_lettering_on_message_expiration=True, + dead_lettering_on_filter_evaluation_exceptions=True, + enable_batched_operations=False, + user_metadata=f"Subscription for reply to {operator.reply_correlation_id} for task ID {context['task'].task_id}", + auto_delete_on_idle="PT6H", # 6 hours + filter_rule_name=operator.subscription_name + operator.REPLY_RULE_SUFFIX, + filter_rule=mock.ANY, + ) + + mock_admin_hook.get_conn.return_value.__enter__.return_value.delete_rule.assert_not_called() + mock_admin_hook.get_conn.return_value.__enter__.return_value.create_rule.assert_not_called() + + @mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") + def test_create_reply_subscription_for_correlation_id_delete_rule_not_found(self, mock_admin_hook): + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name="test_queue", + request_body_generator=lambda: "test_body", + reply_topic_name="reply-topic-name", + ) + + context = mock.MagicMock() + context["task"] = mock.MagicMock() + context["task"].task_id = 789 + + mock_admin_hook.get_conn.return_value.__enter__.return_value.delete_rule.side_effect = ( + ResourceNotFoundError + ) + + operator._create_reply_subscription_for_correlation_id(mock_admin_hook, context) + + mock_admin_hook.create_subscription.assert_called_once_with( + topic_name="reply-topic-name", + subscription_name=operator.subscription_name, + default_message_time_to_live="PT1H", # 1 hour + dead_lettering_on_message_expiration=True, + dead_lettering_on_filter_evaluation_exceptions=True, + enable_batched_operations=False, + user_metadata=f"Subscription for reply to {operator.reply_correlation_id} for task ID {context['task'].task_id}", + auto_delete_on_idle="PT6H", # 6 hours + filter_rule_name=operator.subscription_name + operator.REPLY_RULE_SUFFIX, + filter_rule=mock.ANY, + ) + + @mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") + @mock.patch("airflow.providers.microsoft.azure.operators.asb.MessageHook") + def test_execute(self, mock_message_hook, mock_admin_hook): + request_queue_name = "test_queue" + test_body = "test_body" + reply_topic_name = "reply-topic-name" + + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name=request_queue_name, + request_body_generator=lambda context: test_body, + reply_topic_name="reply-topic-name", + ) + + context = mock.MagicMock() + context["task"] = mock.MagicMock() + context["task"].task_id = 837 + + mock_message_hook_instance = mock_message_hook.return_value + mock_admin_hook_instance = mock_admin_hook.return_value + + operator.execute(context) + + # Check if the reply subscription was created + mock_admin_hook_instance.create_subscription.assert_called_once_with( + topic_name="reply-topic-name", + subscription_name=operator.subscription_name, + default_message_time_to_live="PT1H", # 1 hour + dead_lettering_on_message_expiration=True, + dead_lettering_on_filter_evaluation_exceptions=True, + enable_batched_operations=False, + user_metadata=f"Subscription for reply to {operator.reply_correlation_id} for task ID {context['task'].task_id}", + auto_delete_on_idle="PT6H", # 6 hours + filter_rule_name=operator.subscription_name + operator.REPLY_RULE_SUFFIX, + filter_rule=mock.ANY, + ) + + # Check if the request message was sent + mock_message_hook_instance.send_message.assert_called_once_with( + request_queue_name, + test_body, + message_id=operator.reply_correlation_id, + reply_to=reply_topic_name, + message_headers={"reply_type": "topic"}, + ) + + # Check if the reply message was received + mock_message_hook_instance.receive_subscription_message.assert_called_once_with( + "reply-topic-name", + operator.subscription_name, + context, + max_message_count=1, + max_wait_time=60, + message_callback=None, + ) + + # Check if the reply subscription was removed + mock_admin_hook_instance.delete_subscription.assert_called_once_with( + operator.subscription_name, operator.reply_topic_name + ) + + @mock.patch("airflow.providers.microsoft.azure.operators.asb.AdminClientHook") + @mock.patch("airflow.providers.microsoft.azure.operators.asb.MessageHook") + def test_execute_with_exception(self, mock_message_hook, mock_admin_hook): + operator = AzureServiceBusRequestReplyOperator( + task_id="test_task", + request_queue_name="test_queue", + request_body_generator=lambda context: "test_body", + reply_topic_name="reply-topic-name", + ) + + context = mock.MagicMock() + context["task"] = mock.MagicMock() + context["task"].task_id = 123 + + mock_message_hook_instance = mock_message_hook.return_value + mock_admin_hook_instance = mock_admin_hook.return_value + + # Simulate an exception during message sending + mock_message_hook_instance.send_message.side_effect = Exception("Test exception") + + with pytest.raises(Exception, match="Test exception"): + operator.execute(context) + + # Check if the reply subscription was still removed despite the exception + mock_admin_hook_instance.delete_subscription.assert_called_once_with( + operator.subscription_name, operator.reply_topic_name + )