Fix GET step run endpoint to return unhydrated response if requested (#…

…3240) * Fix GET step run endpoint to return unhydrated response if requested * More hydration fixes * Typo
zenml-io · Dec 4, 2024 · e72aef7 · e72aef7
1 parent 3812ca5
commit e72aef7
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 8 deletions.
diff --git 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)
 
 

diff --git 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 want 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
@@ -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