From e2141d4db2abf1fa8cff99487c4fb498415ee8fa Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 18 Sep 2026 11:51:55 +0100 Subject: [PATCH 1/2] 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. --- airflow-core/src/airflow/api_fastapi/app.py | 3 +- .../airflow/api_fastapi/execution_api/app.py | 141 ++--------------- .../api_fastapi/execution_api/in_process.py | 148 ++++++++++++++++++ .../src/airflow/dag_processing/manager.py | 14 +- .../src/airflow/dag_processing/processor.py | 8 - .../src/airflow/jobs/triggerer_job_runner.py | 11 +- .../api_fastapi/execution_api/test_app.py | 45 ++++++ .../execution_api/test_in_process.py | 56 +++++++ .../tests/unit/api_fastapi/test_app.py | 2 +- .../unit/dag_processing/test_processor.py | 2 +- .../ci/prek/generate_execution_api_schema.py | 2 + 11 files changed, 274 insertions(+), 158 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/in_process.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py diff --git a/airflow-core/src/airflow/api_fastapi/app.py b/airflow-core/src/airflow/api_fastapi/app.py index 65121ab53d04b..974fb2cee7524 100644 --- a/airflow-core/src/airflow/api_fastapi/app.py +++ b/airflow-core/src/airflow/api_fastapi/app.py @@ -142,8 +142,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: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 29b6793c6da6f..8370b0ae9828e 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/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, @@ -38,16 +33,16 @@ from opentelemetry import context as otel_context, propagate as otel_propagate from starlette.middleware.base import BaseHTTPMiddleware -from airflow import settings from airflow.api_fastapi.auth.tokens import ( JWTGenerator, JWTValidator, 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 @@ -55,6 +50,7 @@ logger = structlog.get_logger(logger_name=__name__) __all__ = [ + "InProcessExecutionAPI", "create_task_execution_api_app", "lifespan", "CorrelationIdMiddleware", @@ -286,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 @@ -327,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 @@ -354,119 +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 _IN_PROCESS_NON_TI_CALLER, _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") - or request.headers.get("X-Airflow-In-Process-Attempt-Id") - or "00000000-0000-0000-0000-000000000000" - ) - # Watchers and other trusted in-process callers may have no task identity. - request.scope[_IN_PROCESS_NON_TI_CALLER] = ( - "task_instance_id" not in request.path_params - and "X-Airflow-In-Process-Attempt-Id" not in request.headers - ) - 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): - cm.push_async_callback(settings.dispose_async_engine) - 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) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py new file mode 100644 index 0000000000000..63f0ce78e713b --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/in_process.py @@ -0,0 +1,148 @@ +# 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 _IN_PROCESS_NON_TI_CALLER, _jwt_bearer + + async def always_allow(request: Request): + from uuid import UUID + + ti_id = UUID( + request.path_params.get("task_instance_id") + or request.headers.get("X-Airflow-In-Process-Attempt-Id") + or "00000000-0000-0000-0000-000000000000" + ) + # Watchers and other trusted in-process callers may have no task identity. + request.scope[_IN_PROCESS_NON_TI_CALLER] = ( + "task_instance_id" not in request.path_params + and "X-Airflow-In-Process-Attempt-Id" not in request.headers + ) + 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 + + from airflow import settings + + # 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): + cm.push_async_callback(settings.dispose_async_engine) + 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) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 4d525a40febaa..343d72c5e002d 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -47,6 +47,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.callbacks.callback_requests import DagCallbackRequest from airflow.configuration import conf from airflow.dag_processing.bundles.base import ( @@ -97,22 +98,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 BundleState(NamedTuple): """Persisted refresh state for a DAG bundle.""" @@ -304,7 +294,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): diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index fb3dd55f43832..c3feb3f4959a1 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -87,7 +87,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 @@ -568,13 +567,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 BaseDagFileProcessorProcess(WatchedSubprocess, LoggingMixin): """ diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index ede83fe6c393a..dbef9d2951ab7 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -49,6 +49,7 @@ from airflow._shared.module_loading import import_string from airflow._shared.observability.metrics import stats from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI from airflow.configuration import conf from airflow.exceptions import TaskNotFound from airflow.executors import workloads @@ -140,7 +141,6 @@ from sqlalchemy.orm import Session from structlog.typing import FilteringBoundLogger, WrappedLogger - from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.jobs.job import Job from airflow.sdk.api.client import Client from airflow.sdk.definitions.context import Context @@ -469,13 +469,6 @@ def upload_to_remote(self): upload_to_remote(self.bound_logger, self.ti) -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 TriggerRunnerSupervisor(WatchedSubprocess): """ @@ -584,7 +577,7 @@ def make_client(self) -> Client: """ from airflow.sdk.api.client import Client - client = Client(base_url=None, token="", dry_run=True, transport=in_process_api_server().transport) + client = Client(base_url=None, token="", dry_run=True, transport=InProcessExecutionAPI().transport) # Mypy is wrong -- the setter accepts a string on the property setter! `URLType = URL | str` client.base_url = "http://in-process.invalid./" return client diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py index 1cbb610a418a9..a57607f3b91eb 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py @@ -34,9 +34,11 @@ from sqlalchemy.exc import SQLAlchemyError from airflow import settings +from airflow.api_fastapi.app import create_app from airflow.api_fastapi.execution_api.app import ( InProcessExecutionAPI, _extract_w3c_trace_context, + _jwt_generator, create_task_execution_api_app, ) from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance @@ -51,6 +53,49 @@ pytestmark = pytest.mark.db_test +@pytest.mark.parametrize("supply_dag_bag", [False, True]) +@mock.patch("airflow.api_fastapi.common.dagbag.create_dag_bag", autospec=True) +def test_create_app_initializes_dag_bag(mock_create_dag_bag, supply_dag_bag): + dag_bag = mock.sentinel.dag_bag if supply_dag_bag else None + app = create_task_execution_api_app(dag_bag=dag_bag) + + if supply_dag_bag: + assert app.state.dag_bag is dag_bag + mock_create_dag_bag.assert_not_called() + else: + assert app.state.dag_bag is mock_create_dag_bag.return_value + mock_create_dag_bag.assert_called_once_with() + + +@pytest.mark.parametrize("in_process_first", [False, True]) +@conf_vars({("api_auth", "jwt_secret"): "execution-api-hosting-test-secret"}) +def test_server_and_in_process_apps_isolate_auth(in_process_first, monkeypatch): + monkeypatch.setenv("AIRFLOW_VAR_EXECUTION_API_HOSTING", "shared-value") + if in_process_first: + api = InProcessExecutionAPI() + assert api.app is not None + server_app = create_app(apps="execution") + else: + server_app = create_app(apps="execution") + api = InProcessExecutionAPI() + + with ( + TestClient(server_app) as server, + httpx.Client(transport=api.transport, base_url="http://in-process.invalid") as local, + ): + response = local.get("/variables/execution_api_hosting") + assert response.status_code == 200 + assert response.json() == {"key": "execution_api_hosting", "value": "shared-value"} + assert server.get("/execution/variables/execution_api_hosting").status_code == 401 + token = _jwt_generator().generate({"sub": "00000000-0000-0000-0000-000000000000"}) + authenticated_response = server.get( + "/execution/variables/execution_api_hosting", headers={"Authorization": f"Bearer {token}"} + ) + # Missing credentials bypass JWTValidator; only an authenticated request detects a leaked stub. + assert authenticated_response.status_code == 200 + assert authenticated_response.json() == response.json() + + def test_custom_openapi_includes_extra_schemas(client): """Test to ensure that extra schemas are correctly included in the OpenAPI schema.""" response = client.get("/execution/openapi.json?version=2025-04-28") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py new file mode 100644 index 0000000000000..1686fd2078aa9 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py @@ -0,0 +1,56 @@ +# 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 subprocess +import sys + +import pytest + + +@pytest.mark.parametrize( + "create_instance", + [ + "from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI; " + "api = InProcessExecutionAPI(); ", + "from airflow.dag_processing.manager import DagFileProcessorManager; " + "manager = DagFileProcessorManager(max_runs=1); ", + ], + ids=["in_process_api", "dag_processor_manager"], +) +def test_construction_defers_execution_api_imports(create_instance): + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; " + create_instance + "modules = {'svcs', 'cadwyn', 'fastapi', 'aiohttp', " + "'airflow.api_fastapi.execution_api.app'}; " + "loaded = modules.intersection(sys.modules); " + "assert not loaded, f'{loaded} imported eagerly'", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_app_reexports_in_process_api(): + from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI as legacy_api + from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI + + assert legacy_api is InProcessExecutionAPI diff --git a/airflow-core/tests/unit/api_fastapi/test_app.py b/airflow-core/tests/unit/api_fastapi/test_app.py index de74dbc8d7c8e..dcbbc0a593a1e 100644 --- a/airflow-core/tests/unit/api_fastapi/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/test_app.py @@ -174,7 +174,7 @@ def test_all_apps(mock_create_task_exec_api, mock_init_plugins, mock_init_views) mock_init_plugins.assert_called_once_with(test_app) # Assert that execution-related functions were also called - mock_create_task_exec_api.assert_called_once_with() + mock_create_task_exec_api.assert_called_once_with(dag_bag=test_app.state.dag_bag) @pytest.mark.parametrize("apps", ["all", "core", "execution"]) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 76a5ea85cde09..f5abf5f8b52b6 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -38,12 +38,12 @@ from structlog.typing import FilteringBoundLogger from airflow._shared.timezones import timezone -from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( DagRun as DRDataModel, TaskInstance as TIDataModel, TIRunContext, ) +from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI from airflow.callbacks.callback_requests import ( CallbackRequest, DagCallbackRequest, diff --git a/scripts/ci/prek/generate_execution_api_schema.py b/scripts/ci/prek/generate_execution_api_schema.py index 2efe713ea2269..7b34e0be48487 100644 --- a/scripts/ci/prek/generate_execution_api_schema.py +++ b/scripts/ci/prek/generate_execution_api_schema.py @@ -29,6 +29,8 @@ import httpx +# Keep this import compatible with the target branch: the version check runs this +# generator against both checkouts. from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI app = InProcessExecutionAPI() From cc8c268179cc474ada97a699b453fb5c2c3bdcdf Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 9 Oct 2026 11:15:44 +0100 Subject: [PATCH 2/2] Reduce startup overhead in lazy import tests --- .../execution_api/test_in_process.py | 38 ++++++++++--------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py index 1686fd2078aa9..3f09372162010 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py @@ -16,32 +16,36 @@ # under the License. from __future__ import annotations +import os import subprocess import sys +from textwrap import dedent -import pytest - -@pytest.mark.parametrize( - "create_instance", - [ - "from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI; " - "api = InProcessExecutionAPI(); ", - "from airflow.dag_processing.manager import DagFileProcessorManager; " - "manager = DagFileProcessorManager(max_runs=1); ", - ], - ids=["in_process_api", "dag_processor_manager"], -) -def test_construction_defers_execution_api_imports(create_instance): +def test_construction_defers_execution_api_imports(): result = subprocess.run( [ sys.executable, "-c", - "import sys; " + create_instance + "modules = {'svcs', 'cadwyn', 'fastapi', 'aiohttp', " - "'airflow.api_fastapi.execution_api.app'}; " - "loaded = modules.intersection(sys.modules); " - "assert not loaded, f'{loaded} imported eagerly'", + dedent(""" + import sys + + modules = {'svcs', 'cadwyn', 'fastapi', 'aiohttp', 'airflow.api_fastapi.execution_api.app'} + + from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI + + api = InProcessExecutionAPI() + loaded = modules.intersection(sys.modules) + assert not loaded, f'{loaded} imported eagerly by InProcessExecutionAPI' + + from airflow.dag_processing.manager import DagFileProcessorManager + + manager = DagFileProcessorManager(max_runs=1) + loaded = modules.intersection(sys.modules) + assert not loaded, f'{loaded} imported eagerly by DagFileProcessorManager' + """), ], + env={**os.environ, "_AIRFLOW__AS_LIBRARY": "1"}, capture_output=True, text=True, check=False,