Repository: apache/airflow

Issue:
Use enum for task_states/dag_states
Hello,

We use Python 3+, so we should consider using an enumerated type for objects that are buckets for contents  This will give us better type hints.
https://github.com/apache/airflow/blob/master/airflow/utils/state.py

Best regards,
Kamil Breguła


Discussion:
CC: @turbaszek 
Hi,
I'd like to help out with this issue.
@muyajil should I assign you?
@muyajil I assigned you to this ticket. I look forward to your contributions 😻

@turbaszek I think we can assign people without confirmation. The assignment of people is important because it is visible on the ticket list. This makes finding free tickets easier.
Is this still an open issue?
@chutcheson Yes it is, unfortunately I never got the chance to really get started

I'll take this up, if someone wants to assign me.
I would like to ask, though - do we just want the type hinting here (in which case we can just declare a nice Literal type for the state tuples), or take on the problems that changing a tuple to an enum will cause? There's a few places in Airflow, and presumably in some third-party code, that presume they can do operations on `State.task_states` as if it were a sliceable/iterable.

(and also SQLAlchemy will not like the Enum values, resulting in a lot of casting back to string)
I briefly read related source, and it seems to me the main advantage would be to make sure the right values are passed the right value to SQLAlchemy, so `Literal` feels like a better solution. `Enum` may even be a bit awkward in some places, e.g.

https://github.com/apache/airflow/blob/44a6648fd7482f61ca59153f0caaf0b019c5b3fd/airflow/models/dagrun.py#L490
Yeah, after having done some exploratory coding here, it's basically going to require a big separate Literal type that probably duplicates a lot of the values. I'll draw it up, but we can see if it's worth it in the PR discussion.

Resolving patch:
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -49,7 +49,7 @@
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import provide_session
 from airflow.utils.sqlalchemy import UtcDateTime, nulls_first, skip_locked, with_row_locks
-from airflow.utils.state import State
+from airflow.utils.state import State, TaskInstanceState
 from airflow.utils.types import DagRunType
 
 if TYPE_CHECKING:
@@ -314,7 +314,9 @@ def generate_run_id(run_type: DagRunType, execution_date: datetime) -> str:
         return f"{run_type}__{execution_date.isoformat()}"
 
     @provide_session
