✨ Node block saving in the frontend (#1037)

* ✨ Node styling

* ✨ Remove node_ref

* ✨ Fix issues

* ✨ Fix edge

* ✨ Fix mypy

next/prisma/schema.prisma

@@ -178,15 +178,14 @@ model WorkflowNode {
 }
 
 model NodeBlock {
-  id       String @id @default(cuid())
-  node_ref String @db.VarChar(12)
-  node_id  String
+  id      String @id @default(cuid())
+  node_id String
 
   type  String
   input Json
 
   workflow_node WorkflowNode @relation(fields: [node_id], references: [id], onDelete: Cascade)
 
-  @@unique([node_id, node_ref])
+  @@index([node_id])
   @@map("node_block")
 }

next/src/components/drawer/WorkflowSidebar.tsx

@@ -54,7 +54,7 @@ const NodeBlock = ({ definition, createNode }: NodeBlockProps) => {
   return (
     <div
       className="flex cursor-pointer flex-row gap-2 rounded-md border border-white/20 p-2 hover:bg-white/10"
-      onClick={() => createNode(definition)}
+      onClick={() => createNode({ input: {}, type: definition.type })}
     >
       <div className="h-[30px] w-[30px]">
         <img src={definition.image_url} alt={definition.type} width={30} />

next/src/components/workflow/BasicNode.tsx

@@ -2,38 +2,41 @@ import React, { memo } from "react";
 import { Handle, type NodeProps, Position } from "reactflow";
 import clsx from "clsx";
 import type { WorkflowNode } from "../../types/workflow";
+import { getNodeBlockDefinitions } from "../../services/workflow/node-block-definitions";
 
 function BasicNode({ data }: NodeProps<WorkflowNode>) {
+  const definition = getNodeBlockDefinitions().find((d) => d.type === data.block.type);
+
   return (
     <div
       className={clsx(
-        "border-translucent rounded-md bg-white p-3 shadow-xl shadow-gray-300 dark:shadow-2xl dark:shadow-black",
-        "dark:bg-stone-900 dark:text-white dark:shadow-stone-800",
+        "border-translucent rounded-md p-3 shadow-2xl shadow-black",
+        "bg-stone-900 text-white shadow-stone-800",
         "transition-colors duration-500",
-        data.status === "running" && "border-2 border-amber-500",
-        data.status === "success" && "border-2 border-green-500",
-        !data.status && "border-2 border-gray-500"
+        data.status === "running" && "border border-amber-500",
+        data.status === "success" && "border border-green-500",
+        !data.status && "border border-gray-500"
       )}
     >
       <div className="flex items-center">
         <div className="ml-2">
-          <div className="text-lg font-bold dark:text-gray-100">{data.block.type}</div>
-          <div className="text-md font-thin">{data.block.description}</div>
+          <div className="text-lg font-bold text-gray-100">{definition?.name}</div>
+          <div className="text-md text-sm font-thin">{definition?.description}</div>
         </div>
       </div>
 
       {/* TODO ENABLE THIS BY BLOCK */}
       <Handle
         type="target"
         position={Position.Top}
-        className="bg-black dark:bg-white"
+        className="bg-black"
         style={{ width: "0.5em", height: "0.5em" }}
       />
 
       <Handle
         type="source"
         position={Position.Bottom}
-        className="bg-black dark:bg-white"
+        className="bg-black"
         style={{ width: "0.5em", height: "0.5em" }}
       />
     </div>

next/src/hooks/useWorkflow.ts

@@ -2,12 +2,11 @@ import type { Edge, Node } from "reactflow";
 import { useEffect, useState } from "react";
 import { nanoid } from "nanoid";
 import { useMutation, useQuery } from "@tanstack/react-query";
-import type { Workflow, WorkflowEdge, WorkflowNode } from "../types/workflow";
+import type { NodeBlock, Workflow, WorkflowEdge, WorkflowNode } from "../types/workflow";
 import { toReactFlowEdge, toReactFlowNode } from "../types/workflow";
 import WorkflowApi from "../services/workflow/workflowApi";
 import useSocket from "./useSocket";
 import { z } from "zod";
-import type { NodeBlockDefinition } from "../services/workflow/node-block-definitions";
 import type { Session } from "next-auth";
 
 const eventSchema = z.object({
@@ -77,7 +76,7 @@ export const useWorkflow = (workflowId: string, session: Session | null) => {
     );
   });
 
-  const createNode: createNodeType = (block: NodeBlockDefinition) => {
+  const createNode: createNodeType = (block: NodeBlock) => {
     const ref = nanoid(11);
 
     setNodes((nodes) => [
@@ -125,4 +124,4 @@ export const useWorkflow = (workflowId: string, session: Session | null) => {
   };
 };
 
-export type createNodeType = (block: NodeBlockDefinition) => void;
+export type createNodeType = (block: NodeBlock) => void;

next/src/services/workflow/node-block-definitions.ts

@@ -1,6 +1,7 @@
 import { z } from "zod";
 
 export const NodeBlockDefinitionSchema = z.object({
+  name: z.string(),
   type: z.string(),
   description: z.string(),
   image_url: z.string(),
@@ -9,6 +10,7 @@ export const NodeBlockDefinitionSchema = z.object({
 export type NodeBlockDefinition = z.infer<typeof NodeBlockDefinitionSchema>;
 
 const UrlStatusCheckBlockDefinition: NodeBlockDefinition = {
+  name: "URL Status Check",
   type: "UrlStatusCheck",
   description: "Check the status of a URL",
   image_url: "/tools/web.png",

next/src/types/workflow.ts

@@ -1,10 +1,13 @@
 import { z } from "zod";
 import type { Edge, Node } from "reactflow";
 import type { Dispatch, SetStateAction } from "react";
-import {
-  NodeBlockDefinition,
-  NodeBlockDefinitionSchema,
-} from "../services/workflow/node-block-definitions";
+
+const NodeBlockSchema = z.object({
+  type: z.string(),
+  input: z.any(),
+});
+
+export type NodeBlock = z.infer<typeof NodeBlockSchema>;
 
 type Model<T> = [T, Dispatch<SetStateAction<T>>];
 
@@ -14,7 +17,7 @@ const WorkflowNodeSchema = z.object({
   pos_x: z.number(),
   pos_y: z.number(),
   status: z.enum(["running", "success", "failure"]).optional(),
-  block: NodeBlockDefinitionSchema,
+  block: NodeBlockSchema,
 });
 
 const WorkflowEdgeSchema = z.object({

platform/reworkd_platform/db/crud/node.py

@@ -24,16 +24,16 @@ async def get_nodes(self, workflow_id: str) -> Dict[str, WorkflowNodeModel]:
             for node in (await self.session.execute(query)).scalars().all()
         }
 
-    async def get_node_blocks(self, node_refs: List[str]) -> Dict[str, NodeBlockModel]:
+    async def get_node_blocks(self, node_ids: List[str]) -> Dict[str, NodeBlockModel]:
         """
-        Returns an object mapping node_ref to NodeBlockModel
+        Returns an object mapping node_id to NodeBlockModel
         """
         query = select(NodeBlockModel).where(
-            NodeBlockModel.node_ref.in_(node_refs),
+            NodeBlockModel.node_id.in_(node_ids),
         )
 
         return {
-            block.node_ref: block
+            block.node_id: block
             for block in (await self.session.execute(query)).scalars().all()
         }
 
@@ -49,10 +49,9 @@ async def create_node_with_block(
         ).save(self.session)
 
         await NodeBlockModel(
-            node_ref=node.ref,
             node_id=node.id,
             type=n.block.type,
-            input=n.block.input,
+            input=n.block.input.dict(),
         ).save(self.session)
 
         return node
@@ -62,14 +61,15 @@ async def update_node_with_block(
         node: NodeUpsert,
         existing_node: WorkflowNodeModel,
         existing_block: NodeBlockModel,
-    ) -> None:
+    ) -> WorkflowNodeModel:
         existing_node.pos_x = node.pos_x
         existing_node.pos_y = node.pos_y
         await existing_node.save(self.session)
 
         existing_block.type = node.block.type
-        existing_block.input = node.block.input
+        existing_block.input = node.block.input.dict()
         await existing_block.save(self.session)
+        return existing_node
 
     async def mark_old_nodes_deleted(
         self,

platform/reworkd_platform/db/crud/workflow.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import List
+from typing import Dict, List
 
 from fastapi import Depends
 from sqlalchemy import select
@@ -50,10 +50,10 @@ async def get(self, workflow_id: str) -> WorkflowFull:
 
         # get node blocks
         blocks = await self.node_service.get_node_blocks(
-            [node.ref for node in nodes.values()]
+            [node.id for node in nodes.values()]
         )
 
-        node_block_pairs = [(node, blocks.get(node.ref)) for node in nodes.values()]
+        node_block_pairs = [(node, blocks.get(node.id)) for node in nodes.values()]
 
         return WorkflowFull(
             **workflow.to_schema().dict(),
@@ -77,25 +77,37 @@ async def update(self, workflow_id: str, workflow_update: WorkflowUpdate) -> str
             self.node_service.get_nodes(workflow_id),
             self.edge_service.get_edges(workflow_id),
             self.node_service.get_node_blocks(
-                [node.ref for node in workflow_update.nodes]
+                [node.id for node in workflow_update.nodes if node.id]
             ),
         )
 
         # TODO: use co-routines to make this faster
+        # Map from ref to id so edges can use id's instead of refs
+        ref_to_id: Dict[str, str] = {}
         for n in workflow_update.nodes:
             if not n.id:
-                await self.node_service.create_node_with_block(n, workflow_id)
+                node = await self.node_service.create_node_with_block(n, workflow_id)
             elif n.id not in all_nodes:
                 raise Exception("Node not found")
             else:
-                await self.node_service.update_node_with_block(
-                    n, all_nodes[n.id], all_blocks[n.ref]
+                node = await self.node_service.update_node_with_block(
+                    n, all_nodes[n.id], all_blocks[n.id]
                 )
+            ref_to_id[node.ref] = node.id
 
         # Delete nodes
         await self.node_service.mark_old_nodes_deleted(workflow_update.nodes, all_nodes)
 
-        # Update edges
+        # Mark edges as added
+        for edge_model in all_edges.values():
+            self.edge_service.add_edge(edge_model.source, edge_model.target)
+
+        # Modify the edges' source and target to their corresponding IDs
+        for e in workflow_update.edges:
+            e.source = ref_to_id.get(e.source, e.source)
+            e.target = ref_to_id.get(e.target, e.target)
+
+        # update edges
         for e in workflow_update.edges:
             if self.edge_service.add_edge(e.source, e.target):
                 await self.edge_service.create_edge(e, workflow_id)

platform/reworkd_platform/db/models/workflow.py

@@ -63,14 +63,10 @@ def to_schema(self) -> Edge:
 
 class NodeBlockModel(Base):
     __tablename__ = "node_block"
-    id = mapped_column(String, primary_key=True)
-    node_ref = mapped_column(String, ForeignKey("workflow_node.ref"))
     node_id = mapped_column(String, ForeignKey("workflow_node.id"))
 
     type = mapped_column(String)
     input = mapped_column(JSON)
 
     def to_schema(self) -> Block:
-        return Block(
-            id=self.id, node_ref=self.node_ref, type=self.type, input=self.input
-        )
+        return Block(id=self.id, type=self.type, input=self.input)

platform/reworkd_platform/web/api/workflow/schemas.py

@@ -15,14 +15,12 @@ class Config:
 
 class BlockUpsert(BaseModel):
     id: Optional[str]
-    node_ref: str = Field(description="Reference ID generated by the frontend")
     type: str
     input: BlockIOBase
 
 
 class Block(BaseModel):
     id: str
-    node_ref: str = Field(description="Reference ID generated by the frontend")
     type: str
     input: BlockIOBase