From 42fc70bea34386262160099f2b73737825e62c01 Mon Sep 17 00:00:00 2001
From: Michael Schuster <michael.schuster.ffb@googlemail.com>
Date: Wed, 4 Dec 2024 09:44:50 +0100
Subject: [PATCH 1/3] Fix GET step run endpoint to return unhydrated response
 if requested

---
 src/zenml/zen_server/routers/steps_endpoints.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/zenml/zen_server/routers/steps_endpoints.py b/src/zenml/zen_server/routers/steps_endpoints.py
index 1fa820265b9..aca76f9af91 100644
--- a/src/zenml/zen_server/routers/steps_endpoints.py
+++ b/src/zenml/zen_server/routers/steps_endpoints.py
@@ -142,10 +142,16 @@ def get_step(
     Returns:
         The step.
     """
-    step = zen_store().get_run_step(step_id, hydrate=hydrate)
+    # We always fetch the step hydrated because we need the pipeline_run_id
+    # for the permission checks. If the user requested an unhydrated response,
+    # we later remove the metadata
+    step = zen_store().get_run_step(step_id, hydrate=True)
     pipeline_run = zen_store().get_run(step.pipeline_run_id)
     verify_permission_for_model(pipeline_run, action=Action.READ)
 
+    if hydrate is False:
+        step.metadata = None
+
     return dehydrate_response_model(step)
 
 

From 2dc501c10f72cb621b8e75b403ac18ecd6206104 Mon Sep 17 00:00:00 2001
From: Michael Schuster <michael.schuster.ffb@googlemail.com>
Date: Wed, 4 Dec 2024 11:22:15 +0100
Subject: [PATCH 2/3] More hydration fixes

---
 .../zen_stores/schemas/pipeline_run_schemas.py | 18 ++++++++++++------
 src/zenml/zen_stores/sql_zen_store.py          |  6 +++++-
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
index d0af218b629..d92de22fc27 100644
--- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
+++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
@@ -298,7 +298,7 @@ def to_model(
         )
 
         if self.deployment is not None:
-            deployment = self.deployment.to_model()
+            deployment = self.deployment.to_model(include_metadata=True)
 
             config = deployment.pipeline_configuration
             new_substitutions = config._get_full_substitutions(self.start_time)
@@ -365,12 +365,18 @@ def to_model(
             ):
                 is_templatable = True
 
-            steps = {step.name: step.to_model() for step in self.step_runs}
-
-            step_substitutions = {
-                step_name: step.config.substitutions
-                for step_name, step in steps.items()
+            steps = {
+                step.name: step.to_model(include_metadata=True)
+                for step in self.step_runs
             }
+
+            step_substitutions = {}
+            for step_name, step in steps.items():
+                step_substitutions[step_name] = step.config.substitutions
+                # We fetch the steps hydrated before, but won't them unhydrated
+                # in the response -> We need to reset the metadata here
+                step.metadata = None
+
             metadata = PipelineRunResponseMetadata(
                 workspace=self.workspace.to_model(),
                 run_metadata=self.fetch_metadata(),
diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py
index 5f44873e87b..bb3a77befbd 100644
--- a/src/zenml/zen_stores/sql_zen_store.py
+++ b/src/zenml/zen_stores/sql_zen_store.py
@@ -8633,7 +8633,11 @@ def _update_pipeline_run_status(
 
         # Deployment always exists for pipeline runs of newer versions
         assert pipeline_run.deployment
-        num_steps = len(pipeline_run.deployment.to_model().step_configurations)
+        num_steps = len(
+            pipeline_run.deployment.to_model(
+                include_metadata=True
+            ).step_configurations
+        )
         new_status = get_pipeline_run_status(
             step_statuses=[
                 ExecutionStatus(step_run.status) for step_run in step_runs

From bee43d37dba6f5d42d13a0c61cd674dfd04f654d Mon Sep 17 00:00:00 2001
From: Michael Schuster <michael.schuster.ffb@googlemail.com>
Date: Wed, 4 Dec 2024 11:25:55 +0100
Subject: [PATCH 3/3] Typo

---
 src/zenml/zen_stores/schemas/pipeline_run_schemas.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
index d92de22fc27..67236f0ab7d 100644
--- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
+++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
@@ -373,7 +373,7 @@ def to_model(
             step_substitutions = {}
             for step_name, step in steps.items():
                 step_substitutions[step_name] = step.config.substitutions
-                # We fetch the steps hydrated before, but won't them unhydrated
+                # We fetch the steps hydrated before, but want them unhydrated
                 # in the response -> We need to reset the metadata here
                 step.metadata = None