feat: enhance convertDBStructureToNodes to support show mode

liam-hq · Dec 12, 2024 · a3b0c91 · a3b0c91
1 parent 8ac6a9b
commit a3b0c91
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 74 deletions.
diff --git a/...c/components/ERDRenderer/ERDContent/TableNode/TableDetail/RelatedTables/RelatedTables.tsx b/...c/components/ERDRenderer/ERDContent/TableNode/TableDetail/RelatedTables/RelatedTables.tsx
@@ -15,7 +15,10 @@ type Props = {
 export const RelatedTables: FC<Props> = ({ table }) => {
   const dbStructure = useDBStructureStore()
   const extractedDBStructure = extractDBStructureForTable(table, dbStructure)
-  const { nodes, edges } = convertDBStructureToNodes(extractedDBStructure)
+  const { nodes, edges } = convertDBStructureToNodes({
+    dbStructure: extractedDBStructure,
+    showMode: 'TABLE_NAME',
+  })
 
   return (
     <div className={styles.wrapper}>

diff --git a/frontend/packages/erd-core/src/components/ERDRenderer/ERDRenderer.tsx b/frontend/packages/erd-core/src/components/ERDRenderer/ERDRenderer.tsx
@@ -11,14 +11,18 @@ import { ERDContent } from './ERDContent'
 import styles from './ERDRenderer.module.css'
 import { LeftPane } from './LeftPane'
 import '@/styles/globals.css'
-import { useDBStructureStore } from '@/stores'
+import { useDBStructureStore, useUserEditingStore } from '@/stores'
 import { TableDetailDrawer, TableDetailDrawerRoot } from './TableDetailDrawer'
 import { convertDBStructureToNodes } from './convertDBStructureToNodes'
 
 export const ERDRenderer: FC = () => {
   const defaultOpen = getSidebarStateFromCookie()
+  const { showMode } = useUserEditingStore()
   const dbStructure = useDBStructureStore()
-  const { nodes, edges } = convertDBStructureToNodes(dbStructure)
+  const { nodes, edges } = convertDBStructureToNodes({
+    dbStructure,
+    showMode,
+  })
 
   return (
     <div className={styles.wrapper}>

diff --git a/frontend/packages/erd-core/src/components/ERDRenderer/convertDBStructureToNodes.ts b/frontend/packages/erd-core/src/components/ERDRenderer/convertDBStructureToNodes.ts
@@ -1,76 +1,23 @@
-import type { DBStructure, Table } from '@liam-hq/db-structure'
+import type { ShowMode } from '@/schemas/showMode'
+import type { DBStructure } from '@liam-hq/db-structure'
 import type { Edge, Node } from '@xyflow/react'
 
-type Data = {
-  table: Table
+type Params = {
+  dbStructure: DBStructure
+  showMode: ShowMode
 }
 
-type TableNodeType = Node<Data, 'table'>
-
-const getCategory = (
-  score: { primary: number; foreign: number } | undefined,
-): number => {
-  if (!score) return 4
-  if (score.foreign === 0 && score.primary > 0) return 1
-  if (score.primary > 0) return 2
-  if (score.foreign > 0) return 3
-  return 4
-}
-
-const sortNodes = (
-  nodes: TableNodeType[],
-  dbStructure: DBStructure,
-): TableNodeType[] => {
-  const relationships = Object.values(dbStructure.relationships)
-  const tableScore: Record<string, { primary: number; foreign: number }> = {}
-
-  for (const node of nodes) {
-    tableScore[node.data.table.name] = { primary: 0, foreign: 0 }
-  }
-
-  for (const { primaryTableName, foreignTableName } of relationships) {
-    if (tableScore[primaryTableName]) {
-      tableScore[primaryTableName].primary++
-    }
-    if (tableScore[foreignTableName]) {
-      tableScore[foreignTableName].foreign++
-    }
-  }
-
-  return nodes.sort((a, b) => {
-    const scoreA = tableScore[a.data.table.name]
-    const scoreB = tableScore[b.data.table.name]
-
-    const categoryA = getCategory(scoreA)
-    const categoryB = getCategory(scoreB)
-
-    if (categoryA !== categoryB) {
-      return categoryA - categoryB
-    }
-
-    const primaryA = scoreA?.primary ?? 0
-    const primaryB = scoreB?.primary ?? 0
-    if (primaryA !== primaryB) {
-      return primaryB - primaryA
-    }
-
-    const foreignA = scoreA?.foreign ?? 0
-    const foreignB = scoreB?.foreign ?? 0
-    if (foreignA !== foreignB) {
-      return foreignA - foreignB
-    }
-
-    return a.data.table.name.localeCompare(b.data.table.name)
-  })
-}
-
-export const convertDBStructureToNodes = (
-  dbStructure: DBStructure,
-): { nodes: Node[]; edges: Edge[] } => {
+export const convertDBStructureToNodes = ({
+  dbStructure,
+  showMode,
+}: Params): {
+  nodes: Node[]
+  edges: Edge[]
+} => {
   const tables = Object.values(dbStructure.tables)
   const relationships = Object.values(dbStructure.relationships)
 
-  const nodes: TableNodeType[] = tables.map((table) => {
+  const nodes: Node[] = tables.map((table) => {
     return {
       id: table.name,
       type: 'table',
@@ -83,20 +30,25 @@ export const convertDBStructureToNodes = (
       },
     }
   })
-  const sortedNodes = sortNodes(nodes, dbStructure)
 
-  const edges = relationships.map((rel) => ({
+  const edges: Edge[] = relationships.map((rel) => ({
     id: rel.name,
     type: 'relationship',
     source: rel.primaryTableName,
     target: rel.foreignTableName,
-    sourceHandle: `${rel.primaryTableName}-${rel.primaryColumnName}`,
-    targetHandle: `${rel.foreignTableName}-${rel.foreignColumnName}`,
+    sourceHandle:
+      showMode === 'TABLE_NAME'
+        ? null
+        : `${rel.primaryTableName}-${rel.primaryColumnName}`,
+    targetHandle:
+      showMode === 'TABLE_NAME'
+        ? null
+        : `${rel.foreignTableName}-${rel.foreignColumnName}`,
     data: { relationship: rel },
     style: {
       opacity: 0,
     },
   }))
 
-  return { nodes: sortedNodes, edges }
+  return { nodes, edges }
 }