From 8c3c41d6251fa9a6f704536561cc20b8dbcd6f89 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Thu, 6 Aug 2026 00:51:56 +0800 Subject: [PATCH] Move task instance state enums out of airflow.utils.state airflow.utils has become a dumping ground, and untangling it is a prerequisite for the client-server separation work. TerminalTIState and IntermediateTIState describe the Execution API's task lifecycle protocol and belong next to TriggerRule and WeightRule, which moved to airflow.task for the same reason. The old import paths keep working through the same backwards-compatible shim that the sibling JobState move already established in this module. --- .../airflow/api_fastapi/execution_api/app.py | 3 +- .../execution_api/datamodels/taskinstance.py | 3 +- .../execution_api/routes/task_instances.py | 3 +- airflow-core/src/airflow/task/state.py | 48 ++++++++++ airflow-core/src/airflow/utils/state.py | 94 ++++++++----------- .../versions/head/test_task_instances.py | 3 +- airflow-core/tests/unit/utils/test_state.py | 21 ++++- .../google/cloud/triggers/cloud_composer.py | 2 +- 8 files changed, 116 insertions(+), 61 deletions(-) create mode 100644 airflow-core/src/airflow/task/state.py 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 f9e7cb1725000..675b3f214a52e 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -334,9 +334,10 @@ def get_extra_schemas() -> dict[str, dict]: 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.state import TerminalTIState from airflow.task.trigger_rule import TriggerRule from airflow.task.weight_rule import WeightRule - from airflow.utils.state import TaskInstanceState, TerminalTIState + from airflow.utils.state import TaskInstanceState return { "TaskInstance": TaskInstance.model_json_schema(), diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ad051b3e6d340..a025576fd3663 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -37,11 +37,10 @@ from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile from airflow.api_fastapi.execution_api.datamodels.connection import ConnectionResponse from airflow.api_fastapi.execution_api.datamodels.variable import VariableResponse +from airflow.task.state import IntermediateTIState, TerminalTIState from airflow.utils.state import ( DagRunState, - IntermediateTIState, TaskInstanceState as TIState, - TerminalTIState, ) from airflow.utils.types import DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 41ecf49b053fb..53e823bd1a212 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -90,9 +90,10 @@ from airflow.models.xcom import XComModel from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey from airflow.state import get_state_backend +from airflow.task.state import TerminalTIState from airflow.triggers.base import TriggerEvent from airflow.utils.sqlalchemy import get_dialect_name -from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState +from airflow.utils.state import DagRunState, TaskInstanceState if TYPE_CHECKING: from sqlalchemy.sql.dml import Update diff --git a/airflow-core/src/airflow/task/state.py b/airflow-core/src/airflow/task/state.py new file mode 100644 index 0000000000000..969e730d08d09 --- /dev/null +++ b/airflow-core/src/airflow/task/state.py @@ -0,0 +1,48 @@ +# +# 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 + +from enum import Enum + + +class TerminalTIState(str, Enum): + """States that a Task Instance can be in that indicate it has reached a terminal state.""" + + SUCCESS = "success" + FAILED = "failed" + SKIPPED = "skipped" # A user can raise a AirflowSkipException from a task & it will be marked as skipped + UPSTREAM_FAILED = "upstream_failed" + REMOVED = "removed" + + def __str__(self) -> str: + return self.value + + +class IntermediateTIState(str, Enum): + """States that a Task Instance can be in that indicate it is not yet in a terminal or running state.""" + + SCHEDULED = "scheduled" + QUEUED = "queued" + RESTARTING = "restarting" + UP_FOR_RETRY = "up_for_retry" + UP_FOR_RESCHEDULE = "up_for_reschedule" + DEFERRED = "deferred" + AWAITING_INPUT = "awaiting_input" + + def __str__(self) -> str: + return self.value diff --git a/airflow-core/src/airflow/utils/state.py b/airflow-core/src/airflow/utils/state.py index 89c8efcaa3e02..8d088cf18fcd0 100644 --- a/airflow-core/src/airflow/utils/state.py +++ b/airflow-core/src/airflow/utils/state.py @@ -17,7 +17,16 @@ # under the License. from __future__ import annotations +import warnings from enum import Enum +from importlib import import_module + +# Aliased so the public names stay out of this module's namespace and keep resolving +# through the deprecation shim in ``__getattr__`` below. +from airflow.task.state import ( + IntermediateTIState as _IntermediateTIState, + TerminalTIState as _TerminalTIState, +) class CallbackState(str, Enum): @@ -34,34 +43,6 @@ def __str__(self) -> str: return self.value -class TerminalTIState(str, Enum): - """States that a Task Instance can be in that indicate it has reached a terminal state.""" - - SUCCESS = "success" - FAILED = "failed" - SKIPPED = "skipped" # A user can raise a AirflowSkipException from a task & it will be marked as skipped - UPSTREAM_FAILED = "upstream_failed" - REMOVED = "removed" - - def __str__(self) -> str: - return self.value - - -class IntermediateTIState(str, Enum): - """States that a Task Instance can be in that indicate it is not yet in a terminal or running state.""" - - SCHEDULED = "scheduled" - QUEUED = "queued" - RESTARTING = "restarting" - UP_FOR_RETRY = "up_for_retry" - UP_FOR_RESCHEDULE = "up_for_reschedule" - DEFERRED = "deferred" - AWAITING_INPUT = "awaiting_input" - - def __str__(self) -> str: - return self.value - - class TaskInstanceState(str, Enum): """ All possible states that a Task Instance can be in. @@ -74,21 +55,21 @@ class TaskInstanceState(str, Enum): # Use None instead if need this state. # Set by the scheduler - REMOVED = TerminalTIState.REMOVED # Task vanished from DAG before it ran - SCHEDULED = IntermediateTIState.SCHEDULED # Task should run and will be handed to executor soon + REMOVED = _TerminalTIState.REMOVED # Task vanished from DAG before it ran + SCHEDULED = _IntermediateTIState.SCHEDULED # Task should run and will be handed to executor soon # Set by the task instance itself - QUEUED = IntermediateTIState.QUEUED # Executor has enqueued the task + QUEUED = _IntermediateTIState.QUEUED # Executor has enqueued the task RUNNING = "running" # Task is executing - SUCCESS = TerminalTIState.SUCCESS # Task completed - RESTARTING = IntermediateTIState.RESTARTING # External request to restart (e.g. cleared when running) - FAILED = TerminalTIState.FAILED # Task errored out - UP_FOR_RETRY = IntermediateTIState.UP_FOR_RETRY # Task failed but has retries left - UP_FOR_RESCHEDULE = IntermediateTIState.UP_FOR_RESCHEDULE # A waiting `reschedule` sensor - UPSTREAM_FAILED = TerminalTIState.UPSTREAM_FAILED # One or more upstream deps failed - SKIPPED = TerminalTIState.SKIPPED # Skipped by branching or some other mechanism - DEFERRED = IntermediateTIState.DEFERRED # Deferrable operator waiting on a trigger - AWAITING_INPUT = IntermediateTIState.AWAITING_INPUT # Parked waiting for human input (HITL) + SUCCESS = _TerminalTIState.SUCCESS # Task completed + RESTARTING = _IntermediateTIState.RESTARTING # External request to restart (e.g. cleared when running) + FAILED = _TerminalTIState.FAILED # Task errored out + UP_FOR_RETRY = _IntermediateTIState.UP_FOR_RETRY # Task failed but has retries left + UP_FOR_RESCHEDULE = _IntermediateTIState.UP_FOR_RESCHEDULE # A waiting `reschedule` sensor + UPSTREAM_FAILED = _TerminalTIState.UPSTREAM_FAILED # One or more upstream deps failed + SKIPPED = _TerminalTIState.SKIPPED # Skipped by branching or some other mechanism + DEFERRED = _IntermediateTIState.DEFERRED # Deferrable operator waiting on a trigger + AWAITING_INPUT = _IntermediateTIState.AWAITING_INPUT # Parked waiting for human input (HITL) def __str__(self) -> str: return self.value @@ -235,19 +216,24 @@ def color_fg(cls, state): """ -def __getattr__(name: str): - """Provide backward compatibility for moved classes.""" - if name == "JobState": - import warnings - - from airflow.jobs.job import JobState +_MOVED_ATTRIBUTES = { + "IntermediateTIState": "airflow.task.state.IntermediateTIState", + "JobState": "airflow.jobs.job.JobState", + "TerminalTIState": "airflow.task.state.TerminalTIState", +} - warnings.warn( - "The `airflow.utils.state.JobState` attribute is deprecated and will be removed in a future version. " - "Please use `airflow.jobs.job.JobState` instead.", - DeprecationWarning, - stacklevel=2, - ) - return JobState - raise AttributeError(f"module 'airflow.utils.state' has no attribute '{name}'") +def __getattr__(name: str): + """Provide backward compatibility for moved classes.""" + target = _MOVED_ATTRIBUTES.get(name) + if target is None: + raise AttributeError(f"module 'airflow.utils.state' has no attribute '{name}'") + + warnings.warn( + f"The `airflow.utils.state.{name}` attribute is deprecated and will be removed in a future " + f"version. Please use `{target}` instead.", + DeprecationWarning, + stacklevel=2, + ) + module_name, attribute_name = target.rsplit(".", 1) + return getattr(import_module(module_name), attribute_name) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index bb3c0f7e5a785..74f0c9bc109d4 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -55,7 +55,8 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group from airflow.state.metastore import MetastoreBackend -from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState +from airflow.task.state import TerminalTIState +from airflow.utils.state import DagRunState, State, TaskInstanceState from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import ( diff --git a/airflow-core/tests/unit/utils/test_state.py b/airflow-core/tests/unit/utils/test_state.py index 88a1925842e8b..d93185cf5a06c 100644 --- a/airflow-core/tests/unit/utils/test_state.py +++ b/airflow-core/tests/unit/utils/test_state.py @@ -17,14 +17,17 @@ from __future__ import annotations from datetime import timedelta +from importlib import import_module import pytest from sqlalchemy import select +import airflow.utils.state from airflow.models.dagrun import DagRun from airflow.sdk import DAG +from airflow.task.state import IntermediateTIState, TerminalTIState from airflow.utils.session import create_session -from airflow.utils.state import DagRunState, IntermediateTIState, State, TaskInstanceState, TerminalTIState +from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.test_utils.dag import sync_dag_to_db @@ -139,3 +142,19 @@ def test_all_terminal_states_are_either_failed_or_success(self): f"All terminal states ({all_terminal_states}) except excluded ones ({excluded_states}) " f"should be classified as either failed or success ({classified_states})" ) + + +@pytest.mark.parametrize( + ("name", "target"), + [ + ("IntermediateTIState", "airflow.task.state.IntermediateTIState"), + ("JobState", "airflow.jobs.job.JobState"), + ("TerminalTIState", "airflow.task.state.TerminalTIState"), + ], +) +def test_moved_attributes_are_still_importable(name, target): + module_name, attribute_name = target.rsplit(".", 1) + expected = getattr(import_module(module_name), attribute_name) + + with pytest.warns(DeprecationWarning, match=rf"`airflow\.utils\.state\.{name}` attribute is deprecated"): + assert getattr(airflow.utils.state, name) is expected diff --git a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py index 7269a4a88a2f4..b7746570d9cc0 100644 --- a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py +++ b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py @@ -33,7 +33,7 @@ from airflow.triggers.base import BaseTrigger, TriggerEvent if TYPE_CHECKING: - from airflow.utils.state import TerminalTIState + from airflow.task.state import TerminalTIState class CloudComposerExecutionTrigger(BaseTrigger):