diff --git a/components/job-orchestration/job_orchestration/scheduler/constants.py b/components/job-orchestration/job_orchestration/scheduler/constants.py index cd85016a3e..74cc2bfdcc 100644 --- a/components/job-orchestration/job_orchestration/scheduler/constants.py +++ b/components/job-orchestration/job_orchestration/scheduler/constants.py @@ -11,19 +11,32 @@ class QueueName: QUERY = "query" -class CompressionJobStatus(IntEnum): +class StatusIntEnum(IntEnum): + """ + Delegates __str__ to int.__str__, matching the behavior of IntEnum in Python 3.11+. + TODO: Remove this when our minimum supported Python version is 3.11+. + """ + + def __str__(self) -> str: + return str(self.value) + + def to_str(self) -> str: + return self.name + + +class CompressionJobStatus(StatusIntEnum): PENDING = 0 RUNNING = auto() SUCCEEDED = auto() FAILED = auto() -class CompressionJobCompletionStatus(IntEnum): +class CompressionJobCompletionStatus(StatusIntEnum): SUCCEEDED = 0 FAILED = auto() -class CompressionTaskStatus(IntEnum): +class CompressionTaskStatus(StatusIntEnum): PENDING = 0 RUNNING = auto() SUCCEEDED = auto() @@ -32,7 +45,7 @@ class CompressionTaskStatus(IntEnum): # When adding new states always add them to the end of this enum # and make necessary changes in the UI, Query Scheduler, and Reducer -class QueryJobStatus(IntEnum): +class QueryJobStatus(StatusIntEnum): PENDING = 0 RUNNING = auto() SUCCEEDED = auto() @@ -44,14 +57,8 @@ class QueryJobStatus(IntEnum): def from_str(label: str) -> QueryJobStatus: return QueryJobStatus[label.upper()] - def __str__(self) -> str: - return str(self.value) - - def to_str(self) -> str: - return str(self.name) - -class QueryTaskStatus(IntEnum): +class QueryTaskStatus(StatusIntEnum): PENDING = 0 RUNNING = auto() SUCCEEDED = auto() @@ -62,20 +69,8 @@ class QueryTaskStatus(IntEnum): def from_str(label: str) -> QueryTaskStatus: return QueryTaskStatus[label.upper()] - def __str__(self) -> str: - return str(self.value) - - def to_str(self) -> str: - return str(self.name) - -class QueryJobType(IntEnum): +class QueryJobType(StatusIntEnum): SEARCH_OR_AGGREGATION = 0 EXTRACT_IR = auto() EXTRACT_JSON = auto() - - def __str__(self) -> str: - return str(self.value) - - def to_str(self) -> str: - return str(self.name)