From ba29cf187dd3d43ea2c35bab6c77ea154e64eca7 Mon Sep 17 00:00:00 2001
From: Jordi Enric <37541088+jordienr@users.noreply.github.com>
Date: Thu, 29 May 2025 18:00:37 +0200
Subject: [PATCH] Add copy schema to schema vizz (#34806)

* copy state

* add tooltip and better sql generation

* fix

* fix tooltip
---
 .../Database/Schemas/SchemaGraph.tsx          | 98 ++++++++++++++++++-
 apps/studio/next-env.d.ts                     |  2 +-
 2 files changed, 97 insertions(+), 3 deletions(-)

diff --git a/apps/studio/components/interfaces/Database/Schemas/SchemaGraph.tsx b/apps/studio/components/interfaces/Database/Schemas/SchemaGraph.tsx
index 6a8dc55769..ae81db255c 100644
--- a/apps/studio/components/interfaces/Database/Schemas/SchemaGraph.tsx
+++ b/apps/studio/components/interfaces/Database/Schemas/SchemaGraph.tsx
@@ -1,6 +1,6 @@
 import type { PostgresSchema } from '@supabase/postgres-meta'
 import { toPng, toSvg } from 'html-to-image'
-import { Download, Loader2 } from 'lucide-react'
+import { Check, Download, Loader2, Clipboard, Info } from 'lucide-react'
 import { useTheme } from 'next-themes'
 import { useEffect, useMemo, useState } from 'react'
 import ReactFlow, { Background, BackgroundVariant, MiniMap, useReactFlow } from 'reactflow'
@@ -21,7 +21,7 @@ import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigge
 import { SchemaGraphLegend } from './SchemaGraphLegend'
 import { getGraphDataFromTables, getLayoutedElementsViaDagre } from './Schemas.utils'
 import { TableNode } from './SchemaTableNode'
-
+import { copyToClipboard } from 'ui'
 // [Joshen] Persisting logic: Only save positions to local storage WHEN a node is moved OR when explicitly clicked to reset layout
 
 export const SchemaGraph = () => {
@@ -30,6 +30,13 @@ export const SchemaGraph = () => {
   const { project } = useProjectContext()
   const { selectedSchema, setSelectedSchema } = useQuerySchemaState()
 
+  const [copied, setCopied] = useState(false)
+  useEffect(() => {
+    if (copied) {
+      setTimeout(() => setCopied(false), 2000)
+    }
+  }, [copied])
+
   const [isDownloading, setIsDownloading] = useState(false)
 
   const miniMapNodeColor = '#111318'
@@ -160,6 +167,59 @@ export const SchemaGraph = () => {
     }
   }
 
+  function tablesToSQL(t: typeof tables) {
+    if (!Array.isArray(t)) return ''
+    const warning =
+      '-- WARNING: This schema is for context only and is not meant to be run.\n-- Table order and constraints may not be valid for execution.\n\n'
+    const sql = t
+      .map((table) => {
+        if (!table || !Array.isArray((table as any).columns)) return ''
+
+        const columns = (table as { columns?: any[] }).columns ?? []
+        const columnLines = columns.map((c) => {
+          let line = `  ${c.name} ${c.data_type}`
+          if (c.is_identity) {
+            line += ' GENERATED ALWAYS AS IDENTITY'
+          }
+          if (c.is_nullable === false) {
+            line += ' NOT NULL'
+          }
+          if (c.default_value !== null && c.default_value !== undefined) {
+            line += ` DEFAULT ${c.default_value}`
+          }
+          if (c.is_unique) {
+            line += ' UNIQUE'
+          }
+          if (c.check) {
+            line += ` CHECK (${c.check})`
+          }
+          return line
+        })
+
+        const constraints: string[] = []
+
+        if (Array.isArray(table.primary_keys) && table.primary_keys.length > 0) {
+          const pkCols = table.primary_keys.map((pk) => pk.name).join(', ')
+          constraints.push(`  CONSTRAINT ${table.name}_pkey PRIMARY KEY (${pkCols})`)
+        }
+
+        if (Array.isArray(table.relationships)) {
+          table.relationships.forEach((rel) => {
+            if (rel && rel.source_table_name === table.name) {
+              constraints.push(
+                `  CONSTRAINT ${rel.constraint_name} FOREIGN KEY (${rel.source_column_name}) REFERENCES ${rel.target_table_schema}.${rel.target_table_name}(${rel.target_column_name})`
+              )
+            }
+          })
+        }
+
+        const allLines = [...columnLines, ...constraints]
+        return `CREATE TABLE ${table.schema}.${table.name} (\n${allLines.join(',\n')}\n);`
+      })
+      .join('\n')
+    return warning + sql
+  }
+
   useEffect(() => {
     if (isSuccessTables && isSuccessSchemas && tables.length > 0) {
       const schema = schemas.find((s) => s.name === selectedSchema) as PostgresSchema
@@ -192,6 +252,40 @@ export const SchemaGraph = () => {
               onSelectSchema={setSelectedSchema}
             />
             <div className="flex items-center gap-x-2">
+              <ButtonTooltip
+                type="outline"
+                icon={copied ? <Check /> : <Clipboard />}
+                onClick={() => {
+                  if (tables) {
+                    copyToClipboard(tablesToSQL(tables))
+                    setCopied(true)
+                  }
+                }}
+                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>
+              <ButtonTooltip
+                type="default"
+                loading={isDownloading}
+                className="px-1.5"
+                icon={<Download />}
+                onClick={() => downloadImage('png')}
+                tooltip={{ content: { side: 'bottom', text: 'Download current view as PNG' } }}
+              />
               <DropdownMenu>
                 <DropdownMenuTrigger asChild>
                   <ButtonTooltip
diff --git a/apps/studio/next-env.d.ts b/apps/studio/next-env.d.ts
index 52e831b434..a4a7b3f5cf 100644
--- a/apps/studio/next-env.d.ts
+++ b/apps/studio/next-env.d.ts
@@ -2,4 +2,4 @@
 /// <reference types="next/image-types/global" />
 
 // NOTE: This file should not be edited
-// see https://nextjs.org/docs/pages/api-reference/config/typescript for more information.
+// see https://nextjs.org/docs/pages/building-your-application/configuring/typescript for more information.