Visitar URL original
Correctly shutdown async DB engine on exit. · apache/airflow@d69e81c · GitHub
Skip to content

Commit d69e81c

Browse files
ashbashb
committed
Correctly shutdown async DB engine on exit.
This isn't "a problem" per se, as the process is about to exit anyway, but this does end up with a confusing/scary looking message in the API server logs of: Traceback (most recent call last): File "/usr/python/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 375, in _close_connection self._dialect.do_close(connection) File "/usr/python/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 721, in do_close dbapi_connection.close() File "/usr/python/lib/python3.12/site-packages/sqlalchemy/dialects/sqlite/aiosqlite.py", line 362, in close self._handle_exception(error) File "/usr/python/lib/python3.12/site-packages/sqlalchemy/dialects/sqlite/aiosqlite.py", line 373, in _handle_exception raise error File "/usr/python/lib/python3.12/site-packages/sqlalchemy/dialects/sqlite/aiosqlite.py", line 350, in close self.await_(self._connection.close()) File "/usr/python/lib/python3.12/site-packages/sqlalchemy/util/_concurrency_py3k.py", line 123, in await_only raise exc.MissingGreenlet( sqlalchemy.exc.MissingGreenlet: greenlet_spawn has not been called; can't call await_only() here. Was IO attempted in an unexpected place? (Background on this error at: https://sqlalche.me/e/20/xd2s) This also appears in some unit tests (though only visible when a test fails, as otherwise logs don't get shown), hence the unit test fixture changes
1 parent 489316f commit d69e81c

15 files changed

Lines changed: 418 additions & 213 deletions

File tree

‎airflow-core/src/airflow/api_fastapi/app.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from fastapi.routing import Mount
2828
from starlette.middleware import Middleware
2929

30+
from airflow import settings
3031
from airflow.api_fastapi.common.dagbag import create_dag_bag
3132
from airflow.api_fastapi.common.exceptions import init_error_handlers
3233
from airflow.api_fastapi.common.http_access_log import HttpAccessLogMiddleware
@@ -100,6 +101,9 @@ def _initialize_api_server_stats() -> None:
100101
async def lifespan(app: FastAPI):
101102
_initialize_api_server_stats()
102103
async with AsyncExitStack() as stack:
104+
stack.push_async_callback(settings.dispose_async_orm)
105+
if settings.async_engine is None:
106+
settings._configure_async_session()
103107
for route in app.routes:
104108
if isinstance(route, Mount) and isinstance(route.app, FastAPI):
105109
await stack.enter_async_context(

‎airflow-core/src/airflow/settings.py‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -427,13 +427,7 @@ def create_async_metadata_engine(
427427

428428

429429
def _configure_async_session() -> None:
430-
"""
431-
Configure async SQLAlchemy session.
432-
433-
This exists so tests can reconfigure the session. How SQLAlchemy configures
434-
this does not work well with Pytest and you can end up with issues when the
435-
session and runs in a different event loop from the test itself.
436-
"""
430+
"""Configure the async engine and session factory."""
437431
global AsyncSession, async_engine
438432

439433
if not SQL_ALCHEMY_CONN_ASYNC:
@@ -670,6 +664,16 @@ def dispose_orm(do_log: bool = True):
670664
AsyncSession = None
671665

672666

667+
async def dispose_async_orm() -> None:
668+
"""Dispose the async pool on the event loop that owns its connections."""
669+
global async_engine, AsyncSession
670+
671+
if async_engine is not None:
672+
await async_engine.dispose()
673+
async_engine = None
674+
AsyncSession = None
675+
676+
673677
def reconfigure_orm(disable_connection_pool=False, pool_class=None):
674678
"""Properly close database connections and re-configure ORM."""
675679
dispose_orm()

‎airflow-core/tests/unit/api_fastapi/auth/managers/simple/conftest.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,4 +75,5 @@ def test_client():
7575
): "airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager"
7676
}
7777
):
78-
return TestClient(create_app("core"))
78+
with TestClient(create_app("core")) as client:
79+
yield client

‎airflow-core/tests/unit/api_fastapi/conftest.py‎

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from fastapi.routing import Mount
2828
from fastapi.testclient import TestClient
2929

30+
from airflow import settings
3031
from airflow.api_fastapi.app import create_app
3132
from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser
3233
from airflow.dag_processing.bundles.manager import DagBundlesManager
@@ -48,6 +49,16 @@
4849
BASE_URL = "http://testserver"
4950

5051

