import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { ReactFlow, MiniMap, Controls, Background, useNodesState, useEdgesState, getNodesBounds, getViewportForBounds, type Node, type Edge, } from "@xyflow/react"; import "@xyflow/react/dist/style.css"; import dagre from "dagre"; import { RotateCcw, ChevronUp, ChevronDown, Download, Loader2 } from "lucide-react"; import { toPng, toJpeg, toSvg } from "html-to-image"; import { save } from "@tauri-apps/plugin-dialog"; import { writeFile } from "@tauri-apps/plugin-fs"; import { CrowsFootEdge } from "./CrowsFootEdge"; import { SchemaVisualizerNode } from "./SchemaVisualizerNode"; import { LEGEND_ITEMS } from "./legendHelpers"; import { SelectDropdown } from "../ui/SelectDropdown"; import { getSchemaGraph } from "../../lib/commands"; import { useDbViewerStore } from "../../stores/dbViewerStore"; import { useNotificationStore } from "../../stores/notificationStore"; import type { SchemaGraph, TableNode as TableNodeType } from "../../lib/types"; const nodeTypes = { tableNode: SchemaVisualizerNode }; const edgeTypes = { crowsfoot: CrowsFootEdge }; const CARD_WIDTH = 240; const ROW_HEIGHT = 28; const HEADER_HEIGHT = 32; // Export size for "Entire Schema" renders const EXPORT_WIDTH = 1600; const EXPORT_HEIGHT = 1000; /** * Decode an html-to-image data URL (base64 or URL-encoded) into bytes so it * can be written to disk via the Tauri fs plugin. */ function dataUrlToBytes(dataUrl: string): Uint8Array { const [meta, payload] = dataUrl.split(","); const raw = /;base64/i.test(meta) ? atob(payload) : decodeURIComponent(payload); return Uint8Array.from(raw, (c) => c.charCodeAt(0)); } function getNodeHeight(colCount: number): number { return HEADER_HEIGHT + colCount * ROW_HEIGHT + 4; } function layoutGraph( tables: TableNodeType[], relationships: { source_table: string; target_table: string; source_column: string; target_column: string; cardinality: string }[], ): { nodes: Node[]; edges: Edge[] } { const g = new dagre.graphlib.Graph(); g.setDefaultEdgeLabel(() => ({})); g.setGraph({ rankdir: "TB", nodesep: 60, ranksep: 100, marginx: 40, marginy: 40 }); const cardinalityMap = new Map(); for (const rel of relationships) { cardinalityMap.set( `${rel.source_table}.${rel.source_column}->${rel.target_table}.${rel.target_column}`, rel.cardinality, ); } const nodes: Node[] = []; const edges: Edge[] = []; for (const table of tables) { const height = getNodeHeight(table.columns.length); g.setNode(table.name, { width: CARD_WIDTH, height }); nodes.push({ id: table.name, type: "tableNode", position: { x: 0, y: 0 }, data: { table, isExternal: false }, style: { width: CARD_WIDTH }, width: CARD_WIDTH, height, }); } for (const table of tables) { for (const col of table.columns) { if (col.fk_ref) { const [refSchema, refTable, refColumn] = col.fk_ref; if (tables.some((t) => t.name === refTable && t.schema === refSchema)) { const edgeKey = `${table.name}.${col.name}->${refTable}.${refColumn}`; const cardinality = cardinalityMap.get(edgeKey) ?? "1:N"; const markers = getEdgeMarkers(cardinality); g.setEdge(table.name, refTable, {}); edges.push({ id: edgeKey, source: table.name, target: refTable, sourceHandle: `fk-${col.name}`, targetHandle: `pk-${refColumn}`, type: "crowsfoot", label: cardinality, data: { cardinality, startMarker: markers.markerStart, endMarker: markers.markerEnd, origRight: true }, style: { stroke: "#3b82f6", strokeWidth: 1.5 }, labelStyle: { fill: "#9ca3af", fontSize: 9 }, labelBgStyle: { fill: "#1f2937", fillOpacity: 0.85 }, labelBgPadding: [3, 1], labelBorderRadius: 0, }); } } } } dagre.layout(g); for (const node of nodes) { const dagreNode = g.node(node.id); if (dagreNode) { node.position = { x: dagreNode.x - CARD_WIDTH / 2, y: dagreNode.y - (dagreNode as any).height / 2, }; } } return { nodes, edges }; } function getEdgeMarkers(cardinality: string): { markerStart: string; markerEnd: string } { switch (cardinality) { case "1:1": return { markerStart: "one", markerEnd: "one" }; case "0..1:0..1": return { markerStart: "one", markerEnd: "one" }; case "1:N": return { markerStart: "many", markerEnd: "one" }; case "0..N": return { markerStart: "many", markerEnd: "one" }; case "N:M": return { markerStart: "many", markerEnd: "many" }; default: return { markerStart: "", markerEnd: "" }; } } export interface SchemaVisualizerPageProps { connectionId: string; onSchemaChange?: (schema: string) => void; } export function SchemaVisualizerPage({ connectionId, onSchemaChange, }: SchemaVisualizerPageProps) { const databases = useDbViewerStore((s) => s.databases); const schemas = useDbViewerStore((s) => s.schemas); const currentDatabase = useDbViewerStore((s) => s.currentDatabase); const currentSchema = useDbViewerStore((s) => s.currentSchema); const setCurrentDatabase = useDbViewerStore((s) => s.setCurrentDatabase); const setCurrentSchema = useDbViewerStore((s) => s.setCurrentSchema); const [nodes, setNodes, onNodesChange] = useNodesState([]); const [edges, setEdges, onEdgesChange] = useEdgesState([]); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); const [tableCount, setTableCount] = useState(0); const [legendOpen, setLegendOpen] = useState(true); const [highlightedEdge, setHighlightedEdge] = useState(null); // Export state const containerRef = useRef(null); const exportMenuRef = useRef(null); const exportButtonRef = useRef(null); const [exportOpen, setExportOpen] = useState(false); const [exportScope, setExportScope] = useState<"schema" | "viewport">( "schema", ); const [exporting, setExporting] = useState(false); const [exportError, setExportError] = useState(null); const [exportBackground, setExportBackground] = useState< "opaque" | "transparent" >("opaque"); const transparent = exportBackground === "transparent"; const notify = useNotificationStore((s) => s.notify); // Close the export menu on outside click (ignoring the trigger button) useEffect(() => { if (!exportOpen) return; const close = (e: MouseEvent) => { const target = e.target as Element | null; if (exportButtonRef.current?.contains(target)) return; if (exportMenuRef.current?.contains(target)) return; setExportOpen(false); }; document.addEventListener("mousedown", close); return () => document.removeEventListener("mousedown", close); }, [exportOpen]); const handleExport = useCallback( async (scope: "schema" | "viewport", format: "png" | "jpeg" | "svg") => { const element = document.querySelector( ".react-flow__viewport", ); if (!element) return; setExporting(true); setExportError(null); try { let width: number; let height: number; let style: Partial | undefined; if (scope === "viewport") { const container = containerRef.current; width = container?.clientWidth || 1024; height = container?.clientHeight || 768; } else { width = EXPORT_WIDTH; height = EXPORT_HEIGHT; const bounds = getNodesBounds(nodes); if (bounds.width === 0 && bounds.height === 0) { throw new Error("Nothing to export"); } const viewport = getViewportForBounds( bounds, width, height, 0.5, 2, 0.05, ); style = { width: `${width}px`, height: `${height}px`, transform: `translate(${viewport.x}px, ${viewport.y}px) scale(${viewport.zoom})`, }; } const options = { // JPEG has no alpha channel; transparency only applies to PNG/SVG ...(transparent && format !== "jpeg" ? {} : { backgroundColor: "#0a0a0b" }), width, height, style, pixelRatio: 2, }; const dataUrl = format === "png" ? await toPng(element, options) : format === "jpeg" ? await toJpeg(element, { ...options, quality: 0.95 }) : await toSvg(element, options); // Filename: -. const dbName = currentDatabase ?? currentSchema ?? "schema"; const timestamp = new Date() .toLocaleString() .replace(/[\\/:*?"<>|]/g, "-") .replace(/\s+/g, "-"); const ext = format === "jpeg" ? "jpg" : format; const filename = `${dbName}-${timestamp}.${ext}`; const bytes = dataUrlToBytes(dataUrl); let savedPath: string | null = null; try { const path = await save({ defaultPath: filename, filters: [ { name: format.toUpperCase(), extensions: [ext] }, ], }); if (path) { await writeFile(path, bytes); savedPath = path; } } catch { // Not running in Tauri (e.g. plain browser dev): fall back to the // webview's default download handler. const a = document.createElement("a"); a.href = dataUrl; a.download = filename; a.click(); } if (savedPath) { notify(`Schema exported to ${savedPath}`, "success"); } } catch (e) { const msg = e instanceof Error ? e.message : String(e); setExportError(`Export failed: ${msg}`); } finally { setExporting(false); setExportOpen(false); } }, [nodes, currentSchema, currentDatabase, exportBackground, notify], ); const fetchGraph = useCallback(async () => { if (!currentSchema) return; setLoading(true); setError(null); try { const graph: SchemaGraph = await getSchemaGraph(connectionId, currentSchema); if (graph.tables.length === 0) { setNodes([]); setEdges([]); setTableCount(0); } else { const { nodes: layoutedNodes, edges: layoutedEdges } = layoutGraph(graph.tables, graph.relationships); setNodes(layoutedNodes); setEdges(layoutedEdges); if (graph.tables.length > 200) { const proceed = window.confirm( `This schema has ${graph.tables.length} tables. Rendering the full diagram may be slow. Continue?`, ); if (!proceed) { setLoading(false); return; } } setTableCount(graph.tables.length); } } catch (e) { const msg = e instanceof Error ? e.message : String(e); setError(msg); } finally { setLoading(false); } }, [connectionId, currentSchema, setNodes, setEdges]); useEffect(() => { fetchGraph(); }, [fetchGraph]); const handleResetLayout = useCallback(() => { fetchGraph(); setHighlightedEdge(null); }, [fetchGraph]); const handleEdgeClick = useCallback( (_event: React.MouseEvent, edge: Edge) => { setHighlightedEdge(edge.id === highlightedEdge ? null : edge.id); }, [highlightedEdge], ); const handlePaneClick = useCallback(() => { setHighlightedEdge(null); }, []); // Derive edges with highlighting applied const displayEdges = useMemo(() => { if (!highlightedEdge) return edges; return edges.map((e) => { if (e.id === highlightedEdge) { return { ...e, zIndex: 1000, style: { ...e.style, stroke: "#f59e0b", strokeWidth: 2.5, opacity: 1 }, labelStyle: { ...e.labelStyle, fill: "#f59e0b" }, labelBgStyle: { ...e.labelBgStyle, fill: "#1f2937", fillOpacity: 0.95 }, }; } return { ...e, style: { ...e.style, opacity: 0.15 } }; }); }, [edges, highlightedEdge]); const highlightedCardinality = useMemo(() => { if (!highlightedEdge) return null; const edge = edges.find((e) => e.id === highlightedEdge); return (edge?.data as any)?.cardinality as string | null; }, [edges, highlightedEdge]); const handleSchemaChange = useCallback( (schema: string) => { setCurrentSchema(schema); onSchemaChange?.(schema); }, [setCurrentSchema, onSchemaChange], ); const schemaOptions = useMemo( () => schemas.map((s) => ({ value: s, label: s })), [schemas], ); return (
{/* Toolbar */}
{databases.length > 1 && ( ({ value: d, label: d }))} placeholder="Select database" variant="ghost" /> )} {databases.length > 1 && schemas.length > 0 && ( | )} {schemas.length > 0 && ( )}
{tableCount} {tableCount === 1 ? "table" : "tables"} {/* Export */}
{exportError && ( {exportError} )}
{exportOpen && !exporting && (
Scope
setExportScope(v as "schema" | "viewport") } options={[ { value: "schema", label: "Entire Schema" }, { value: "viewport", label: "Viewport" }, ]} aria-label="Export scope" variant="pill" />
Background
setExportBackground(v as "opaque" | "transparent") } options={[ { value: "opaque", label: "Opaque" }, { value: "transparent", label: "Transparent" }, ]} aria-label="Export background" variant="pill" />
{[ { format: "png" as const, label: "PNG" }, { format: "jpeg" as const, label: "JPEG" }, { format: "svg" as const, label: "SVG" }, ] .filter((f) => !(transparent && f.format === "jpeg")) .map(({ format, label }) => ( ))}
)}
{/* Canvas */}
{loading && (

Loading schema...

)} {error && (

Failed to load schema: {error}

)} {!loading && !error && tableCount === 0 && (

No tables found in schema "{currentSchema}"

)} {/* Legend */}
{legendOpen && (
{LEGEND_ITEMS.map((item) => { const isActive = highlightedCardinality === item.cardinality; return (
{/* Start marker */} {item.markerStart === "one" ? ( ) : ( <> )} {/* End marker */} {item.markerEnd === "one" ? ( ) : ( <> )} {item.label}
)})}
)}
{/* Powered by React Flow */}
Powered by React Flow
); }