Fix DatabricksWorkflowTaskGroup leaking TaskGroupContext on internal exception#66582
Merged
eladkal merged 6 commits intoMay 8, 2026
Conversation
The previous test raised an AirflowException from inside the workflow body, but super().__exit__ is reached even on the unfixed code in that path (the for-loop over tasks is empty), so the test passed without the fix. Replace it with a regression test that triggers the raise inside DatabricksWorkflowTaskGroup.__exit__ itself (via an EmptyOperator that does not implement _convert_to_databricks_workflow_task) and asserts TaskGroupContext is empty afterward — fails on unfixed code, passes on fixed code. Also drops two unused imports introduced earlier.
2 tasks
arpitrathore
pushed a commit
to arpitrathore/airflow
that referenced
this pull request
May 9, 2026
…l exception (apache#66582) * Removing resources properly for DatabricksWorkflowGroup * Add the unit test to check that super().__exit()__ called even exception happen * remove unnecessary comment * Fix unit tests * Fix unit tests * Strengthen test to actually exercise the bug The previous test raised an AirflowException from inside the workflow body, but super().__exit__ is reached even on the unfixed code in that path (the for-loop over tasks is empty), so the test passed without the fix. Replace it with a regression test that triggers the raise inside DatabricksWorkflowTaskGroup.__exit__ itself (via an EmptyOperator that does not implement _convert_to_databricks_workflow_task) and asserts TaskGroupContext is empty afterward — fails on unfixed code, passes on fixed code. Also drops two unused imports introduced earlier. --------- Co-authored-by: artsiomyudovin <a.yudovin6695@gmail.com>
jason810496
pushed a commit
to jason810496/airflow
that referenced
this pull request
May 11, 2026
…l exception (apache#66582) * Removing resources properly for DatabricksWorkflowGroup * Add the unit test to check that super().__exit()__ called even exception happen * remove unnecessary comment * Fix unit tests * Fix unit tests * Strengthen test to actually exercise the bug The previous test raised an AirflowException from inside the workflow body, but super().__exit__ is reached even on the unfixed code in that path (the for-loop over tasks is empty), so the test passed without the fix. Replace it with a regression test that triggers the raise inside DatabricksWorkflowTaskGroup.__exit__ itself (via an EmptyOperator that does not implement _convert_to_databricks_workflow_task) and asserts TaskGroupContext is empty afterward — fails on unfixed code, passes on fixed code. Also drops two unused imports introduced earlier. --------- Co-authored-by: artsiomyudovin <a.yudovin6695@gmail.com>
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
3 participants
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.
DatabricksWorkflowTaskGroup.exit can raise mid-execution (for example when a child task does not implement
_convert_to_databricks_workflow_task). Without try/finally, super().exit never runs and the workflow group is left on the global TaskGroupContext stack, breaking later DAGs withCannot mix TaskGroups from different DAGs.Wrap the body of
__exit__in try/finally so super().exit is always invoked and the context stack is cleaned up.Picks up @ayudovin's earlier work in #55891 (closed by author with explicit permission to take over) and adds a regression test that exercises the bug path directly: an EmptyOperator (which lacks
_convert_to_databricks_workflow_task) is added inside the workflow body, which makes__exit__raise mid-way; the test then asserts thatTaskGroupContext._contextis empty, which fails on unfixed code and passes on fixed code.closes: #42164
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 4.7) following the guidelines
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.