Add cleanup to Kafka AwaitMessageTrigger for consumer management

providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/await_message.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from collections.abc import Sequence
 from functools import partial
 from typing import Any
@@ -34,6 +35,8 @@
 else:
     from airflow.triggers.base import BaseTrigger as BaseEventTrigger  # type: ignore
 
+log = logging.getLogger(__name__)
+
 
 class AwaitMessageTrigger(BaseEventTrigger):
     """
@@ -87,6 +90,7 @@ def __init__(
         self.poll_timeout = poll_timeout
         self.poll_interval = poll_interval
         self.commit_offset = commit_offset
+        self._consumer = None
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
@@ -107,10 +111,10 @@ async def run(self):
         consumer_hook = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id)
 
         async_get_consumer = sync_to_async(consumer_hook.get_consumer)
-        consumer = await async_get_consumer()
+        self._consumer = await async_get_consumer()
 
-        async_poll = sync_to_async(consumer.poll)
-        async_commit = sync_to_async(consumer.commit)
+        async_poll = sync_to_async(self._consumer.poll)
+        async_commit = sync_to_async(self._consumer.commit)
 
         async_message_process = None
         if self.apply_function:
@@ -141,3 +145,12 @@ async def run(self):
                     if self.commit_offset:
                         await async_commit(message=message, asynchronous=False)
                     await asyncio.sleep(self.poll_interval)
+
+    async def cleanup(self) -> None:
+        consumer = self._consumer
+        if consumer is not None:
+            self._consumer = None
+            try:
+                await sync_to_async(consumer.close)()
+            except Exception:
+                log.warning("Failed to close Kafka consumer", exc_info=True)

providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py

@@ -18,6 +18,7 @@
 
 import asyncio
 import json
+from unittest.mock import call
 
 import pytest
 
@@ -59,6 +60,9 @@ def poll(*args, **kwargs):
     def commit(*args, **kwargs):
         return True
 
+    def close(*args, **kwargs):
+        return None
+
 
 class TestTrigger:
     @pytest.fixture(autouse=True)
@@ -142,6 +146,40 @@ async def test_trigger_run_bad(self, mocker):
         assert task.done() is False
         asyncio.get_event_loop().stop()
 
+    @pytest.mark.asyncio
+    async def test_cleanup_closes_consumer(self, mocker):
+        consumer = MockedConsumer()
+        close_mock = mocker.patch.object(consumer, "close")
+
+        mocker.patch.object(KafkaConsumerHook, "get_consumer", return_value=consumer)
+
+        trigger = AwaitMessageTrigger(
+            kafka_config_id="kafka_d",
+            apply_function="unit.apache.kafka.triggers.test_await_message.apply_function_true",
+            topics=["noop"],
+            poll_timeout=0.0001,
+            poll_interval=5,
+        )
+
+        generator = trigger.run()
+        await generator.__anext__()
+        await trigger.cleanup()
+        await generator.aclose()
+
+        assert close_mock.mock_calls == [call()]
+
+    @pytest.mark.asyncio
+    async def test_cleanup_does_not_raise_without_consumer(self):
+        trigger = AwaitMessageTrigger(
+            kafka_config_id="kafka_d",
+            apply_function="unit.apache.kafka.triggers.test_await_message.apply_function_true",
+            topics=["noop"],
+            poll_timeout=0.0001,
+            poll_interval=5,
+        )
+
+        await trigger.cleanup()
+
 
 @mark_common_msg_queue_test
 class TestMessageQueueTrigger: