Add per-route request metrics to the API server

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)

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)

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