52+
@pytest.fixture(autouse=True)
53+
def isolate_async_orm(monkeypatch, pytestconfig):
54+
monkeypatch.setattr(settings, "async_engine", None)
55+
monkeypatch.setattr(settings, "AsyncSession", None)
56+
if pytestconfig.getoption("skip_db_tests"):
57+
monkeypatch.setattr(
58+
settings, "_configure_async_session", mock.create_autospec(settings._configure_async_session)
59+
)
60+
61+
5162
def get_api_path(request):
5263
"""Determine the API path based on the test's subdirectory."""
5364
test_dir = os.path.dirname(request.path)
@@ -133,11 +144,12 @@ def _authed_test_client(app: FastAPI, request):
133144
),
134145
)
135146
with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False):
136-
yield TestClient(
147+
with TestClient(
137148
app,
138149
headers={"Authorization": f"Bearer {token}"},
139150
base_url=f"{BASE_URL}{get_api_path(request)}",
140-
)
151+
) as test_client:
152+
yield test_client
141153

142154

143155
@pytest.fixture
@@ -167,22 +179,28 @@ def fresh_test_client(request):
167179

168180
@pytest.fixture
169181
def unauthenticated_test_client(request, _isolated_shared_app):
170-
return TestClient(_isolated_shared_app, base_url=f"{BASE_URL}{get_api_path(request)}")
182+
with TestClient(_isolated_shared_app, base_url=f"{BASE_URL}{get_api_path(request)}") as test_client:
183+
yield test_client
171184

172185

173186
@pytest.fixture
174-
def unauthorized_test_client(request, _isolated_shared_app):
175-
app = _isolated_shared_app
176-
auth_manager: SimpleAuthManager = app.state.auth_manager
187+
def unauthorized_headers(_isolated_shared_app):
188+
auth_manager: SimpleAuthManager = _isolated_shared_app.state.auth_manager
177189
token = auth_manager._get_token_signer().generate(
178190
auth_manager.serialize_user(SimpleAuthManagerUser(username="dummy", role=None))
179191
)
192+
return {"Authorization": f"Bearer {token}"}
193+
194+
195+
@pytest.fixture
196+
def unauthorized_test_client(request, _isolated_shared_app, unauthorized_headers):
180197
with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False):
181-
yield TestClient(
182-
app,
183-
headers={"Authorization": f"Bearer {token}"},
198+
with TestClient(
199+
_isolated_shared_app,
200+
headers=unauthorized_headers,
184201
base_url=f"{BASE_URL}{get_api_path(request)}",
185-
)
202+
) as test_client:
203+
yield test_client
186204

187205

188206
@pytest.fixture

‎airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import jwt
2424
import pytest
25+
from fastapi.testclient import TestClient
2526

2627
from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
2728
from airflow.models.revoked_token import RevokedToken
@@ -169,10 +170,7 @@ def cleanup_revoked_tokens(self):
169170
clear_db_revoked_tokens()
170171

171172
@pytest.fixture
172-
def logout_client(self):
173-
"""A test client without the is_revoked mock so revocation tests hit the real DB."""
174-
from fastapi.testclient import TestClient
175-
173+
def logout_app(self):
176174
from airflow.api_fastapi.app import create_app
177175

178176
with conf_vars(
@@ -183,8 +181,13 @@ def logout_client(self):
183181
): "airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager"
184182
}
185183
):
186-
app = create_app()
187-
yield TestClient(app, base_url="http://testserver/api/v2")
184+
yield create_app()
185+
186+
@pytest.fixture
187+
def logout_client(self, logout_app):
188+
"""A test client without the is_revoked mock so revocation tests hit the real DB."""
189+
with TestClient(logout_app, base_url="http://testserver/api/v2") as client:
190+
yield client
188191

189192
def test_logout_revokes_token(self, logout_client):
190193
"""Test that logout revokes the JWT token and persists it in the database."""
@@ -282,7 +285,7 @@ def test_logout_revokes_every_credential_presented(self, logout_client):
282285
assert RevokedToken.is_revoked("test-jti-both-bearer") is True
283286
assert RevokedToken.is_revoked("test-jti-both-cookie") is True
284287

285-
def test_logout_revokes_both_even_when_a_trusted_user_is_cached(self, logout_client):
288+
def test_logout_revokes_both_even_when_a_trusted_user_is_cached(self, logout_app):
286289
"""The trusted-middleware shortcut must not change what logout revokes.
287290
288291
On protected routes `get_user()` can return a user cached by JWTRefreshMiddleware
@@ -291,7 +294,7 @@ def test_logout_revokes_both_even_when_a_trusted_user_is_cached(self, logout_cli
291294
"""
292295
from airflow.api_fastapi.core_api.security import USER_INJECTED_BY_TRUSTED_MIDDLEWARE
293296

294-
auth_manager = logout_client.app.state.auth_manager
297+
auth_manager = logout_app.state.auth_manager
295298
bearer_token = self._mint(auth_manager, "test-jti-trusted-bearer")
296299
cookie_token = self._mint(auth_manager, "test-jti-trusted-cookie")
297300

