supabase/apps/studio/components/interfaces/Database/Schemas/SchemaGraph.tsx

import type { PGSchema } from '@supabase/pg-meta'
import { PermissionAction } from '@supabase/shared-types/out/constants'
import {
  Background,
  BackgroundVariant,
  ColorMode,
  Edge,
  MiniMap,
  Node,
  OnSelectionChangeParams,
  Panel,
  ReactFlow,
  useReactFlow,
} from '@xyflow/react'
import { Check, ChevronDown, Copy, Download, Loader2, Plus } from 'lucide-react'
import { useTheme } from 'next-themes'
import Link from 'next/link'
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import { toast } from 'sonner'

import '@xyflow/react/dist/style.css'

import { LOCAL_STORAGE_KEYS, useParams } from 'common'
import {
  AlertDialog,
  AlertDialogAction,
  AlertDialogCancel,
  AlertDialogContent,
  AlertDialogDescription,
  AlertDialogFooter,
  AlertDialogHeader,
  AlertDialogTitle,
  AlertDialogTrigger,
  Button,
  copyToClipboard,
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuTrigger,
} from 'ui'
import { Admonition } from 'ui-patterns/Admonition'

import { SidePanelEditor } from '../../TableGridEditor/SidePanelEditor/SidePanelEditor'
import { DefaultEdge } from './DefaultEdge'
import { FindTableSelector } from './FindTableSelector'
import { SchemaGraphContextProvider, SchemaGraphContextType } from './SchemaGraphContext'
import { SchemaGraphLegend } from './SchemaGraphLegend'
import { EdgeData, TableNodeData } from './Schemas.constants'
import {
  getEnumsAsMarkdown,
  getGraphDataFromTables,
  getLayoutedElementsViaDagre,
  getPoliciesAsMarkdown,
  getSchemaAsMarkdown,
} from './Schemas.utils'
import { TableNode } from './SchemaTableNode'
import { useExportSchemaToImage } from './useExportSchemaToImage'
import { AlertError } from '@/components/ui/AlertError'
import { ButtonTooltip } from '@/components/ui/ButtonTooltip'
import { SchemaSelector } from '@/components/ui/SchemaSelector'
import { Shortcut } from '@/components/ui/Shortcut'
import { useDatabasePoliciesQuery } from '@/data/database-policies/database-policies-query'
import { useSchemasQuery } from '@/data/database/schemas-query'
import { useEnumeratedTypesQuery } from '@/data/enumerated-types/enumerated-types-query'
import { useInfiniteTablesQuery } 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 type { SafePostgresTable } from '@/lib/postgres-types'
import { SHORTCUT_IDS } from '@/state/shortcuts/registry'
import { useShortcut } from '@/state/shortcuts/useShortcut'
import { useTableEditorStateSnapshot } from '@/state/table-editor'

