This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 0c031565d85 [v3-4-test] Keep the audit log off the API server event
loop (#76416) (#76431)
0c031565d85 is described below
commit 0c031565d8580174677a24f231a76443c1f43f85
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):