@@ -300,9 +303,12 @@ async def _inject(request, call_next):
300303
request.state.user_authenticated_via = USER_INJECTED_BY_TRUSTED_MIDDLEWARE
301304
return await call_next(request)
302305

303-
logout_client.app.middleware("http")(_inject)
304-
logout_client.cookies.set(COOKIE_NAME_JWT_TOKEN, cookie_token)
305-
with patch.object(auth_manager, "get_url_logout", return_value=None):
306+
logout_app.middleware("http")(_inject)
307+
with (
308+
TestClient(logout_app, base_url="http://testserver/api/v2") as logout_client,
309+
patch.object(auth_manager, "get_url_logout", return_value=None),
310+
):
311+
logout_client.cookies.set(COOKIE_NAME_JWT_TOKEN, cookie_token)
306312
response = logout_client.get(
307313
"/auth/logout",
308314
headers={"Authorization": f"Bearer {bearer_token}"},

‎airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323

2424
import pendulum
2525
import pytest
26-
from fastapi.testclient import TestClient
2726
from sqlalchemy import and_, func, select
2827
from sqlalchemy.exc import OperationalError, ProgrammingError
2928

@@ -83,18 +82,13 @@ def clean_db():
8382

8483

8584
@pytest.fixture
86-
def dag_reader_test_client(test_client):
85+
def dag_reader_headers(test_client):
8786
"""A caller who may read the Dags but not write them: viewer is below the role edits require."""
8887
auth_manager = test_client.app.state.auth_manager
8988
token = auth_manager._get_token_signer().generate(
9089
auth_manager.serialize_user(SimpleAuthManagerUser(username="reader", role="viewer"))
9190
)
92-
with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False):
93-
yield TestClient(
94-
test_client.app,
95-
headers={"Authorization": f"Bearer {token}"},
96-
base_url=str(test_client.base_url),
97-
)
91+
return {"Authorization": f"Bearer {token}"}
9892

9993

10094
def make_dags():
@@ -1599,22 +1593,24 @@ def test_pause_backfill_404_when_the_dag_is_unreadable(self, session, unauthoriz
15991593
response = unauthorized_test_client.put(f"/backfills/{backfill.id}/pause")
16001594
assert response.status_code == 404
16011595

1602-
def test_pause_backfill_403(self, session, dag_reader_test_client):
1596+
def test_pause_backfill_403(self, session, dag_reader_headers, test_client):
16031597
(dag,) = self._create_dag_models()
16041598
from_date = timezone.utcnow()
16051599
to_date = timezone.utcnow()
16061600
backfill = Backfill(dag_id=dag.dag_id, from_date=from_date, to_date=to_date)
16071601
session.add(backfill)
16081602
session.commit()
1609-
response = dag_reader_test_client.put(f"/backfills/{backfill.id}/pause")
1603+
response = test_client.put(f"/backfills/{backfill.id}/pause", headers=dag_reader_headers)
16101604
assert response.status_code == 403
16111605

16121606
def test_pause_backfill_unknown_id_is_not_authorized_by_a_body_dag_id(
1613-
self, session, dag_reader_test_client
1607+
self, session, dag_reader_headers, test_client
16141608
):
16151609
(dag,) = self._create_dag_models()
16161610
session.commit()
1617-
response = dag_reader_test_client.put(f"/backfills/{231984098}/pause", json={"dag_id": dag.dag_id})
1611+
response = test_client.put(
1612+
f"/backfills/{231984098}/pause", json={"dag_id": dag.dag_id}, headers=dag_reader_headers
1613+
)
16181614
assert response.status_code == 404
16191615
assert response.json().get("detail") == "Backfill not found"
16201616

‎airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1732,16 +1732,15 @@ def test_get_status_requires_token_header(self, test_client):
17321732
assert response.status_code == 422
17331733

17341734
@mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
1735-
def test_get_status_unauthorized_user_does_not_leak_row(
1736-
self, test_client, unauthorized_test_client, session
1737-
):
1735+
def test_get_status_unauthorized_user_does_not_leak_row(self, test_client, unauthorized_headers, session):
17381736
"""A user without rights on the conn_id never sees the row payload via GET-by-token."""
17391737
post_response = test_client.post("/connections/enqueue-test", json=self.TEST_REQUEST_BODY)
17401738
assert post_response.status_code == 202
17411739
token = post_response.json()["token"]
17421740

1743-
response = unauthorized_test_client.get(
1744-
"/connections/enqueue-test", headers={"Airflow-Connection-Test-Token": token}
1741+
response = test_client.get(
1742+
"/connections/enqueue-test",
1743+
headers={**unauthorized_headers, "Airflow-Connection-Test-Token": token},
17451744
)
17461745
assert response.status_code in (401, 403, 404)
17471746
body = (

0 commit comments

Comments
 (0)