Remove SHUTDOWN status (#2027)

# Description Remove the task instance state which no longer exists. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Ephraim Anierobi <[email protected]>
astronomer · Aug 30, 2023 · d32832c · d32832c
1 parent 7583814
commit d32832c
Showing 1 changed file with 7 additions and 9 deletions.
diff --git a/python-sdk/src/astro/sql/operators/cleanup.py b/python-sdk/src/astro/sql/operators/cleanup.py
@@ -19,8 +19,7 @@
     MappedOperator = None
 
 
-from airflow.models.taskinstance import TaskInstance
-from airflow.utils.state import State
+from airflow.models.taskinstance import TaskInstance, TaskInstanceState
 
 from astro.databases import create_database
 from astro.sql.operators.base_decorator import BaseSQLDecoratedOperator
@@ -134,7 +133,7 @@ def _has_task_failed(self, task_instances: list[TaskInstance]) -> bool:
         failed_tasks = [
             (ti.task_id, ti.state)
             for ti in task_instances
-            if ti.task_id != self.task_id and ti.state == State.FAILED
+            if ti.task_id != self.task_id and ti.state == TaskInstanceState.FAILED
         ]
         if failed_tasks:
             self.log.info(
@@ -159,12 +158,11 @@ def _is_dag_running(self, task_instances: list[TaskInstance]) -> bool:
             if ti.task_id != self.task_id
             and ti.state
             not in [
-                State.SUCCESS,
-                State.FAILED,
-                State.SKIPPED,
-                State.UPSTREAM_FAILED,
-                State.REMOVED,
-                State.SHUTDOWN,
+                TaskInstanceState.SUCCESS,
+                TaskInstanceState.FAILED,
+                TaskInstanceState.SKIPPED,
+                TaskInstanceState.UPSTREAM_FAILED,
+                TaskInstanceState.REMOVED,
             ]
         ]
         if running_tasks: