From 30e062a66b2651cda80a3af728960c8b1d339c43 Mon Sep 17 00:00:00 2001
From: Tyler Ohlsen <ohltyler@amazon.com>
Date: Tue, 17 Oct 2023 19:31:21 -0700
Subject: [PATCH] Update schema on component deletion

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
---
 .../workspace/resizable_workspace.tsx         | 25 +++++++++++++++++++
 .../workflow_detail/workspace/workspace.tsx   |  1 +
 .../workspace_component.tsx                   |  6 ++++-
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/public/pages/workflow_detail/workspace/resizable_workspace.tsx b/public/pages/workflow_detail/workspace/resizable_workspace.tsx
index 2c5f8223..ecabbdf4 100644
--- a/public/pages/workflow_detail/workspace/resizable_workspace.tsx
+++ b/public/pages/workflow_detail/workspace/resizable_workspace.tsx
@@ -6,6 +6,7 @@
 import React, { useRef, useState, useEffect } from 'react';
 import { ReactFlowProvider } from 'reactflow';
 import { Form, Formik } from 'formik';
+import { ObjectSchema } from 'yup';
 import * as yup from 'yup';
 import { EuiButton, EuiResizableContainer } from '@elastic/eui';
 import {
@@ -59,6 +60,25 @@ export function ResizableWorkspace(props: ResizableWorkspaceProps) {
     }
   }, [props.workflow]);
 
+  // If a component is added or deleted in the workspace, we update the values.
+  // However, we cannot dynamically update the schema due to this known
+  // formik bug: https://github.com/jaredpalmer/formik/issues/3335
+  // So, we trigger this during validation when there is a mismatch
+  // in the count of component ids
+  function onComponentChange(values: WorkspaceFormValues): void {
+    const updatedComponentIds = Object.keys(values);
+    const newSchemaObj = {} as WorkspaceSchemaObj;
+
+    Object.keys(formSchema.fields).forEach((componentId) => {
+      if (updatedComponentIds.includes(componentId)) {
+        newSchemaObj[componentId] = formSchema.fields[
+          `${componentId}`
+        ] as ObjectSchema<any, any, any>;
+      }
+    });
+    setFormSchema(yup.object(newSchemaObj));
+  }
+
   return (
     <Formik
       enableReinitialize={true}
@@ -69,6 +89,11 @@ export function ResizableWorkspace(props: ResizableWorkspaceProps) {
       }}
       validate={(values) => {
         console.log('values on validate: ', values);
+        const componentCountValues = Object.keys(values).length;
+        const componentCountSchema = Object.keys(formSchema.fields).length;
+        if (componentCountSchema !== componentCountValues) {
+          onComponentChange(values);
+        }
       }}
     >
       {(formikProps) => (
diff --git a/public/pages/workflow_detail/workspace/workspace.tsx b/public/pages/workflow_detail/workspace/workspace.tsx
index 6762da3d..4f004ea9 100644
--- a/public/pages/workflow_detail/workspace/workspace.tsx
+++ b/public/pages/workflow_detail/workspace/workspace.tsx
@@ -97,6 +97,7 @@ export function Workspace(props: WorkspaceProps) {
 
       // TODO: on node addition, need to update the formik values to include
       // a new property with this id.
+      // TODO: on node addition, update yup schema
       setNodes((nds) => nds.concat(newNode));
       dispatch(setDirty());
     },
diff --git a/public/pages/workflow_detail/workspace_component/workspace_component.tsx b/public/pages/workflow_detail/workspace_component/workspace_component.tsx
index d25ce46a..17f0b8b7 100644
--- a/public/pages/workflow_detail/workspace_component/workspace_component.tsx
+++ b/public/pages/workflow_detail/workspace_component/workspace_component.tsx
@@ -30,7 +30,7 @@ interface WorkspaceComponentProps {
 export function WorkspaceComponent(props: WorkspaceComponentProps) {
   const component = props.data;
   const { deleteNode } = useContext(rfContext);
-  const { values } = useFormikContext<WorkspaceFormValues>();
+  const { values, validateForm } = useFormikContext<WorkspaceFormValues>();
 
   return (
     <EuiCard
@@ -48,6 +48,10 @@ export function WorkspaceComponent(props: WorkspaceComponentProps) {
               onClick={() => {
                 deleteNode(component.id);
                 delete values[`${component.id}`];
+                validateForm();
+                // TODO: use below way instead of hacky change via validation.
+                // Formik bug reference: https://github.com/jaredpalmer/formik/issues/3335
+                // delete validationSchema[`${component.id}`];
               }}
               aria-label="Delete"
             />