This is an automated email from the ASF dual-hosted git repository.
rahulvats pushed a commit to branch v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-2-test by this push:
new b7b4ac6dcd1 Fix dispose_orm() not disposing async engine on shutdown
(#65274) (#65284)
b7b4ac6dcd1 is described below
commit b7b4ac6dcd1f24bfc5e577d6021c3348dda298b1
Author: Rahul Vats <[email protected]>
AuthorDate: Wed Apr 15 13:34:22 2026 +0530
Fix dispose_orm() not disposing async engine on shutdown (#65274) (#65284)
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.
(cherry picked from commit 69a88bfecc63ff336070529c8443c61315e2bdef)
Co-authored-by: Kaxil Naik <[email protected]>
---
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 223ec291b4d..b9889657f29 100644
--- a/airflow-core/src/airflow/settings.py
+++ b/airflow-core/src/airflow/settings.py
@@ -531,10 +531,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:
@@ -552,6 +556,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 24a5485402d..16f28c38cb2 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_POLICY = """
@@ -228,3 +231,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()