-    def get_task_instances(self, state=None, session=None) -> Iterable[TI]:
+    def get_task_instances(
+        self, state: Optional[Iterable[TaskInstanceState]] = None, session=None
+    ) -> Iterable[TI]:
         """Returns the task instances for this dag run"""
         tis = session.query(TI).filter(
             TI.dag_id == self.dag_id,
diff --git a/airflow/typing_compat.py b/airflow/typing_compat.py
--- a/airflow/typing_compat.py
+++ b/airflow/typing_compat.py
@@ -22,7 +22,7 @@
 """
 
 try:
-    # Protocol and TypedDict are only added to typing module starting from
+    # Literal, Protocol and TypedDict are only added to typing module starting from
     # python 3.8 we can safely remove this shim import after Airflow drops
     # support for <3.8
     from typing import Literal, Protocol, TypedDict, runtime_checkable  # type: ignore
diff --git a/airflow/utils/state.py b/airflow/utils/state.py
--- a/airflow/utils/state.py
+++ b/airflow/utils/state.py
@@ -16,70 +16,101 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from enum import Enum
+from typing import Dict, FrozenSet, Tuple
+
 from airflow.settings import STATE_COLORS
+from airflow.utils.types import Optional
 
 
-class State:
+class TaskInstanceState(str, Enum):
     """
-    Static class with task instance states constants and color method to
-    avoid hardcoding.
+    Enum that represents all possible states that a Task Instance can be in.
+
+    Note that None is also allowed, so always use this in a type hint with Optional.
     """
 
-    # scheduler
-    NONE = None  # type: None
-    REMOVED = "removed"
-    SCHEDULED = "scheduled"
+    # Set by the scheduler
+    # None - Task is created but should not run yet
+    REMOVED = "removed"  # Task vanished from DAG before it ran
+    SCHEDULED = "scheduled"  # Task should run and will be handed to executor soon
+
+    # Set by the task instance itself
+    QUEUED = "queued"  # Executor has enqueued the task
+    RUNNING = "running"  # Task is executing
+    SUCCESS = "success"  # Task completed
+    SHUTDOWN = "shutdown"  # External request to shut down
+    FAILED = "failed"  # Task errored out
+    UP_FOR_RETRY = "up_for_retry"  # Task failed but has retries left
+    UP_FOR_RESCHEDULE = "up_for_reschedule"  # A waiting `reschedule` sensor
+    UPSTREAM_FAILED = "upstream_failed"  # One or more upstream deps failed
+    SKIPPED = "skipped"  # Skipped by branching or some other mechanism
+    SENSING = "sensing"  # Smart sensor offloaded to the sensor DAG
 
-    # set by the executor (t.b.d.)
-    # LAUNCHED = "launched"
+    def __str__(self) -> str:  # pylint: disable=invalid-str-returned
+        return self.value
+
+
+class DagRunState(str, Enum):
+    """
+    Enum that represents all possible states that a DagRun can be in.
+
+    These are "shared" with TaskInstanceState in some parts of the code,
+    so please ensure that their values always match the ones with the
+    same name in TaskInstanceState.
+    """
 
-    # set by a task
-    QUEUED = "queued"
     RUNNING = "running"
     SUCCESS = "success"
-    SHUTDOWN = "shutdown"  # External request to shut down
     FAILED = "failed"
-    UP_FOR_RETRY = "up_for_retry"
-    UP_FOR_RESCHEDULE = "up_for_reschedule"
-    UPSTREAM_FAILED = "upstream_failed"
-    SKIPPED = "skipped"
-    SENSING = "sensing"
-
-    task_states = (
-        SUCCESS,
-        RUNNING,
-        FAILED,
-        UPSTREAM_FAILED,
-        SKIPPED,
-        UP_FOR_RETRY,
-        UP_FOR_RESCHEDULE,
-        QUEUED,
-        NONE,
-        SCHEDULED,
-        SENSING,
-        REMOVED,
-    )
 
-    dag_states = (
-        SUCCESS,
-        RUNNING,
-        FAILED,
+
+class State:
+    """
+    Static class with task instance state constants and color methods to
+    avoid hardcoding.
+    """
+
+    # Backwards-compat constants for code that does not yet use the enum
+    # These first three are shared by DagState and TaskState
+    SUCCESS = TaskInstanceState.SUCCESS
+    RUNNING = TaskInstanceState.RUNNING
+    FAILED = TaskInstanceState.FAILED
+
+    # These are TaskState only
+    NONE = None
+    REMOVED = TaskInstanceState.REMOVED
+    SCHEDULED = TaskInstanceState.SCHEDULED
+    QUEUED = TaskInstanceState.QUEUED
+    SHUTDOWN = TaskInstanceState.SHUTDOWN
+    UP_FOR_RETRY = TaskInstanceState.UP_FOR_RETRY
+    UP_FOR_RESCHEDULE = TaskInstanceState.UP_FOR_RESCHEDULE
+    UPSTREAM_FAILED = TaskInstanceState.UPSTREAM_FAILED
+    SKIPPED = TaskInstanceState.SKIPPED
+    SENSING = TaskInstanceState.SENSING
+
+    task_states: Tuple[Optional[TaskInstanceState], ...] = (None,) + tuple(TaskInstanceState)
+
+    dag_states: Tuple[DagRunState, ...] = (
+        DagRunState.SUCCESS,
+        DagRunState.RUNNING,
+        DagRunState.FAILED,
     )
 
-    state_color = {
-        QUEUED: 'gray',
-        RUNNING: 'lime',
-        SUCCESS: 'green',
-        SHUTDOWN: 'blue',
-        FAILED: 'red',
-        UP_FOR_RETRY: 'gold',
-        UP_FOR_RESCHEDULE: 'turquoise',
-        UPSTREAM_FAILED: 'orange',
-        SKIPPED: 'pink',
-        REMOVED: 'lightgrey',
-        SCHEDULED: 'tan',
-        NONE: 'lightblue',
-        SENSING: 'lightseagreen',
+    state_color: Dict[Optional[TaskInstanceState], str] = {
+        None: 'lightblue',
+        TaskInstanceState.QUEUED: 'gray',
+        TaskInstanceState.RUNNING: 'lime',
+        TaskInstanceState.SUCCESS: 'green',
+        TaskInstanceState.SHUTDOWN: 'blue',
+        TaskInstanceState.FAILED: 'red',
+        TaskInstanceState.UP_FOR_RETRY: 'gold',
+        TaskInstanceState.UP_FOR_RESCHEDULE: 'turquoise',
+        TaskInstanceState.UPSTREAM_FAILED: 'orange',
+        TaskInstanceState.SKIPPED: 'pink',
+        TaskInstanceState.REMOVED: 'lightgrey',
+        TaskInstanceState.SCHEDULED: 'tan',
+        TaskInstanceState.SENSING: 'lightseagreen',
     }
     state_color.update(STATE_COLORS)  # type: ignore
 
@@ -96,17 +127,17 @@ def color_fg(cls, state):
             return 'white'
         return 'black'
 
-    running = frozenset([RUNNING, SENSING])
+    running: FrozenSet[TaskInstanceState] = frozenset([TaskInstanceState.RUNNING, TaskInstanceState.SENSING])
     """
     A list of states indicating that a task is being executed.
     """
 
-    finished = frozenset(
+    finished: FrozenSet[TaskInstanceState] = frozenset(
         [
-            SUCCESS,
-            FAILED,
-            SKIPPED,
-            UPSTREAM_FAILED,
+            TaskInstanceState.SUCCESS,
+            TaskInstanceState.FAILED,
+            TaskInstanceState.SKIPPED,
+            TaskInstanceState.UPSTREAM_FAILED,
         ]
     )
     """
@@ -118,16 +149,16 @@ def color_fg(cls, state):
     case, it is no longer running.
     """
 
-    unfinished = frozenset(
+    unfinished: FrozenSet[Optional[TaskInstanceState]] = frozenset(
         [
-            NONE,
-            SCHEDULED,
-            QUEUED,
-            RUNNING,
-            SENSING,
-            SHUTDOWN,
-            UP_FOR_RETRY,
-            UP_FOR_RESCHEDULE,
+            None,
+            TaskInstanceState.SCHEDULED,
+            TaskInstanceState.QUEUED,
+            TaskInstanceState.RUNNING,
+            TaskInstanceState.SENSING,
+            TaskInstanceState.SHUTDOWN,
+            TaskInstanceState.UP_FOR_RETRY,
+            TaskInstanceState.UP_FOR_RESCHEDULE,
         ]
     )
     """
@@ -135,12 +166,16 @@ def color_fg(cls, state):
     a run or has not even started.
     """
 
-    failed_states = frozenset([FAILED, UPSTREAM_FAILED])
+    failed_states: FrozenSet[TaskInstanceState] = frozenset(
+        [TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED]
+    )
     """
     A list of states indicating that a task or dag is a failed state.
     """
 
-    success_states = frozenset([SUCCESS, SKIPPED])
+    success_states: FrozenSet[TaskInstanceState] = frozenset(
+        [TaskInstanceState.SUCCESS, TaskInstanceState.SKIPPED]
+    )
     """
     A list of states indicating that a task or dag is a success state.
     """


Test patch: