From d2b33f6a2d4e3180aea2b166c0d5ce40fc615b87 Mon Sep 17 00:00:00 2001
From: Tyler Ohlsen <ohltyler@amazon.com>
Date: Thu, 5 Oct 2023 10:44:50 -0700
Subject: [PATCH] Add edge deletion and custom edge styling (#52)

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
---
 .../workspace/reactflow-styles.scss           | 20 +++++-
 .../workflow_detail/workspace/workspace.tsx   | 15 +++--
 .../workspace_component/input_handle.tsx      |  3 -
 .../workspace_component/output_handle.tsx     |  3 -
 .../workspace_component.tsx                   |  2 +-
 .../workspace_edge/deletable-edge-styles.scss | 14 ++++
 .../workspace_edge/deletable_edge.tsx         | 66 +++++++++++++++++++
 .../workflow_detail/workspace_edge/index.ts   |  6 ++
 .../context/react_flow_context_provider.tsx   | 10 ++-
 9 files changed, 124 insertions(+), 15 deletions(-)
 create mode 100644 public/pages/workflow_detail/workspace_edge/deletable-edge-styles.scss
 create mode 100644 public/pages/workflow_detail/workspace_edge/deletable_edge.tsx
 create mode 100644 public/pages/workflow_detail/workspace_edge/index.ts

diff --git a/public/pages/workflow_detail/workspace/reactflow-styles.scss b/public/pages/workflow_detail/workspace/reactflow-styles.scss
index 9c027fd5..6c5293e0 100644
--- a/public/pages/workflow_detail/workspace/reactflow-styles.scss
+++ b/public/pages/workflow_detail/workspace/reactflow-styles.scss
@@ -1,3 +1,7 @@
+$handle-color: #5a5a5a;
+$handle-color-valid: #55dd99;
+$handle-color-invalid: #ff6060;
+
 .reactflow-parent-wrapper {
   display: flex;
   flex-grow: 1;
@@ -15,6 +19,20 @@
   padding: 0;
 }
 
-.workspace-component {
+.reactflow-workspace .react-flow__node {
   width: 300px;
 }
+
+.reactflow-workspace .react-flow__handle {
+  height: 10px;
+  width: 10px;
+  background: $handle-color;
+}
+
+.reactflow-workspace .react-flow__handle-connecting {
+  background: $handle-color-invalid;
+}
+
+.reactflow-workspace .react-flow__handle-valid {
+  background: $handle-color-valid;
+}
diff --git a/public/pages/workflow_detail/workspace/workspace.tsx b/public/pages/workflow_detail/workspace/workspace.tsx
index 1acd61a2..561cb2f6 100644
--- a/public/pages/workflow_detail/workspace/workspace.tsx
+++ b/public/pages/workflow_detail/workspace/workspace.tsx
@@ -17,17 +17,19 @@ import { IComponent, Workflow } from '../../../../common';
 import { generateId } from '../../../utils';
 import { getCore } from '../../../services';
 import { WorkspaceComponent } from '../workspace_component';
+import { DeletableEdge } from '../workspace_edge';
 
 // styling
 import 'reactflow/dist/style.css';
 import './reactflow-styles.scss';
+import '../workspace_edge/deletable-edge-styles.scss';
 
 interface WorkspaceProps {
   workflow?: Workflow;
 }
 
 const nodeTypes = { customComponent: WorkspaceComponent };
-// TODO: probably have custom edge types here too
+const edgeTypes = { customEdge: DeletableEdge };
 
 export function Workspace(props: WorkspaceProps) {
   const reactFlowWrapper = useRef(null);
@@ -38,11 +40,12 @@ export function Workspace(props: WorkspaceProps) {
 
   const onConnect = useCallback(
     (params) => {
-      setEdges((eds) => addEdge(params, eds));
+      const edge = {
+        ...params,
+        type: 'customEdge',
+      };
+      setEdges((eds) => addEdge(edge, eds));
     },
-    // TODO: add customized logic to prevent connections based on the node's
-    // allowed inputs. If allowed, update that node state as well with the added
-    // connection details.
     [setEdges]
   );
 
@@ -132,12 +135,14 @@ export function Workspace(props: WorkspaceProps) {
                 nodes={nodes}
                 edges={edges}
                 nodeTypes={nodeTypes}
+                edgeTypes={edgeTypes}
                 onNodesChange={onNodesChange}
                 onEdgesChange={onEdgesChange}
                 onConnect={onConnect}
                 onInit={setReactFlowInstance}
                 onDrop={onDrop}
                 onDragOver={onDragOver}
+                className="reactflow-workspace"
                 fitView
               >
                 <Controls />
diff --git a/public/pages/workflow_detail/workspace_component/input_handle.tsx b/public/pages/workflow_detail/workspace_component/input_handle.tsx
index d4f87444..33407d9e 100644
--- a/public/pages/workflow_detail/workspace_component/input_handle.tsx
+++ b/public/pages/workflow_detail/workspace_component/input_handle.tsx
@@ -37,9 +37,6 @@ export function InputHandle(props: InputHandleProps) {
             isValidConnection(connection, reactFlowInstance)
           }
           style={{
-            height: 10,
-            width: 10,
-            backgroundColor: 'black',
             top: position,
           }}
         />
diff --git a/public/pages/workflow_detail/workspace_component/output_handle.tsx b/public/pages/workflow_detail/workspace_component/output_handle.tsx
index 92d8ef68..b19d2d61 100644
--- a/public/pages/workflow_detail/workspace_component/output_handle.tsx
+++ b/public/pages/workflow_detail/workspace_component/output_handle.tsx
@@ -38,9 +38,6 @@ export function OutputHandle(props: OutputHandleProps) {
             isValidConnection(connection, reactFlowInstance)
           }
           style={{
-            height: 10,
-            width: 10,
-            backgroundColor: 'black',
             top: position,
           }}
         />
diff --git a/public/pages/workflow_detail/workspace_component/workspace_component.tsx b/public/pages/workflow_detail/workspace_component/workspace_component.tsx
index 7e634b30..f0b640db 100644
--- a/public/pages/workflow_detail/workspace_component/workspace_component.tsx
+++ b/public/pages/workflow_detail/workspace_component/workspace_component.tsx
@@ -31,7 +31,7 @@ export function WorkspaceComponent(props: WorkspaceComponentProps) {
     : component.fields;
 
   return (
-    <EuiCard title={component.label} className="workspace-component">
+    <EuiCard title={component.label}>
       <EuiFlexGroup direction="column">
         {/* <EuiFlexItem>
           {component.allowsCreation ? (
diff --git a/public/pages/workflow_detail/workspace_edge/deletable-edge-styles.scss b/public/pages/workflow_detail/workspace_edge/deletable-edge-styles.scss
new file mode 100644
index 00000000..a249a641
--- /dev/null
+++ b/public/pages/workflow_detail/workspace_edge/deletable-edge-styles.scss
@@ -0,0 +1,14 @@
+.delete-edge-button {
+  width: 20px;
+  height: 20px;
+  background: #eee;
+  border: 1px solid #fff;
+  cursor: pointer;
+  border-radius: 50%;
+  font-size: 12px;
+  line-height: 1;
+}
+
+.delete-edge-button:hover {
+  box-shadow: 0 0 6px 2px rgba(0, 0, 0, 0.08);
+}
diff --git a/public/pages/workflow_detail/workspace_edge/deletable_edge.tsx b/public/pages/workflow_detail/workspace_edge/deletable_edge.tsx
new file mode 100644
index 00000000..0b46f47f
--- /dev/null
+++ b/public/pages/workflow_detail/workspace_edge/deletable_edge.tsx
@@ -0,0 +1,66 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React, { useContext } from 'react';
+import {
+  BaseEdge,
+  EdgeLabelRenderer,
+  EdgeProps,
+  getBezierPath,
+} from 'reactflow';
+import { rfContext } from '../../../store';
+
+// styling
+import './deletable-edge-styles.scss';
+
+type DeletableEdgeProps = EdgeProps;
+
+/**
+ * A custom deletable edge. Renders a delete button in the center of the edge once connected.
+ * Using bezier path by default. For all edge types,
+ * see https://reactflow.dev/docs/examples/edges/edge-types/
+ */
+export function DeletableEdge(props: DeletableEdgeProps) {
+  const [edgePath, labelX, labelY] = getBezierPath({
+    sourceX: props.sourceX,
+    sourceY: props.sourceY,
+    sourcePosition: props.sourcePosition,
+    targetX: props.targetX,
+    targetY: props.targetY,
+    targetPosition: props.targetPosition,
+  });
+
+  const { deleteEdge } = useContext(rfContext);
+
+  const onEdgeClick = (event: any, edgeId: string) => {
+    event.stopPropagation();
+    deleteEdge(edgeId);
+  };
+
+  return (
+    <>
+      <BaseEdge path={edgePath} markerEnd={props.markerEnd} />
+      <EdgeLabelRenderer>
+        {/** Using in-line styling since scss can't support dynamic values*/}
+        <div
+          style={{
+            position: 'absolute',
+            transform: `translate(-50%, -50%) translate(${labelX}px,${labelY}px)`,
+            fontSize: 12,
+            pointerEvents: 'all',
+          }}
+          className="nodrag nopan"
+        >
+          <button
+            className="delete-edge-button"
+            onClick={(event) => onEdgeClick(event, props.id)}
+          >
+            ×
+          </button>
+        </div>
+      </EdgeLabelRenderer>
+    </>
+  );
+}
diff --git a/public/pages/workflow_detail/workspace_edge/index.ts b/public/pages/workflow_detail/workspace_edge/index.ts
new file mode 100644
index 00000000..a1b34598
--- /dev/null
+++ b/public/pages/workflow_detail/workspace_edge/index.ts
@@ -0,0 +1,6 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export { DeletableEdge } from './deletable_edge';
diff --git a/public/store/context/react_flow_context_provider.tsx b/public/store/context/react_flow_context_provider.tsx
index 8434e358..73ef1cce 100644
--- a/public/store/context/react_flow_context_provider.tsx
+++ b/public/store/context/react_flow_context_provider.tsx
@@ -4,6 +4,9 @@
  */
 
 import React, { createContext, useState } from 'react';
+import { useDispatch } from 'react-redux';
+import { Edge } from 'reactflow';
+import { setDirty } from '../reducers';
 
 const initialValues = {
   reactFlowInstance: null,
@@ -23,6 +26,7 @@ export const rfContext = createContext(initialValues);
  * nested child components.
  */
 export function ReactFlowContextProvider({ children }: any) {
+  const dispatch = useDispatch();
   const [reactFlowInstance, setReactFlowInstance] = useState(null);
 
   const deleteNode = (nodeId: string) => {
@@ -31,8 +35,10 @@ export function ReactFlowContextProvider({ children }: any) {
   };
 
   const deleteEdge = (edgeId: string) => {
-    // TODO: implement edge deletion
-    // reactFlowInstance.setEdges(...)
+    reactFlowInstance.setEdges(
+      reactFlowInstance.getEdges().filter((edge: Edge) => edge.id !== edgeId)
+    );
+    dispatch(setDirty());
   };
 
   return (