Remove state dependency from airflow core in sdk (apache#55292)

Remove state dependency from airflow core in sdk

Author: sunank200

task-sdk/src/airflow/sdk/definitions/dag.py

@@ -43,7 +43,7 @@
     RemovedInAirflow4Warning,
     TaskNotFound,
 )
-from airflow.sdk import TriggerRule
+from airflow.sdk import TaskInstanceState, TriggerRule
 from airflow.sdk.bases.operator import BaseOperator
 from airflow.sdk.definitions._internal.node import validate_key
 from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
@@ -85,6 +85,15 @@
     "dag",
 ]
 
+FINISHED_STATES = frozenset(
+    [
+        TaskInstanceState.SUCCESS,
+        TaskInstanceState.FAILED,
+        TaskInstanceState.SKIPPED,
+        TaskInstanceState.UPSTREAM_FAILED,
+        TaskInstanceState.REMOVED,
+    ]
+)
 
 DagStateChangeCallback = Callable[[Context], None]
 ScheduleInterval = None | str | timedelta | relativedelta
@@ -1166,10 +1175,9 @@ def test(
         from airflow import settings
         from airflow.configuration import secrets_backend_list
         from airflow.models.dagrun import DagRun, get_or_create_dagrun
-        from airflow.sdk import DagRunState, TaskInstanceState, timezone
+        from airflow.sdk import DagRunState, timezone
         from airflow.secrets.local_filesystem import LocalFilesystemBackend
         from airflow.serialization.serialized_objects import SerializedDAG
-        from airflow.utils.state import State
         from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
         exit_stack = ExitStack()
@@ -1291,7 +1299,7 @@ def test(
                 # triggerer may mark tasks scheduled so we read from DB
                 all_tis = set(dr.get_task_instances(session=session))
                 scheduled_tis = {x for x in all_tis if x.state == TaskInstanceState.SCHEDULED}
-                ids_unrunnable = {x for x in all_tis if x.state not in State.finished} - scheduled_tis
+                ids_unrunnable = {x for x in all_tis if x.state not in FINISHED_STATES} - scheduled_tis
                 if not scheduled_tis and ids_unrunnable:
                     log.warning("No tasks to run. unrunnable tasks: %s", ids_unrunnable)
                     time.sleep(1)
@@ -1331,7 +1339,7 @@ def test(
                         # Run the task locally
                         try:
                             if mark_success:
-                                ti.set_state(State.SUCCESS)
+                                ti.set_state(TaskInstanceState.SUCCESS)
                                 log.info("[DAG TEST] Marking success for %s on %s", task, ti.logical_date)
                             else:
                                 _run_task(ti=ti, task=task, run_triggerer=True)
@@ -1363,7 +1371,6 @@ def _run_task(
     possible.  This function is only meant for the `dag.test` function as a helper function.
     """
     from airflow.sdk.module_loading import import_string
-    from airflow.utils.state import State
 
     taskrun_result: TaskRunResult | None
     log.info("[DAG TEST] starting task_id=%s map_index=%s", ti.task_id, ti.map_index)
@@ -1378,7 +1385,7 @@ def _run_task(
 
             # The API Server expects the task instance to be in QUEUED state before
             # it is run.
-            ti.set_state(State.QUEUED)
+            ti.set_state(TaskInstanceState.QUEUED)
             task_sdk_ti = TaskInstanceSDK(
                 id=ti.id,
                 task_id=ti.task_id,
@@ -1394,12 +1401,12 @@ def _run_task(
             ti.set_state(taskrun_result.ti.state)
             ti.task = create_scheduler_operator(taskrun_result.ti.task)
 
-            if ti.state == State.DEFERRED and isinstance(msg, DeferTask) and run_triggerer:
+            if ti.state == TaskInstanceState.DEFERRED and isinstance(msg, DeferTask) and run_triggerer:
                 from airflow.utils.session import create_session
 
                 # API Server expects the task instance to be in QUEUED state before
                 # resuming from deferral.
-                ti.set_state(State.QUEUED)
+                ti.set_state(TaskInstanceState.QUEUED)
 
                 log.info("[DAG TEST] running trigger in line")
                 trigger = import_string(msg.classpath)(**msg.trigger_kwargs)
@@ -1410,15 +1417,15 @@ def _run_task(
 
                 # Set the state to SCHEDULED so that the task can be resumed.
                 with create_session() as session:
-                    ti.state = State.SCHEDULED
+                    ti.state = TaskInstanceState.SCHEDULED
                     session.add(ti)
                 continue
 
             break
         except Exception:
             log.exception("[DAG TEST] Error running task %s", ti)
-            if ti.state not in State.finished:
-                ti.set_state(State.FAILED)
+            if ti.state not in FINISHED_STATES:
+                ti.set_state(TaskInstanceState.FAILED)
                 taskrun_result = None
                 break
             raise

task-sdk/tests/task_sdk/api/test_client.py

@@ -39,6 +39,7 @@
     DagRunStateResponse,
     HITLDetailResponse,
     HITLUser,
+    TerminalTIState,
     VariableResponse,
     XComResponse,
 )
@@ -52,7 +53,6 @@
     RescheduleTask,
     TaskRescheduleStartDate,
 )
-from airflow.utils.state import TerminalTIState
 
 if TYPE_CHECKING:
     from time_machine import TimeMachineFixture

task-sdk/tests/task_sdk/bases/test_sensor.py

@@ -34,12 +34,11 @@
 )
 from airflow.models.trigger import TriggerFailureReason
 from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.sdk import timezone
+from airflow.sdk import TaskInstanceState, timezone
 from airflow.sdk.bases.sensor import BaseSensorOperator, PokeReturnValue, poke_mode_only
 from airflow.sdk.definitions.dag import DAG
 from airflow.sdk.execution_time.comms import RescheduleTask, TaskRescheduleStartDate
 from airflow.sdk.timezone import datetime
-from airflow.utils.state import State
 
 if TYPE_CHECKING:
     from airflow.sdk.definitions.context import Context
@@ -172,22 +171,22 @@ def test_ok_with_reschedule(self, run_task, make_sensor, time_machine):
 
         state, msg, _ = run_task(task=sensor)
 
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
         assert msg.reschedule_date == date1 + timedelta(seconds=sensor.poke_interval)
 
         # second poke returns False and task is re-scheduled
         time_machine.coordinates.shift(sensor.poke_interval)
         date2 = date1 + timedelta(seconds=sensor.poke_interval)
         state, msg, _ = run_task(task=sensor)
 
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
         assert msg.reschedule_date == date2 + timedelta(seconds=sensor.poke_interval)
 
         # third poke returns True and task succeeds
         time_machine.coordinates.shift(sensor.poke_interval)
         state, _, _ = run_task(task=sensor)
 
-        assert state == State.SUCCESS
+        assert state == TaskInstanceState.SUCCESS
 
     def test_fail_with_reschedule(self, run_task, make_sensor, time_machine, mock_supervisor_comms):
         sensor = make_sensor(return_value=False, poke_interval=10, timeout=5, mode="reschedule")
@@ -198,7 +197,7 @@ def test_fail_with_reschedule(self, run_task, make_sensor, time_machine, mock_su
 
         state, msg, _ = run_task(task=sensor)
 
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
         assert msg.reschedule_date == date1 + timedelta(seconds=sensor.poke_interval)
 
         # second poke returns False, timeout occurs
@@ -208,7 +207,7 @@ def test_fail_with_reschedule(self, run_task, make_sensor, time_machine, mock_su
         mock_supervisor_comms.send.return_value = TaskRescheduleStartDate(start_date=date1)
         state, msg, error = run_task(task=sensor, context_update={"task_reschedule_count": 1})
 
-        assert state == State.FAILED
+        assert state == TaskInstanceState.FAILED
         assert isinstance(error, AirflowSensorTimeout)
 
     def test_soft_fail_with_reschedule(self, run_task, make_sensor, time_machine, mock_supervisor_comms):
@@ -221,15 +220,15 @@ def test_soft_fail_with_reschedule(self, run_task, make_sensor, time_machine, mo
         time_machine.move_to(date1, tick=False)
 
         state, msg, _ = run_task(task=sensor)
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
 
         # second poke returns False, timeout occurs
         time_machine.coordinates.shift(sensor.poke_interval)
 
         # Mocking values from DB/API-server
         mock_supervisor_comms.send.return_value = TaskRescheduleStartDate(start_date=date1)
         state, msg, _ = run_task(task=sensor, context_update={"task_reschedule_count": 1})
-        assert state == State.SKIPPED
+        assert state == TaskInstanceState.SKIPPED
 
     def test_ok_with_reschedule_and_exponential_backoff(
         self, run_task, make_sensor, time_machine, mock_supervisor_comms
@@ -261,7 +260,7 @@ def run_duration():
             curr_date = curr_date + timedelta(seconds=new_interval)
             time_machine.coordinates.shift(new_interval)
             state, msg, _ = run_task(sensor, context_update={"task_reschedule_count": _poke_count})
-            assert state == State.UP_FOR_RESCHEDULE
+            assert state == TaskInstanceState.UP_FOR_RESCHEDULE
             old_interval = new_interval
             new_interval = sensor._get_next_poke_interval(task_start_date, run_duration, _poke_count)
             assert old_interval < new_interval  # actual test
@@ -272,7 +271,7 @@ def run_duration():
         time_machine.coordinates.shift(new_interval)
 
         state, msg, _ = run_task(sensor, context_update={"task_reschedule_count": false_count + 1})
-        assert state == State.SUCCESS
+        assert state == TaskInstanceState.SUCCESS
 
     def test_invalid_mode(self):
         with pytest.raises(AirflowException):
@@ -291,22 +290,22 @@ def test_ok_with_custom_reschedule_exception(self, make_sensor, run_task):
         with time_machine.travel(date1, tick=False):
             state, msg, error = run_task(sensor)
 
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
         assert isinstance(msg, RescheduleTask)
         assert msg.reschedule_date == date2
 
         # second poke returns False and task is re-scheduled
         with time_machine.travel(date2, tick=False):
             state, msg, error = run_task(sensor)
 
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
         assert isinstance(msg, RescheduleTask)
         assert msg.reschedule_date == date3
 
         # third poke returns True and task succeeds
         with time_machine.travel(date3, tick=False):
             state, _, _ = run_task(sensor)
-        assert state == State.SUCCESS
+        assert state == TaskInstanceState.SUCCESS
 
     def test_sensor_with_invalid_poke_interval(self):
         negative_poke_interval = -10
@@ -523,36 +522,36 @@ def _run_task():
                 context_update={"task_reschedule_count": test_state["task_reschedule_count"]},
             )
 
-            if state == State.UP_FOR_RESCHEDULE:
+            if state == TaskInstanceState.UP_FOR_RESCHEDULE:
                 test_state["task_reschedule_count"] += 1
                 # Only set first_reschedule_date on the first successful reschedule
                 if test_state["first_reschedule_date"] is None:
                     test_state["first_reschedule_date"] = test_state["current_time"]
-            elif state == State.UP_FOR_RETRY:
+            elif state == TaskInstanceState.UP_FOR_RETRY:
                 test_state["try_number"] += 1
             return state, msg, error
 
         # Phase 1: Initial execution until failure
         # First poke - should reschedule
         state, _, _ = _run_task()
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
 
         # Second poke - should raise RuntimeError and retry
         test_state["current_time"] += timedelta(seconds=sensor.poke_interval)
         state, _, error = _run_task()
-        assert state == State.UP_FOR_RETRY
+        assert state == TaskInstanceState.UP_FOR_RETRY
         assert isinstance(error, RuntimeError)
 
         # Third poke - should reschedule again
         test_state["current_time"] += sensor.retry_delay + timedelta(seconds=1)
         state, _, _ = _run_task()
-        assert state == State.UP_FOR_RESCHEDULE
+        assert state == TaskInstanceState.UP_FOR_RESCHEDULE
 
         # Fourth poke - should timeout
         test_state["current_time"] += timedelta(seconds=sensor.poke_interval)
         state, _, error = _run_task()
         assert isinstance(error, AirflowSensorTimeout)
-        assert state == State.FAILED
+        assert state == TaskInstanceState.FAILED
 
         # Phase 2: After clearing the failed sensor
         # Reset supervisor comms to return None, simulating a fresh start after clearing
@@ -564,13 +563,13 @@ def _run_task():
         for _ in range(3):
             test_state["current_time"] += timedelta(seconds=sensor.poke_interval)
             state, _, _ = _run_task()
-            assert state == State.UP_FOR_RESCHEDULE
+            assert state == TaskInstanceState.UP_FOR_RESCHEDULE
 
         # Final poke - should timeout
         test_state["current_time"] += timedelta(seconds=sensor.poke_interval)
         state, _, error = _run_task()
         assert isinstance(error, AirflowSensorTimeout)
-        assert state == State.FAILED
+        assert state == TaskInstanceState.FAILED
 
     def test_sensor_with_xcom(self, make_sensor):
         xcom_value = "TestValue"
@@ -615,7 +614,7 @@ def timeout():
         state, _, error = run_task(task=task, dag_id=f"test_sensor_timeout_{mode}_{retries}")
 
         assert isinstance(error, AirflowSensorTimeout)
-        assert state == State.FAILED
+        assert state == TaskInstanceState.FAILED
 
 
 @poke_mode_only