Move task instance state enums out of airflow.utils.state

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(),

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
 

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

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

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)

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 (

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

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):