Add cleanup to Kafka AwaitMessageTrigger for consumer management#64612
Open
jason810496 wants to merge 2 commits intoapache:mainfrom
Open
Add cleanup to Kafka AwaitMessageTrigger for consumer management#64612jason810496 wants to merge 2 commits intoapache:mainfrom
jason810496 wants to merge 2 commits intoapache:mainfrom
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds explicit Kafka consumer cleanup to AwaitMessageTrigger to prevent lingering consumers (and related max.poll.interval.ms issues) when triggers finish or are removed by the triggerer.
Changes:
- Store the created Kafka consumer on the trigger instance for later cleanup.
- Add a
cleanup()implementation that closes the consumer. - Add unit tests covering consumer close on cleanup and no-op cleanup when no consumer exists.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/await_message.py | Persist consumer reference and add cleanup() to close it. |
| providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py | Add unit tests validating cleanup behavior and extend mocked consumer with close(). |
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/await_message.py
Outdated
Show resolved
Hide resolved
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/await_message.py
Show resolved
Hide resolved
Lee-W
reviewed
Apr 2, 2026
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/await_message.py
Outdated
Show resolved
Hide resolved
providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py
Outdated
Show resolved
Hide resolved
providers/apache/kafka/tests/unit/apache/kafka/triggers/test_await_message.py
Outdated
Show resolved
Hide resolved
jason810496
force-pushed
the
providers/apache-kafka/add-proper-cleanup-for-kafka-await-message
branch
from
a3fd51c to
ab34e72
Compare
Lee-W
reviewed
Apr 7, 2026
| self._consumer = None | ||
| try: | ||
| await sync_to_async(consumer.close)() | ||
| except Exception: |
Member
There was a problem hiding this comment.
is it possible to narrow the exception?
Comment on lines
+114
to
+117
| 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) |
Comment on lines
+149
to
+156
| 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) |
| try: | ||
| await sync_to_async(consumer.close)() | ||
| except Exception: | ||
| log.warning("Failed to close Kafka consumer", exc_info=True) |
Comment on lines
+164
to
+167
| generator = trigger.run() | ||
| await generator.__anext__() | ||
| await trigger.cleanup() | ||
| await generator.aclose() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
When testing out Kafka external-event driven Dag for ##64205 scenario, I encounter the following error with only enabling one consumer Dag
What
I verify locally that adding proper resource cleanup for Kafka Consumer can avoid the above error.
Was generative AI tooling used to co-author this PR?