(airflow) branch main updated: Fix dispose_orm() not disposing async engine on shutdown (#65274)

[ra...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

rahulvats pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 69a88bfecc6 Fix dispose_orm() not disposing async engine on shutdown (#65274)
69a88bfecc6 is described below

commit 69a88bfecc63ff336070529c8443c61315e2bdef
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Wed Apr 15 07:00:52 2026 +0100

    Fix dispose_orm() not disposing async engine on shutdown (#65274)
    
    When _configure_async_session() was extracted from configure_orm() in
    PR #51920, the corresponding cleanup in dispose_orm() was not updated.
    This left async_engine connections abandoned on process exit, gunicorn
    worker restarts, and atexit -- gradually exhausting PostgreSQL's
    max_connections.
    
    Dispose async_engine.sync_engine (the synchronous path, matching
    the existing clean_in_fork pattern) and clear both async_engine and
    AsyncSession references.
---
 airflow-core/src/airflow/settings.py          | 13 ++++-
 airflow-core/tests/unit/core/test_settings.py | 71 +++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/airflow-core/src/airflow/settings.py b/airflow-core/src/airflow/settings.py
index a00339603be..bf47f86a2c2 100644
--- a/airflow-core/src/airflow/settings.py
+++ b/airflow-core/src/airflow/settings.py
@@ -567,10 +567,14 @@ def prepare_engine_args(disable_connection_pool=False, pool_class=None):
 
 def dispose_orm(do_log: bool = True):
     """Properly close pooled database connections."""
-    global Session, engine, NonScopedSession
+    global Session, engine, NonScopedSession, async_engine, AsyncSession
 
     _globals = globals()
-    if _globals.get("engine") is None and _globals.get("Session") is None:
+    if (
+        _globals.get("engine") is None
+        and _globals.get("Session") is None
+        and _globals.get("async_engine") is None
+    ):
         return
 
     if do_log:
@@ -588,6 +592,11 @@ def dispose_orm(do_log: bool = True):
         engine.dispose()
         engine = None
 
+    if "async_engine" in _globals and async_engine is not None:
+        async_engine.sync_engine.dispose()
+        async_engine = None
+        AsyncSession = None
+
 
 def reconfigure_orm(disable_connection_pool=False, pool_class=None):
     """Properly close database connections and re-configure ORM."""
diff --git a/airflow-core/tests/unit/core/test_settings.py b/airflow-core/tests/unit/core/test_settings.py
index a865094a574..22cb00b36bb 100644
--- a/airflow-core/tests/unit/core/test_settings.py
+++ b/airflow-core/tests/unit/core/test_settings.py
@@ -25,7 +25,10 @@ from unittest import mock
 from unittest.mock import MagicMock, call, patch
 
 import pytest
+from sqlalchemy.engine import Engine
+from sqlalchemy.ext.asyncio import AsyncEngine
 
+from airflow import settings
 from airflow.exceptions import AirflowClusterPolicyViolation, AirflowConfigException
 
 SETTINGS_FILE_CUSTOM_ENGINE = """
@@ -370,3 +373,71 @@ def test_sqlite_relative_path(value, expectation):
     ):
         with expectation:
             settings.configure_orm()
+
+
+class TestDisposeOrm:
+    """Tests for dispose_orm() async engine disposal."""
+
+    def setup_method(self):
+        self._orig = {
+            "engine": settings.engine,
+            "Session": settings.Session,
+            "NonScopedSession": settings.NonScopedSession,
+            "async_engine": settings.async_engine,
+            "AsyncSession": settings.AsyncSession,
+        }
+
+    def teardown_method(self):
+        for attr, val in self._orig.items():
+            setattr(settings, attr, val)
+
+    def test_disposes_async_engine(self):
+        """dispose_orm() must dispose the async engine and clear AsyncSession."""
+        mock_sync_engine = MagicMock(spec=Engine)
+        mock_async_engine = MagicMock(spec=AsyncEngine)
+        mock_async_engine.sync_engine = mock_sync_engine
+
+        settings.engine = None
+        settings.Session = None
+        settings.NonScopedSession = None
+        settings.async_engine = mock_async_engine
+        settings.AsyncSession = MagicMock()
+
+        settings.dispose_orm(do_log=False)
+
+        mock_sync_engine.dispose.assert_called_once()
+        assert settings.async_engine is None
+        assert settings.AsyncSession is None
+
+    def test_disposes_both_sync_and_async(self):
+        """dispose_orm() disposes both engines when both are set."""
+        mock_engine = MagicMock(spec=Engine)
+        mock_sync_engine = MagicMock(spec=Engine)
+        mock_async_engine = MagicMock(spec=AsyncEngine)
+        mock_async_engine.sync_engine = mock_sync_engine
+
+        settings.engine = mock_engine
+        settings.Session = MagicMock()
+        settings.NonScopedSession = MagicMock()
+        settings.async_engine = mock_async_engine
+        settings.AsyncSession = MagicMock()
+
+        with patch("sqlalchemy.orm.session.close_all_sessions"):
+            settings.dispose_orm(do_log=False)
+
+        mock_engine.dispose.assert_called_once()
+        mock_sync_engine.dispose.assert_called_once()
+        assert settings.engine is None
+        assert settings.async_engine is None
+        assert settings.AsyncSession is None
+
+    def test_early_return_when_all_none(self):
+        """dispose_orm() returns early without side effects when nothing is configured."""
+        settings.engine = None
+        settings.Session = None
+        settings.async_engine = None
+
+        with patch("sqlalchemy.orm.session.close_all_sessions") as mock_close:
+            settings.dispose_orm(do_log=False)
+
+        mock_close.assert_not_called()