From 46bcefffb07a39d231590648fb2f2a9cfdec8427 Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Tue, 17 Dec 2024 15:02:06 +1000 Subject: [PATCH] Add inline completions to SQL Editor (#30706) * monaco complete * progress * just add panel to sqleditor * update panel and prompt * adjust prompt * clean up * clean up and improve prompt * update prompt * refine prompt * remove package * use constants file * revert v3 prompt * combine diff action bars * use widget * refactor * Some updates * fix: useCompletion base path * fix: typo on monaco action label * feat: use entire selected line(s) for inline completions * fix widget consistency * widget styling * Fix TS * feat: lock llm model version * fix enter key issue * move complete endpoint * update complete prompt * remove file * placeholder text --------- Co-authored-by: Joshen Lim Co-authored-by: Greg Richardson --- .../interfaces/SQLEditor/AskAIWidget.tsx | 79 +++ .../interfaces/SQLEditor/InlineWidget.tsx | 108 ++++ .../interfaces/SQLEditor/MonacoEditor.tsx | 58 +- .../SQLEditor/RunQueryWarningModal.tsx | 66 ++ .../interfaces/SQLEditor/SQLEditor.tsx | 564 ++++++++++-------- .../interfaces/SQLEditor/SQLEditor.utils.ts | 8 +- .../SQLEditor/UtilityPanel/UtilityPanel.tsx | 2 +- .../UtilityPanel/UtilityTabResults.tsx | 2 +- .../components/interfaces/SQLEditor/hooks.ts | 97 ++- apps/studio/middleware.ts | 1 + apps/studio/package.json | 2 +- apps/studio/pages/api/ai/sql/complete.ts | 148 +++++ package-lock.json | 1 - 13 files changed, 860 insertions(+), 276 deletions(-) create mode 100644 apps/studio/components/interfaces/SQLEditor/AskAIWidget.tsx create mode 100644 apps/studio/components/interfaces/SQLEditor/InlineWidget.tsx create mode 100644 apps/studio/components/interfaces/SQLEditor/RunQueryWarningModal.tsx create mode 100644 apps/studio/pages/api/ai/sql/complete.ts diff --git a/apps/studio/components/interfaces/SQLEditor/AskAIWidget.tsx b/apps/studio/components/interfaces/SQLEditor/AskAIWidget.tsx new file mode 100644 index 00000000000..7030c9e2561 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/AskAIWidget.tsx @@ -0,0 +1,79 @@ +import { detectOS } from 'lib/helpers' +import { ArrowDownLeft, Loader2, Wand } from 'lucide-react' +import { useCallback } from 'react' +import { Input, Button } from 'ui' + +interface AskAIWidgetProps { + value: string + onChange: (value: string) => void + onSubmit: (prompt: string) => void + onAccept?: () => void + onReject?: () => void + isDiffVisible: boolean + isLoading?: boolean +} + +export const AskAIWidget = ({ + value, + onChange, + onSubmit, + onAccept, + onReject, + isDiffVisible, + isLoading = false, +}: AskAIWidgetProps) => { + const os = detectOS() + + // Auto focus input + const inputRef = useCallback((input: HTMLInputElement | null) => { + setTimeout(() => { + input?.focus() + }, 0) + }, []) + + const handleSubmit = () => { + if (value.trim() && !isLoading) { + onSubmit(value) + } + } + + return ( +
+ onChange(e.target.value)} + onKeyDown={(e) => { + if (e.key === 'Enter' && !e.metaKey && !e.ctrlKey) { + handleSubmit() + } + }} + disabled={isLoading} + /> + {isDiffVisible && ( +
+ + +
+ )} +
+ ) +} diff --git a/apps/studio/components/interfaces/SQLEditor/InlineWidget.tsx b/apps/studio/components/interfaces/SQLEditor/InlineWidget.tsx new file mode 100644 index 00000000000..afa44a8ba67 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/InlineWidget.tsx @@ -0,0 +1,108 @@ +import { editor } from 'monaco-editor' +import { PropsWithChildren, useEffect, useMemo } from 'react' +import { createPortal } from 'react-dom' + +export interface InlineWidgetProps { + /** + * The Monaco editor instance. + */ + editor: editor.IStandaloneCodeEditor | editor.IStandaloneDiffEditor + + /** + * ID used by Monaco to reference this widget + */ + id: string + + /** + * The line number after or before which this zone should appear. + * Use 0 to place a view zone before the first line number. + */ + afterLineNumber: number + beforeLineNumber?: number + + /** + * The height in lines of the view zone. + * @default 1 + */ + heightInLines?: number +} + +/** + * Adds an inline widget to a Monaco editor instance. Acts as a + * container for custom widgets that are rendered inline. + * + * Implemented using the same techniques VS Code uses for their inline + * widgets (such as "Go to References"), but built for React. + * Uses a combination of a view zone and overlay widget. + */ +const InlineWidget = ({ + children, + editor, + id, + beforeLineNumber, + afterLineNumber, + heightInLines = 1, +}: PropsWithChildren) => { + const lineNumber = beforeLineNumber ?? afterLineNumber + const key = `${id}-${lineNumber.toString()}` + + const containerElement = useMemo(() => document.createElement('div'), []) + + // Get the appropriate editor instance for diff editor + const targetEditor = 'getModifiedEditor' in editor ? editor.getModifiedEditor() : editor + + useEffect(() => { + let zoneId: string + let viewZoneTop = 0 + let viewZoneHeight = 0 + + const overlayWidget: editor.IOverlayWidget = { + getId: () => id, + getDomNode: () => containerElement, + getPosition: () => null, + } + + const recalculateLayout = () => { + const layoutInfo = targetEditor.getLayoutInfo() + + if (!layoutInfo) { + return + } + + containerElement.style.left = `${layoutInfo.contentLeft}px` + containerElement.style.top = `${viewZoneTop}px` + containerElement.style.width = `${layoutInfo.width - layoutInfo.contentLeft}px` + containerElement.style.height = `${viewZoneHeight}px` + } + + targetEditor.changeViewZones((accessor) => { + zoneId = accessor.addZone({ + afterLineNumber: beforeLineNumber ?? afterLineNumber, + heightInLines, + domNode: document.createElement('div'), + onDomNodeTop: (top) => { + viewZoneTop = top + recalculateLayout() + }, + onComputedHeight: (height) => { + viewZoneHeight = height + recalculateLayout() + }, + }) + + targetEditor.addOverlayWidget(overlayWidget) + }) + + // Remove the view zone & overlay widget on unmount + return () => { + targetEditor.changeViewZones((accessor) => { + accessor.removeZone(zoneId) + targetEditor.removeOverlayWidget(overlayWidget) + }) + } + }, [targetEditor, id, beforeLineNumber, afterLineNumber, heightInLines, containerElement]) + + return createPortal(children, containerElement, key) +} + +export default InlineWidget diff --git a/apps/studio/components/interfaces/SQLEditor/MonacoEditor.tsx b/apps/studio/components/interfaces/SQLEditor/MonacoEditor.tsx index 4b05881848d..8a4900bf0b3 100644 --- a/apps/studio/components/interfaces/SQLEditor/MonacoEditor.tsx +++ b/apps/studio/components/interfaces/SQLEditor/MonacoEditor.tsx @@ -24,6 +24,13 @@ export type MonacoEditorProps = { autoFocus?: boolean executeQuery: () => void onHasSelection: (value: boolean) => void + onPrompt?: (value: { + selection: string + beforeSelection: string + afterSelection: string + startLineNumber: number + endLineNumber: number + }) => void } const MonacoEditor = ({ @@ -34,6 +41,7 @@ const MonacoEditor = ({ className, executeQuery, onHasSelection, + onPrompt, }: MonacoEditorProps) => { const router = useRouter() const { profile } = useProfile() @@ -92,6 +100,36 @@ const MonacoEditor = ({ }, }) + if (onPrompt) { + editor.addAction({ + id: 'generate-sql', + label: 'Generate SQL', + keybindings: [monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyK], + run: () => { + const selection = editor.getSelection() + const model = editor.getModel() + if (!model || !selection) return + + const allLines = model.getLinesContent() + + const startLineIndex = selection.startLineNumber - 1 + const endLineIndex = selection.endLineNumber + + const beforeSelection = allLines.slice(0, startLineIndex).join('\n') + '\n' + const selectedText = allLines.slice(startLineIndex, endLineIndex).join('\n') + const afterSelection = '\n' + allLines.slice(endLineIndex).join('\n') + + onPrompt({ + selection: selectedText, + beforeSelection, + afterSelection, + startLineNumber: selection?.startLineNumber ?? 0, + endLineNumber: selection?.endLineNumber ?? 0, + }) + }, + }) + } + editor.onDidChangeCursorSelection(({ selection }) => { const noSelection = selection.startLineNumber === selection.endLineNumber && @@ -99,25 +137,13 @@ const MonacoEditor = ({ onHasSelection(!noSelection) }) - // add margin above first line - editorRef.current.changeViewZones((accessor) => { - accessor.addZone({ - afterLineNumber: 0, - heightInPx: 4, - domNode: document.createElement('div'), - }) - }) - if (autoFocus) { if (editor.getValue().length === 1) editor.setPosition({ lineNumber: 1, column: 2 }) editor.focus() } } - // [Joshen] Also needs updating here - const debouncedSetSql = debounce((id, value) => { - snapV2.setSql(id, value) - }, 1000) + const debouncedSetSql = debounce((id, value) => snapV2.setSql(id, value), 1000) function handleEditorChange(value: string | undefined) { const snippetCheck = snapV2.snippets[id] @@ -145,9 +171,7 @@ const MonacoEditor = ({ // if an SQL query is passed by the content parameter, set the editor value to its content. This // is usually used for sending the user to SQL editor from other pages with SQL. useEffect(() => { - if (content && content.length > 0) { - handleEditorChange(content) - } + if (content && content.length > 0) handleEditorChange(content) }, []) return ( @@ -171,9 +195,11 @@ const MonacoEditor = ({ options={{ tabSize: 2, fontSize: 13, + lineDecorationsWidth: 0, readOnly: disableEdit, minimap: { enabled: false }, wordWrap: 'on', + padding: { top: 4 }, // [Joshen] Commenting the following out as it causes the autocomplete suggestion popover // to be positioned wrongly somehow. I'm not sure if this affects anything though, but leaving // comment just in case anyone might be wondering. Relevant issues: diff --git a/apps/studio/components/interfaces/SQLEditor/RunQueryWarningModal.tsx b/apps/studio/components/interfaces/SQLEditor/RunQueryWarningModal.tsx new file mode 100644 index 00000000000..4a37804be1c --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/RunQueryWarningModal.tsx @@ -0,0 +1,66 @@ +import { Separator } from 'ui' +import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' + +interface RunQueryWarningModalProps { + visible: boolean + hasDestructiveOperations: boolean + hasUpdateWithoutWhere: boolean + onCancel: () => void + onConfirm: () => void +} + +export const RunQueryWarningModal = ({ + visible, + hasDestructiveOperations, + hasUpdateWithoutWhere, + onCancel, + onConfirm, +}: RunQueryWarningModalProps) => { + return ( + +
+
    + {hasDestructiveOperations && ( +
  • + Query has destructive operation + + Make sure you are not accidentally removing something important. + +
    • + )} + {hasDestructiveOperations && hasUpdateWithoutWhere && } + {hasUpdateWithoutWhere && ( +
  • + Query uses update without a where clause + + Without a where clause, this could update all rows + in the table. + +
    • + )} +
+
+

+ Please confirm that you would like to execute this query. +

+ + ) +} diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx b/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx index 3fda4b48718..5d3e5f58021 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx @@ -1,18 +1,19 @@ import type { Monaco } from '@monaco-editor/react' import { useQueryClient } from '@tanstack/react-query' -import { motion } from 'framer-motion' -import { ChevronUp, Loader2 } from 'lucide-react' +import { AnimatePresence, motion } from 'framer-motion' +import { ChevronUp, Command, Loader2 } from 'lucide-react' import dynamic from 'next/dynamic' import { useRouter } from 'next/router' import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { toast } from 'sonner' -import { Separator } from '@ui/components/SidePanel/SidePanel' +import { useCompletion } from 'ai/react' import { useParams } from 'common' import { GridFooter } from 'components/ui/GridFooter' import { useSqlDebugMutation } from 'data/ai/sql-debug-mutation' import { useSqlTitleGenerateMutation } from 'data/ai/sql-title-mutation' import { useEntityDefinitionsQuery } from 'data/database/entity-definitions-query' +import { constructHeaders } from 'data/fetchers' import { lintKeys } from 'data/lint/keys' import { useReadReplicasQuery } from 'data/read-replicas/replicas-query' import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' @@ -20,16 +21,16 @@ import { useFormatQueryMutation } from 'data/sql/format-sql-query' import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query' import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import { isError } from 'data/utils/error-check' -import { useLocalStorageQuery } from 'hooks/misc/useLocalStorage' import { useOrgOptedIntoAi } from 'hooks/misc/useOrgOptedIntoAi' import { useSchemasForAi } from 'hooks/misc/useSchemasForAi' import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization' import { useSelectedProject } from 'hooks/misc/useSelectedProject' -import { IS_PLATFORM, LOCAL_STORAGE_KEYS } from 'lib/constants' +import { BASE_PATH, IS_PLATFORM, LOCAL_STORAGE_KEYS } from 'lib/constants' +import { detectOS, uuidv4 } from 'lib/helpers' import { TelemetryActions } from 'lib/constants/telemetry' -import { uuidv4 } from 'lib/helpers' import { useProfile } from 'lib/profile' import { wrapWithRoleImpersonation } from 'lib/role-impersonation' +import { format } from 'sql-formatter' import { useAppStateSnapshot } from 'state/app-state' import { useDatabaseSelectorStateSnapshot } from 'state/database-selector' import { isRoleImpersonationEnabled, useGetImpersonatedRole } from 'state/role-impersonation-state' @@ -47,28 +48,21 @@ import { TooltipContent_Shadcn_, TooltipTrigger_Shadcn_, Tooltip_Shadcn_, - cn, } from 'ui' -import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' import { subscriptionHasHipaaAddon } from '../Billing/Subscription/Subscription.utils' -import { DiffActionBar } from './DiffActionBar' +import { AskAIWidget } from './AskAIWidget' +import { useSqlEditorDiff, useSqlEditorPrompt } from './hooks' +import InlineWidget from './InlineWidget' +import { RunQueryWarningModal } from './RunQueryWarningModal' import { ROWS_PER_PAGE_OPTIONS, sqlAiDisclaimerComment, untitledSnippetTitle, } from './SQLEditor.constants' -import { - ContentDiff, - DiffType, - IStandaloneCodeEditor, - IStandaloneDiffEditor, -} from './SQLEditor.types' +import { DiffType, IStandaloneCodeEditor, IStandaloneDiffEditor } from './SQLEditor.types' import { checkDestructiveQuery, checkIfAppendLimitRequired, - compareAsAddition, - compareAsModification, - compareAsNewSnippet, createSqlSnippetSkeletonV2, isUpdateWithoutWhere, suffixWithLimit, @@ -83,57 +77,69 @@ const DiffEditor = dynamic( ) export const SQLEditor = () => { + const os = detectOS() const router = useRouter() const { ref, id: urlId } = useParams() - - // generate an id to be used for new snippets. The dependency on urlId is to avoid a bug which - // shows up when clicking on the SQL Editor while being in the SQL editor on a random snippet. - const generatedId = useMemo(() => uuidv4(), [urlId]) - // the id is stable across renders - it depends either on the url or on the memoized generated id - const id = !urlId || urlId === 'new' ? generatedId : urlId - const { profile } = useProfile() + const queryClient = useQueryClient() const project = useSelectedProject() const organization = useSelectedOrganization() const appSnap = useAppStateSnapshot() const snapV2 = useSqlEditorV2StateSnapshot() const getImpersonatedRole = useGetImpersonatedRole() const databaseSelectorState = useDatabaseSelectorStateSnapshot() - const queryClient = useQueryClient() + const isOptedInToAI = useOrgOptedIntoAi() + const [selectedSchemas] = useSchemasForAi(project?.ref!) + const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM - const { mutate: formatQuery } = useFormatQueryMutation() - const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation() - const { mutateAsync: debugSql, isLoading: isDebugSqlLoading } = useSqlDebugMutation() - - const [sourceSqlDiff, setSourceSqlDiff] = useState() - const [pendingTitle, setPendingTitle] = useState() - const [hasSelection, setHasSelection] = useState(false) + const { + sourceSqlDiff, + setSourceSqlDiff, + selectedDiffType, + setSelectedDiffType, + pendingTitle, + setPendingTitle, + isAcceptDiffLoading, + setIsAcceptDiffLoading, + isDiffOpen, + defaultSqlDiff, + closeDiff, + } = useSqlEditorDiff() + const { promptState, setPromptState, promptInput, setPromptInput, resetPrompt } = + useSqlEditorPrompt() const editorRef = useRef(null) const monacoRef = useRef(null) const diffEditorRef = useRef(null) - - const { data: subscription } = useOrgSubscriptionQuery({ orgSlug: organization?.slug }) - const { data: databases, isSuccess: isSuccessReadReplicas } = useReadReplicasQuery({ - projectRef: ref, - }) - + const [hasSelection, setHasSelection] = useState(false) + const [lineHighlights, setLineHighlights] = useState([]) + const [isDiffEditorMounted, setIsDiffEditorMounted] = useState(false) const [showPotentialIssuesModal, setShowPotentialIssuesModal] = useState(false) const [queryHasDestructiveOperations, setQueryHasDestructiveOperations] = useState(false) const [queryHasUpdateWithoutWhere, setQueryHasUpdateWithoutWhere] = useState(false) + const [showWidget, setShowWidget] = useState(false) - const isOptedInToAI = useOrgOptedIntoAi() - const [selectedSchemas] = useSchemasForAi(project?.ref!) - // Customers on HIPAA plans should not have access to Supabase AI + // generate an id to be used for new snippets. The dependency on urlId is to avoid a bug which + // shows up when clicking on the SQL Editor while being in the SQL editor on a random snippet. + // eslint-disable-next-line react-hooks/exhaustive-deps + const generatedId = useMemo(() => uuidv4(), [urlId]) + // the id is stable across renders - it depends either on the url or on the memoized generated id + const id = !urlId || urlId === 'new' ? generatedId : urlId + + const limit = snapV2.limit + const results = snapV2.results[id]?.[0] + const snippetIsLoading = !( + id in snapV2.snippets && snapV2.snippets[id].snippet.content !== undefined + ) + const isLoading = urlId === 'new' ? false : snippetIsLoading + + /** React query data fetching */ + const { data: subscription } = useOrgSubscriptionQuery({ orgSlug: organization?.slug }) const hasHipaaAddon = subscriptionHasHipaaAddon(subscription) - const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM - const [isAcceptDiffLoading, setIsAcceptDiffLoading] = useState(false) - const [, setAiQueryCount] = useLocalStorageQuery('supabase_sql-editor-ai-query-count', 0) - - const [selectedDiffType, setSelectedDiffType] = useState(undefined) - const [isFirstRender, setIsFirstRender] = useState(true) - const [lineHighlights, setLineHighlights] = useState([]) + const { data: databases, isSuccess: isSuccessReadReplicas } = useReadReplicasQuery({ + projectRef: ref, + }) const { data, refetch: refetchEntityDefinitions } = useEntityDefinitionsQuery( { @@ -143,19 +149,13 @@ export const SQLEditor = () => { }, { enabled: includeSchemaMetadata } ) - const entityDefinitions = includeSchemaMetadata ? data?.map((def) => def.sql.trim()) : undefined - const isDiffOpen = !!sourceSqlDiff - - const limit = snapV2.limit - const results = snapV2.results[id]?.[0] - const snippetIsLoading = !( - id in snapV2.snippets && snapV2.snippets[id].snippet.content !== undefined - ) - const isLoading = urlId === 'new' ? false : snippetIsLoading + /* React query mutations */ + const { mutate: formatQuery } = useFormatQueryMutation() + const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation() + const { mutateAsync: debugSql, isLoading: isDebugSqlLoading } = useSqlDebugMutation() const { mutate: sendEvent } = useSendEventMutation() - const { mutate: execute, isLoading: isExecuting } = useExecuteSqlMutation({ onSuccess(data, vars) { if (id) snapV2.addResult(id, data.result, vars.autoLimit) @@ -327,6 +327,7 @@ export const SQLEditor = () => { }) } }, + // eslint-disable-next-line react-hooks/exhaustive-deps [ isDiffOpen, id, @@ -363,35 +364,10 @@ export const SQLEditor = () => { toast.error(`Failed to create new query: ${error.message}`) } }, + // eslint-disable-next-line react-hooks/exhaustive-deps [profile?.id, project?.id, ref, router, snapV2] ) - const updateEditorWithCheckForDiff = useCallback( - ({ diffType, sql }: { diffType: DiffType; sql: string }) => { - const editorModel = editorRef.current?.getModel() - if (!editorModel) return - - setAiQueryCount((count) => count + 1) - - const existingValue = editorRef.current?.getValue() ?? '' - if (existingValue.length === 0) { - // if the editor is empty, just copy over the code - editorRef.current?.executeEdits('apply-ai-message', [ - { - text: `${sqlAiDisclaimerComment}\n\n${sql}`, - range: editorModel.getFullModelRange(), - }, - ]) - } else { - const currentSql = editorRef.current?.getValue() - const diff = { original: currentSql || '', modified: sql } - setSourceSqlDiff(diff) - setSelectedDiffType(diffType) - } - }, - [setAiQueryCount] - ) - const onDebug = useCallback(async () => { try { const snippet = snapV2.snippets[id] @@ -411,27 +387,20 @@ export const SQLEditor = () => { ) } } + // eslint-disable-next-line react-hooks/exhaustive-deps }, [debugSql, entityDefinitions, id, snapV2.results, snapV2.snippets]) const acceptAiHandler = useCallback(async () => { try { setIsAcceptDiffLoading(true) - if (!sourceSqlDiff) { - return - } - // TODO: show error if undefined - if (!editorRef.current || !diffEditorRef.current) { - return - } + if (!sourceSqlDiff || !editorRef.current || !diffEditorRef.current) return const editorModel = editorRef.current.getModel() const diffModel = diffEditorRef.current.getModel() - if (!editorModel || !diffModel) { - return - } + if (!editorModel || !diffModel) return const sql = diffModel.modified.getValue() @@ -454,11 +423,12 @@ export const SQLEditor = () => { sendEvent({ action: TelemetryActions.ASSISTANT_SUGGESTION_ACCEPTED }) setSelectedDiffType(DiffType.Modification) - setSourceSqlDiff(undefined) - setPendingTitle(undefined) + resetPrompt() + closeDiff() } finally { setIsAcceptDiffLoading(false) } + // eslint-disable-next-line react-hooks/exhaustive-deps }, [ sourceSqlDiff, selectedDiffType, @@ -472,129 +442,201 @@ export const SQLEditor = () => { const discardAiHandler = useCallback(() => { sendEvent({ action: TelemetryActions.ASSISTANT_SUGGESTION_REJECTED }) - setSourceSqlDiff(undefined) - setPendingTitle(undefined) - }, [router]) + resetPrompt() + closeDiff() + }, [closeDiff, resetPrompt, sendEvent]) + + const { + complete, + completion, + isLoading: isCompletionLoading, + } = useCompletion({ + api: `${BASE_PATH}/api/ai/sql/complete`, + body: { + projectRef: project?.ref, + connectionString: project?.connectionString, + includeSchemaMetadata, + }, + onResponse: (response) => { + if (!response.ok) throw new Error('Failed to generate completion') + }, + onError: (error) => { + toast.error(`Failed to generate SQL: ${error.message}`) + }, + }) + + const handlePrompt = async ( + prompt: string, + context: { + beforeSelection: string + selection: string + afterSelection: string + } + ) => { + try { + setPromptState((prev) => ({ + ...prev, + selection: context.selection, + beforeSelection: context.beforeSelection, + afterSelection: context.afterSelection, + })) + const headerData = await constructHeaders() + + await complete(prompt, { + headers: { Authorization: headerData.get('Authorization') ?? '' }, + body: { + completionMetadata: { + textBeforeCursor: context.beforeSelection, + textAfterCursor: context.afterSelection, + language: 'pgsql', + prompt, + selection: context.selection, + }, + }, + }) + } catch (error) { + setPromptState((prev) => ({ ...prev, isLoading: false })) + } + } + + /** All useEffects are at the bottom before returning the TSX */ + + useEffect(() => { + if (id) { + closeDiff() + setPromptState((prev) => ({ ...prev, isOpen: false })) + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [closeDiff, id]) useEffect(() => { const handler = (e: KeyboardEvent) => { - if (!isDiffOpen) { - return - } + if (!isDiffOpen && !promptState.isOpen) return switch (e.key) { case 'Enter': - acceptAiHandler() + if ((os === 'macos' ? e.metaKey : e.ctrlKey) && isDiffOpen) { + acceptAiHandler() + resetPrompt() + } return case 'Escape': - discardAiHandler() + if (isDiffOpen) discardAiHandler() + resetPrompt() + editorRef.current?.focus() return } } - window.addEventListener('keydown', handler) - return () => window.removeEventListener('keydown', handler) - }, [isDiffOpen, acceptAiHandler, discardAiHandler]) + }, [os, isDiffOpen, promptState.isOpen, acceptAiHandler, discardAiHandler, resetPrompt]) useEffect(() => { - const applyDiff = ({ original, modified }: { original: string; modified: string }) => { - const model = diffEditorRef.current?.getModel() + if (isDiffOpen) { + const diffEditor = diffEditorRef.current + const model = diffEditor?.getModel() if (model && model.original && model.modified) { - model.original.setValue(original) - model.modified.setValue(modified) + model.original.setValue(defaultSqlDiff.original) + model.modified.setValue(defaultSqlDiff.modified) + // scroll to the start line of the modification + const modifiedEditor = diffEditor!.getModifiedEditor() + const startLine = promptState.startLineNumber + modifiedEditor.revealLineInCenter(startLine) } } - - const model = diffEditorRef.current?.getModel() - try { - if (model?.original && model.modified && sourceSqlDiff) { - switch (selectedDiffType) { - case DiffType.Modification: { - const transformedDiff = compareAsModification(sourceSqlDiff) - applyDiff(transformedDiff) - return - } - - case DiffType.Addition: { - const transformedDiff = compareAsAddition(sourceSqlDiff) - applyDiff(transformedDiff) - return - } - - case DiffType.NewSnippet: { - const transformedDiff = compareAsNewSnippet(sourceSqlDiff) - applyDiff(transformedDiff) - return - } - - default: - throw new Error(`Unknown diff type '${selectedDiffType}'`) - } - } - } catch (e) { - console.log(e) - } + // eslint-disable-next-line react-hooks/exhaustive-deps }, [selectedDiffType, sourceSqlDiff]) - // Used for cleaner framer motion transitions - useEffect(() => { - setIsFirstRender(false) - }, []) - useEffect(() => { if (isSuccessReadReplicas) { const primaryDatabase = databases.find((db) => db.identifier === ref) databaseSelectorState.setSelectedDatabaseId(primaryDatabase?.identifier) } + // eslint-disable-next-line react-hooks/exhaustive-deps }, [isSuccessReadReplicas, databases, ref]) useEffect(() => { if (snapV2.diffContent !== undefined) { - updateEditorWithCheckForDiff(snapV2.diffContent) + const { diffType, sql }: { diffType: DiffType; sql: string } = snapV2.diffContent + const editorModel = editorRef.current?.getModel() + if (!editorModel) return + + const existingValue = editorRef.current?.getValue() ?? '' + if (existingValue.length === 0) { + // if the editor is empty, just copy over the code + editorRef.current?.executeEdits('apply-ai-message', [ + { + text: `${sql}`, + range: editorModel.getFullModelRange(), + }, + ]) + } else { + const currentSql = editorRef.current?.getValue() + const diff = { original: currentSql || '', modified: sql } + setSourceSqlDiff(diff) + setSelectedDiffType(diffType) + } } + // eslint-disable-next-line react-hooks/exhaustive-deps }, [snapV2.diffContent]) - const defaultSqlDiff = useMemo(() => { - if (!sourceSqlDiff) { - return { original: '', modified: '' } + useEffect(() => { + if (!completion) { + return } - switch (selectedDiffType) { - case DiffType.Modification: { - return compareAsModification(sourceSqlDiff) - } - case DiffType.Addition: { - return compareAsAddition(sourceSqlDiff) - } + const original = + promptState.beforeSelection + promptState.selection + promptState.afterSelection + const modified = promptState.beforeSelection + completion + promptState.afterSelection - case DiffType.NewSnippet: { - return compareAsNewSnippet(sourceSqlDiff) - } + if (isCompletionLoading) { + let formattedModified = modified - default: - return { original: '', modified: '' } + // Attempt to format the modified SQL in case the LLM left out indentation, etc + try { + formattedModified = format( + promptState.beforeSelection + completion + promptState.afterSelection, + { + language: 'postgresql', + keywordCase: 'lower', + } + ) + } catch (error) {} + + setSourceSqlDiff({ + original, + modified: formattedModified, + }) + setSelectedDiffType(DiffType.Modification) + setPromptState((prev) => ({ ...prev, isLoading: false })) } - }, [selectedDiffType, sourceSqlDiff]) + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [ + completion, + promptState.beforeSelection, + promptState.selection, + promptState.afterSelection, + isCompletionLoading, + ]) + + // We want to check if the diff editor is mounted and if it is, we want to show the widget + // We also want to cleanup the widget when the diff editor is closed + useEffect(() => { + if (!isDiffOpen) { + setIsDiffEditorMounted(false) + setShowWidget(false) + } else if (diffEditorRef.current && isDiffEditorMounted) { + setShowWidget(true) + return () => setShowWidget(false) + } + }, [isDiffOpen, isDiffEditorMounted]) return ( <> - { setShowPotentialIssuesModal(false) setQueryHasDestructiveOperations(false) @@ -605,33 +647,7 @@ export const SQLEditor = () => { setShowPotentialIssuesModal(false) executeQuery(true) }} - > -
-
    - {queryHasDestructiveOperations && ( -
  • - Query has destructive operation - - Make sure you are not accidentally removing something important. - -
    • - )} - {queryHasDestructiveOperations && queryHasUpdateWithoutWhere && } - {queryHasUpdateWithoutWhere && ( -
  • - Query uses update without a where clause - - Without a where clause, this could update all - rows in the table. - -
    • - )} -
-
-

- Please confirm that you would like to execute this query. -

- + /> { direction="vertical" autoSaveId={LOCAL_STORAGE_KEYS.SQL_EDITOR_SPLIT_SIZE} > - {!hasHipaaAddon && isDiffOpen && ( - - setSelectedDiffType(diffType)} - onAccept={acceptAiHandler} - onCancel={discardAiHandler} - /> - - )}
{isLoading ? ( @@ -674,15 +669,7 @@ export const SQLEditor = () => { ) : ( <> {isDiffOpen && ( - +
{ modified={defaultSqlDiff.modified} onMount={(editor) => { diffEditorRef.current = editor + setIsDiffEditorMounted(true) + }} + options={{ + fontSize: 13, + renderSideBySide: false, + minimap: { enabled: false }, + wordWrap: 'on', + lineNumbers: 'on', + folding: false, + padding: { top: 4 }, + lineNumbersMinChars: 3, }} - options={{ fontSize: 13 }} /> - + {showWidget && ( + + { + handlePrompt(prompt, { + beforeSelection: promptState.beforeSelection, + selection: promptState.selection || defaultSqlDiff.modified, + afterSelection: promptState.afterSelection, + }) + }} + value={promptInput} + onChange={setPromptInput} + onAccept={acceptAiHandler} + onReject={discardAiHandler} + isDiffVisible={true} + isLoading={isCompletionLoading} + /> + + )} +
)} - +
{ monacoRef={monacoRef} executeQuery={executeQuery} onHasSelection={setHasSelection} + onPrompt={({ + selection, + beforeSelection, + afterSelection, + startLineNumber, + endLineNumber, + }) => { + setPromptState((prev) => ({ + ...prev, + isOpen: true, + selection, + beforeSelection, + afterSelection, + startLineNumber, + endLineNumber, + })) + }} /> - + {editorRef.current && promptState.isOpen && !isDiffOpen && ( + + { + handlePrompt(prompt, { + beforeSelection: promptState.beforeSelection, + selection: promptState.selection, + afterSelection: promptState.afterSelection, + }) + }} + isDiffVisible={false} + isLoading={isCompletionLoading} + /> + + )} + + {!promptState.isOpen && !editorRef.current?.getValue() && ( + + Hit {os === 'macos' ? : `CTRL+`}K to edit with the + Assistant + + )} + +
)}
diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts index 4e98e2bedd5..f3fbeac8b7a 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts @@ -137,18 +137,14 @@ export const compareAsModification = (sqlDiff: ContentDiff) => { return { original: sqlDiff.original, - modified: `${sqlAiDisclaimerComment}\n\n${formattedModified}`, + modified: `${formattedModified}`, } } export const compareAsAddition = (sqlDiff: ContentDiff) => { const formattedOriginal = sqlDiff.original.replace(sqlAiDisclaimerComment, '').trim() const formattedModified = sqlDiff.modified.replace(sqlAiDisclaimerComment, '').trim() - const newModified = - sqlAiDisclaimerComment + - '\n\n' + - (formattedOriginal ? formattedOriginal + '\n\n' : '') + - formattedModified + const newModified = (formattedOriginal ? formattedOriginal + '\n\n' : '') + formattedModified return { original: sqlDiff.original, diff --git a/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityPanel.tsx b/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityPanel.tsx index 1bd4af220ff..b3c4c9901a6 100644 --- a/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityPanel.tsx +++ b/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityPanel.tsx @@ -113,7 +113,7 @@ const UtilityPanel = ({ return ( - +
Results diff --git a/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityTabResults.tsx b/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityTabResults.tsx index 9931c8faab9..084f9a0b6fb 100644 --- a/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityTabResults.tsx +++ b/apps/studio/components/interfaces/SQLEditor/UtilityPanel/UtilityTabResults.tsx @@ -156,7 +156,7 @@ const UtilityTabResults = ({ } else if (!result) { return (
-

+

Click Run to execute your query.

diff --git a/apps/studio/components/interfaces/SQLEditor/hooks.ts b/apps/studio/components/interfaces/SQLEditor/hooks.ts index f3a06b50be5..265461e982a 100644 --- a/apps/studio/components/interfaces/SQLEditor/hooks.ts +++ b/apps/studio/components/interfaces/SQLEditor/hooks.ts @@ -1,5 +1,6 @@ import { PermissionAction } from '@supabase/shared-types/out/constants' import { useRouter } from 'next/router' +import { useCallback, useEffect, useMemo, useState } from 'react' import { toast } from 'sonner' import { useParams } from 'common' @@ -8,7 +9,13 @@ import { useCheckPermissions } from 'hooks/misc/useCheckPermissions' import { uuidv4 } from 'lib/helpers' import { useProfile } from 'lib/profile' import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2' -import { createSqlSnippetSkeletonV2 } from './SQLEditor.utils' +import { ContentDiff, DiffType } from './SQLEditor.types' +import { + compareAsAddition, + compareAsModification, + compareAsNewSnippet, + createSqlSnippetSkeletonV2, +} from './SQLEditor.utils' export const useNewQuery = () => { const router = useRouter() @@ -50,4 +57,92 @@ export const useNewQuery = () => { return { newQuery } } +export function useSqlEditorDiff() { + const [sourceSqlDiff, setSourceSqlDiff] = useState() + const [selectedDiffType, setSelectedDiffType] = useState() + const [pendingTitle, setPendingTitle] = useState() + const [isAcceptDiffLoading, setIsAcceptDiffLoading] = useState(false) + + const isDiffOpen = !!sourceSqlDiff + + const defaultSqlDiff = useMemo(() => { + if (!sourceSqlDiff) { + return { original: '', modified: '' } + } + + switch (selectedDiffType) { + case DiffType.Modification: + return compareAsModification(sourceSqlDiff) + case DiffType.Addition: + return compareAsAddition(sourceSqlDiff) + case DiffType.NewSnippet: + return compareAsNewSnippet(sourceSqlDiff) + default: + return { original: '', modified: '' } + } + }, [selectedDiffType, sourceSqlDiff]) + + const closeDiff = useCallback(() => { + setSourceSqlDiff(undefined) + setPendingTitle(undefined) + setSelectedDiffType(undefined) + }, []) + + return { + sourceSqlDiff, + setSourceSqlDiff, + selectedDiffType, + setSelectedDiffType, + pendingTitle, + setPendingTitle, + isAcceptDiffLoading, + setIsAcceptDiffLoading, + isDiffOpen, + defaultSqlDiff, + closeDiff, + } +} + +interface PromptState { + isOpen: boolean + selection: string + beforeSelection: string + afterSelection: string + startLineNumber: number + endLineNumber: number +} + +const initialPromptState: PromptState = { + isOpen: false, + selection: '', + beforeSelection: '', + afterSelection: '', + startLineNumber: 0, + endLineNumber: 0, +} + +export function useSqlEditorPrompt() { + const [promptState, setPromptState] = useState(initialPromptState) + const [promptInput, setPromptInput] = useState('') + + useEffect(() => { + if (!promptState.isOpen) { + setPromptInput('') + } + }, [promptState.isOpen]) + + const resetPrompt = () => { + setPromptState(initialPromptState) + setPromptInput('') + } + + return { + promptState, + setPromptState, + promptInput, + setPromptInput, + resetPrompt, + } +} + export default useNewQuery diff --git a/apps/studio/middleware.ts b/apps/studio/middleware.ts index eb510b1988a..ee838fa8116 100644 --- a/apps/studio/middleware.ts +++ b/apps/studio/middleware.ts @@ -10,6 +10,7 @@ const HOSTED_SUPPORTED_API_URLS = [ '/ai/sql/suggest', '/ai/sql/generate-v2', '/ai/sql/generate-v3', + '/ai/sql/complete', '/ai/sql/title', '/ai/sql/debug', '/ai/sql/cron', diff --git a/apps/studio/package.json b/apps/studio/package.json index 6466ef888b1..0cb3c602bd4 100644 --- a/apps/studio/package.json +++ b/apps/studio/package.json @@ -30,6 +30,7 @@ "@hookform/resolvers": "^3.1.1", "@monaco-editor/react": "^4.6.0", "@next/bundle-analyzer": "^14.2.3", + "@number-flow/react": "^0.3.2", "@radix-ui/react-popover": "^1.0.7", "@radix-ui/react-tooltip": "^1.0.7", "@sentry/nextjs": "^8.34.0", @@ -77,7 +78,6 @@ "mobx": "^6.10.2", "mobx-react-lite": "^4.0.5", "monaco-editor": "0.33.0", - "@number-flow/react": "^0.3.2", "next": "^14.2.10", "next-themes": "^0.3.0", "nuqs": "^1.17.8", diff --git a/apps/studio/pages/api/ai/sql/complete.ts b/apps/studio/pages/api/ai/sql/complete.ts new file mode 100644 index 00000000000..29d6d63bf49 --- /dev/null +++ b/apps/studio/pages/api/ai/sql/complete.ts @@ -0,0 +1,148 @@ +import { streamText } from 'ai' +import { openai } from '@ai-sdk/openai' +import { NextApiRequest, NextApiResponse } from 'next' +import { getTools } from '../sql/tools' +import { executeSql } from 'data/sql/execute-sql-query' +import pgMeta from '@supabase/pg-meta' + +export const maxDuration = 30 +const openAiKey = process.env.OPENAI_API_KEY +const pgMetaSchemasList = pgMeta.schemas.list() + +export default async function handler(req: NextApiRequest, res: NextApiResponse) { + if (!openAiKey) { + return new Response( + JSON.stringify({ + error: 'No OPENAI_API_KEY set. Create this environment variable to use AI features.', + }), + { + status: 500, + headers: { 'Content-Type': 'application/json' }, + } + ) + } + + if (req.method !== 'POST') { + return new Response( + JSON.stringify({ data: null, error: { message: `Method ${req.method} Not Allowed` } }), + { + status: 405, + headers: { 'Content-Type': 'application/json', Allow: 'POST' }, + } + ) + } + + try { + const { completionMetadata, projectRef, connectionString, includeSchemaMetadata } = req.body + const { textBeforeCursor, textAfterCursor, language, prompt, selection } = completionMetadata + + if (!projectRef) { + return res.status(400).json({ + error: 'Missing project_ref in request body', + }) + } + + const authorization = req.headers.authorization + + const { result: schemas } = includeSchemaMetadata + ? await executeSql( + { + projectRef, + connectionString, + sql: pgMetaSchemasList.sql, + }, + undefined, + { + 'Content-Type': 'application/json', + ...(authorization && { Authorization: authorization }), + } + ) + : { result: [] } + + const result = await streamText({ + model: openai('gpt-4o-mini-2024-07-18'), + maxSteps: 5, + tools: getTools({ projectRef, connectionString, authorization, includeSchemaMetadata }), + system: ` + You are a Supabase Postgres expert who can do the following things. + + # You generate and debug SQL + The generated SQL (must be valid SQL), and must adhere to the following: + - Always use double apostrophe in SQL strings (eg. 'Night''s watch') + - Always use semicolons + - Use vector(384) data type for any embedding/vector related query + - When debugging, retrieve sql schema details to ensure sql is correct + - In Supabase, the auth schema already has a users table which is used to store users. It is common practice to create a profiles table in the public schema that links to auth.users to store user information instead. You don't need to create a new users table. + - Never suggest creating a view to retrieve information from the users table of the auth schema. This is against our best practices. + + When generating tables, do the following: + - Ensure that all tables always have a primary key + - Ensure that all tables have RLS enabled. Inform the user that they will need to create RLS policies before being able to read or write to the table over Supabase APIs. + - For primary keys, always use "id bigint primary key generated always as identity" (not serial) + - Prefer creating foreign key references in the create statement + - Prefer 'text' over 'varchar' + - Prefer 'timestamp with time zone' over 'date' + - If a foreign key is created, include a statement to create a database index on the columns involved in the foreign key. This is a best practice as it is common to use foreign key columns in join conditions when querying the database, hence the index will improve the performance of those joins and reduce database resource consumption. + - If the table is a foreign table, create it in a new schema called private. If the private schema doesnt exist, create it. + - Explain to the user that we create it outside of the public schema as it can present a security risk as foreign tables views cannot be configured to respect RLS policies. Provide this link for users to view more details about the security risk: https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0017_foreign_table_in_api + + When generating views, do the following: + - All views should include 'with (security_invoker=on)' clause in the SQL statement for creating views. + - Place the 'with (security_invoker=on)' immediately after the CREATE VIEW statement, before AS + - If the view is a materialized view, create it in a new schema called private. If the private schema doesnt exist, create it. + - Explain to the user that we create it outside of the public schema as it can present a security risk as materialized views cannot be configured to respect RLS policies of the underlying tables they are built upon, nor can they be secured with RLS directly. Provide this link for users to view more details about the security risk: https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0016_materialized_view_in_api + + Feel free to suggest corrections for suspected typos. + + # You write row level security policies. + + Your purpose is to generate a policy with the constraints given by the user. + - First, use getSchema to retrieve more information about a schema or schemas that will contain policies, usually the public schema. + - Then retrieve existing RLS policies and guidelines on how to write policies using the getRlsKnowledge tool . + - Then write new policies or update existing policies based on the prompt + - When asked to suggest policies, either alter existing policies or add new ones to the public schema. + - When writing policies that use a function from the auth schema, ensure that the calls are wrapped with parentheses e.g select auth.uid() should be written as (select auth.uid()) instead + + # You write database functions + Your purpose is to generate a database function with the constraints given by the user. The output may also include a database trigger + if the function returns a type of trigger. When generating functions, do the following: + - If the function returns a trigger type, ensure that it uses security definer, otherwise default to security invoker. Include this in the create functions SQL statement. + - Ensure to set the search_path configuration parameter as '', include this in the create functions SQL statement. + - Default to create or replace whenever possible for updating an existing function, otherwise use the alter function statement + Please make sure that all queries are valid Postgres SQL queries + + # For all your abilities, follow these instructions: + - First look at the list of provided schemas and if needed, get more information about a schema. You will almost always need to retrieve information about the public schema before answering a question. + - If the question is about users or involves creating a users table, also retrieve the auth schema. + + Here are the existing database schema names you can retrieve: ${schemas} + `, + messages: [ + { + role: 'user', + content: `You are helping me edit some pgsql code. + Here is the context: + ${textBeforeCursor}${selection}${textAfterCursor} + + Instructions: + 1. Only modify the selected text based on this prompt: ${prompt} + 2. Your response should be ONLY the modified selection text, nothing else. Remove selected text if needed. + 3. Do not wrap in code blocks or markdown + 4. You can respond with one word or multiple words + 5. Ensure the modified text flows naturally within the current line + 6. Avoid duplicating SQL keywords (SELECT, FROM, WHERE, etc) when considering the full statement + 7. If there is no surrounding context (before or after), make sure your response is a complete valid SQL statement that can be run and resolves the prompt. + + Modify the selected text now:`, + }, + ], + }) + + return result.pipeDataStreamToResponse(res) + } catch (error) { + console.error('Completion error:', error) + return res.status(500).json({ + error: 'Failed to generate completion', + }) + } +} diff --git a/package-lock.json b/package-lock.json index f62d2332c9b..42529488a4e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13868,7 +13868,6 @@ "version": "0.1.6", "resolved": "https://registry.npmjs.org/@supabase/sql-to-rest/-/sql-to-rest-0.1.6.tgz", "integrity": "sha512-06KgjeINtc6405XQvfnchBE1azEsU8G2NElfadmvVHKmHa5l2bFzjbtFbpaYgpgTzccHlcDmBaCgedVf2Gyl8Q==", - "license": "MIT", "dependencies": { "@babel/parser": "^7.24.5", "libpg-query": "^15.1.0",