From 2bdc05f52f20116cf56ee53438ae72951be4ccb6 Mon Sep 17 00:00:00 2001 From: Jayachandra Kasarla Date: Mon, 6 Jul 2026 14:01:04 +0530 Subject: [PATCH 1/2] Added request metrics middleware to export fastapi requests metrics --- .../api_fastapi/common/request_metrics.py | 85 +++++++++++++++++++ .../src/airflow/api_fastapi/core_api/app.py | 12 +++ 2 files changed, 97 insertions(+) create mode 100644 airflow-core/src/airflow/api_fastapi/common/request_metrics.py diff --git a/airflow-core/src/airflow/api_fastapi/common/request_metrics.py b/airflow-core/src/airflow/api_fastapi/common/request_metrics.py new file mode 100644 index 0000000000000..a57a0e90f6bcc --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/common/request_metrics.py @@ -0,0 +1,85 @@ +# 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. +"""Per-route request metrics middleware for the API server.""" + +from __future__ import annotations + +import contextlib +import time +from typing import TYPE_CHECKING + +from airflow.observability.stats import Stats + +if TYPE_CHECKING: + from starlette.types import ASGIApp, Message, Receive, Scope, Send + +REQUEST_COUNT_METRIC = "api_server.request.count" +REQUEST_DURATION_METRIC = "api_server.request.duration" + +# Emitted when a request does not resolve to a known route. Using a constant keeps the +# "route" tag cardinality bounded — the raw path would let unmatched URLs explode it. +_UNMATCHED_ROUTE = "__unmatched__" + + +class RequestMetricsMiddleware: + """ + Emit per-route request count and latency metrics through the "Stats" backend. + + Records one counter increment and one timing observation per completed request, + tagged with the templated route, HTTP method, and response status. The templated + route (e.g. "/api/v2/dags/{dag_id}/dagRuns") is used rather than the concrete + path so that per-"dag_id" requests do not each become a distinct metric series. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + start = time.monotonic_ns() + response: Message | None = None + + async def capture_send(message: Message) -> None: + nonlocal response + if message["type"] == "http.response.start": + response = message + await send(message) + + try: + await self.app(scope, receive, capture_send) + except Exception: + if response is None: + response = {"status": 500} + raise + finally: + duration_ms = (time.monotonic_ns() - start) / 1_000_000 + status = response["status"] if response is not None else 0 + route = scope.get("route") + tags = { + "method": scope.get("method", ""), + "status": str(status), + "route": getattr(route, "path", _UNMATCHED_ROUTE), + } + + # Never let a metrics failure replace an application exception propagating + # through this "finally" (see HttpAccessLogMiddleware for the rationale). + with contextlib.suppress(Exception): + Stats.incr(REQUEST_COUNT_METRIC, tags=tags) + Stats.timing(REQUEST_DURATION_METRIC, duration_ms, tags=tags) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/app.py b/airflow-core/src/airflow/api_fastapi/core_api/app.py index 27213a4c0a2e8..fa0ca3829943c 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/app.py @@ -181,6 +181,7 @@ def init_middlewares(app: FastAPI) -> None: from airflow.api_fastapi.app import get_auth_manager from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware from airflow.api_fastapi.common.http_access_log import HttpAccessLogMiddleware + from airflow.configuration import conf app.add_middleware(JWTRefreshMiddleware) @@ -194,3 +195,14 @@ def init_middlewares(app: FastAPI) -> None: # HttpAccessLogMiddleware must be outermost (added last) so it times the full # request lifecycle including all inner middleware. app.add_middleware(HttpAccessLogMiddleware) + + # Only wire per-route request metrics when a metrics backend is configured; otherwise + # the Stats singleton is a no-op and the middleware would add cost for nothing. + if ( + conf.getboolean("metrics", "otel_on") + or conf.getboolean("metrics", "statsd_on") + or conf.getboolean("metrics", "statsd_datadog_enabled") + ): + from airflow.api_fastapi.common.request_metrics import RequestMetricsMiddleware + + app.add_middleware(RequestMetricsMiddleware) From fc7634eaa911f69bcdaa415b0521a4d8d452e977 Mon Sep 17 00:00:00 2001 From: Jayachandra Kasarla Date: Mon, 6 Jul 2026 14:18:59 +0530 Subject: [PATCH 2/2] Added tests for request_metrics.py --- .../common/test_request_metrics.py | 169 ++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 airflow-core/tests/unit/api_fastapi/common/test_request_metrics.py diff --git a/airflow-core/tests/unit/api_fastapi/common/test_request_metrics.py b/airflow-core/tests/unit/api_fastapi/common/test_request_metrics.py new file mode 100644 index 0000000000000..4d71618933f81 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/common/test_request_metrics.py @@ -0,0 +1,169 @@ +# 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 asyncio +from unittest import mock + +import pytest +from fastapi import FastAPI +from fastapi.responses import PlainTextResponse +from starlette.testclient import TestClient + +from airflow.api_fastapi.common.request_metrics import ( + _UNMATCHED_ROUTE, + REQUEST_COUNT_METRIC, + REQUEST_DURATION_METRIC, + RequestMetricsMiddleware, +) + +from tests_common.test_utils.config import conf_vars + +STATS_TARGET = "airflow.api_fastapi.common.request_metrics.Stats" + + +def _make_app(*, raise_exc: bool = False, status_code: int = 200) -> FastAPI: + app = FastAPI() + + @app.get("/items/{item_id}") + def items(item_id: str): + if raise_exc: + raise RuntimeError("boom") + return PlainTextResponse("ok", status_code=status_code) + + app.add_middleware(RequestMetricsMiddleware) + return app + + +@mock.patch(STATS_TARGET, autospec=True) +def test_records_count_and_duration_with_templated_route(mock_stats): + client = TestClient(_make_app(), raise_server_exceptions=False) + client.get("/items/42") + + expected_tags = {"method": "GET", "status": "200", "route": "/items/{item_id}"} + mock_stats.incr.assert_called_once_with(REQUEST_COUNT_METRIC, tags=expected_tags) + + mock_stats.timing.assert_called_once() + name, duration = mock_stats.timing.call_args.args + assert name == REQUEST_DURATION_METRIC + assert duration >= 0 + assert mock_stats.timing.call_args.kwargs["tags"] == expected_tags + + +@mock.patch(STATS_TARGET, autospec=True) +def test_records_response_status(mock_stats): + client = TestClient(_make_app(status_code=418), raise_server_exceptions=False) + client.get("/items/1") + + assert mock_stats.incr.call_args.kwargs["tags"]["status"] == "418" + + +@mock.patch(STATS_TARGET, autospec=True) +def test_unmatched_route_uses_constant_tag(mock_stats): + client = TestClient(_make_app(), raise_server_exceptions=False) + client.get("/does-not-exist") + + tags = mock_stats.incr.call_args.kwargs["tags"] + assert tags["route"] == _UNMATCHED_ROUTE + assert tags["status"] == "404" + + +@mock.patch(STATS_TARGET, autospec=True) +def test_records_status_500_on_app_exception(mock_stats): + client = TestClient(_make_app(raise_exc=True), raise_server_exceptions=False) + client.get("/items/1") + + assert mock_stats.incr.call_args.kwargs["tags"]["status"] == "500" + + +@mock.patch(STATS_TARGET, autospec=True) +def test_non_http_scope_is_passed_through(mock_stats): + called = False + + async def downstream(scope, receive, send): + nonlocal called + called = True + + middleware = RequestMetricsMiddleware(downstream) + asyncio.run(middleware({"type": "lifespan"}, None, None)) + + assert called + mock_stats.incr.assert_not_called() + mock_stats.timing.assert_not_called() + + +@mock.patch(STATS_TARGET, autospec=True) +def test_stats_failure_does_not_mask_app_exception(mock_stats): + mock_stats.incr.side_effect = RuntimeError("stats broken") + + async def raising_app(scope, receive, send): + await send({"type": "http.response.start", "status": 503, "headers": []}) + raise RuntimeError("app exception") + + middleware = RequestMetricsMiddleware(raising_app) + scope = {"type": "http", "method": "GET", "path": "/boom", "headers": []} + + async def receive(): + return {"type": "http.request", "body": b""} + + async def send(_message): + return None + + with pytest.raises(RuntimeError, match="app exception"): + asyncio.run(middleware(scope, receive, send)) + + +@mock.patch(STATS_TARGET, autospec=True) +def test_stats_failure_swallowed_on_clean_request(mock_stats): + mock_stats.incr.side_effect = RuntimeError("stats broken") + + client = TestClient(_make_app(), raise_server_exceptions=False) + response = client.get("/items/1") + + assert response.status_code == 200 + + +@pytest.mark.parametrize( + ("flag", "enabled", "should_register"), + [ + ("otel_on", "True", True), + ("statsd_on", "True", True), + ("statsd_datadog_enabled", "True", True), + ("otel_on", "False", False), + ], +) +@mock.patch("airflow.api_fastapi.app.get_auth_manager") +def test_middleware_registered_only_when_metrics_backend_enabled( + mock_get_auth_manager, flag, enabled, should_register +): + from airflow.api_fastapi.core_api.app import init_middlewares + + mock_get_auth_manager.return_value.get_fastapi_middlewares.return_value = [] + + app = FastAPI() + with conf_vars( + { + ("metrics", "otel_on"): "False", + ("metrics", "statsd_on"): "False", + ("metrics", "statsd_datadog_enabled"): "False", + ("metrics", flag): enabled, + } + ): + init_middlewares(app) + + registered = {m.cls for m in app.user_middleware} + assert (RequestMetricsMiddleware in registered) is should_register