import type { PostgresSchema, PostgresTable } from '@supabase/postgres-meta' import { PermissionAction } from '@supabase/shared-types/out/constants' import { toPng, toSvg } from 'html-to-image' import { Check, Copy, Download, Loader2, Plus } from 'lucide-react' import { useTheme } from 'next-themes' import Link from 'next/link' import { useEffect, useMemo, useRef, useState } from 'react' import ReactFlow, { Background, BackgroundVariant, MiniMap, useReactFlow } from 'reactflow' import 'reactflow/dist/style.css' import { LOCAL_STORAGE_KEYS, useParams } from 'common' import AlertError from 'components/ui/AlertError' import { ButtonTooltip } from 'components/ui/ButtonTooltip' import SchemaSelector from 'components/ui/SchemaSelector' import { useSchemasQuery } from 'data/database/schemas-query' import { useTablesQuery } from 'data/tables/tables-query' import { useAsyncCheckPermissions } from 'hooks/misc/useCheckPermissions' import { useLocalStorage } from 'hooks/misc/useLocalStorage' import { useQuerySchemaState } from 'hooks/misc/useSchemaQueryState' import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject' import { useIsProtectedSchema } from 'hooks/useProtectedSchemas' import { tablesToSQL } from 'lib/helpers' import { toast } from 'sonner' import { useTableEditorStateSnapshot } from 'state/table-editor' import { Button, copyToClipboard, DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger, } from 'ui' import { Admonition } from 'ui-patterns/admonition' import { SidePanelEditor } from '../../TableGridEditor/SidePanelEditor/SidePanelEditor' import { ColumnEditionContextProvider, ColumnEditionContextType } from './ColumnEditionContext' import { SchemaGraphLegend } from './SchemaGraphLegend' import { getGraphDataFromTables, getLayoutedElementsViaDagre } from './Schemas.utils' import { TableNode } from './SchemaTableNode' // [Joshen] Persisting logic: Only save positions to local storage WHEN a node is moved OR when explicitly clicked to reset layout export const SchemaGraph = () => { const { ref } = useParams() const { resolvedTheme } = useTheme() const { data: project } = useSelectedProjectQuery() const { selectedSchema, setSelectedSchema } = useQuerySchemaState() const [selectedTable, setSelectedTable] = useState(null) const snap = useTableEditorStateSnapshot() const [copied, setCopied] = useState(false) useEffect(() => { if (copied) { setTimeout(() => setCopied(false), 2000) } }, [copied]) const [isDownloading, setIsDownloading] = useState(false) const miniMapNodeColor = '#111318' const miniMapMaskColor = resolvedTheme?.includes('dark') ? 'rgb(17, 19, 24, .8)' : 'rgb(237, 237, 237, .8)' const reactFlowInstance = useReactFlow() const nodeTypes = useMemo( () => ({ table: TableNode, }), [] ) const { data: schemas, error: errorSchemas, isSuccess: isSuccessSchemas, isPending: isLoadingSchemas, isError: isErrorSchemas, } = useSchemasQuery({ projectRef: project?.ref, connectionString: project?.connectionString, }) const { data: tables = [], error: errorTables, isSuccess: isSuccessTables, isPending: isLoadingTables, isError: isErrorTables, } = useTablesQuery({ projectRef: project?.ref, connectionString: project?.connectionString, schema: selectedSchema, includeColumns: true, }) const hasNoTables = isSuccessSchemas && tables.length === 0 const schema = (schemas ?? []).find((s) => s.name === selectedSchema) const [, setStoredPositions] = useLocalStorage( LOCAL_STORAGE_KEYS.SCHEMA_VISUALIZER_POSITIONS(ref as string, schema?.id ?? 0), {} ) const { can: canUpdateTables } = useAsyncCheckPermissions( PermissionAction.TENANT_SQL_ADMIN_WRITE, 'tables' ) const { isSchemaLocked } = useIsProtectedSchema({ schema: selectedSchema }) const canAddTables = canUpdateTables && !isSchemaLocked const resetLayout = () => { const nodes = reactFlowInstance.getNodes() const edges = reactFlowInstance.getEdges() getLayoutedElementsViaDagre(nodes, edges) reactFlowInstance.setNodes(nodes) reactFlowInstance.setEdges(edges) setTimeout(() => reactFlowInstance.fitView({})) saveNodePositions() } const saveNodePositions = () => { if (schema === undefined) return console.error('Schema is required') const nodes = reactFlowInstance.getNodes() if (nodes.length > 0) { const nodesPositionData = nodes.reduce((a, b) => { return { ...a, [b.id]: b.position } }, {}) setStoredPositions(nodesPositionData) } } const downloadImage = (format: 'png' | 'svg') => { const reactflowViewport = document.querySelector('.react-flow__viewport') as HTMLElement if (!reactflowViewport) return setIsDownloading(true) const width = reactflowViewport.clientWidth const height = reactflowViewport.clientHeight const { x, y, zoom } = reactFlowInstance.getViewport() if (format === 'svg') { toSvg(reactflowViewport, { backgroundColor: 'white', width, height, style: { width: width.toString(), height: height.toString(), transform: `translate(${x}px, ${y}px) scale(${zoom})`, }, }) .then((data) => { const a = document.createElement('a') a.setAttribute('download', `supabase-schema-${ref}.svg`) a.setAttribute('href', data) a.click() toast.success('Successfully downloaded as SVG') }) .catch((error) => { console.error('Failed to download:', error) toast.error('Failed to download current view:', error.message) }) .finally(() => { setIsDownloading(false) }) } else if (format === 'png') { toPng(reactflowViewport, { backgroundColor: 'white', width, height, style: { width: width.toString(), height: height.toString(), transform: `translate(${x}px, ${y}px) scale(${zoom})`, }, }) .then((data) => { const a = document.createElement('a') a.setAttribute('download', `supabase-schema-${ref}.png`) a.setAttribute('href', data) a.click() toast.success('Successfully downloaded as PNG') }) .catch((error) => { console.error('Failed to download:', error) toast.error('Failed to download current view:', error.message) }) .finally(() => { setIsDownloading(false) }) } } const isFirstLoad = useRef(true) useEffect(() => { if (isSuccessTables && isSuccessSchemas && tables.length > 0) { const schema = schemas.find((s) => s.name === selectedSchema) as PostgresSchema getGraphDataFromTables(ref as string, schema, tables).then(({ nodes, edges }) => { reactFlowInstance.setNodes(nodes) reactFlowInstance.setEdges(edges) // Prevent resetting a view after first load to avoid layout changes after editing a column if (isFirstLoad.current) { isFirstLoad.current = false setTimeout(() => reactFlowInstance.fitView({})) // it needs to happen during next event tick } }) } }, [isSuccessTables, isSuccessSchemas, tables, resolvedTheme]) const columnEditionContext = useMemo( () => ({ onEditColumn: (tableId, columnId) => { const table = tables.find((table) => table.id === tableId) if (!table || table.columns == null) return const column = table.columns.find((column) => column.id === columnId) if (!column) return setSelectedTable(table) snap.onEditColumn(column) }, }), [tables, snap] ) return ( <>
{isLoadingSchemas && (
)} {isErrorSchemas && ( )} {isSuccessSchemas && ( <> {!hasNoTables && (
: } onClick={() => { if (tables) { copyToClipboard(tablesToSQL(tables)) setCopied(true) } }} tooltip={{ content: { side: 'bottom', text: (

Note

This schema is for context or debugging only. Table order and constraints may be invalid. Not meant to be run as-is.

), }, }} > Copy as SQL
} tooltip={{ content: { side: 'bottom', text: 'Download current view' } }} /> downloadImage('png')}> Download as PNG downloadImage('svg')}> Download as SVG Auto layout
)} )}
{isLoadingTables && (

Loading tables

)} {isErrorTables && (
)} {isSuccessTables && ( <> {hasNoTables ? (
{canAddTables && ( )}
) : (
saveNodePositions()} >
)} )} ) }