konnosDashboard
SQLEditor.tsx1099 lines · main
1// @ts-nocheck
2import type { Monaco } from '@monaco-editor/react'
3import {
4 acceptUntrustedSql,
5 rawSql,
6 safeSql,
7 type SafeSqlFragment,
8 type UntrustedSqlFragment,
9} from '@supabase/pg-meta'
10import { wrapWithRollback } from '@supabase/pg-meta/src/query'
11import { useQueryClient } from '@tanstack/react-query'
12import { IS_PLATFORM, LOCAL_STORAGE_KEYS, useFlag, useParams } from 'common'
13import { ChevronUp, Loader2 } from 'lucide-react'
14import dynamic from 'next/dynamic'
15import { useRouter } from 'next/router'
16import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
17import { toast } from 'sonner'
18import {
19 Button,
20 cn,
21 DropdownMenu,
22 DropdownMenuContent,
23 DropdownMenuRadioGroup,
24 DropdownMenuRadioItem,
25 DropdownMenuTrigger,
26 ResizableHandle,
27 ResizablePanel,
28 ResizablePanelGroup,
29 Tooltip,
30 TooltipContent,
31 TooltipTrigger,
32} from 'ui'
33
34import { useSqlEditorDiff, useSqlEditorPrompt } from './hooks'
35import { RunQueryWarningModal } from './RunQueryWarningModal'
36import {
37 generateSnippetTitle,
38 ROWS_PER_PAGE_OPTIONS,
39 sqlAiDisclaimerComment,
40 untitledSnippetTitle,
41} from './SQLEditor.constants'
42import {
43 DiffType,
44 IStandaloneCodeEditor,
45 IStandaloneDiffEditor,
46 type PotentialIssues,
47} from './SQLEditor.types'
48import {
49 appendEnableRLSStatements,
50 checkAlterDatabaseConnection,
51 checkDestructiveQuery,
52 checkIfAppendLimitRequired,
53 createSqlSnippetSkeletonV2,
54 filterTablesCoveredByEnsureRLSTrigger,
55 getCreateTablesMissingRLS,
56 hasActiveEnsureRLSTrigger,
57 isUpdateWithoutWhere,
58 suffixWithLimit,
59} from './SQLEditor.utils'
60import { useAddDefinitions } from './useAddDefinitions'
61import { UtilityPanel } from './UtilityPanel/UtilityPanel'
62import {
63 isExplainQuery,
64 isExplainSql,
65 splitSqlStatements,
66} from '@/components/interfaces/ExplainVisualizer/ExplainVisualizer.utils'
67import { SIDEBAR_KEYS } from '@/components/layouts/ProjectLayout/LayoutSidebar/LayoutSidebarProvider'
68import ResizableAIWidget from '@/components/ui/AIEditor/ResizableAIWidget'
69import { GridFooter } from '@/components/ui/GridFooter'
70import { useSqlTitleGenerateMutation } from '@/data/ai/sql-title-mutation'
71import { useDatabaseEventTriggersQuery } from '@/data/database-event-triggers/database-event-triggers-query'
72import { constructHeaders, isValidConnString } from '@/data/fetchers'
73import { lintKeys } from '@/data/lint/keys'
74import { useReadReplicasQuery } from '@/data/read-replicas/replicas-query'
75import { useExecuteSqlMutation } from '@/data/sql/execute-sql-mutation'
76import { useSendEventMutation } from '@/data/telemetry/send-event-mutation'
77import { isError } from '@/data/utils/error-check'
78import { useOrgAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
79import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
80import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
81import { generateUuid } from '@/lib/api/snippets.browser'
82import { BASE_PATH } from '@/lib/constants'
83import { formatSql } from '@/lib/formatSql'
84import { detectOS } from '@/lib/helpers'
85import { useProfile } from '@/lib/profile'
86import { wrapWithRoleImpersonation } from '@/lib/role-impersonation'
87import { useAiAssistantStateSnapshot } from '@/state/ai-assistant-state'
88import { useDatabaseSelectorStateSnapshot } from '@/state/database-selector'
89import {
90 isRoleImpersonationEnabled,
91 useGetImpersonatedRoleState,
92} from '@/state/role-impersonation-state'
93import { SHORTCUT_IDS } from '@/state/shortcuts/registry'
94import { useShortcut } from '@/state/shortcuts/useShortcut'
95import { useSidebarManagerSnapshot } from '@/state/sidebar-manager-state'
96import { getSqlEditorV2StateSnapshot, useSqlEditorV2StateSnapshot } from '@/state/sql-editor-v2'
97import { createTabId, useTabsStateSnapshot } from '@/state/tabs'
98
99// Load the monaco editor client-side only (does not behave well server-side)
100const MonacoEditor = dynamic(() => import('./MonacoEditor'), { ssr: false })
101const DiffEditor = dynamic(
102 () => import('../../ui/DiffEditor').then(({ DiffEditor }) => DiffEditor),
103 { ssr: false }
104)
105
106export const SQLEditor = () => {
107 const os = detectOS()
108 const router = useRouter()
109 const { ref, id: urlId } = useParams()
110
111 const { profile } = useProfile()
112 const { data: project } = useSelectedProjectQuery()
113 const { data: org } = useSelectedOrganizationQuery()
114
115 const queryClient = useQueryClient()
116 const tabs = useTabsStateSnapshot()
117 const aiSnap = useAiAssistantStateSnapshot()
118 const { openSidebar } = useSidebarManagerSnapshot()
119 const snapV2 = useSqlEditorV2StateSnapshot()
120 const getImpersonatedRoleState = useGetImpersonatedRoleState()
121 const databaseSelectorState = useDatabaseSelectorStateSnapshot()
122 const { isHipaaProjectDisallowed } = useOrgAiOptInLevel()
123 const showPrettyExplain = useFlag('ShowPrettyExplain')
124
125 const {
126 sourceSqlDiff,
127 setSourceSqlDiff,
128 selectedDiffType,
129 setSelectedDiffType,
130 setIsAcceptDiffLoading,
131 isDiffOpen,
132 defaultSqlDiff,
133 closeDiff,
134 } = useSqlEditorDiff()
135 const { promptState, setPromptState, promptInput, setPromptInput, resetPrompt } =
136 useSqlEditorPrompt()
137
138 const editorRef = useRef<IStandaloneCodeEditor | null>(null)
139 const monacoRef = useRef<Monaco | null>(null)
140 const diffEditorRef = useRef<IStandaloneDiffEditor | null>(null)
141 const scrollTopRef = useRef<number>(0)
142 const shouldRefocusAfterRunRef = useRef(false)
143
144 const [hasSelection, setHasSelection] = useState<boolean>(false)
145 const [lineHighlights, setLineHighlights] = useState<string[]>([])
146 const [isDiffEditorMounted, setIsDiffEditorMounted] = useState(false)
147 const [potentialIssues, setPotentialIssues] = useState<PotentialIssues>()
148
149 const [showWidget, setShowWidget] = useState(false)
150 const [activeUtilityTab, setActiveUtilityTab] = useState<string>('results')
151
152 const refocusEditor = useCallback(() => {
153 requestAnimationFrame(() => {
154 setTimeout(() => editorRef.current?.focus(), 0)
155 })
156 }, [])
157
158 useShortcut(SHORTCUT_IDS.SQL_EDITOR_FOCUS_EDITOR, refocusEditor, {
159 registerInCommandMenu: true,
160 })
161
162 const openNewSnippet = useCallback(() => {
163 if (!ref) return
164 // skip=true bypasses the "load last visited snippet" redirect on /sql/new.
165 // Without it, the effect in pages/project/[ref]/sql/[id].tsx bounces back
166 // to the previous snippet.
167 router.push(`/project/${ref}/sql/new?skip=true`)
168 }, [ref, router])
169
170 useShortcut(SHORTCUT_IDS.SQL_EDITOR_NEW_SNIPPET, openNewSnippet, {
171 registerInCommandMenu: true,
172 })
173
174 const clearPendingRunRefocus = useCallback(() => {
175 shouldRefocusAfterRunRef.current = false
176 }, [])
177
178 const refocusEditorAfterRunIfNeeded = useCallback(() => {
179 if (!shouldRefocusAfterRunRef.current) return
180
181 shouldRefocusAfterRunRef.current = false
182 refocusEditor()
183 }, [refocusEditor])
184
185 // generate a new snippet title and an id to be used for new snippets. The dependency on urlId is to avoid a bug which
186 // shows up when clicking on the SQL Editor while being in the SQL editor on a random snippet.
187 const [generatedNewSnippetName, generatedId] = useMemo(() => {
188 const name = generateSnippetTitle()
189 return [name, generateUuid([`${name}.sql`])]
190 }, [urlId])
191
192 // the id is stable across renders - it depends either on the url or on the memoized generated id
193 const id = !urlId || urlId === 'new' ? generatedId : urlId
194
195 const limit = snapV2.limit
196 const results = snapV2.results[id]?.[0]
197 const snippetIsLoading = !(
198 id in snapV2.snippets && snapV2.snippets[id].snippet.content !== undefined
199 )
200 const isLoading = urlId === 'new' ? false : snippetIsLoading
201
202 useAddDefinitions(id, monacoRef.current)
203
204 const { data: databases, isSuccess: isSuccessReadReplicas } = useReadReplicasQuery(
205 {
206 projectRef: ref,
207 },
208 { enabled: isValidConnString(project?.connectionString) }
209 )
210
211 const { data: eventTriggers } = useDatabaseEventTriggersQuery(
212 {
213 projectRef: project?.ref,
214 connectionString: project?.connectionString,
215 },
216 { enabled: isValidConnString(project?.connectionString) }
217 )
218
219 /* React query mutations */
220 const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation()
221 const { mutate: sendEvent } = useSendEventMutation()
222 const { mutate: execute, isPending: isExecuting } = useExecuteSqlMutation({
223 onSuccess(data, vars) {
224 if (id) {
225 snapV2.addResult(id, data.result, vars.autoLimit)
226
227 if (showPrettyExplain && isExplainQuery(data.result)) {
228 snapV2.addExplainResult(id, data.result)
229 setActiveUtilityTab('explain')
230 } else if (activeUtilityTab === 'explain') {
231 // If on Explain tab but ran a non-EXPLAIN query, switch to Results tab
232 setActiveUtilityTab('results')
233 }
234 }
235
236 // revalidate lint query
237 queryClient.invalidateQueries({ queryKey: lintKeys.lint(ref) })
238 refocusEditorAfterRunIfNeeded()
239 },
240 onError(error: any, vars) {
241 if (id) {
242 if (error.position && monacoRef.current) {
243 const editor = editorRef.current
244 const monaco = monacoRef.current
245
246 const startLineNumber = hasSelection ? (editor?.getSelection()?.startLineNumber ?? 0) : 0
247
248 const formattedError = error.formattedError ?? ''
249 const lineError = formattedError.slice(formattedError.indexOf('LINE'))
250 const line =
251 startLineNumber + Number(lineError.slice(0, lineError.indexOf(':')).split(' ')[1])
252
253 if (!isNaN(line)) {
254 const decorations = editor?.deltaDecorations(
255 [],
256 [
257 {
258 range: new monaco.Range(line, 1, line, 20),
259 options: {
260 isWholeLine: true,
261 inlineClassName: 'bg-warning-400',
262 },
263 },
264 ]
265 )
266 if (decorations) {
267 editor?.revealLineInCenter(line)
268 setLineHighlights(decorations)
269 }
270 }
271 }
272
273 snapV2.addResultError(id, error, vars.autoLimit)
274 }
275
276 refocusEditorAfterRunIfNeeded()
277 },
278 })
279
280 const { mutate: executeExplain, isPending: isExplainExecuting } = useExecuteSqlMutation({
281 onSuccess(data) {
282 if (id) {
283 snapV2.addExplainResult(id, data.result)
284 setActiveUtilityTab('explain')
285 }
286 },
287 onError(error) {
288 if (id) {
289 snapV2.addExplainResultError(id, error)
290 setActiveUtilityTab('explain')
291 }
292 },
293 })
294
295 const setAiTitle = useCallback(
296 async (id: string, sql: string) => {
297 try {
298 const { title: name } = await generateSqlTitle({ sql })
299 snapV2.updateSnippet({ id, snippet: { name } })
300 snapV2.addNeedsSaving(id)
301 const tabId = createTabId('sql', { id })
302 tabs.updateTab(tabId, { label: name })
303 } catch (error) {
304 // [Joshen] No error handler required as this happens in the background and not necessary to ping the user
305 }
306 },
307 [generateSqlTitle, snapV2]
308 )
309
310 const prettifyQuery = useCallback(async () => {
311 if (isDiffOpen) return
312
313 // use the latest state
314 const state = getSqlEditorV2StateSnapshot()
315 const snippet = state.snippets[id]
316
317 if (editorRef.current && project) {
318 const editor = editorRef.current
319 const selection = editor.getSelection()
320 const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined
321 const sql = snippet
322 ? ((selectedValue || editorRef.current?.getValue()) ??
323 snippet.snippet.content?.unchecked_sql)
324 : selectedValue || editorRef.current?.getValue()
325 const formattedSql = formatSql(sql)
326
327 const editorModel = editorRef?.current?.getModel()
328 if (editorRef.current && editorModel) {
329 editorRef.current.executeEdits('apply-prettify-edit', [
330 {
331 text: formattedSql,
332 range: editorModel.getFullModelRange(),
333 },
334 ])
335 snapV2.setSql({ id, sql: formattedSql })
336 }
337 }
338 }, [id, isDiffOpen, project, snapV2])
339
340 useShortcut(SHORTCUT_IDS.SQL_EDITOR_FORMAT, prettifyQuery, {
341 registerInCommandMenu: true,
342 })
343
344 const executeQuery = useCallback(
345 async (force: boolean = false, sqlOverride?: SafeSqlFragment) => {
346 if (isDiffOpen) {
347 clearPendingRunRefocus()
348 return
349 }
350
351 // use the latest state
352 const state = getSqlEditorV2StateSnapshot()
353 const snippet = state.snippets[id]
354
355 if (editorRef.current === null || isExecuting || project === undefined) {
356 clearPendingRunRefocus()
357 return
358 }
359
360 const editor = editorRef.current
361 const selection = editor.getSelection()
362 const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined
363
364 const editorSql = snippet
365 ? ((selectedValue || editorRef.current?.getValue()) ??
366 snippet.snippet.content?.unchecked_sql)
367 : selectedValue || editorRef.current?.getValue()
368 const sql = sqlOverride ?? editorSql
369
370 const hasDestructiveOperations = checkDestructiveQuery(sql)
371 const hasUpdateWithoutWhere = isUpdateWithoutWhere(sql)
372 const hasAlterDatabasePreventConnection = checkAlterDatabaseConnection(sql)
373 const createTablesMissingRLS = filterTablesCoveredByEnsureRLSTrigger(
374 getCreateTablesMissingRLS(sql),
375 hasActiveEnsureRLSTrigger(eventTriggers)
376 )
377
378 const queryHasIssues =
379 !force &&
380 (hasDestructiveOperations ||
381 hasUpdateWithoutWhere ||
382 hasAlterDatabasePreventConnection ||
383 createTablesMissingRLS.length > 0)
384
385 if (queryHasIssues) {
386 setPotentialIssues({
387 hasDestructiveOperations,
388 hasUpdateWithoutWhere,
389 hasAlterDatabasePreventConnection,
390 createTablesMissingRLS,
391 })
392 return
393 }
394
395 if (
396 !isHipaaProjectDisallowed &&
397 snippet?.snippet.name.startsWith(untitledSnippetTitle) &&
398 IS_PLATFORM
399 ) {
400 // Intentionally don't await title gen (lazy)
401 setAiTitle(id, sql)
402 }
403
404 if (lineHighlights.length > 0) {
405 editor?.deltaDecorations(lineHighlights, [])
406 setLineHighlights([])
407 }
408
409 const impersonatedRoleState = getImpersonatedRoleState()
410 const connectionString = databases?.find(
411 (db) => db.identifier === databaseSelectorState.selectedDatabaseId
412 )?.connectionString
413 if (!isValidConnString(connectionString)) {
414 clearPendingRunRefocus()
415 return toast.error('Unable to run query: Connection string is missing')
416 }
417
418 const userSql = rawSql(sql)
419 const { appendAutoLimit } = checkIfAppendLimitRequired(userSql, limit)
420 const formattedSql = suffixWithLimit(userSql, limit)
421
422 execute({
423 projectRef: project.ref,
424 connectionString: connectionString,
425 sql: wrapWithRoleImpersonation(formattedSql, impersonatedRoleState),
426 autoLimit: appendAutoLimit ? limit : undefined,
427 isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRoleState.role),
428 isStatementTimeoutDisabled: true,
429 contextualInvalidation: true,
430 handleError: (error) => {
431 throw error
432 },
433 })
434
435 sendEvent({
436 action: 'sql_editor_query_run_button_clicked',
437 groups: { project: ref ?? 'Unknown', organization: org?.slug ?? 'Unknown' },
438 })
439 },
440 // eslint-disable-next-line react-hooks/exhaustive-deps
441 [
442 clearPendingRunRefocus,
443 isDiffOpen,
444 id,
445 isExecuting,
446 project,
447 isHipaaProjectDisallowed,
448 execute,
449 getImpersonatedRoleState,
450 setAiTitle,
451 databaseSelectorState.selectedDatabaseId,
452 databases,
453 eventTriggers,
454 limit,
455 ]
456 )
457
458 const executeQueryFromButton = useCallback(() => {
459 shouldRefocusAfterRunRef.current = true
460 refocusEditor()
461 void executeQuery()
462 }, [executeQuery, refocusEditor])
463
464 const executeExplainQuery = useCallback(async () => {
465 if (isDiffOpen) return
466
467 // use the latest state
468 const state = getSqlEditorV2StateSnapshot()
469 const snippet = state.snippets[id]
470
471 if (editorRef.current !== null && !isExplainExecuting && project !== undefined) {
472 const editor = editorRef.current
473 const selection = editor.getSelection()
474 const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined
475
476 const sql = snippet
477 ? ((selectedValue || editorRef.current?.getValue()) ??
478 snippet.snippet.content?.unchecked_sql)
479 : selectedValue || editorRef.current?.getValue()
480
481 // Check for multiple statements - EXPLAIN only works on a single statement
482 const statements = splitSqlStatements(sql)
483 if (statements.length > 1) {
484 snapV2.addExplainResultError(id, {
485 message:
486 'EXPLAIN only works on a single SQL statement. Please select just one query to analyze.',
487 })
488 setActiveUtilityTab('explain')
489 return
490 }
491
492 if (lineHighlights.length > 0) {
493 editor?.deltaDecorations(lineHighlights, [])
494 setLineHighlights([])
495 }
496
497 const impersonatedRoleState = getImpersonatedRoleState()
498 const connectionString = databases?.find(
499 (db) => db.identifier === databaseSelectorState.selectedDatabaseId
500 )?.connectionString
501 if (!isValidConnString(connectionString)) {
502 return toast.error('Unable to run query: Connection string is missing')
503 }
504
505 // Wrap the query with EXPLAIN ANALYZE only if it's not already an EXPLAIN query
506 const userSql = rawSql(sql ?? '')
507 const explainSql = isExplainSql(sql) ? userSql : safeSql`EXPLAIN ANALYZE ${userSql}`
508
509 // Wrap EXPLAIN queries in a transaction with rollback to prevent data modifications
510 // This ensures EXPLAIN ANALYZE INSERT/UPDATE/DELETE queries don't actually modify data
511 const explainSqlWithTransaction = wrapWithRollback(
512 wrapWithRoleImpersonation(explainSql, impersonatedRoleState)
513 )
514
515 executeExplain({
516 projectRef: project.ref,
517 connectionString: connectionString,
518 sql: explainSqlWithTransaction,
519 isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRoleState.role),
520 handleError: (error) => {
521 throw error
522 },
523 })
524 }
525 }, [
526 isDiffOpen,
527 id,
528 isExplainExecuting,
529 project,
530 executeExplain,
531 getImpersonatedRoleState,
532 databaseSelectorState.selectedDatabaseId,
533 databases,
534 lineHighlights,
535 snapV2,
536 ])
537
538 useShortcut(SHORTCUT_IDS.SQL_EDITOR_EXPLAIN, executeExplainQuery, {
539 registerInCommandMenu: true,
540 })
541
542 const handleNewQuery = useCallback(
543 async (sql: string, name: string) => {
544 if (!ref) return console.error('Project ref is required')
545 if (!profile) return console.error('Profile is required')
546 if (!project) return console.error('Project is required')
547
548 try {
549 const snippet = createSqlSnippetSkeletonV2({
550 name,
551 sql,
552 owner_id: profile.id,
553 project_id: project.id,
554 })
555 snapV2.addSnippet({ projectRef: ref, snippet })
556 snapV2.addNeedsSaving(snippet.id!)
557 router.push(`/project/${ref}/sql/${snippet.id}`)
558 } catch (error: any) {
559 toast.error(`Failed to create new query: ${error.message}`)
560 }
561 },
562 // eslint-disable-next-line react-hooks/exhaustive-deps
563 [profile?.id, project?.id, ref, router, snapV2]
564 )
565
566 const onMount = (editor: IStandaloneCodeEditor) => {
567 const tabId = createTabId('sql', { id })
568 const tabData = tabs.tabsMap[tabId]
569
570 // [Joshen] Tiny timeout to give a bit of time for the content to load before scrolling
571 setTimeout(() => {
572 if (tabData?.metadata?.scrollTop) {
573 editor.setScrollTop(tabData.metadata.scrollTop)
574 }
575 }, 20)
576 editor.onDidScrollChange((e) => (scrollTopRef.current = e.scrollTop))
577 }
578
579 const buildDebugPrompt = useCallback(() => {
580 const snippet = snapV2.snippets[id]
581 const result = snapV2.results[id]?.[0]
582 const sql = (snippet?.snippet.content?.unchecked_sql ?? '')
583 .replace(sqlAiDisclaimerComment, '')
584 .trim()
585 const errorMessage = result?.error?.message ?? 'Unknown error'
586 const prompt = `Help me to debug the attached sql snippet which gives the following error: \n\n${errorMessage}`
587
588 return `${prompt}\n\nSQL Query:\n\`\`\`sql\n${sql}\n\`\`\``
589 }, [id, snapV2.results, snapV2.snippets])
590
591 const onDebug = useCallback(async () => {
592 try {
593 const snippet = snapV2.snippets[id]
594 const result = snapV2.results[id]?.[0]
595 openSidebar(SIDEBAR_KEYS.AI_ASSISTANT)
596 aiSnap.newChat({
597 name: 'Debug SQL snippet',
598 sqlSnippets: [
599 (snippet.snippet.content?.unchecked_sql ?? '').replace(sqlAiDisclaimerComment, '').trim(),
600 ],
601 initialInput: `Help me to debug the attached sql snippet which gives the following error: \n\n${result.error.message}`,
602 })
603 } catch (error: unknown) {
604 // [Joshen] There's a tendency for the SQL debug to chuck a lengthy error message
605 // that's not relevant for the user - so we prettify it here by avoiding to return the
606 // entire error body from the assistant
607 if (isError(error)) {
608 toast.error(
609 `Sorry, the assistant failed to debug your query! Please try again with a different one.`
610 )
611 }
612 }
613 // eslint-disable-next-line react-hooks/exhaustive-deps
614 }, [id, snapV2.results, snapV2.snippets])
615
616 const acceptAiHandler = useCallback(async () => {
617 try {
618 setIsAcceptDiffLoading(true)
619
620 // TODO: show error if undefined
621 if (!sourceSqlDiff || !editorRef.current || !diffEditorRef.current) return
622
623 const editorModel = editorRef.current.getModel()
624 const diffModel = diffEditorRef.current.getModel()
625
626 if (!editorModel || !diffModel) return
627
628 const sql = diffModel.modified.getValue()
629
630 if (selectedDiffType === DiffType.NewSnippet) {
631 const { title } = await generateSqlTitle({ sql })
632 await handleNewQuery(sql, title)
633 } else {
634 editorRef.current.executeEdits('apply-ai-edit', [
635 {
636 text: sql,
637 range: editorModel.getFullModelRange(),
638 },
639 ])
640 }
641
642 sendEvent({
643 action: 'assistant_sql_diff_handler_evaluated',
644 properties: { handlerAccepted: true },
645 groups: { project: ref ?? 'Unknown', organization: org?.slug ?? 'Unknown' },
646 })
647
648 setSelectedDiffType(DiffType.Modification)
649 resetPrompt()
650 closeDiff()
651 } finally {
652 setIsAcceptDiffLoading(false)
653 }
654 // eslint-disable-next-line react-hooks/exhaustive-deps
655 }, [sourceSqlDiff, selectedDiffType, handleNewQuery, generateSqlTitle, router, id, snapV2])
656
657 const discardAiHandler = useCallback(() => {
658 sendEvent({
659 action: 'assistant_sql_diff_handler_evaluated',
660 properties: { handlerAccepted: false },
661 groups: { project: ref ?? 'Unknown', organization: org?.slug ?? 'Unknown' },
662 })
663 resetPrompt()
664 closeDiff()
665 }, [closeDiff, resetPrompt, sendEvent])
666
667 const [isCompletionLoading, setIsCompletionLoading] = useState<boolean>(false)
668
669 const complete = useCallback(
670 async (
671 _prompt: string,
672 options?: {
673 headers?: Record<string, string>
674 body?: { completionMetadata?: any }
675 }
676 ) => {
677 try {
678 setIsCompletionLoading(true)
679
680 const response = await fetch(`${BASE_PATH}/api/ai/code/complete`, {
681 method: 'POST',
682 headers: {
683 'Content-Type': 'application/json',
684 ...(options?.headers ?? {}),
685 },
686 body: JSON.stringify({
687 projectRef: project?.ref,
688 connectionString: project?.connectionString,
689 language: 'sql',
690 orgSlug: org?.slug,
691 ...(options?.body ?? {}),
692 }),
693 })
694
695 if (!response.ok) {
696 const errorText = await response.text()
697 throw new Error(errorText || 'Failed to generate completion')
698 }
699
700 // API returns a JSON-encoded string
701 const text: string = await response.json()
702
703 const meta = options?.body?.completionMetadata ?? {}
704 const beforeSelection: string = meta.textBeforeCursor ?? ''
705 const afterSelection: string = meta.textAfterCursor ?? ''
706 const selection: string = meta.selection ?? ''
707
708 const original = beforeSelection + selection + afterSelection
709 const modified = beforeSelection + text + afterSelection
710
711 const formattedModified = formatSql(modified)
712 setSourceSqlDiff({ original, modified: formattedModified })
713 setSelectedDiffType(DiffType.Modification)
714 setPromptState((prev) => ({ ...prev, isLoading: false }))
715 setIsCompletionLoading(false)
716 } catch (error: any) {
717 toast.error(`Failed to generate SQL: ${error?.message ?? 'Unknown error'}`)
718 setIsCompletionLoading(false)
719 throw error
720 }
721 },
722 [
723 org?.slug,
724 project?.connectionString,
725 project?.ref,
726 setPromptState,
727 setSelectedDiffType,
728 setSourceSqlDiff,
729 ]
730 )
731
732 const handlePrompt = async (
733 prompt: string,
734 context: {
735 beforeSelection: string
736 selection: string
737 afterSelection: string
738 }
739 ) => {
740 try {
741 setPromptState((prev) => ({
742 ...prev,
743 selection: context.selection,
744 beforeSelection: context.beforeSelection,
745 afterSelection: context.afterSelection,
746 }))
747 const headerData = await constructHeaders()
748
749 const authorizationHeader = headerData.get('Authorization')
750
751 await complete(prompt, {
752 ...(authorizationHeader ? { headers: { Authorization: authorizationHeader } } : undefined),
753 body: {
754 completionMetadata: {
755 textBeforeCursor: context.beforeSelection,
756 textAfterCursor: context.afterSelection,
757 language: 'pgsql',
758 prompt,
759 selection: context.selection,
760 },
761 },
762 })
763 } catch (error) {
764 setPromptState((prev) => ({ ...prev, isLoading: false }))
765 }
766 }
767
768 /** All useEffects are at the bottom before returning the TSX */
769
770 useEffect(() => {
771 if (id) {
772 closeDiff()
773 setPromptState((prev) => ({ ...prev, isOpen: false }))
774 }
775 return () => {
776 if (ref) {
777 const tabId = createTabId('sql', { id })
778 tabs.updateTab(tabId, { scrollTop: scrollTopRef.current })
779 }
780 }
781 // eslint-disable-next-line react-hooks/exhaustive-deps
782 }, [closeDiff, id])
783
784 useEffect(() => {
785 const handler = (e: KeyboardEvent) => {
786 if (!isDiffOpen && !promptState.isOpen) return
787
788 switch (e.key) {
789 case 'Enter':
790 if ((os === 'macos' ? e.metaKey : e.ctrlKey) && isDiffOpen) {
791 acceptAiHandler()
792 resetPrompt()
793 }
794 return
795 case 'Escape':
796 if (isDiffOpen) discardAiHandler()
797 resetPrompt()
798 editorRef.current?.focus()
799 return
800 }
801 }
802 window.addEventListener('keydown', handler)
803 return () => window.removeEventListener('keydown', handler)
804 }, [os, isDiffOpen, promptState.isOpen, acceptAiHandler, discardAiHandler, resetPrompt])
805
806 useEffect(() => {
807 if (isDiffOpen) {
808 const diffEditor = diffEditorRef.current
809 const model = diffEditor?.getModel()
810 if (model && model.original && model.modified) {
811 model.original.setValue(defaultSqlDiff.original)
812 model.modified.setValue(defaultSqlDiff.modified)
813 // scroll to the start line of the modification
814 const modifiedEditor = diffEditor!.getModifiedEditor()
815 const startLine = promptState.startLineNumber
816 modifiedEditor.revealLineInCenter(startLine)
817 }
818 }
819 // eslint-disable-next-line react-hooks/exhaustive-deps
820 }, [selectedDiffType, sourceSqlDiff])
821
822 useEffect(() => {
823 if (isSuccessReadReplicas) {
824 const primaryDatabase = databases.find((db) => db.identifier === ref)
825 databaseSelectorState.setSelectedDatabaseId(primaryDatabase?.identifier)
826 }
827 // eslint-disable-next-line react-hooks/exhaustive-deps
828 }, [isSuccessReadReplicas, databases, ref])
829
830 useEffect(() => {
831 if (snapV2.diffContent !== undefined) {
832 const { diffType, sql }: { diffType: DiffType; sql: string } = snapV2.diffContent
833 const editorModel = editorRef.current?.getModel()
834 if (!editorModel) return
835
836 const existingValue = editorRef.current?.getValue() ?? ''
837 if (existingValue.length === 0) {
838 // if the editor is empty, just copy over the code
839 editorRef.current?.executeEdits('apply-ai-message', [
840 {
841 text: `${sql}`,
842 range: editorModel.getFullModelRange(),
843 },
844 ])
845 } else {
846 const currentSql = editorRef.current?.getValue()
847 const diff = { original: currentSql || '', modified: sql }
848 setSourceSqlDiff(diff)
849 setSelectedDiffType(diffType)
850 }
851 }
852 // eslint-disable-next-line react-hooks/exhaustive-deps
853 }, [snapV2.diffContent])
854
855 // We want to check if the diff editor is mounted and if it is, we want to show the widget
856 // We also want to cleanup the widget when the diff editor is closed
857 useEffect(() => {
858 if (!isDiffOpen) {
859 setIsDiffEditorMounted(false)
860 setShowWidget(false)
861 } else if (diffEditorRef.current && isDiffEditorMounted) {
862 setShowWidget(true)
863 return () => setShowWidget(false)
864 }
865 }, [isDiffOpen, isDiffEditorMounted])
866
867 return (
868 <>
869 <RunQueryWarningModal
870 visible={!!potentialIssues}
871 potentialIssues={potentialIssues}
872 onCancel={() => {
873 clearPendingRunRefocus()
874 setPotentialIssues(undefined)
875 refocusEditor()
876 }}
877 onConfirm={() => {
878 shouldRefocusAfterRunRef.current = true
879 setPotentialIssues(undefined)
880 refocusEditor()
881 void executeQuery(true)
882 }}
883 onConfirmWithRLS={() => {
884 const tables = potentialIssues?.createTablesMissingRLS ?? []
885 if (tables.length === 0) return
886 const editor = editorRef.current
887 const selection = editor?.getSelection()
888 const selectedValue = selection
889 ? editor?.getModel()?.getValueInRange(selection)
890 : undefined
891 const baseSql = selectedValue || editor?.getValue() || ''
892 const rewrittenSql = appendEnableRLSStatements(baseSql, tables)
893 shouldRefocusAfterRunRef.current = true
894 setPotentialIssues(undefined)
895 refocusEditor()
896 void executeQuery(true, acceptUntrustedSql(rewrittenSql as UntrustedSqlFragment))
897 }}
898 />
899
900 <div className="flex h-full">
901 <ResizablePanelGroup
902 className="relative"
903 orientation="vertical"
904 autoSaveId={LOCAL_STORAGE_KEYS.SQL_EDITOR_SPLIT_SIZE}
905 >
906 <ResizablePanel defaultSize="50" maxSize="70">
907 <div className="grow overflow-y-auto border-b h-full">
908 {isLoading ? (
909 <div className="flex h-full w-full items-center justify-center">
910 <Loader2 className="animate-spin text-brand" />
911 </div>
912 ) : (
913 <>
914 {isDiffOpen && (
915 <div className="w-full h-full">
916 <DiffEditor
917 language="pgsql"
918 original={defaultSqlDiff.original}
919 modified={defaultSqlDiff.modified}
920 onMount={(editor) => {
921 diffEditorRef.current = editor
922 setIsDiffEditorMounted(true)
923 }}
924 />
925 {showWidget && (
926 <ResizableAIWidget
927 editor={diffEditorRef.current!}
928 id="ask-ai-diff"
929 value={promptInput}
930 onChange={setPromptInput}
931 onSubmit={(prompt: string) => {
932 handlePrompt(prompt, {
933 beforeSelection: promptState.beforeSelection,
934 selection: promptState.selection || defaultSqlDiff.modified,
935 afterSelection: promptState.afterSelection,
936 })
937 }}
938 onAccept={acceptAiHandler}
939 onReject={discardAiHandler}
940 onCancel={resetPrompt}
941 isDiffVisible={true}
942 isLoading={isCompletionLoading}
943 startLineNumber={Math.max(0, promptState.startLineNumber)}
944 endLineNumber={promptState.endLineNumber}
945 />
946 )}
947 </div>
948 )}
949 <div key={id} className="w-full h-full relative">
950 <MonacoEditor
951 autoFocus
952 placeholder={
953 !promptState.isOpen && !editorRef.current?.getValue()
954 ? 'Hit ' +
955 (os === 'macos' ? 'CMD+SHIFT+K' : `CTRL+SHIFT+K`) +
956 ' to generate query or just start typing'
957 : ''
958 }
959 id={id}
960 snippetName={
961 urlId === 'new'
962 ? generatedNewSnippetName
963 : (snapV2.snippets[id]?.snippet.name ?? generatedNewSnippetName)
964 }
965 className={cn(isDiffOpen && 'hidden')}
966 editorRef={editorRef}
967 monacoRef={monacoRef}
968 executeQuery={executeQuery}
969 executeExplainQuery={executeExplainQuery}
970 prettifyQuery={prettifyQuery}
971 onHasSelection={setHasSelection}
972 onMount={onMount}
973 onPrompt={({
974 selection,
975 beforeSelection,
976 afterSelection,
977 startLineNumber,
978 endLineNumber,
979 }) => {
980 setPromptState((prev) => ({
981 ...prev,
982 isOpen: true,
983 selection,
984 beforeSelection,
985 afterSelection,
986 startLineNumber,
987 endLineNumber,
988 }))
989 }}
990 />
991 {editorRef.current && promptState.isOpen && !isDiffOpen && (
992 <ResizableAIWidget
993 editor={editorRef.current}
994 id="ask-ai"
995 value={promptInput}
996 onChange={setPromptInput}
997 onSubmit={(prompt: string) => {
998 handlePrompt(prompt, {
999 beforeSelection: promptState.beforeSelection,
1000 selection: promptState.selection,
1001 afterSelection: promptState.afterSelection,
1002 })
1003 }}
1004 onCancel={resetPrompt}
1005 isDiffVisible={false}
1006 isLoading={isCompletionLoading}
1007 startLineNumber={Math.max(0, promptState.startLineNumber)}
1008 endLineNumber={promptState.endLineNumber}
1009 />
1010 )}
1011 </div>
1012 </>
1013 )}
1014 </div>
1015 </ResizablePanel>
1016
1017 <ResizableHandle withHandle />
1018
1019 <ResizablePanel defaultSize="50" maxSize="70">
1020 {isLoading ? (
1021 <div className="flex h-full w-full items-center justify-center">
1022 <Loader2 className="animate-spin text-brand" />
1023 </div>
1024 ) : (
1025 <UtilityPanel
1026 id={id}
1027 isExecuting={isExecuting}
1028 isExplainExecuting={isExplainExecuting}
1029 isDisabled={isDiffOpen}
1030 hasSelection={hasSelection}
1031 prettifyQuery={prettifyQuery}
1032 executeQuery={executeQueryFromButton}
1033 executeExplainQuery={executeExplainQuery}
1034 onDebug={onDebug}
1035 buildDebugPrompt={buildDebugPrompt}
1036 activeTab={activeUtilityTab}
1037 onActiveTabChange={setActiveUtilityTab}
1038 />
1039 )}
1040 </ResizablePanel>
1041
1042 <div className="h-9">
1043 {results?.rows !== undefined && !isExecuting && (
1044 <GridFooter className="flex items-center justify-between gap-2">
1045 <Tooltip>
1046 <TooltipTrigger>
1047 <p className="text-xs">
1048 <span className="text-foreground">
1049 {results.rows.length} row{results.rows.length > 1 ? 's' : ''}
1050 </span>
1051 <span className="text-foreground-lighter ml-1">
1052 {results.autoLimit !== undefined &&
1053 ` (Limited to only ${results.autoLimit} rows)`}
1054 </span>
1055 </p>
1056 </TooltipTrigger>
1057 <TooltipContent className="max-w-xs">
1058 <p className="flex flex-col gap-y-1">
1059 <span>
1060 Results are automatically limited to preserve browser performance, in
1061 particular if your query returns an exceptionally large number of rows.
1062 </span>
1063
1064 <span className="text-foreground-light">
1065 You may change or remove this limit from the dropdown on the right
1066 </span>
1067 </p>
1068 </TooltipContent>
1069 </Tooltip>
1070 {results.autoLimit !== undefined && (
1071 <DropdownMenu>
1072 <DropdownMenuTrigger asChild>
1073 <Button type="default" iconRight={<ChevronUp size={14} />}>
1074 Limit results to:{' '}
1075 {ROWS_PER_PAGE_OPTIONS.find((opt) => opt.value === snapV2.limit)?.label}
1076 </Button>
1077 </DropdownMenuTrigger>
1078 <DropdownMenuContent className="w-40" align="end">
1079 <DropdownMenuRadioGroup
1080 value={snapV2.limit.toString()}
1081 onValueChange={(val) => snapV2.setLimit(Number(val))}
1082 >
1083 {ROWS_PER_PAGE_OPTIONS.map((option) => (
1084 <DropdownMenuRadioItem key={option.label} value={option.value.toString()}>
1085 {option.label}
1086 </DropdownMenuRadioItem>
1087 ))}
1088 </DropdownMenuRadioGroup>
1089 </DropdownMenuContent>
1090 </DropdownMenu>
1091 )}
1092 </GridFooter>
1093 )}
1094 </div>
1095 </ResizablePanelGroup>
1096 </div>
1097 </>
1098 )
1099}