Lee-W commented on code in PR #41601:
URL: https://github.com/apache/airflow/pull/41601#discussion_r1738376359
##########
airflow/providers/microsoft/azure/hooks/asb.py:
##########
@@ -326,5 +336,27 @@ def receive_subscription_message(
max_message_count=max_message_count,
max_wait_time=max_wait_time
)
for msg in received_msgs:
- self.log.info(msg)
- subscription_receiver.complete_message(msg)
+ self._process_message(msg, message_callback,
subscription_receiver)
+
+ def _process_message(self, msg, message_callback, receiver):
Review Comment:
Could we add type annotation here? Thanks!
##########
airflow/providers/microsoft/azure/hooks/asb.py:
##########
@@ -270,14 +273,21 @@ def send_batch_message(sender: ServiceBusSender,
messages: list[str]):
sender.send_messages(batch_message)
def receive_message(
- self, queue_name, max_message_count: int | None = 1, max_wait_time:
float | None = None
+ self,
+ queue_name,
Review Comment:
I guess this should be a str?
```suggestion
queue_name: str,
```
##########
airflow/providers/microsoft/azure/hooks/asb.py:
##########
@@ -326,5 +336,27 @@ def receive_subscription_message(
max_message_count=max_message_count,
max_wait_time=max_wait_time
)
for msg in received_msgs:
- self.log.info(msg)
- subscription_receiver.complete_message(msg)
+ self._process_message(msg, message_callback,
subscription_receiver)
+
+ def _process_message(self, msg, message_callback, receiver):
+ """
+ Process the message by calling the message_callback or logging the
message.
+
+ :param msg: The message to process.
+ :param message_callback: Optional callback to process each message. If
not provided, then
+ the message will be logged and completed. If provided, and throws
an exception, the
+ message will be abandoned for future redelivery.
+ :param receiver: The receiver that received the message.
+ """
+ print("message_callback:", message_callback)
+ if message_callback is None:
+ self.log.info(msg)
+ receiver.complete_message(msg)
+ else:
+ try:
+ message_callback(msg)
+ receiver.complete_message(msg)
+ except Exception as e:
+ self.log.error("Error processing message: %s", e)
+ receiver.abandon_message(msg)
+ raise e
Review Comment:
I guess we don't expect `message_callback` to raise an exception here
```suggestion
try:
message_callback(msg)
except Exception as e:
self.log.error("Error processing message: %s", e)
receiver.abandon_message(msg)
raise e
else:
receiver.complete_message(msg)
```
##########
airflow/providers/microsoft/azure/hooks/asb.py:
##########
@@ -326,5 +336,27 @@ def receive_subscription_message(
max_message_count=max_message_count,
max_wait_time=max_wait_time
)
for msg in received_msgs:
- self.log.info(msg)
- subscription_receiver.complete_message(msg)
+ self._process_message(msg, message_callback,
subscription_receiver)
+
+ def _process_message(self, msg, message_callback, receiver):
+ """
+ Process the message by calling the message_callback or logging the
message.
+
+ :param msg: The message to process.
+ :param message_callback: Optional callback to process each message. If
not provided, then
+ the message will be logged and completed. If provided, and throws
an exception, the
+ message will be abandoned for future redelivery.
+ :param receiver: The receiver that received the message.
+ """
+ print("message_callback:", message_callback)
Review Comment:
We probably would like to use logging instead
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
