From d74717e5b6980661f1e1f31da31a98a50d2ac400 Mon Sep 17 00:00:00 2001
From: Victor Bustamante <victor@systeminit.com>
Date: Wed, 11 Dec 2024 17:28:52 -0300
Subject: [PATCH] Fix: Restore geometries too

Signed-off-by: Victor Bustamante <victor@systeminit.com>
---
 lib/dal/src/component.rs        |  4 ++++
 lib/dal/src/diagram/geometry.rs | 37 +++++++++++++++++++++++++++++++--
 lib/dal/src/diagram/view.rs     | 23 ++++++++++----------
 3 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/lib/dal/src/component.rs b/lib/dal/src/component.rs
index 3aa4cce195..fe0d0c7106 100644
--- a/lib/dal/src/component.rs
+++ b/lib/dal/src/component.rs
@@ -3332,6 +3332,10 @@ impl Component {
         ctx.add_dependent_values_and_enqueue(component.input_socket_attribute_values(ctx).await?)
             .await?;
 
+        Geometry::restore_all_for_component_id(ctx, component_id)
+            .await
+            .map_err(|e| ComponentError::Diagram(Box::new(e)))?;
+
         Ok(())
     }
 
diff --git a/lib/dal/src/diagram/geometry.rs b/lib/dal/src/diagram/geometry.rs
index 9dbfb1e120..bf4fd834e7 100644
--- a/lib/dal/src/diagram/geometry.rs
+++ b/lib/dal/src/diagram/geometry.rs
@@ -234,8 +234,8 @@ impl Geometry {
         Ok(geo_represents)
     }
 
-    pub async fn get_by_id(ctx: &DalContext, component_id: GeometryId) -> DiagramResult<Self> {
-        let (node_weight, content) = Self::get_node_weight_and_content(ctx, component_id).await?;
+    pub async fn get_by_id(ctx: &DalContext, geometry_id: GeometryId) -> DiagramResult<Self> {
+        let (node_weight, content) = Self::get_node_weight_and_content(ctx, geometry_id).await?;
         Ok(Self::assemble(node_weight, content))
     }
 
@@ -528,4 +528,37 @@ impl Geometry {
 
         Ok(())
     }
+
+    pub async fn restore_all_for_component_id(
+        ctx: &DalContext,
+        component_id: ComponentId,
+    ) -> DiagramResult<()> {
+        let base_change_set_ctx = ctx.clone_with_base().await?;
+
+        for geo_id in Self::list_ids_by_component(&base_change_set_ctx, component_id).await? {
+            let view_id = Self::get_view_id_by_id(&base_change_set_ctx, geo_id).await?;
+
+            // Check if view exists on this changeset
+            if View::try_get_by_id(ctx, view_id).await?.is_none() {
+                continue;
+            };
+
+            let head_geometry = Self::get_by_id(&base_change_set_ctx, geo_id).await?;
+
+            Self::new_for_component(ctx, component_id, view_id)
+                .await?
+                .update(
+                    ctx,
+                    RawGeometry {
+                        x: head_geometry.x(),
+                        y: head_geometry.y(),
+                        width: head_geometry.width(),
+                        height: head_geometry.height(),
+                    },
+                )
+                .await?;
+        }
+
+        Ok(())
+    }
 }
diff --git a/lib/dal/src/diagram/view.rs b/lib/dal/src/diagram/view.rs
index 422426101c..278d0df9a9 100644
--- a/lib/dal/src/diagram/view.rs
+++ b/lib/dal/src/diagram/view.rs
@@ -116,8 +116,18 @@ impl View {
     }
 
     pub async fn get_by_id(ctx: &DalContext, view_id: ViewId) -> DiagramResult<Self> {
-        let (node_weight, content) = Self::get_node_weight_and_content(ctx, view_id).await?;
-        Ok(Self::assemble(node_weight, content))
+        Self::try_get_by_id(ctx, view_id)
+            .await?
+            .ok_or(DiagramError::ViewNotFound(view_id))
+    }
+
+    pub async fn try_get_by_id(ctx: &DalContext, view_id: ViewId) -> DiagramResult<Option<Self>> {
+        let Some((node_weight, content)) =
+            Self::try_get_node_weight_and_content(ctx, view_id).await?
+        else {
+            return Ok(None);
+        };
+        Ok(Some(Self::assemble(node_weight, content)))
     }
 
     pub async fn find_by_name(ctx: &DalContext, name: &str) -> DiagramResult<Option<Self>> {
@@ -206,15 +216,6 @@ impl View {
         Ok(default_view.into())
     }
 
-    async fn get_node_weight_and_content(
-        ctx: &DalContext,
-        view_id: ViewId,
-    ) -> DiagramResult<(ViewNodeWeight, ViewContent)> {
-        Self::try_get_node_weight_and_content(ctx, view_id)
-            .await?
-            .ok_or(DiagramError::ViewNotFound(view_id))
-    }
-
     async fn try_get_node_weight_and_content(
         ctx: &DalContext,
         view_id: ViewId,