Consolidate in-process Execution API hosting

Dag parsing, trigger execution, local Dag tests and schema generation
need consistent Execution API setup. Keeping in-process authentication
and transport lifecycle together prevents setup drift while preserving
isolation from network-served apps.

Manager construction should not load the API server dependency stack.
Existing Task SDK consumers and schema generation against older
checkouts also need the established app.py import path to remain valid.

Author: ephraimbuddy

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

@@ -140,8 +140,7 @@ def create_app(apps: str = "all") -> FastAPI:
     dag_bag = create_dag_bag()
 
     if "all" in apps_list or "execution" in apps_list:
-        task_exec_api_app = create_task_execution_api_app()
-        task_exec_api_app.state.dag_bag = dag_bag
+        task_exec_api_app = create_task_execution_api_app(dag_bag=dag_bag)
         app.mount("/execution", task_exec_api_app)
 
     if "all" in apps_list or "core" in apps_list:

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

@@ -19,14 +19,9 @@
 
 import asyncio
 import json
-import threading
 import time
-import weakref
-from contextlib import AsyncExitStack
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
 
-import attrs
 import svcs
 from cadwyn import (
     Cadwyn,
@@ -44,16 +39,18 @@
     get_sig_validation_args,
     get_signing_args,
 )
+from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI
 
 if TYPE_CHECKING:
-    import httpx
+    from airflow.models.dagbag import DBDagBag
 
 import structlog
 from structlog.contextvars import bind_contextvars
 
 logger = structlog.get_logger(logger_name=__name__)
 
 __all__ = [
+    "InProcessExecutionAPI",
     "create_task_execution_api_app",
     "lifespan",
     "CorrelationIdMiddleware",
@@ -285,8 +282,18 @@ def _inject_trace_context_dep(routes, mode: str) -> None:
                     route.dependencies.append(dep)
 
 
-def create_task_execution_api_app(lifespan: svcs.fastapi.lifespan = lifespan) -> FastAPI:
-    """Create FastAPI app for task execution API."""
+def create_task_execution_api_app(
+    lifespan: svcs.fastapi.lifespan = lifespan,
+    *,
+    dag_bag: DBDagBag | None = None,
+) -> FastAPI:
+    """
+    Create the Execution API app.
+
+    :param lifespan: Lifespan manager that registers the app's services.
+    :param dag_bag: Dag bag to use; a new one is created when omitted.
+    """
+    from airflow.api_fastapi.common.dagbag import create_dag_bag
     from airflow.api_fastapi.common.exceptions import init_error_handlers
     from airflow.api_fastapi.execution_api.routes import execution_api_router
     from airflow.api_fastapi.execution_api.versions import bundle
@@ -326,6 +333,7 @@ def handle_exceptions(request: Request, exc: Exception):
             content["correlation-id"] = correlation_id
         return JSONResponse(status_code=500, content=content)
 
+    app.state.dag_bag = dag_bag if dag_bag is not None else create_dag_bag()
     return app
 
 
@@ -353,111 +361,3 @@ def get_extra_schemas() -> dict[str, dict]:
             "x-enum-varnames": [DagAttributeTypes.OP.name, DagAttributeTypes.TASK_GROUP.name],
         },
     }
