From 0fb5faf1145fb81522ffdc413cf1d32421b40d28 Mon Sep 17 00:00:00 2001
From: Greg Richardson <greg.nmr@gmail.com>
Date: Thu, 13 Jun 2024 11:25:13 -0600
Subject: [PATCH] feat(db-playground): animated edge path length

---
 .../components/schema/table-graph.tsx         | 43 ++++++++----
 .../components/schema/table-node.tsx          | 69 ++++++++++++++++++-
 2 files changed, 96 insertions(+), 16 deletions(-)

diff --git a/apps/database-playground/components/schema/table-graph.tsx b/apps/database-playground/components/schema/table-graph.tsx
index b9af2442f56..d01bf3ed821 100644
--- a/apps/database-playground/components/schema/table-graph.tsx
+++ b/apps/database-playground/components/schema/table-graph.tsx
@@ -12,7 +12,7 @@ import ReactFlow, {
   useReactFlow,
 } from 'reactflow'
 import SchemaGraphLegend from './legend'
-import { TABLE_NODE_ROW_HEIGHT, TABLE_NODE_WIDTH, TableNode } from './table-node'
+import { TABLE_NODE_ROW_HEIGHT, TABLE_NODE_WIDTH, TableEdge, TableNode } from './table-node'
 
 export default function TablesGraph({ tables }: { tables: PostgresTable[] }) {
   const { resolvedTheme } = useTheme()
@@ -24,6 +24,12 @@ export default function TablesGraph({ tables }: { tables: PostgresTable[] }) {
     }),
     []
   )
+  const edgeTypes = useMemo(
+    () => ({
+      table: TableEdge,
+    }),
+    []
+  )
 
   useEffect(() => {
     getGraphDataFromTables(tables).then(({ nodes, edges }) => {
@@ -49,14 +55,17 @@ export default function TablesGraph({ tables }: { tables: PostgresTable[] }) {
         defaultEdges={[]}
         defaultEdgeOptions={{
           type: 'smoothstep',
-          animated: true,
           deletable: false,
           style: {
             stroke: 'hsl(var(--border-stronger))',
-            strokeWidth: 0.5,
+            strokeWidth: 1,
+            strokeDasharray: 5,
+            // Manually create animation so that it doesn't interfere with our custom edge component
+            animation: 'dashdraw 0.5s linear infinite',
           },
         }}
         nodeTypes={nodeTypes}
+        edgeTypes={edgeTypes}
         minZoom={0.8}
         maxZoom={1.8}
         proOptions={{ hideAttribution: true }}
@@ -103,18 +112,20 @@ async function getGraphDataFromTables(tables: PostgresTable[]): Promise<{
   }
 
   const nodes = tables.map((table) => {
-    const columns = (table.columns || []).map((column) => {
-      return {
-        id: column.id,
-        isPrimary: table.primary_keys.some((pk) => pk.name === column.name),
-        name: column.name,
-        format: column.format,
-        isNullable: column.is_nullable,
-        isUnique: column.is_unique,
-        isUpdateable: column.is_updatable,
-        isIdentity: column.is_identity,
-      }
-    })
+    const columns = (table.columns || [])
+      .sort((a, b) => a.ordinal_position - b.ordinal_position)
+      .map((column) => {
+        return {
+          id: column.id,
+          isPrimary: table.primary_keys.some((pk) => pk.name === column.name),
+          name: column.name,
+          format: column.format,
+          isNullable: column.is_nullable,
+          isUnique: column.is_unique,
+          isUpdateable: column.is_updatable,
+          isIdentity: column.is_identity,
+        }
+      })
 
     return {
       id: `${table.id}`,
@@ -163,6 +174,7 @@ async function getGraphDataFromTables(tables: PostgresTable[]): Promise<{
       if (source) {
         edges.push({
           id: String(rel.id),
+          type: 'table',
           source,
           sourceHandle,
           target: rel.constraint_name,
@@ -188,6 +200,7 @@ async function getGraphDataFromTables(tables: PostgresTable[]): Promise<{
     if (source && target) {
       edges.push({
         id: String(rel.id),
+        type: 'table',
         source,
         sourceHandle,
         target,
diff --git a/apps/database-playground/components/schema/table-node.tsx b/apps/database-playground/components/schema/table-node.tsx
index 2ac2bdfc37a..e39b58718ab 100644
--- a/apps/database-playground/components/schema/table-node.tsx
+++ b/apps/database-playground/components/schema/table-node.tsx
@@ -1,7 +1,14 @@
 import { AnimatePresence, m } from 'framer-motion'
 import { DiamondIcon, Fingerprint, Hash, Key, Table2 } from 'lucide-react'
 import { useState } from 'react'
-import { Handle, NodeProps, useUpdateNodeInternals } from 'reactflow'
+import {
+  EdgeProps,
+  Handle,
+  NodeProps,
+  Position,
+  getSmoothStepPath,
+  useUpdateNodeInternals,
+} from 'reactflow'
 import { cn } from 'ui'
 
 // ReactFlow is scaling everything by the factor of 2
@@ -33,6 +40,9 @@ const inOutTop = {
   },
 }
 
+/**
+ * Custom node to display database tables.
+ */
 export const TableNode = ({
   id,
   data,
@@ -258,3 +268,60 @@ export const TableNode = ({
     </m.div>
   )
 }
+
+/**
+ * Custom edge that animates its path length.
+ */
+export function TableEdge({
+  id,
+  sourceX,
+  sourceY,
+  targetX,
+  targetY,
+  style,
+  sourcePosition = Position.Bottom,
+  targetPosition = Position.Top,
+  markerEnd,
+  markerStart,
+  pathOptions,
+}: EdgeProps) {
+  const [path] = getSmoothStepPath({
+    sourceX,
+    sourceY,
+    sourcePosition,
+    targetX,
+    targetY,
+    targetPosition,
+    borderRadius: pathOptions?.borderRadius,
+    offset: pathOptions?.offset,
+  })
+
+  return (
+    <>
+      <defs>
+        {/* Create a mask with the same shape that animates its path length */}
+        <mask id={`mask-${id}`}>
+          <m.path
+            d={path}
+            fill="none"
+            stroke="white"
+            strokeWidth={10}
+            initial={{ pathLength: 0 }}
+            animate={{ pathLength: 1 }}
+            transition={{ duration: 0.25 }}
+          />
+        </mask>
+      </defs>
+      <path
+        id={id}
+        d={path}
+        style={style}
+        className={cn(['react-flow__edge-path'])}
+        fill="none"
+        mask={`url(#mask-${id})`}
+        markerEnd={markerEnd}
+        markerStart={markerStart}
+      />
+    </>
+  )
+}