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..3f09372162010 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_in_process.py @@ -0,0 +1,60 @@ +# 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 os +import subprocess +import sys +from textwrap import dedent + + +def test_construction_defers_execution_api_imports(): + result = subprocess.run( + [ + sys.executable, + "-c", + 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, + ) + 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()