Add request-reply operator to Microsoft Azure provider

docs/spelling_wordlist.txt

@@ -790,6 +790,7 @@ hiveserver
 hmsclient
 Hoc
 hoc
+Hohpe
 homebrew
 honoured
 hookable

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(

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<howto/connection:azure_service_bus>`.
+    :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,21 +123,29 @@ 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)
         self.queue_name = queue_name
         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."""
         # Create the hook
         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<howto/connection:azure_service_bus>`.
+    :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)