-
-
-# Note: _shutdown_loop is used as a finalizer for the WSGI transport returned by
-# ``InProcessExecutionAPI.transport``. As such, its arguments must not directly or indirectly reference that
-# transport, as this would prevent the transport from being garbage collected.
-def _shutdown_loop(
-    loop: asyncio.AbstractEventLoop,
-    thread: threading.Thread,
-    cm: AsyncExitStack,
-) -> None:
-    """Close the FastAPI lifespan and stop the background event loop + thread."""
-    try:
-        asyncio.run_coroutine_threadsafe(cm.aclose(), loop).result(timeout=5)
-    except Exception:
-        logger.exception("Error while closing in-process execution API lifespan")
-    loop.call_soon_threadsafe(loop.stop)
-    thread.join(timeout=5)
-
-
-@attrs.define()
-class InProcessExecutionAPI:
-    """
-    A helper class to make it possible to run the ExecutionAPI "in-process".
-
-    The sync version of this makes use of a2wsgi which runs the async loop in a separate thread. This is
-    needed so that we can use the sync httpx client
-    """
-
-    _app: FastAPI | None = None
-
-    @cached_property
-    def app(self):
-        if not self._app:
-            from airflow.api_fastapi.common.dagbag import create_dag_bag
-            from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
-            from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
-            from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
-            from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access
-            from airflow.api_fastapi.execution_api.security import _jwt_bearer
-
-            # Give this app its own lifespan + services registry so that stubbing services
-            # (e.g. JWTValidator) doesn't affect the module-level ``lifespan.registry``.
-            registry = svcs.Registry()
-            private_lifespan = attrs.evolve(lifespan, registry=registry)
-            self._app = create_task_execution_api_app(lifespan=private_lifespan)
-
-            # In-process callers don't need a real JWTValidator: auth is bypassed below via
-            # ``dependency_overrides``.
-            registry.register_value(JWTValidator, None)
-
-            # Set up dag_bag in app state for dependency injection
-            self._app.state.dag_bag = create_dag_bag()
-
-            async def always_allow(request: Request):
-                from uuid import UUID
-
-                ti_id = UUID(
-                    request.path_params.get("task_instance_id", "00000000-0000-0000-0000-000000000000")
-                )
-                claims = TIClaims(scope="execution")
-                return TIToken(id=ti_id, claims=claims)
-
-            self._app.dependency_overrides[_jwt_bearer] = always_allow
-            self._app.dependency_overrides[has_connection_access] = always_allow
-            self._app.dependency_overrides[has_variable_access] = always_allow
-            self._app.dependency_overrides[has_xcom_access] = always_allow
-
-        return self._app
-
-    @cached_property
-    def transport(self) -> httpx.WSGITransport:
-        import httpx
-        from a2wsgi import ASGIMiddleware
-
-        # We choose to own the event loop + executor thread here so that we can have explicit control over
-        # their lifecycle.
-        loop = asyncio.new_event_loop()
-        thread = threading.Thread(target=loop.run_forever, name="InProcessExecutionAPI-loop", daemon=True)
-        thread.start()
-
-        middleware = ASGIMiddleware(self.app, loop=loop)
-
-        # https://github.com/abersheeran/a2wsgi/discussions/64
-        async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
-            await cm.enter_async_context(app.router.lifespan_context(app))
-
-        cm = AsyncExitStack()
-
-        # Wait for lifespan startup to complete so callers see a ready app and so the finalizer can
-        # safely aclose() a context whose __aenter__ has actually run.
-        asyncio.run_coroutine_threadsafe(start_lifespan(cm, self.app), loop).result()
-
-        transport = httpx.WSGITransport(app=middleware)  # type: ignore[arg-type]
-
-        # Stop the loop + thread and unwind the lifespan when the *transport* is garbage collected, not
-        # this InProcessExecutionAPI instance. Callers commonly build a Client from ``.transport`` and drop
-        # the factory object (e.g. ``Client(transport=InProcessExecutionAPI().transport)``); finalizing on
-        # ``self`` would stop the loop while the transport is still in use, so every later request would
-        # hang on the now-dead loop.
-        weakref.finalize(transport, _shutdown_loop, loop, thread, cm)
-
-        return transport
-
-    @cached_property
-    def atransport(self) -> httpx.ASGITransport:
-        import httpx
-
-        return httpx.ASGITransport(app=self.app)

airflow-core/src/airflow/api_fastapi/execution_api/in_process.py

