ERDiagram.tsx 10.9 KB
import { useMemo, useEffect } from 'react'
import {
  ReactFlow,
  Background,
  Controls,
  MiniMap,
  Handle,
  Position,
  useNodesState,
  useEdgesState,
  MarkerType,
  BaseEdge,
  EdgeLabelRenderer,
  getSmoothStepPath,
  type Node,
  type Edge,
  type NodeProps,
  type EdgeProps,
} from '@xyflow/react'
import dagre from '@dagrejs/dagre'
import '@xyflow/react/dist/style.css'
import type { Ontology, OntoObject, OntoAttribute, RelKind } from '../ontology/types'

// 交互式 ER 图:表 = React Flow 自定义节点(字段级 handle),外键 = 节点间连线(蓝色基数样式);
// 组合(composition)= 父对象 → 子实体的紫色含义连线(含菱形标记,读作 containment)。
// 拖拽 / 缩放 / 平移 / 小地图均由 React Flow 内置。首屏用 dagre 做有向布局(按依赖左右排开)。

const NODE_W = 232
const HDL: React.CSSProperties = { width: 9, height: 9, background: '#5b7fe0', border: '2px solid #fff' }

// ---- 节点/边 data 负载类型 ----------------------------------------------------
interface TableAttr {
  name: string
  type: string
  pk: boolean
  fk: boolean
}
interface TableNodeData {
  name: string
  alias: string
  attributes: TableAttr[]
  [key: string]: unknown
}
interface CardinalityEdgeData {
  srcCard: string
  tgtCard: string
  name: string
  [key: string]: unknown
}
interface CompositionEdgeData {
  name: string
  [key: string]: unknown
}

type TableFlowNode = Node<TableNodeData, 'table'>
type CardinalityFlowEdge = Edge<CardinalityEdgeData, 'cardinality'>
type CompositionFlowEdge = Edge<CompositionEdgeData, 'composition'>

// 自定义节点:一张表卡片,每行字段左右各一个 handle(外键行出右侧 source,主键行出左侧 target)
function TableNode({ data }: NodeProps<TableFlowNode>) {
  return (
    <div style={{ width: NODE_W, borderRadius: 8, background: '#fff', boxShadow: '0 2px 12px rgba(0,0,0,.10)' }}>
      <div style={{ background: '#5b7fe0', color: '#fff', fontWeight: 700, fontSize: 13, textAlign: 'center', padding: '8px', borderRadius: '8px 8px 0 0' }}>
        {data.name} <span style={{ fontWeight: 500, opacity: 0.85 }}>({data.alias})</span>
      </div>
      <div style={{ border: '1px solid #e6ecf8', borderTop: 0, borderRadius: '0 0 8px 8px' }}>
        {data.attributes.map((a) => (
          <div key={a.name} style={{ position: 'relative', display: 'flex', justifyContent: 'space-between', alignItems: 'center', padding: '6px 12px', fontSize: 11.5, borderBottom: '1px solid #f4f6fb' }}>
            <Handle type="target" position={Position.Left} id={`${a.name}__target`} style={{ ...HDL, opacity: a.pk ? 1 : 0 }} isConnectable={false} />
            <span className="om-mono" style={{ color: a.pk ? '#d99a1a' : a.fk ? '#c2528a' : '#3f3f46', whiteSpace: 'nowrap' }}>{a.pk ? '🔑 ' : a.fk ? '🔗 ' : ''}{a.name}</span>
            <span style={{ color: '#b4b4ba', marginLeft: 8 }}>{a.type}</span>
            <Handle type="source" position={Position.Right} id={`${a.name}__source`} style={{ ...HDL, opacity: a.fk ? 1 : 0 }} isConnectable={false} />
          </div>
        ))}
      </div>
    </div>
  )
}

const nodeTypes = { table: TableNode }

// 关系型基数:ONE_TO_MANY 的 FK 一般在“多”侧,故按 源→目标 方向给两端基数
const CARD: Record<RelKind, [string, string]> = {
  ONE_TO_ONE: ['1', '1'],
  ONE_TO_MANY: ['1', 'N'],
  MANY_TO_ONE: ['N', '1'],
  MANY_TO_MANY: ['N', 'N'], // 关系库里无直接 N-N(经中间表实现),此处仅兜底
}