// [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<SafePostgresTable | null>(null)
  const snap = useTableEditorStateSnapshot()
  const { isDownloading, exportSchemaToImage } = useExportSchemaToImage()

  const [copied, setCopied] = useState(false)
  useEffect(() => {
    if (copied) {
      setTimeout(() => setCopied(false), 2000)
    }
  }, [copied])

  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 edgeTypes = useMemo(
    () => ({
      default: DefaultEdge,
    }),
    []
  )

  const {
    data: schemas,
    error: errorSchemas,
    isSuccess: isSuccessSchemas,
    isPending: isLoadingSchemas,
    isError: isErrorSchemas,
  } = useSchemasQuery({
    projectRef: project?.ref,
    connectionString: project?.connectionString,
  })

  const {
    data: tablesData,
    error: errorTables,
    isSuccess: isSuccessTables,
    isPending: isLoadingTables,
    isError: isErrorTables,
    hasNextPage,
    isFetchingNextPage,
    fetchNextPage,
  } = useInfiniteTablesQuery({
    projectRef: project?.ref,
    connectionString: project?.connectionString,
    schema: selectedSchema,
    includeColumns: true,
    pageSize: 100,
  })
  const tables = useMemo(() => tablesData?.pages.flat() ?? [], [tablesData])
  const hasNoTables = isSuccessTables && isSuccessSchemas && tables.length === 0 && !hasNextPage

  const { data: enumeratedTypes = [], isPending: isLoadingEnumeratedTypes } =
    useEnumeratedTypesQuery({
      projectRef: project?.ref,
      connectionString: project?.connectionString,
    })

  const { data: policies = [], isPending: isLoadingPolicies } = useDatabasePoliciesQuery({
    projectRef: project?.ref,
    connectionString: project?.connectionString,
    schemas: [selectedSchema],
  })

  const isMarkdownDataLoading = isLoadingEnumeratedTypes || isLoadingPolicies

  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 = async () => {
    const nodes = reactFlowInstance.getNodes()
    const edges = reactFlowInstance.getEdges()

    getLayoutedElementsViaDagre(
      nodes.filter((item) => item.type === 'table') as Node<TableNodeData>[],
      edges
    )
    reactFlowInstance.setNodes(nodes)
    reactFlowInstance.setEdges(edges)
    await new Promise<void>((resolve) =>
      setTimeout(async () => {
        await reactFlowInstance.fitView({})
        resolve()
      })
    )
    saveNodePositions()
  }

  const saveNodePositions = useCallback(() => {
    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)
    }
  }, [schema, reactFlowInstance, setStoredPositions])

  const [selectedEdge, setSelectedEdge] = useState<Edge | undefined>(undefined)
  const handleSelectionChange = useCallback(
    (params: OnSelectionChangeParams<Node<TableNodeData>, Edge<EdgeData>>) => {
      if (params.edges.length === 1) {
        setSelectedEdge(params.edges[0])
      } else {
        setSelectedEdge(undefined)
      }

      const selectedNodeIds = new Set(params.nodes.map((n) => n.id))
      const currentEdges = reactFlowInstance.getEdges()
      let hasChanges = false
      const nextEdges = currentEdges.map((edge) => {
        const shouldAnimate =
          selectedNodeIds.size > 0 &&
          (selectedNodeIds.has(edge.source) || selectedNodeIds.has(edge.target))
        if (edge.animated === shouldAnimate) return edge
        hasChanges = true
        return { ...edge, animated: shouldAnimate }
      })
      if (hasChanges) reactFlowInstance.setEdges(nextEdges)
    },
    [reactFlowInstance, setSelectedEdge]
  )

  const downloadImage = async (format: 'png' | 'svg') => {
    const reactflowViewport = document.querySelector('.react-flow__viewport') as HTMLElement
    if (!reactflowViewport) return
    if (!ref) return
    const { x, y, zoom } = reactFlowInstance.getViewport()
    exportSchemaToImage({ element: reactflowViewport, format, x, y, zoom, projectRef: ref })
  }

  const copyAsSQL = () => {
    if (!tables) return
    copyToClipboard(tablesToSQL(tables))
    setCopied(true)
    toast.success('Successfully copied as SQL')
  }

  const copyAsMarkdown = () => {
    if (isMarkdownDataLoading) return

    const tableNodes = reactFlowInstance
      .getNodes()
      .filter((node) => node.type === 'table')
      .map((node) => node.data as TableNodeData)

    let markdown = getSchemaAsMarkdown(selectedSchema, tableNodes)

    const enumsMarkdown = getEnumsAsMarkdown(
      selectedSchema,
      enumeratedTypes.map((e) => ({ name: e.name, schema: e.schema, enums: e.enums }))
    )
    if (enumsMarkdown) markdown += enumsMarkdown

    const policiesMarkdown = getPoliciesAsMarkdown(
      selectedSchema,
      policies.map((p) => ({
        name: p.name,
        schema: p.schema,
        table: p.table,
        command: p.command,
        roles: p.roles,
        action: p.action,
        definition: p.definition ? String(p.definition) : null,
        check: p.check ? String(p.check) : null,
      }))
    )
    if (policiesMarkdown) markdown += policiesMarkdown

    copyToClipboard(markdown)
    setCopied(true)
    toast.success('Successfully copied as Markdown')
  }

  const [schemaSelectorOpen, setSchemaSelectorOpen] = useState(false)
  const [findTableOpen, setFindTableOpen] = useState(false)
  const [autoLayoutDialogOpen, setAutoLayoutDialogOpen] = useState(false)

  const handleSelectSchema = (name: string) => {
    setFindTableOpen(false)
    setSelectedSchema(name)
  }

  const shortcutsEnabled = isSuccessSchemas && !hasNoTables

  useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_COPY_SQL, copyAsSQL, { enabled: shortcutsEnabled })
  useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_COPY_MARKDOWN, copyAsMarkdown, {
    enabled: shortcutsEnabled && !isMarkdownDataLoading,
  })
  useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_DOWNLOAD_PNG, () => downloadImage('png'), {
    enabled: shortcutsEnabled,
  })
  useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_DOWNLOAD_SVG, () => downloadImage('svg'), {
    enabled: shortcutsEnabled,
  })
  useShortcut(SHORTCUT_IDS.SCHEMA_VISUALIZER_FIND_TABLE, () => setFindTableOpen(true), {
    enabled: shortcutsEnabled,
  })

  const isFirstLoad = useRef(true)
  const fitViewOnNextLayout = useRef(false)
  const pendingFocusTableIdRef = useRef<string | null>(null)
  useEffect(() => {
    if (isSuccessTables && isSuccessSchemas && tables.length > 0) {
      const schema = schemas.find((s) => s.name === selectedSchema) as PGSchema
      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 || fitViewOnNextLayout.current) {
          isFirstLoad.current = false
          fitViewOnNextLayout.current = false
          setTimeout(() => reactFlowInstance.fitView({})) // it needs to happen during next event tick
        }
        const pendingId = pendingFocusTableIdRef.current
        if (pendingId !== null && nodes.some((n) => n.id === pendingId)) {
          pendingFocusTableIdRef.current = null
          setTimeout(() =>
            reactFlowInstance.fitView({
              nodes: [{ id: pendingId }],
              duration: 300,
              maxZoom: 1.5,
            })
          )
        }
      })
    }
  }, [isSuccessTables, isSuccessSchemas, tables, reactFlowInstance, ref, schemas, selectedSchema])

  const handleFindTableSelect = async (table: SafePostgresTable) => {
    const targetId = String(table.id)
    if (reactFlowInstance.getNode(targetId)) {
      reactFlowInstance.fitView({
        nodes: [{ id: targetId }],
        duration: 300,
        maxZoom: 1.5,
      })
      return
    }
    // Selected table isn't loaded yet — queue the fitView and pull pages until
    // it shows up. The build-effect above will consume the pending id once the
    // node is mounted.
    pendingFocusTableIdRef.current = targetId
    let result = await fetchNextPage()
    while (
      result.hasNextPage &&
      !result.data?.pages.some((page) => page.some((t) => t.id === table.id))
    ) {
      result = await fetchNextPage()
    }
  }

  const schemaGraphContext = useMemo<SchemaGraphContextType>(
    () => ({
      selectedEdge,
      isDownloading,
      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)
      },
      onEditTable: (tableId) => {
        const table = tables.find((table) => table.id === tableId)
        if (!table || table.columns == null) return

        setSelectedTable(table)
        snap.onEditTable()
      },
    }),
    [tables, snap, isDownloading, selectedEdge]
  )

  return (
    <>
      <div className="flex items-center justify-between p-4 border-b border-muted h-(--header-height)">
        {isLoadingSchemas && (
          <div className="h-[34px] w-[260px] bg-foreground-lighter rounded-sm shimmering-loader" />
        )}

        {isErrorSchemas && <AlertError error={errorSchemas} subject="Failed to retrieve schemas" />}

        {isSuccessSchemas && (
          <>
            <div className="flex items-center gap-x-2">
              <Shortcut
                id={SHORTCUT_IDS.SCHEMA_VISUALIZER_FOCUS_SCHEMA}
                onTrigger={() => setSchemaSelectorOpen(true)}
                options={{ enabled: isSuccessSchemas }}
                side="bottom"
                tooltipOpen={schemaSelectorOpen ? false : undefined}
              >
                <SchemaSelector
                  className="w-[180px]"
                  size="tiny"
                  showError={false}
                  selectedSchemaName={selectedSchema}
                  onSelectSchema={handleSelectSchema}
                  open={schemaSelectorOpen}
                  onOpenChange={setSchemaSelectorOpen}
                />
              </Shortcut>
              {!hasNoTables && (
                <Shortcut
                  id={SHORTCUT_IDS.SCHEMA_VISUALIZER_FIND_TABLE}
                  onTrigger={() => setFindTableOpen(true)}
                  options={{ enabled: shortcutsEnabled }}
                  side="bottom"
                  tooltipOpen={findTableOpen ? false : undefined}
                >
                  <FindTableSelector
                    projectRef={project?.ref}
                    connectionString={project?.connectionString}
                    schema={selectedSchema}
                    open={findTableOpen}
                    onOpenChange={setFindTableOpen}
                    onSelect={handleFindTableSelect}
                  />
                </Shortcut>
              )}
            </div>
            {!hasNoTables && (
              <div className="flex items-center gap-x-2">
                <div className="flex items-center gap-0">
                  <ButtonTooltip
                    variant="default"
                    className="rounded-r-none hover:z-10 focus-visible:z-10 focus-visible:rounded-r-sm"
                    icon={copied ? <Check data-testid="copy-sql-ready" /> : <Copy />}
                    onClick={copyAsSQL}
                    tooltip={{
                      content: {
                        side: 'bottom',
                        text: (
                          <div className="max-w-[180px] space-y-2 text-foreground-light">
                            <p className="text-foreground">Note</p>
                            <p>
                              This schema is for context or debugging only. Table order and
                              constraints may be invalid. Not meant to be run as-is.
                            </p>
                          </div>
                        ),
                      },
                    }}
                  >
                    Copy as SQL
                  </ButtonTooltip>
                  <DropdownMenu>
                    <DropdownMenuTrigger asChild>
                      <Button
                        variant="default"
                        size="tiny"
                        aria-label="Export options"
                        className="shrink-0 rounded-l-none px-[4px] py-[5px] -ml-px focus-visible:z-10 focus-visible:rounded-l-sm"
                        icon={<ChevronDown size={12} />}
                      />
                    </DropdownMenuTrigger>
                    <DropdownMenuContent align="end" className="w-44">
                      <DropdownMenuItem
                        className="flex items-center space-x-2 whitespace-nowrap"
                        disabled={isMarkdownDataLoading}
                        onClick={(e) => {
                          e.stopPropagation()
                          copyAsMarkdown()
                        }}
                      >
                        {isMarkdownDataLoading ? (
                          <Loader2 size={12} className="animate-spin" />
                        ) : (
                          <Copy size={12} />
                        )}
                        <span>Copy as Markdown</span>
                      </DropdownMenuItem>
                      <DropdownMenuItem
                        className="flex items-center space-x-2 whitespace-nowrap"
                        onClick={(e) => {
                          e.stopPropagation()
                          downloadImage('png')
                        }}
                      >
                        <Download size={12} />
                        <span>Download as PNG</span>
                      </DropdownMenuItem>
                      <DropdownMenuItem
                        className="flex items-center space-x-2 whitespace-nowrap"
                        onClick={(e) => {
                          e.stopPropagation()
                          downloadImage('svg')
                        }}
                      >
                        <Download size={12} />
                        <span>Download as SVG</span>
                      </DropdownMenuItem>
                    </DropdownMenuContent>
                  </DropdownMenu>
                </div>
                <AlertDialog open={autoLayoutDialogOpen} onOpenChange={setAutoLayoutDialogOpen}>
                  <Shortcut
                    id={SHORTCUT_IDS.SCHEMA_VISUALIZER_AUTO_LAYOUT}
                    onTrigger={() => setAutoLayoutDialogOpen(true)}
                    options={{ enabled: shortcutsEnabled }}
                    side="bottom"
                    tooltipOpen={autoLayoutDialogOpen ? false : undefined}
                  >
                    <AlertDialogTrigger asChild>
                      <Button variant="default">Auto layout</Button>
                    </AlertDialogTrigger>
                  </Shortcut>
                  <AlertDialogContent>
                    <AlertDialogHeader>
                      <AlertDialogTitle>Confirm to rearrange all nodes</AlertDialogTitle>
                      <AlertDialogDescription>
                        Auto layout will rearrange all nodes in the graph. This cannot be undone.
                      </AlertDialogDescription>
                    </AlertDialogHeader>
                    <AlertDialogFooter>
                      <AlertDialogCancel>Cancel</AlertDialogCancel>
                      <AlertDialogAction onClick={resetLayout}>Apply</AlertDialogAction>
                    </AlertDialogFooter>
                  </AlertDialogContent>
                </AlertDialog>
              </div>
            )}
          </>
        )}
      </div>
      {isLoadingTables && (
        <div className="w-full h-full flex items-center justify-center gap-x-2">
          <Loader2 className="animate-spin text-foreground-light" size={16} />
          <p className="text-sm text-foreground-light">Loading tables</p>
        </div>
      )}
      {isErrorTables && (
        <div className="w-full h-full flex items-center justify-center px-20">
          <AlertError subject="Failed to retrieve tables" error={errorTables} />
        </div>
      )}
      {isSuccessTables && (
        <>
          {hasNoTables ? (
            <div className="flex items-center justify-center w-full h-full">
              <Admonition
                type="default"
                className="max-w-md"
                title="No tables in schema"
                description={
                  isSchemaLocked
                    ? `The “${selectedSchema}” schema is managed by Supabase and is read-only through
                    the dashboard.`
                    : !canUpdateTables
                      ? 'You need additional permissions to create tables'
                      : `The “${selectedSchema}” schema doesn’t have any tables.`
                }
              >
                {canAddTables && (
                  <Button asChild className="mt-2 w-min" variant="default" icon={<Plus />}>
                    <Link href={`/project/${ref}/editor?create=table`}>New table</Link>
                  </Button>
                )}
              </Admonition>
            </div>
          ) : (
            <SchemaGraphContextProvider value={schemaGraphContext}>
              <div className="w-full h-full">
                <ReactFlow<Node<TableNodeData>, Edge<EdgeData>>
                  // FIXME: https://github.com/xyflow/xyflow/issues/4876
                  colorMode={'' as unknown as ColorMode}
                  defaultNodes={[]}
                  defaultEdges={[]}
                  defaultEdgeOptions={{
                    type: 'default',
                    animated: false,
                    deletable: false,
                  }}
                  nodeTypes={nodeTypes}
                  edgeTypes={edgeTypes}
                  fitView
                  minZoom={0.8}
                  maxZoom={1.8}
                  onlyRenderVisibleElements
                  proOptions={{ hideAttribution: true }}
                  onNodeDragStop={saveNodePositions}
                  onSelectionChange={handleSelectionChange}
                >
                  <Background
                    gap={16}
                    className="*:stroke-foreground-muted opacity-25"
                    variant={BackgroundVariant.Dots}
                    color={'inherit'}
                  />
                  <MiniMap
                    pannable
                    zoomable
                    nodeColor={miniMapNodeColor}
                    maskColor={miniMapMaskColor}
                    className="border rounded-md shadow-xs mb-11!"
                  />
                  <SchemaGraphLegend />
                  {hasNextPage && (
                    <Panel position="bottom-center" className="mb-11!">
                      <Button
                        variant="default"
                        size="tiny"
                        loading={isFetchingNextPage}
                        onClick={() => {
                          fitViewOnNextLayout.current = true
                          fetchNextPage()
                        }}
                      >
                        Load more tables
                      </Button>
                    </Panel>
                  )}
                </ReactFlow>
              </div>
            </SchemaGraphContextProvider>
          )}
        </>
      )}
      <SidePanelEditor selectedTable={selectedTable ?? undefined} includeColumns />
    </>
  )
}