Repository: apache/airflow
File: airflow-core/src/airflow/api_fastapi/execution_api/app.py (Python)

Code before:

    return app


def get_extra_schemas() -> dict[str, dict]:
    """Get all the extra schemas that are not part of the main FastAPI app."""
    from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance
    from airflow.executors.workloads import BundleInfo
    from airflow.task.trigger_rule import TriggerRule
    from airflow.task.weight_rule import WeightRule
    from airflow.utils.state import TaskInstanceState, TerminalTIState

    return {
        "TaskInstance": TaskInstance.model_json_schema(),
        "BundleInfo": BundleInfo.model_json_schema(),
        # Include the combined state enum too. In the datamodels we separate out SUCCESS from the other states
        # as that has different payload requirements
        "TerminalTIState": {"type": "string", "enum": list(TerminalTIState)},
        "TaskInstanceState": {"type": "string", "enum": list(TaskInstanceState)},
        "WeightRule": {"type": "string", "enum": list(WeightRule)},
        "TriggerRule": {"type": "string", "enum": list(TriggerRule)},
    }


@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
    _cm: AsyncExitStack | 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.app import create_task_execution_api_app
            from airflow.api_fastapi.execution_api.deps import (
                JWTBearerDep,
                JWTBearerTIPathDep,
            )
            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

            self._app = create_task_execution_api_app()


Review comment (suggestion):
Can these be dynamic? e.g. `[m.name for m in DagAttributeTypes]` so the overhead of having all the attribute types wouldn’t be that much (just a couple hundred extra bytes in the generated type)

Code after:

    return app


def get_extra_schemas() -> dict[str, dict]:
    """Get all the extra schemas that are not part of the main FastAPI app."""
    from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance
    from airflow.executors.workloads import BundleInfo
    from airflow.serialization.enums import DagAttributeTypes
    from airflow.task.trigger_rule import TriggerRule
    from airflow.task.weight_rule import WeightRule
    from airflow.utils.state import TaskInstanceState, TerminalTIState

    return {
        "TaskInstance": TaskInstance.model_json_schema(),
        "BundleInfo": BundleInfo.model_json_schema(),
        # Include the combined state enum too. In the datamodels we separate out SUCCESS from the other states
        # as that has different payload requirements
        "TerminalTIState": {"type": "string", "enum": list(TerminalTIState)},
        "TaskInstanceState": {"type": "string", "enum": list(TaskInstanceState)},
        "WeightRule": {"type": "string", "enum": list(WeightRule)},
        "TriggerRule": {"type": "string", "enum": list(TriggerRule)},
        "DagAttributeTypes": {
            "type": "string",
            "enum": [DagAttributeTypes.OP.value, DagAttributeTypes.TASK_GROUP.value],
            "x-enum-varnames": [DagAttributeTypes.OP.name, DagAttributeTypes.TASK_GROUP.name],
        },
    }


@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
    _cm: AsyncExitStack | 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.app import create_task_execution_api_app
            from airflow.api_fastapi.execution_api.deps import (
                JWTBearerDep,
                JWTBearerTIPathDep,
            )