// 自定义边:平滑折线 + 两端基数徽标(1 / N)+ 中间关系名(蓝色,reference/外键)
function CardinalityEdge({ id, sourceX, sourceY, targetX, targetY, sourcePosition, targetPosition, markerEnd, style, data }: EdgeProps<CardinalityFlowEdge>) {
  const [path] = getSmoothStepPath({ sourceX, sourceY, sourcePosition, targetX, targetY, targetPosition, borderRadius: 10 })
  const dx = targetX - sourceX, dy = targetY - sourceY
  const len = Math.hypot(dx, dy) || 1
  const ux = dx / len, uy = dy / len
  const badge = (x: number, y: number, text: string, extra?: React.CSSProperties) => (
    <div style={{ position: 'absolute', transform: `translate(-50%,-50%) translate(${x}px,${y}px)`, pointerEvents: 'none', fontSize: 10.5, fontWeight: 700, lineHeight: 1, padding: '2px 5px', borderRadius: 6, background: '#fff', border: '1px solid #cdd6ee', color: '#4a6fd6', ...extra }}>{text}</div>
  )
  return (
    <>
      <BaseEdge id={id} path={path} markerEnd={markerEnd} style={style} />
      <EdgeLabelRenderer>
        {badge(sourceX + ux * 20, sourceY + uy * 20, data?.srcCard ?? '')}
        {badge(targetX - ux * 20, targetY - uy * 20, data?.tgtCard ?? '')}
        {data?.name && badge((sourceX + targetX) / 2, (sourceY + targetY) / 2, data.name, { color: '#8a93ad', fontWeight: 500, borderColor: '#e6e9f2', background: 'rgba(255,255,255,.9)' })}
      </EdgeLabelRenderer>
    </>
  )
}

// 组合边:父 → 子实体(聚合内 composition)。紫色虚线 + 源端菱形(containment)+ “组合 (1:N)”标签。
const COMP_COLOR = '#7c4dcf'
function CompositionEdge({ id, sourceX, sourceY, targetX, targetY, sourcePosition, targetPosition, markerEnd, style, data }: EdgeProps<CompositionFlowEdge>) {
  const [path] = getSmoothStepPath({ sourceX, sourceY, sourcePosition, targetX, targetY, targetPosition, borderRadius: 10 })
  const dx = targetX - sourceX, dy = targetY - sourceY
  const len = Math.hypot(dx, dy) || 1
  const ux = dx / len, uy = dy / len
  return (
    <>
      <BaseEdge id={id} path={path} markerEnd={markerEnd} style={style} />
      <EdgeLabelRenderer>
        {/* 源端(父/整体)菱形标记:旋转 45° 的实心方块,读作 UML 组合 */}
        <div style={{ position: 'absolute', transform: `translate(-50%,-50%) translate(${sourceX + ux * 12}px,${sourceY + uy * 12}px) rotate(45deg)`, width: 11, height: 11, background: COMP_COLOR, border: '2px solid #fff', boxShadow: '0 0 0 1px rgba(124,77,207,.4)', pointerEvents: 'none' }} />
        {/* 关系标签 */}
        <div style={{ position: 'absolute', transform: `translate(-50%,-50%) translate(${(sourceX + targetX) / 2}px,${(sourceY + targetY) / 2}px)`, pointerEvents: 'none', fontSize: 10.5, fontWeight: 700, lineHeight: 1, padding: '2px 6px', borderRadius: 6, background: 'rgba(255,255,255,.92)', border: `1px solid ${COMP_COLOR}`, color: COMP_COLOR, whiteSpace: 'nowrap' }}>
          {data?.name ? `◆ ${data.name}` : '◆ 组合 (1:N)'}
        </div>
      </EdgeLabelRenderer>
    </>
  )
}

const edgeTypes = { cardinality: CardinalityEdge, composition: CompositionEdge }

// 主键 / 首字段的锚定 handle(组合边用节点级锚点,无需专门字段 handle)
function anchorAttr(o: OntoObject): OntoAttribute | undefined {
  return (o.attributes || []).find((a) => a.is_primary_key) || (o.attributes || [])[0]
}

