This is an automated email from the ASF dual-hosted git repository.
kaxil 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 afb1ce5396b Fix ``DbApiHook.inspector`` keeping a database connection
open until process exit (#75327)
afb1ce5396b is described below
commit afb1ce5396b6f78ceac14a67cd098541b4c7ed7c
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 21 22:34:22 2026 +0100
Fix ``DbApiHook.inspector`` keeping a database connection open until
process exit (#75327)
* Fix DbApiHook.inspector keeping a database connection open until process
exit
DbApiHook.inspector is cached for the hook's lifetime, and SQLAlchemy's
default QueuePool keeps the DBAPI connection the inspector reflected with
open in the pool. Nothing disposes that engine, so the connection stays
open until the interpreter exits. Build the inspector's engine with
NullPool so each inspection closes its connection, as run() and
get_records() already do.
* Reword DbApiHook.inspector NullPool comment
---
.../src/airflow/providers/common/sql/hooks/sql.py | 7 ++++-
.../sql/tests/unit/common/sql/hooks/test_sql.py | 33 ++++++++++++++++++++++
2 files changed, 39 insertions(+), 1 deletion(-)
diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
index 3e30501cb8f..4f8ec0748cc 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py
@@ -35,10 +35,12 @@ try:
from sqlalchemy import create_engine, inspect
from sqlalchemy.engine import make_url
from sqlalchemy.exc import ArgumentError, NoSuchModuleError
+ from sqlalchemy.pool import NullPool
except ImportError:
create_engine = None # type: ignore[assignment]
inspect = None # type: ignore[assignment]
make_url = None # type: ignore[assignment]
+ NullPool = None # type: ignore[assignment,misc]
ArgumentError = Exception # type: ignore[misc,assignment]
NoSuchModuleError = Exception # type: ignore[misc,assignment]
@@ -378,7 +380,10 @@ class DbApiHook(BaseHook):
"SQLAlchemy is required for database inspection. "
"Install it with: pip install
'apache-airflow-providers-common-sql[sqlalchemy]'"
)
- return inspect(self.get_sqlalchemy_engine())
+ # The inspector is cached for the hook's lifetime, so a pooled engine
would keep the
+ # connection it reflected with open until the process exits. With
NullPool, no
+ # connection outlives the inspection that opened it.
+ return inspect(self.get_sqlalchemy_engine({"poolclass": NullPool}))
def get_table_schema(self, table_name: str, schema: str | None = None) ->
list[dict[str, str]]:
"""
diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py
b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py
index c2fbfec3cbb..3207721367a 100644
--- a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py
+++ b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py
@@ -20,11 +20,14 @@ from __future__ import annotations
import inspect
import logging
+import sqlite3
+from contextlib import closing
from unittest.mock import MagicMock, PropertyMock, patch
import pandas as pd
import polars as pl
import pytest
+from sqlalchemy import event
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import Connection
@@ -470,3 +473,33 @@ def test_inspector_is_cached():
inspector2 = hook.inspector
assert inspector1 is inspector2
mock_get_engine.assert_called_once()
+
+
+def
test_inspector_closes_its_dbapi_connections_after_each_inspection(tmp_path):
+ """A cached inspector must not hold a pooled DBAPI connection open for the
hook's lifetime."""
+ db_path = tmp_path / "inspect.db"
+ with closing(sqlite3.connect(db_path)) as conn:
+ conn.execute("CREATE TABLE users (id INTEGER, name TEXT)")
+ hook = DBApiHookForTests(conn_id=DEFAULT_CONN_ID)
+ opened, closed = [], []
+ real_get_sqlalchemy_engine = hook.get_sqlalchemy_engine
+
+ def get_sqlalchemy_engine_with_listeners(engine_kwargs=None):
+ engine = real_get_sqlalchemy_engine(engine_kwargs)
+ event.listen(engine, "connect", lambda dbapi_conn, _record:
opened.append(dbapi_conn))
+ event.listen(engine, "close", lambda dbapi_conn, _record:
closed.append(dbapi_conn))
+ return engine
+
+ with (
+ patch.object(hook, "get_uri", autospec=True,
return_value=f"sqlite:///{db_path}"),
+ patch.object(
+ hook, "get_sqlalchemy_engine", autospec=True,
side_effect=get_sqlalchemy_engine_with_listeners
+ ),
+ ):
+ columns = hook.get_table_schema("users")
+ table_names = hook.inspector.get_table_names()
+
+ assert columns == [{"name": "id", "type": "INTEGER"}, {"name": "name",
"type": "TEXT"}]
+ assert table_names == ["users"]
+ assert opened
+ assert closed == opened