@@ -0,0 +1,136 @@
+# 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.
+"""In-process Execution API hosting, authentication overrides, and transport lifecycle."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+import weakref
+from contextlib import AsyncExitStack
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+import attrs
+import structlog
+from starlette.requests import Request
+
+if TYPE_CHECKING:
+    import httpx
+    from fastapi import FastAPI
+
+logger = structlog.get_logger(logger_name=__name__)
+
+
+# Finalizer arguments must not reference the transport, or it cannot be garbage collected.
+def _shutdown_loop(
+    loop: asyncio.AbstractEventLoop,
+    thread: threading.Thread,
+    cm: AsyncExitStack,
+) -> None:
+    """Close the FastAPI lifespan and stop the background event loop + thread."""
+    try:
+        asyncio.run_coroutine_threadsafe(cm.aclose(), loop).result(timeout=5)
+    except Exception:
+        logger.exception("Error while closing in-process execution API lifespan")
+    loop.call_soon_threadsafe(loop.stop)
+    thread.join(timeout=5)
+
+
+def _configure_in_process_auth(app: FastAPI) -> None:
+    from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
+    from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
+    from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
+    from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access
+    from airflow.api_fastapi.execution_api.security import _jwt_bearer
+
+    async def always_allow(request: Request):
+        from uuid import UUID
+
+        ti_id = UUID(request.path_params.get("task_instance_id", "00000000-0000-0000-0000-000000000000"))
+        claims = TIClaims(scope="execution")
+        return TIToken(id=ti_id, claims=claims)
+
+    app.dependency_overrides[_jwt_bearer] = always_allow
+    app.dependency_overrides[has_connection_access] = always_allow
+    app.dependency_overrides[has_variable_access] = always_allow
+    app.dependency_overrides[has_xcom_access] = always_allow
+
+
+@attrs.define()
+class InProcessExecutionAPI:
+    """
+    A helper class to make it possible to run the ExecutionAPI "in-process".
+
+    The sync version of this makes use of a2wsgi which runs the async loop in a separate thread. This is
+    needed so that we can use the sync httpx client
+    """
+
+    _app: FastAPI | None = None
+
+    @cached_property
+    def app(self):
+        import svcs
+
+        from airflow.api_fastapi.auth.tokens import JWTValidator
+        from airflow.api_fastapi.execution_api.app import create_task_execution_api_app, lifespan
+
+        if not self._app:
+            # Keep the stub private: a shared None validator would prevent lifespan() from
+            # registering the real API server's validator.
+            registry = svcs.Registry()
+            self._app = create_task_execution_api_app(lifespan=attrs.evolve(lifespan, registry=registry))
+            # Auth dependency overrides bypass validation, so the stub skips validator setup in lifespan().
+            registry.register_value(JWTValidator, None)
+            _configure_in_process_auth(self._app)
+
+        return self._app
+
+    @cached_property
+    def transport(self) -> httpx.WSGITransport:
+        import httpx
+        from a2wsgi import ASGIMiddleware
+
+        # Own the event loop and thread so the transport controls their lifecycle.
+        loop = asyncio.new_event_loop()
+        thread = threading.Thread(target=loop.run_forever, name="InProcessExecutionAPI-loop", daemon=True)
+        thread.start()
+
+        middleware = ASGIMiddleware(self.app, loop=loop)
+
+        # https://github.com/abersheeran/a2wsgi/discussions/64
+        async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
+            await cm.enter_async_context(app.router.lifespan_context(app))
+
+        cm = AsyncExitStack()
+
+        # Wait for startup so callers see a ready app and the finalizer can close an entered context.
+        asyncio.run_coroutine_threadsafe(start_lifespan(cm, self.app), loop).result()
+
+        transport = httpx.WSGITransport(app=middleware)  # type: ignore[arg-type]
+
+        # Callers can retain the transport after dropping this instance; finalizing on self would
+        # stop the loop while requests still need it.
+        weakref.finalize(transport, _shutdown_loop, loop, thread, cm)
+
+        return transport
+
+    @cached_property
+    def atransport(self) -> httpx.ASGITransport:
+        import httpx
+
+        return httpx.ASGITransport(app=self.app)

airflow-core/src/airflow/dag_processing/manager.py

@@ -48,6 +48,7 @@
 from airflow._shared.observability.metrics import stats
 from airflow._shared.observability.metrics.stats import normalize_name_for_stats
 from airflow._shared.timezones import timezone
+from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI
 from airflow.configuration import conf
 from airflow.dag_processing.bundles.base import (
     BundleUsageTrackingManager,
@@ -90,22 +91,11 @@
     from sqlalchemy.orm import Session
     from sqlalchemy.sql import Select
 
-    from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI
     from airflow.callbacks.callback_requests import CallbackRequest
     from airflow.dag_processing.bundles.base import BaseDagBundle
     from airflow.sdk.api.client import Client
 
 
-def _make_execution_api() -> InProcessExecutionAPI:
-    # This is a seriously weighty import, pulling in svcs, cadwyn, fastapi, aiohttp, etc.
-    #
-    # Defer it so that an import of this module for types (e.g. DagFileStat, DagFileInto) doesn't need to pay
-    # that cost.
-    from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI
-
-    return InProcessExecutionAPI()
-
-
 class DagParsingStat(NamedTuple):
     """Information on processing progress."""
 
@@ -303,7 +293,7 @@ class DagFileProcessorManager(LoggingMixin):
     )
     """Resolved once per process so file discovery and the deactivation scan use the same value."""
 
-    _api_server: InProcessExecutionAPI = attrs.field(init=False, factory=_make_execution_api)
+    _api_server: InProcessExecutionAPI = attrs.field(init=False, factory=InProcessExecutionAPI)
     """API server to interact with Metadata DB"""
 
     def register_exit_signals(self):

airflow-core/src/airflow/dag_processing/processor.py

@@ -101,7 +101,6 @@
 
     from structlog.typing import FilteringBoundLogger
 
-    from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI
     from airflow.sdk.api.client import Client
     from airflow.sdk.bases.operator import BaseOperator
     from airflow.sdk.definitions.context import Context
@@ -567,13 +566,6 @@ def _execute_email_callbacks(dagbag: DagBag, request: EmailRequest, log: Filteri
         )
 
 
-def in_process_api_server() -> InProcessExecutionAPI:
-    from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI
-
-    api = InProcessExecutionAPI()
-    return api
-
-
 @attrs.define(kw_only=True)
 class DagFileProcessorProcess(WatchedSubprocess, LoggingMixin):
     """