This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch changes-3.4.0rc3 in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 440af3f87d23037d731a49b0675fec875eb8d0d1 Author: Jason(Zhe-You) Liu <[email protected]> AuthorDate: Thu Nov 5 00:20:10 2026 +0800 [v3-4-test] Keep the audit log off the API server event loop (#76416) (#76431) (cherry picked from commit 575d8af602f3a885b2f976e88ca9cb15d33da326) --- .../src/airflow/api_fastapi/logging/decorators.py | 36 +++++++++-------- .../tests/unit/api_fastapi/core_api/conftest.py | 38 ++++++++++++++++++ .../core_api/routes/public/test_execution.py | 1 + .../unit/api_fastapi/logging/test_decorators.py | 45 +++++++++++++++++----- .../src/tests_common/test_utils/asserts.py | 12 ++++-- 5 files changed, 105 insertions(+), 27 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/logging/decorators.py b/airflow-core/src/airflow/api_fastapi/logging/decorators.py index 03b01f87d75..3fb0ccfe14a 100644 --- a/airflow-core/src/airflow/api_fastapi/logging/decorators.py +++ b/airflow-core/src/airflow/api_fastapi/logging/decorators.py @@ -29,7 +29,7 @@ from pendulum.parsing.exceptions import ParserError from sqlalchemy import select from airflow._shared.secrets_masker import secrets_masker -from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.db.common import AsyncSessionDep from airflow.api_fastapi.core_api.security import GetUserDep from airflow.configuration import conf from airflow.models import Connection, Log, Pool, Variable @@ -40,7 +40,7 @@ from airflow.models.team import find_invalid_team_names if TYPE_CHECKING: from collections.abc import Mapping - from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession logger = logging.getLogger(__name__) @@ -157,8 +157,8 @@ def _mask_variable_entity(extra_fields): return result -def _find_addressed_attempt( - path_params: Mapping[str, str], params: Mapping, *, session: Session +async def _find_addressed_attempt( + path_params: Mapping[str, str], params: Mapping, *, session: AsyncSession ) -> UUID | None: """ Return the live attempt of the one task instance a request addresses, or None if it names no single one. @@ -176,11 +176,15 @@ def _find_addressed_attempt( region_id=UUID(str(region_id)) if region_id else None, region_index=int(region_index) if region_index is not None else None, ) - attempt_ids = session.scalars( - select(TaskInstance.id) - .where(TaskInstance.dag_id == dag_id, TaskInstance.run_id == run_id, TaskInstance.task_id == task_id) - .where(*filters) - .limit(2) + attempt_ids = ( + await session.scalars( + select(TaskInstance.id) + .where( + TaskInstance.dag_id == dag_id, TaskInstance.run_id == run_id, TaskInstance.task_id == task_id + ) + .where(*filters) + .limit(2) + ) ).all() return attempt_ids[0] if len(attempt_ids) == 1 else None @@ -197,7 +201,7 @@ def _mask_xcom_fields(extra_fields): return {k: "***" if k == "value" else v for k, v in extra_fields.items()} -def _resolve_team_name(params: dict, *, dag_id: str | None, session: Session) -> str | None: +async def _resolve_team_name(params: dict, *, dag_id: str | None, session: AsyncSession) -> str | None: """ Return the team the audited action belongs to, for the resources that own no Dag. @@ -222,14 +226,16 @@ def _resolve_team_name(params: dict, *, dag_id: str | None, session: Session) -> return None for param, (team_column, resource_column) in _TEAM_SCOPED_RESOURCES.items(): if (resource_id := params.get(param)) is not None: - return session.scalar(select(team_column).where(resource_column == resource_id)) + return await session.scalar(select(team_column).where(resource_column == resource_id)) return None def action_logging(event: str | None = None): async def log_action( request: Request, - session: SessionDep, + # Not the route's sync ``SessionDep``: this dependency runs on the event loop, where a sync + # pool checkout blocks the requests holding the pooled connections from releasing them. + session: AsyncSessionDep, user: GetUserDep, ): """Log user actions.""" @@ -320,7 +326,7 @@ def action_logging(event: str | None = None): scope.update(masked_body_json) dag_id = scope.get("dag_id") try: - task_instance_id = _find_addressed_attempt(request.path_params, params, session=session) + task_instance_id = await _find_addressed_attempt(request.path_params, params, session=session) except (TypeError, ValueError): # Malformed coordinates are the route's to reject; the access is still logged. task_instance_id = None @@ -336,7 +342,7 @@ def action_logging(event: str | None = None): task_id=scope.get("task_id"), dag_id=dag_id, run_id=scope.get("run_id") or scope.get("dag_run_id"), - team_name=_resolve_team_name(params, dag_id=dag_id, session=session), + team_name=await _resolve_team_name(params, dag_id=dag_id, session=session), ) if "logical_date" in request.query_params: @@ -357,6 +363,6 @@ def action_logging(event: str | None = None): session.add(log) # Explicit commit to persist the access log independently if the path operation fails or not. # Also it cannot be deferred to a 'function' scoped dependency because of the `request` parameter. - session.commit() + await session.commit() return log_action diff --git a/airflow-core/tests/unit/api_fastapi/core_api/conftest.py b/airflow-core/tests/unit/api_fastapi/core_api/conftest.py new file mode 100644 index 00000000000..d6b881e8cec --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/core_api/conftest.py @@ -0,0 +1,38 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest + +from airflow import settings + +from tests_common.test_utils.config import conf_vars + + [email protected] +def async_engine_without_pool(): + """ + Give the test an unpooled async engine. + + For tests that send write requests through the test client and also call ``dag.test()``: the two + run on different event loops, and a pooled async connection only works on the loop that opened it. + """ + previous = settings.async_engine, settings.AsyncSession + with conf_vars({("database", "sql_alchemy_pool_enabled"): "False"}): + settings._configure_async_session() + yield + settings.async_engine, settings.AsyncSession = previous diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py index dbf7a1b67ac..312a5e0dc8c 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_execution.py @@ -72,6 +72,7 @@ def test_execution_lists_live_work_and_selects_archived_tries(test_client, dag_m assert reader.return_value.log_handler.get_external_log_url.call_args.args[0].id == archived_id [email protected]("async_engine_without_pool") @pytest.mark.parametrize("retain_later", [False, True]) def test_loop_repeated_selective_clears_archive_only_replaced_task_instances( test_client, dag_maker, session, retain_later diff --git a/airflow-core/tests/unit/api_fastapi/logging/test_decorators.py b/airflow-core/tests/unit/api_fastapi/logging/test_decorators.py index d3a25732410..76bb01f9d10 100644 --- a/airflow-core/tests/unit/api_fastapi/logging/test_decorators.py +++ b/airflow-core/tests/unit/api_fastapi/logging/test_decorators.py @@ -25,8 +25,9 @@ import pytest from fastapi import Request from fastapi.dependencies.models import Dependant from sqlalchemy import delete, func, select -from sqlalchemy.orm import Session +from sqlalchemy.ext.asyncio import AsyncSession +from airflow import settings from airflow.api_fastapi.auth.managers.models.base_user import BaseUser from airflow.api_fastapi.logging.decorators import ( _mask_connection_fields, @@ -40,6 +41,7 @@ from airflow.models.dag import DagModel, clear_team_name_cache from airflow.models.dagbundle import DagBundleModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.utils.session import create_session_async from tests_common.test_utils.api_fastapi import iter_api_routes from tests_common.test_utils.config import conf_vars @@ -56,6 +58,20 @@ from tests_common.test_utils.db import ( from tests_common.test_utils.mock_operators import MockOperator +def _run_log_action(event: str, request: Request) -> None: + """Run ``log_action`` on a real async session, as the API server does.""" + + async def _run(): + try: + async with create_session_async() as session: + await action_logging(event=event)(request=request, session=session, user=None) + finally: + # Pooled async connections are bound to this event loop, which ``asyncio.run`` closes. + await settings.dispose_async_engine() + + asyncio.run(_run()) + + class TestSanitizeForStdlibLog: """User input passed to stdlib ``%s``-style logging must have CR/LF stripped. @@ -333,7 +349,7 @@ class TestActionLoggingUserFields: return "Jane Doe" request = Request({"type": "http", "method": "GET", "headers": [], "query_string": b""}) - session = MagicMock(spec=Session) + session = MagicMock(spec=AsyncSession) asyncio.run(action_logging(event="test_event")(request=request, session=session, user=FakeUser())) @@ -342,6 +358,17 @@ class TestActionLoggingUserFields: assert logged.owner_display_name == "Jane Doe" +class TestActionLoggingCommit: + def test_the_commit_is_awaited(self): + """A blocking commit waits for a pooled connection on the event loop, stalling every request.""" + request = Request({"type": "http", "method": "POST", "headers": [], "query_string": b""}) + session = MagicMock(spec=AsyncSession) + + asyncio.run(action_logging(event="test_event")(request=request, session=session, user=None)) + + session.commit.assert_awaited_once() + + class TestActionLoggingTeamName: """A team-scoped resource that owns no Dag records its team on the Log row itself, since there is no ``dag_id`` to resolve one through.""" @@ -355,7 +382,7 @@ class TestActionLoggingTeamName: ) def test_team_name_is_taken_from_the_request(self, query_string, expected_team_name): request = Request({"type": "http", "method": "POST", "headers": [], "query_string": query_string}) - session = MagicMock(spec=Session) + session = MagicMock(spec=AsyncSession) asyncio.run(action_logging(event="post_pool")(request=request, session=session, user=None)) @@ -373,7 +400,7 @@ class TestActionLoggingTeamName: """The row is committed before the endpoint validates the request, so a name the column cannot hold would fail the insert instead of the request.""" request = Request({"type": "http", "method": "POST", "headers": [], "query_string": query_string}) - session = MagicMock(spec=Session) + session = MagicMock(spec=AsyncSession) asyncio.run(action_logging(event="post_pool")(request=request, session=session, user=None)) @@ -416,7 +443,7 @@ class TestActionLoggingResourceTeamName: "path_params": path_params, } ) - asyncio.run(action_logging(event="delete_resource")(request=request, session=session, user=None)) + _run_log_action("delete_resource", request) return session.scalar(select(Log).order_by(Log.id.desc())) @pytest.mark.parametrize( @@ -477,7 +504,7 @@ class TestActionLoggingResourceTeamName: } ) - asyncio.run(action_logging(event="delete_connection")(request=request, session=session, user=None)) + _run_log_action("delete_connection", request) log = session.scalar(select(Log).order_by(Log.id.desc())) assert log.dag_id is None @@ -499,7 +526,7 @@ class TestActionLoggingDagScope: request = Request( {"type": "http", "method": "POST", "headers": [], "query_string": query_string, "path_params": {}} ) - session = MagicMock(spec=Session) + session = MagicMock(spec=AsyncSession) asyncio.run(action_logging(event=event)(request=request, session=session, user=None)) (logged,) = session.add.call_args.args return logged @@ -549,7 +576,7 @@ class TestActionLoggingUnparsableBody: ], ) def test_the_access_is_logged_with_no_body(self, body): - session = MagicMock(spec=Session) + session = MagicMock(spec=AsyncSession) asyncio.run( action_logging(event="test_event")( @@ -678,7 +705,7 @@ class TestActionLoggingAddressedAttempt: "path_params": {"dag_id": "audit_dag", "dag_run_id": "audit_run", **path_params}, } ) - asyncio.run(action_logging(event="patch_task_instance")(request=request, session=session, user=None)) + _run_log_action("patch_task_instance", request) return session.scalar(select(Log).order_by(Log.id.desc())) @pytest.mark.parametrize( diff --git a/devel-common/src/tests_common/test_utils/asserts.py b/devel-common/src/tests_common/test_utils/asserts.py index ab79ad768a0..c0eaa8510a6 100644 --- a/devel-common/src/tests_common/test_utils/asserts.py +++ b/devel-common/src/tests_common/test_utils/asserts.py @@ -103,7 +103,7 @@ class QueriesTraceInfo(NamedTuple): class CountQueries: """ - Counts the number of queries sent to Airflow Database in a given context. + Counts the number of queries sent to Airflow Database in a given context, on both the sync and async engines. Does not support multiple processes. When a new process is started in context, its queries will not be included. @@ -125,14 +125,20 @@ class CountQueries: if self.session: event.listen(self.session, "do_orm_execute", self.after_cursor_execute) else: - event.listen(airflow.settings.engine, "after_cursor_execute", self.after_cursor_execute) + self.engines = [airflow.settings.engine] + # Absent before Airflow 3, which provider tests still run against. + if async_engine := getattr(airflow.settings, "async_engine", None): + self.engines.append(async_engine.sync_engine) + for engine in self.engines: + event.listen(engine, "after_cursor_execute", self.after_cursor_execute) return self.result def __exit__(self, type_, value, tb): if self.session: event.remove(self.session, "do_orm_execute", self.after_cursor_execute) else: - event.remove(airflow.settings.engine, "after_cursor_execute", self.after_cursor_execute) + for engine in self.engines: + event.remove(engine, "after_cursor_execute", self.after_cursor_execute) log.debug("Queries count: %d", sum(self.result.values())) def after_cursor_execute(self, *args, **kwargs):
