From 4090c421822a853216c181d390876583c27e01c8 Mon Sep 17 00:00:00 2001
From: Michael Schuster <michael.schuster.ffb@googlemail.com>
Date: Wed, 4 Dec 2024 18:12:50 +0100
Subject: [PATCH] Fix workload token expiration for cached steps/runs

---
 src/zenml/zen_server/auth.py                   | 10 ++--------
 src/zenml/zen_server/routers/auth_endpoints.py | 11 ++---------
 2 files changed, 4 insertions(+), 17 deletions(-)

diff --git a/src/zenml/zen_server/auth.py b/src/zenml/zen_server/auth.py
index a7918f3e81c..80290f091cc 100644
--- a/src/zenml/zen_server/auth.py
+++ b/src/zenml/zen_server/auth.py
@@ -413,10 +413,7 @@ def get_pipeline_run_status(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
-            if pipeline_run_status in [
-                ExecutionStatus.FAILED,
-                ExecutionStatus.COMPLETED,
-            ]:
+            if pipeline_run_status.is_finished:
                 error = (
                     f"The execution of pipeline run "
                     f"{decoded_token.pipeline_run_id} has already concluded and "
@@ -461,10 +458,7 @@ def get_step_run_status(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
-            if step_run_status in [
-                ExecutionStatus.FAILED,
-                ExecutionStatus.COMPLETED,
-            ]:
+            if step_run_status.is_finished:
                 error = (
                     f"The execution of step run "
                     f"{decoded_token.step_run_id} has already concluded and "
diff --git a/src/zenml/zen_server/routers/auth_endpoints.py b/src/zenml/zen_server/routers/auth_endpoints.py
index e970ba535e7..a1339c10bf3 100644
--- a/src/zenml/zen_server/routers/auth_endpoints.py
+++ b/src/zenml/zen_server/routers/auth_endpoints.py
@@ -41,7 +41,6 @@
 from zenml.enums import (
     APITokenType,
     AuthScheme,
-    ExecutionStatus,
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
@@ -589,10 +588,7 @@ def api_token(
                 "security reasons."
             )
 
-        if pipeline_run.status in [
-            ExecutionStatus.FAILED,
-            ExecutionStatus.COMPLETED,
-        ]:
+        if pipeline_run.status.is_finished:
             raise ValueError(
                 f"The execution of pipeline run {pipeline_run_id} has already "
                 "concluded and API tokens can no longer be generated for it "
@@ -609,10 +605,7 @@ def api_token(
                 "be generated for non-existent step runs for security reasons."
             )
 
-        if step_run.status in [
-            ExecutionStatus.FAILED,
-            ExecutionStatus.COMPLETED,
-        ]:
+        if step_run.status.is_finished:
             raise ValueError(
                 f"The execution of step run {step_run_id} has already "
                 "concluded and API tokens can no longer be generated for it "