konnosDashboard
Schemas.utils.ts294 lines · main
1import dagre from '@dagrejs/dagre'
2import type { PGSchema, PGTable } from '@supabase/pg-meta'
3import { Edge, Node, Position } from '@xyflow/react'
4import { uniqBy } from 'lodash'
5
6import '@xyflow/react/dist/style.css'
7
8import { LOCAL_STORAGE_KEYS } from 'common'
9
10import { TableNodeData } from './Schemas.constants'
11import { TABLE_NODE_ROW_HEIGHT, TABLE_NODE_WIDTH } from './SchemaTableNode'
12import { tryParseJson } from '@/lib/helpers'
13
14const NODE_SEP = 25
15const RANK_SEP = 50
16
17export async function getGraphDataFromTables(
18 ref?: string,
19 schema?: PGSchema,
20 tables?: PGTable[]
21): Promise<{
22 nodes: Node<TableNodeData>[]
23 edges: Edge[]
24}> {
25 if (!tables?.length) {
26 return { nodes: [], edges: [] }
27 }
28
29 const nodes = tables.map((table) => {
30 const columns = (table.columns || []).map((column) => {
31 return {
32 id: column.id,
33 isPrimary: table.primary_keys.some((pk) => pk.name === column.name),
34 name: column.name,
35 format: column.format,
36 isNullable: column.is_nullable,
37 isUnique: column.is_unique,
38 isUpdateable: column.is_updatable,
39 isIdentity: column.is_identity,
40 description: column.comment ?? '',
41 }
42 })
43
44 const data: TableNodeData = {
45 ref,
46 id: table.id,
47 name: table.name,
48 description: table.comment ?? '',
49 schema: table.schema,
50 isForeign: false,
51 columns,
52 }
53
54 return {
55 data,
56 id: `${table.id}`,
57 type: 'table',
58 position: { x: 0, y: 0 },
59 }
60 })
61
62 const edges: Edge[] = []
63 const currentSchema = tables[0].schema
64 const uniqueRelationships = uniqBy(
65 tables.flatMap((t) => t.relationships),
66 'id'
67 )
68
69 for (const rel of uniqueRelationships) {
70 // TODO: Support [external->this] relationship?
71 if (rel.source_schema !== currentSchema) {
72 continue
73 }
74
75 // Create additional [this->foreign] node that we can point to on the graph.
76 if (rel.target_table_schema !== currentSchema) {
77 const targetId = `${rel.target_table_schema}.${rel.target_table_name}.${rel.target_column_name}`
78
79 const targetNode = nodes.find((n) => n.id === targetId)
80 if (!targetNode) {
81 const data: TableNodeData = {
82 id: rel.id,
83 ref: ref!,
84 schema: rel.target_table_schema,
85 name: targetId,
86 description: '',
87 isForeign: true,
88 columns: [],
89 }
90
91 nodes.push({
92 id: targetId,
93 type: 'table',
94 data: data,
95 position: { x: 0, y: 0 },
96 })
97 }
98
99 const [source, sourceHandle] = findTablesHandleIds(
100 tables,
101 rel.source_table_name,
102 rel.source_column_name
103 )
104
105 if (source) {
106 edges.push({
107 id: String(rel.id),
108 source,
109 sourceHandle,
110 target: targetId,
111 targetHandle: targetId,
112 deletable: false,
113 data: {
114 sourceName: rel.source_table_name,
115 sourceSchemaName: rel.source_schema,
116 sourceColumnName: rel.source_column_name,
117 targetName: rel.target_table_name,
118 targetSchemaName: rel.target_table_schema,
119 targetColumnName: rel.target_column_name,
120 },
121 })
122 }
123
124 continue
125 }
126
127 const [source, sourceHandle] = findTablesHandleIds(
128 tables,
129 rel.source_table_name,
130 rel.source_column_name
131 )
132 const [target, targetHandle] = findTablesHandleIds(
133 tables,
134 rel.target_table_name,
135 rel.target_column_name
136 )
137
138 // We do not support [external->this] flow currently.
139 if (source && target) {
140 edges.push({
141 id: String(rel.id),
142 source,
143 sourceHandle,
144 target,
145 targetHandle,
146 type: 'default',
147 data: {
148 sourceName: rel.source_table_name,
149 sourceSchemaName: rel.source_schema,
150 sourceColumnName: rel.source_column_name,
151 targetName: rel.target_table_name,
152 targetSchemaName: rel.target_table_schema,
153 targetColumnName: rel.target_column_name,
154 },
155 })
156 }
157 }
158
159 const savedPositionsLocalStorage = localStorage.getItem(
160 LOCAL_STORAGE_KEYS.SCHEMA_VISUALIZER_POSITIONS(ref ?? 'project', schema?.id ?? 0)
161 )
162 const savedPositions = tryParseJson(savedPositionsLocalStorage)
163 return !!savedPositions
164 ? getLayoutedElementsViaLocalStorage(nodes, edges, savedPositions)
165 : getLayoutedElementsViaDagre(nodes, edges)
166}
167
168function findTablesHandleIds(
169 tables: PGTable[],
170 table_name: string,
171 column_name: string
172): [string?, string?] {
173 for (const table of tables) {
174 if (table_name !== table.name) continue
175
176 for (const column of table.columns || []) {
177 if (column_name !== column.name) continue
178
179 return [String(table.id), column.id]
180 }
181 }
182
183 return []
184}
185
186export const getLayoutedElementsViaDagre = (nodes: Node<TableNodeData>[], edges: Edge[]) => {
187 const dagreGraph = new dagre.graphlib.Graph()
188 dagreGraph.setDefaultEdgeLabel(() => ({}))
189 dagreGraph.setGraph({
190 rankdir: 'LR',
191 align: 'UR',
192 nodesep: NODE_SEP,
193 ranksep: RANK_SEP,
194 })
195
196 nodes.forEach((node) => {
197 dagreGraph.setNode(node.id, {
198 width: TABLE_NODE_WIDTH / 2,
199 height: (TABLE_NODE_ROW_HEIGHT / 2) * (node.data.columns.length + 1), // columns + header
200 })
201 })
202
203 edges.forEach((edge) => {
204 dagreGraph.setEdge(edge.source, edge.target)
205 })
206
207 dagre.layout(dagreGraph)
208
209 nodes.forEach((node) => {
210 const nodeWithPosition = dagreGraph.node(node.id)
211 node.targetPosition = Position.Left
212 node.sourcePosition = Position.Right
213 // We are shifting the dagre node position (anchor=center center) to the top left
214 // so it matches the React Flow node anchor point (top left).
215 node.position = {
216 x: nodeWithPosition.x - nodeWithPosition.width / 2,
217 y: nodeWithPosition.y - nodeWithPosition.height / 2,
218 }
219
220 return node
221 })
222
223 return { nodes, edges }
224}
225
226const getLayoutedElementsViaLocalStorage = (
227 nodes: Node<TableNodeData>[],
228 edges: Edge[],
229 positions: { [key: string]: { x: number; y: number } }
230) => {
231 // [Joshen] Potentially look into auto fitting new nodes?
232 // https://github.com/xyflow/xyflow/issues/1113
233
234 const nodesWithNoSavedPositons = nodes.filter((n) => !(n.id in positions))
235 let newNodeCount = 0
236 let basePosition = {
237 x: 0,
238 y: -(NODE_SEP + TABLE_NODE_ROW_HEIGHT + nodesWithNoSavedPositons.length * 10),
239 }
240
241 nodes.forEach((node) => {
242 const existingPosition = positions?.[node.id]
243
244 node.targetPosition = Position.Left
245 node.sourcePosition = Position.Right
246
247 if (existingPosition) {
248 node.position = existingPosition
249 } else {
250 node.position = {
251 x: basePosition.x + newNodeCount * 10,
252 y: basePosition.y + newNodeCount * 10,
253 }
254 newNodeCount += 1
255 }
256 })
257 return { nodes, edges }
258}
259
260export const getTableDefinitionAsMarkdown = (table: TableNodeData) => {
261 let markdown = `## Table \`${escapeForMarkdown(table.name)}\`\n\n`
262 if (table.description) {
263 markdown += `${table.description}\n\n`
264 }
265 markdown += `### Columns\n\n`
266 markdown += `| Name | Type | Constraints |\n`
267 markdown += `|------|------|-------------|\n`
268
269 return table.columns.reduce((current, column) => {
270 current += `| \`${escapeForMarkdown(column.name)}\` | \`${escapeForMarkdown(column.format)}\` | ${column.isPrimary ? 'Primary' : ''}${column.isNullable ? ' Nullable' : ''}${column.isUnique ? ' Unique' : ''}${column.isIdentity ? ' Identity' : ''} |\n`
271 return current
272 }, markdown)
273}
274
275export const getSchemaAsMarkdown = (schema: string, tables: TableNodeData[]) => {
276 return tables.reduce((current, table) => {
277 if (table.schema === schema) {
278 current += `${getTableDefinitionAsMarkdown(table)}\n`
279 }
280 return current
281 }, '')
282}
283
284const escapeForMarkdown = (str: string) => {
285 return (
286 str
287 // Escape backslashes first so later escapes are not ambiguous
288 .replace(/\\/g, '\\\\')
289 // Escape backticks and pipes for markdown tables
290 .replace(/([|`])/g, '\\$1')
291 // Remove new lines
292 .replace(/\n/g, ' ')
293 )
294}