konnosDashboard
SchemaGraph.tsx500 lines · main
1import type { PGSchema } from '@supabase/pg-meta'
2import { PermissionAction } from '@supabase/shared-types/out/constants'
3import {
4 Background,
5 BackgroundVariant,
6 ColorMode,
7 Edge,
8 MiniMap,
9 Node,
10 OnSelectionChangeParams,
11 ReactFlow,
12 useReactFlow,
13} from '@xyflow/react'
14import { Check, ChevronDown, Copy, Download, Loader2, Plus } from 'lucide-react'
15import { useTheme } from 'next-themes'
16import Link from 'next/link'
17import { useEffect, useMemo, useRef, useState } from 'react'
18import { toast } from 'sonner'
19
20import '@xyflow/react/dist/style.css'
21
22import { LOCAL_STORAGE_KEYS, useParams } from 'common'
23import {
24 AlertDialog,
25 AlertDialogAction,
26 AlertDialogCancel,
27 AlertDialogContent,
28 AlertDialogDescription,
29 AlertDialogFooter,
30 AlertDialogHeader,
31 AlertDialogTitle,
32 AlertDialogTrigger,
33 Button,
34 copyToClipboard,
35 DropdownMenu,
36 DropdownMenuContent,
37 DropdownMenuItem,
38 DropdownMenuTrigger,
39} from 'ui'
40import { Admonition } from 'ui-patterns/admonition'
41
42import { SidePanelEditor } from '../../TableGridEditor/SidePanelEditor/SidePanelEditor'
43import { DefaultEdge } from './DefaultEdge'
44import { SchemaGraphContextProvider, SchemaGraphContextType } from './SchemaGraphContext'
45import { SchemaGraphLegend } from './SchemaGraphLegend'
46import { EdgeData, TableNodeData } from './Schemas.constants'
47import {
48 getGraphDataFromTables,
49 getLayoutedElementsViaDagre,
50 getSchemaAsMarkdown,
51} from './Schemas.utils'
52import { TableNode } from './SchemaTableNode'
53import { useExportSchemaToImage } from './useExportSchemaToImage'
54import AlertError from '@/components/ui/AlertError'
55import { ButtonTooltip } from '@/components/ui/ButtonTooltip'
56import SchemaSelector from '@/components/ui/SchemaSelector'
57import { Shortcut } from '@/components/ui/Shortcut'
58import { useSchemasQuery } from '@/data/database/schemas-query'
59import { useTablesQuery } from '@/data/tables/tables-query'
60import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
61import { useLocalStorage } from '@/hooks/misc/useLocalStorage'
62import { useQuerySchemaState } from '@/hooks/misc/useSchemaQueryState'
63import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
64import { useIsProtectedSchema } from '@/hooks/useProtectedSchemas'
65import { useStaticEffectEvent } from '@/hooks/useStaticEffectEvent'
66import { tablesToSQL } from '@/lib/helpers'
67import type { SafePostgresTable } from '@/lib/postgres-types'
68import { SHORTCUT_IDS } from '@/state/shortcuts/registry'
69import { useShortcut } from '@/state/shortcuts/useShortcut'
70import { useTableEditorStateSnapshot } from '@/state/table-editor'
71
72// [Joshen] Persisting logic: Only save positions to local storage WHEN a node is moved OR when explicitly clicked to reset layout
73
74export const SchemaGraph = () => {
75 const { ref } = useParams()
76 const { resolvedTheme } = useTheme()
77 const { data: project } = useSelectedProjectQuery()
78 const { selectedSchema, setSelectedSchema } = useQuerySchemaState()
79 const [selectedTable, setSelectedTable] = useState<SafePostgresTable | null>(null)
80 const snap = useTableEditorStateSnapshot()
81 const { isDownloading, exportSchemaToImage } = useExportSchemaToImage()
82
83 const [copied, setCopied] = useState(false)
84 useEffect(() => {
85 if (copied) {
86 setTimeout(() => setCopied(false), 2000)
87 }
88 }, [copied])
89
90 const miniMapNodeColor = '#111318'
91 const miniMapMaskColor = resolvedTheme?.includes('dark')
92 ? 'rgb(17, 19, 24, .8)'
93 : 'rgb(237, 237, 237, .8)'
94
95 const reactFlowInstance = useReactFlow()
96 const nodeTypes = useMemo(
97 () => ({
98 table: TableNode,
99 }),
100 []
101 )
102 const edgeTypes = useMemo(
103 () => ({
104 default: DefaultEdge,
105 }),
106 []
107 )
108
109 const {
110 data: schemas,
111 error: errorSchemas,
112 isSuccess: isSuccessSchemas,
113 isPending: isLoadingSchemas,
114 isError: isErrorSchemas,
115 } = useSchemasQuery({
116 projectRef: project?.ref,
117 connectionString: project?.connectionString,
118 })
119
120 const {
121 data: tables = [],
122 error: errorTables,
123 isSuccess: isSuccessTables,
124 isPending: isLoadingTables,
125 isError: isErrorTables,
126 } = useTablesQuery({
127 projectRef: project?.ref,
128 connectionString: project?.connectionString,
129 schema: selectedSchema,
130 includeColumns: true,
131 })
132 const hasNoTables = isSuccessSchemas && tables.length === 0
133
134 const schema = (schemas ?? []).find((s) => s.name === selectedSchema)
135 const [, setStoredPositions] = useLocalStorage(
136 LOCAL_STORAGE_KEYS.SCHEMA_VISUALIZER_POSITIONS(ref as string, schema?.id ?? 0),
137 {}
138 )
139
140 const { can: canUpdateTables } = useAsyncCheckPermissions(
141 PermissionAction.TENANT_SQL_ADMIN_WRITE,
142 'tables'
143 )
144
145 const { isSchemaLocked } = useIsProtectedSchema({ schema: selectedSchema })
146
147 const canAddTables = canUpdateTables && !isSchemaLocked
148
149 const resetLayout = async () => {
150 const nodes = reactFlowInstance.getNodes()
151 const edges = reactFlowInstance.getEdges()
152
153 getLayoutedElementsViaDagre(
154 nodes.filter((item) => item.type === 'table') as Node<TableNodeData>[],
155 edges
156 )
157 reactFlowInstance.setNodes(nodes)
158 reactFlowInstance.setEdges(edges)
159 await new Promise<void>((resolve) =>
160 setTimeout(async () => {
161 await reactFlowInstance.fitView({})
162 resolve()
163 })
164 )
165 saveNodePositions()
166 }
167
168 const saveNodePositions = useStaticEffectEvent(() => {
169 if (schema === undefined) return console.error('Schema is required')
170
171 const nodes = reactFlowInstance.getNodes()
172 if (nodes.length > 0) {
173 const nodesPositionData = nodes.reduce((a, b) => {
174 return { ...a, [b.id]: b.position }
175 }, {})
176 setStoredPositions(nodesPositionData)
177 }
178 })
179
180 const [selectedEdge, setSelectedEdge] = useState<Edge | undefined>(undefined)
181 const handleSelectionChange = useStaticEffectEvent(
182 (params: OnSelectionChangeParams<Node<TableNodeData>, Edge<EdgeData>>) => {
183 if (params.edges.length === 1) {
184 setSelectedEdge(params.edges[0])
185 } else {
186 setSelectedEdge(undefined)
187 }
188
189 const selectedNodeIds = new Set(params.nodes.map((n) => n.id))
190 reactFlowInstance.setEdges(
191 reactFlowInstance.getEdges().map((edge) => ({
192 ...edge,
193 animated:
194 selectedNodeIds.size > 0 &&
195 (selectedNodeIds.has(edge.source) || selectedNodeIds.has(edge.target)),
196 }))
197 )
198 }
199 )
200
201 const downloadImage = async (format: 'png' | 'svg') => {
202 const reactflowViewport = document.querySelector('.react-flow__viewport') as HTMLElement
203 if (!reactflowViewport) return
204 if (!ref) return
205 const { x, y, zoom } = reactFlowInstance.getViewport()
206 exportSchemaToImage({ element: reactflowViewport, format, x, y, zoom, projectRef: ref })
207 }
208
209 const copyAsSQL = () => {
210 if (!tables) return
211 copyToClipboard(tablesToSQL(tables))
212 setCopied(true)
213 toast.success('Successfully copied as SQL')
214 }
215
216 const copyAsMarkdown = () => {
217 const tableNodes = reactFlowInstance
218 .getNodes()
219 .filter((node) => node.type === 'table')
220 .map((node) => node.data as TableNodeData)
221 copyToClipboard(getSchemaAsMarkdown(selectedSchema, tableNodes))
222 setCopied(true)
223 toast.success('Successfully copied as Markdown')
224 }
225
226 const [schemaSelectorOpen, setSchemaSelectorOpen] = useState(false)
227 const [autoLayoutDialogOpen, setAutoLayoutDialogOpen] = useState(false)
228
229 const shortcutsEnabled = isSuccessSchemas && !hasNoTables
230
231 useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_COPY_SQL, copyAsSQL, { enabled: shortcutsEnabled })
232 useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_COPY_MARKDOWN, copyAsMarkdown, {
233 enabled: shortcutsEnabled,
234 })
235 useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_DOWNLOAD_PNG, () => downloadImage('png'), {
236 enabled: shortcutsEnabled,
237 })
238 useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_DOWNLOAD_SVG, () => downloadImage('svg'), {
239 enabled: shortcutsEnabled,
240 })
241
242 const isFirstLoad = useRef(true)
243 useEffect(() => {
244 if (isSuccessTables && isSuccessSchemas && tables.length > 0) {
245 const schema = schemas.find((s) => s.name === selectedSchema) as PGSchema
246 getGraphDataFromTables(ref as string, schema, tables).then(({ nodes, edges }) => {
247 reactFlowInstance.setNodes(nodes)
248 reactFlowInstance.setEdges(edges)
249 // Prevent resetting a view after first load to avoid layout changes after editing a column
250 if (isFirstLoad.current) {
251 isFirstLoad.current = false
252 setTimeout(() => reactFlowInstance.fitView({})) // it needs to happen during next event tick
253 }
254 })
255 }
256 }, [
257 isSuccessTables,
258 isSuccessSchemas,
259 tables,
260 reactFlowInstance,
261 ref,
262 resolvedTheme,
263 schemas,
264 selectedSchema,
265 ])
266
267 const schemaGraphContext = useMemo<SchemaGraphContextType>(
268 () => ({
269 selectedEdge,
270 isDownloading,
271 onEditColumn: (tableId, columnId) => {
272 const table = tables.find((table) => table.id === tableId)
273 if (!table || table.columns == null) return
274
275 const column = table.columns.find((column) => column.id === columnId)
276 if (!column) return
277
278 setSelectedTable(table)
279 snap.onEditColumn(column)
280 },
281 onEditTable: (tableId) => {
282 const table = tables.find((table) => table.id === tableId)
283 if (!table || table.columns == null) return
284
285 setSelectedTable(table)
286 snap.onEditTable()
287 },
288 }),
289 [tables, snap, isDownloading, selectedEdge]
290 )
291
292 return (
293 <>
294 <div className="flex items-center justify-between p-4 border-b border-muted h-(--header-height)">
295 {isLoadingSchemas && (
296 <div className="h-[34px] w-[260px] bg-foreground-lighter rounded-sm shimmering-loader" />
297 )}
298
299 {isErrorSchemas && <AlertError error={errorSchemas} subject="Failed to retrieve schemas" />}
300
301 {isSuccessSchemas && (
302 <>
303 <Shortcut
304 id={SHORTCUT_IDS.SCHEMA_VISUALIZER_FOCUS_SCHEMA}
305 onTrigger={() => setSchemaSelectorOpen(true)}
306 options={{ enabled: isSuccessSchemas }}
307 side="bottom"
308 tooltipOpen={schemaSelectorOpen ? false : undefined}
309 >
310 <SchemaSelector
311 className="w-[180px]"
312 size="tiny"
313 showError={false}
314 selectedSchemaName={selectedSchema}
315 onSelectSchema={setSelectedSchema}
316 open={schemaSelectorOpen}
317 onOpenChange={setSchemaSelectorOpen}
318 />
319 </Shortcut>
320 {!hasNoTables && (
321 <div className="flex items-center gap-x-2">
322 <div className="flex items-center gap-0">
323 <ButtonTooltip
324 type="default"
325 className="rounded-r-none border-r-0"
326 icon={copied ? <Check data-testid="copy-sql-ready" /> : <Copy />}
327 onClick={copyAsSQL}
328 tooltip={{
329 content: {
330 side: 'bottom',
331 text: (
332 <div className="max-w-[180px] space-y-2 text-foreground-light">
333 <p className="text-foreground">Note</p>
334 <p>
335 This schema is for context or debugging only. Table order and
336 constraints may be invalid. Not meant to be run as-is.
337 </p>
338 </div>
339 ),
340 },
341 }}
342 >
343 Copy as SQL
344 </ButtonTooltip>
345 <DropdownMenu>
346 <DropdownMenuTrigger asChild>
347 <Button
348 type="default"
349 size="tiny"
350 className="rounded-l-none pl-1 pr-0"
351 icon={<ChevronDown size={12} />}
352 >
353 <span className="sr-only">Export options</span>
354 </Button>
355 </DropdownMenuTrigger>
356 <DropdownMenuContent align="end" className="w-44">
357 <DropdownMenuItem
358 className="flex items-center space-x-2 whitespace-nowrap"
359 onClick={(e) => {
360 e.stopPropagation()
361 copyAsMarkdown()
362 }}
363 >
364 <Copy size={12} />
365 <span>Copy as Markdown</span>
366 </DropdownMenuItem>
367 <DropdownMenuItem
368 className="flex items-center space-x-2 whitespace-nowrap"
369 onClick={(e) => {
370 e.stopPropagation()
371 downloadImage('png')
372 }}
373 >
374 <Download size={12} />
375 <span>Download as PNG</span>
376 </DropdownMenuItem>
377 <DropdownMenuItem
378 className="flex items-center space-x-2 whitespace-nowrap"
379 onClick={(e) => {
380 e.stopPropagation()
381 downloadImage('svg')
382 }}
383 >
384 <Download size={12} />
385 <span>Download as SVG</span>
386 </DropdownMenuItem>
387 </DropdownMenuContent>
388 </DropdownMenu>
389 </div>
390 <AlertDialog open={autoLayoutDialogOpen} onOpenChange={setAutoLayoutDialogOpen}>
391 <Shortcut
392 id={SHORTCUT_IDS.SCHEMA_VISUALIZER_AUTO_LAYOUT}
393 onTrigger={() => setAutoLayoutDialogOpen(true)}
394 options={{ enabled: shortcutsEnabled }}
395 side="bottom"
396 tooltipOpen={autoLayoutDialogOpen ? false : undefined}
397 >
398 <AlertDialogTrigger asChild>
399 <Button type="default">Auto layout</Button>
400 </AlertDialogTrigger>
401 </Shortcut>
402 <AlertDialogContent>
403 <AlertDialogHeader>
404 <AlertDialogTitle>Confirm to rearrange all nodes</AlertDialogTitle>
405 <AlertDialogDescription>
406 Auto layout will rearrange all nodes in the graph. This cannot be undone.
407 </AlertDialogDescription>
408 </AlertDialogHeader>
409 <AlertDialogFooter>
410 <AlertDialogCancel>Cancel</AlertDialogCancel>
411 <AlertDialogAction onClick={resetLayout}>Apply</AlertDialogAction>
412 </AlertDialogFooter>
413 </AlertDialogContent>
414 </AlertDialog>
415 </div>
416 )}
417 </>
418 )}
419 </div>
420 {isLoadingTables && (
421 <div className="w-full h-full flex items-center justify-center gap-x-2">
422 <Loader2 className="animate-spin text-foreground-light" size={16} />
423 <p className="text-sm text-foreground-light">Loading tables</p>
424 </div>
425 )}
426 {isErrorTables && (
427 <div className="w-full h-full flex items-center justify-center px-20">
428 <AlertError subject="Failed to retrieve tables" error={errorTables} />
429 </div>
430 )}
431 {isSuccessTables && (
432 <>
433 {hasNoTables ? (
434 <div className="flex items-center justify-center w-full h-full">
435 <Admonition
436 type="default"
437 className="max-w-md"
438 title="No tables in schema"
439 description={
440 isSchemaLocked
441 ? `The “${selectedSchema}” schema is managed by Briven and is read-only through
442 the dashboard.`
443 : !canUpdateTables
444 ? 'You need additional permissions to create tables'
445 : `The “${selectedSchema}” schema doesn’t have any tables.`
446 }
447 >
448 {canAddTables && (
449 <Button asChild className="mt-2" type="default" icon={<Plus />}>
450 <Link href={`/project/${ref}/editor?create=table`}>New table</Link>
451 </Button>
452 )}
453 </Admonition>
454 </div>
455 ) : (
456 <SchemaGraphContextProvider value={schemaGraphContext}>
457 <div className="w-full h-full">
458 <ReactFlow<Node<TableNodeData>, Edge<EdgeData>>
459 // FIXME: https://github.com/xyflow/xyflow/issues/4876
460 colorMode={'' as unknown as ColorMode}
461 defaultNodes={[]}
462 defaultEdges={[]}
463 defaultEdgeOptions={{
464 type: 'default',
465 animated: false,
466 deletable: false,
467 }}
468 nodeTypes={nodeTypes}
469 edgeTypes={edgeTypes}
470 fitView
471 minZoom={0.8}
472 maxZoom={1.8}
473 proOptions={{ hideAttribution: true }}
474 onNodeDragStop={saveNodePositions}
475 onSelectionChange={handleSelectionChange}
476 >
477 <Background
478 gap={16}
479 className="*:stroke-foreground-muted opacity-25"
480 variant={BackgroundVariant.Dots}
481 color={'inherit'}
482 />
483 <MiniMap
484 pannable
485 zoomable
486 nodeColor={miniMapNodeColor}
487 maskColor={miniMapMaskColor}
488 className="border rounded-md shadow-xs"
489 />
490 <SchemaGraphLegend />
491 </ReactFlow>
492 </div>
493 </SchemaGraphContextProvider>
494 )}
495 </>
496 )}
497 <SidePanelEditor selectedTable={selectedTable ?? undefined} includeColumns />
498 </>
499 )
500}