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