// 由本体对象模型构建 React Flow 的 nodes/edges,并用 dagre 计算初始坐标
function buildGraph(ontology: Ontology): { nodes: TableFlowNode[]; edges: (CardinalityFlowEdge | CompositionFlowEdge)[] } {
  const objects: OntoObject[] = (ontology && ontology.objects) || []
  const byId = (id: string) => objects.find((o) => o.id === id)

  const g = new dagre.graphlib.Graph()
  g.setGraph({ rankdir: 'LR', nodesep: 36, ranksep: 96, marginx: 24, marginy: 24 })
  g.setDefaultEdgeLabel(() => ({}))
  objects.forEach((o) => g.setNode(o.id, { width: NODE_W, height: 40 + (o.attributes || []).length * 25 }))

  const edges: (CardinalityFlowEdge | CompositionFlowEdge)[] = []
  objects.forEach((o) => (o.relationships || []).forEach((r, i) => {
    const tgt = byId(r.target_object_id)
    if (!tgt) return

    if (r.refType === 'composition') {
      // 组合:父对象 → 子实体(无 fk_field,node 级连线,锚在双方 PK/首字段 handle)
      const src = anchorAttr(o)
      const dst = anchorAttr(tgt)
      g.setEdge(o.id, tgt.id)
      edges.push({
        id: `c-${o.id}-${tgt.id}-${i}`,
        source: o.id, sourceHandle: src ? `${src.name}__source` : undefined,
        target: tgt.id, targetHandle: dst ? `${dst.name}__target` : undefined,
        type: 'composition',
        data: { name: r.label || r.name || '组合 (1:N)' },
        style: { stroke: COMP_COLOR, strokeWidth: 1.8, strokeDasharray: '6 4' },
        markerEnd: { type: MarkerType.ArrowClosed, color: COMP_COLOR, width: 15, height: 15 },
      })
      return
    }

    // reference:外键连线(需源侧持有 fk_field 且目标有主键)
    if (!r.fk_field) return
    const hasFk = (o.attributes || []).some((a) => a.name === r.fk_field)
    const pk = (tgt.attributes || []).find((a) => a.is_primary_key)
    if (!hasFk || !pk) return
    g.setEdge(o.id, tgt.id)
    const [srcCard, tgtCard] = CARD[r.kind] || ['N', '1']
    edges.push({
      id: `e-${o.id}-${r.fk_field}-${tgt.id}-${i}`,
      source: o.id, sourceHandle: `${r.fk_field}__source`,
      target: tgt.id, targetHandle: `${pk.name}__target`,
      type: 'cardinality',
      data: { srcCard, tgtCard, name: r.name || '' },
      style: { stroke: '#8aa2e6', strokeWidth: 1.6 },
      markerEnd: { type: MarkerType.ArrowClosed, color: '#8aa2e6', width: 15, height: 15 },
    })
  }))
  dagre.layout(g)

  const nodes: TableFlowNode[] = objects.map((o) => {
    const n = g.node(o.id)
    const fkSet = new Set((o.relationships || []).map((r) => r.fk_field).filter(Boolean) as string[])
    return {
      id: o.id,
      type: 'table',
      position: { x: (n ? n.x : 0) - NODE_W / 2, y: (n ? n.y - n.height / 2 : 0) },
      data: {
        name: o.name,
        alias: o.alias,
        attributes: (o.attributes || []).map((a) => ({
          name: a.name,
          type: a.type,
          pk: !!a.is_primary_key,
          fk: !!a.is_foreign_key || fkSet.has(a.name),
        })),
      },
    }
  })
  return { nodes, edges }
}

export default function ERDiagram({ ontology }: { ontology: Ontology }) {
  const graph = useMemo(() => buildGraph(ontology), [ontology])
  const [nodes, setNodes, onNodesChange] = useNodesState<TableFlowNode>(graph.nodes)
  const [edges, setEdges, onEdgesChange] = useEdgesState<CardinalityFlowEdge | CompositionFlowEdge>(graph.edges)
  // 本体变化时重建(重新布局并复位)
  useEffect(() => { setNodes(graph.nodes); setEdges(graph.edges) }, [graph, setNodes, setEdges])

  return (
    <div style={{ position: 'absolute', inset: 0 }}>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        fitView
        fitViewOptions={{ padding: 0.2 }}
        minZoom={0.2}
        maxZoom={2.5}
        nodesConnectable={false}
        elementsSelectable
        proOptions={{ hideAttribution: false }}
      >
        <Background gap={18} size={1} color="#e9ecf3" />
        <MiniMap pannable zoomable nodeColor="#c7d3f0" nodeStrokeColor="#5b7fe0" maskColor="rgba(240,242,247,.6)" />
        <Controls showInteractive={false} />
      </ReactFlow>
    </div>
  )
}