diff --git a/apps/studio/components/interfaces/SQLEditor/usePrettifyQuery.test.tsx b/apps/studio/components/interfaces/SQLEditor/usePrettifyQuery.test.tsx new file mode 100644 index 00000000000..d96a2d7ba84 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/usePrettifyQuery.test.tsx @@ -0,0 +1,67 @@ +import { act, waitFor } from '@testing-library/react' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' + +import { usePrettifyQuery } from './usePrettifyQuery' +import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' +import { formatSql } from '@/lib/formatSql' +import { sqlEditorState } from '@/state/sql-editor/sql-editor-state' +import { + createInMemoryEditor, + renderSqlEditorHook, + resetSqlEditorStores, + seedSnippet, + setupSqlEditorMocks, +} from '@/tests/lib/sql-editor-test-utils' + +const SNIPPET_ID = 'prettify-snippet' +const MESSY_SQL = 'select id,name from users' + +function usePrettifyHarness({ isDiffOpen }: { isDiffOpen: boolean }) { + const { data: project } = useSelectedProjectQuery() + const prettifyQuery = usePrettifyQuery({ id: SNIPPET_ID, isDiffOpen }) + return { prettifyQuery, isReady: !!project } +} + +beforeEach(() => { + resetSqlEditorStores() + setupSqlEditorMocks() + seedSnippet({ id: SNIPPET_ID, name: 'My query', sql: MESSY_SQL }) +}) + +afterEach(() => { + resetSqlEditorStores() +}) + +describe('usePrettifyQuery', () => { + it('formats the editor SQL in place and writes it back to the snippet store', async () => { + const inMemoryEditor = createInMemoryEditor(MESSY_SQL) + const { result } = renderSqlEditorHook( + (props: { isDiffOpen: boolean }) => usePrettifyHarness(props), + { inMemoryEditor, initialProps: { isDiffOpen: false } } + ) + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.prettifyQuery() + }) + + const expected = formatSql(MESSY_SQL) + expect(inMemoryEditor.editor.getValue()).toBe(expected) + expect(sqlEditorState.snippets[SNIPPET_ID].snippet.content?.unchecked_sql).toBe(expected) + }) + + it('is a no-op while a diff is open', async () => { + const inMemoryEditor = createInMemoryEditor(MESSY_SQL) + const { result } = renderSqlEditorHook( + (props: { isDiffOpen: boolean }) => usePrettifyHarness(props), + { inMemoryEditor, initialProps: { isDiffOpen: true } } + ) + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.prettifyQuery() + }) + + expect(inMemoryEditor.editor.getValue()).toBe(MESSY_SQL) + }) +}) diff --git a/apps/studio/components/interfaces/SQLEditor/useSnippetIdentity.test.tsx b/apps/studio/components/interfaces/SQLEditor/useSnippetIdentity.test.tsx new file mode 100644 index 00000000000..61e8b723917 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/useSnippetIdentity.test.tsx @@ -0,0 +1,42 @@ +import { waitFor } from '@testing-library/react' +import { beforeEach, describe, expect, it } from 'vitest' + +import { useSnippetIdentity } from './useSnippetIdentity' +import { + renderSqlEditorHook, + resetSqlEditorStores, + seedSnippet, +} from '@/tests/lib/sql-editor-test-utils' + +/** + * `useParams` is globally stubbed to `{ ref: 'default' }` (no snippet id), so + * these tests exercise the generated-id branch. The pure id/loading derivation + * itself is proven exhaustively against `deriveSnippetIdentity` in + * `SQLEditor.utils.test`; here we assert the hook wires the store in. + */ +describe('useSnippetIdentity', () => { + beforeEach(() => { + resetSqlEditorStores() + }) + + it('generates a fresh snippet id + name when there is no URL id', () => { + const { result } = renderSqlEditorHook(useSnippetIdentity) + + expect(result.current.urlId).toBeUndefined() + expect(result.current.id).toEqual(expect.any(String)) + expect(result.current.id.length).toBeGreaterThan(0) + expect(result.current.generatedNewSnippetName).toEqual(expect.any(String)) + expect(result.current.generatedNewSnippetName.length).toBeGreaterThan(0) + }) + + it('reports loading until the snippet content is present in the store', async () => { + const { result } = renderSqlEditorHook(useSnippetIdentity) + + // The generated snippet has no content in the store yet. + expect(result.current.isLoading).toBe(true) + + seedSnippet({ id: result.current.id, name: 'My query', sql: 'select 1;' }) + + await waitFor(() => expect(result.current.isLoading).toBe(false)) + }) +}) diff --git a/apps/studio/components/interfaces/SQLEditor/useSnippetTitleGenerator.test.tsx b/apps/studio/components/interfaces/SQLEditor/useSnippetTitleGenerator.test.tsx new file mode 100644 index 00000000000..93012e33a44 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/useSnippetTitleGenerator.test.tsx @@ -0,0 +1,50 @@ +import { act, waitFor } from '@testing-library/react' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' + +import { useSnippetTitleGenerator } from './useSnippetTitleGenerator' +import { sqlEditorState } from '@/state/sql-editor/sql-editor-state' +import { + renderSqlEditorHook, + resetSqlEditorStores, + seedSnippet, + setupSqlEditorMocks, +} from '@/tests/lib/sql-editor-test-utils' + +const SNIPPET_ID = 'title-snippet' + +beforeEach(() => { + resetSqlEditorStores() + // The title endpoint (POST /ai/sql/title-v2) returns { title: 'Generated title' }. + setupSqlEditorMocks() + seedSnippet({ id: SNIPPET_ID, name: 'Untitled query', sql: 'select 1;' }) +}) + +afterEach(() => { + resetSqlEditorStores() +}) + +describe('useSnippetTitleGenerator', () => { + it('names an untitled snippet from the generated title and queues it for saving', async () => { + const { result } = renderSqlEditorHook(useSnippetTitleGenerator) + + await act(async () => { + await result.current.setAiTitle(SNIPPET_ID, 'select 1;') + }) + + await waitFor(() => + expect(sqlEditorState.snippets[SNIPPET_ID].snippet.name).toBe('Generated title') + ) + expect(sqlEditorState.needsSaving.has(SNIPPET_ID)).toBe(true) + }) + + it('returns the generated title from generateSqlTitle', async () => { + const { result } = renderSqlEditorHook(useSnippetTitleGenerator) + + let title: string | undefined + await act(async () => { + title = (await result.current.generateSqlTitle({ sql: 'select 1;' })).title + }) + + expect(title).toBe('Generated title') + }) +}) diff --git a/apps/studio/components/interfaces/SQLEditor/useSqlEditorAi.test.tsx b/apps/studio/components/interfaces/SQLEditor/useSqlEditorAi.test.tsx new file mode 100644 index 00000000000..244f4fc4431 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/useSqlEditorAi.test.tsx @@ -0,0 +1,152 @@ +import { act, waitFor } from '@testing-library/react' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' + +import { useSqlEditorDiff, useSqlEditorPrompt } from './hooks' +import { DiffType } from './SQLEditor.types' +import { useSqlEditorAi } from './useSqlEditorAi' +import { SIDEBAR_KEYS } from '@/components/layouts/ProjectLayout/LayoutSidebar/LayoutSidebarProvider' +import { sidebarManagerState } from '@/state/sidebar-manager-state' +import { sqlEditorDiffRequestState } from '@/state/sql-editor/sql-editor-diff-request' +import { sqlEditorSessionState } from '@/state/sql-editor/sql-editor-session-state' +import { + createInMemoryEditor, + renderSqlEditorHook, + resetSqlEditorStores, + seedSnippet, + setupSqlEditorMocks, +} from '@/tests/lib/sql-editor-test-utils' + +const SNIPPET_ID = 'ai-snippet' + +/** + * Composes the diff + prompt state hooks the AI hook depends on (production + * wires these together in `SQLEditorControllers`), so tests drive the real + * accept/discard/drain flows end to end. + */ +function useAiHarness({ editorMountCount = 1 }: { editorMountCount?: number } = {}) { + const diff = useSqlEditorDiff() + const prompt = useSqlEditorPrompt() + const ai = useSqlEditorAi({ id: SNIPPET_ID, editorMountCount, diff, prompt }) + return { ai, diff, prompt } +} + +beforeEach(() => { + resetSqlEditorStores() + setupSqlEditorMocks() +}) + +afterEach(() => { + resetSqlEditorStores() +}) + +describe('useSqlEditorAi — diff-request drain', () => { + it('copies the requested SQL straight into an empty editor and consumes the request', async () => { + sqlEditorDiffRequestState.requestDiff('select 100;', DiffType.Modification) + const inMemoryEditor = createInMemoryEditor('') + + const { result } = renderSqlEditorHook(useAiHarness, { inMemoryEditor }) + + await waitFor(() => expect(inMemoryEditor.editor.getValue()).toBe('select 100;')) + // One-shot: the request is drained so it can't re-apply to a later editor. + expect(sqlEditorDiffRequestState.pending).toBeUndefined() + expect(result.current.diff.isDiffOpen).toBe(false) + }) + + it('opens a diff between existing and requested SQL when the editor is non-empty', async () => { + sqlEditorDiffRequestState.requestDiff('select 42;', DiffType.Modification) + const inMemoryEditor = createInMemoryEditor('select 1;') + + const { result } = renderSqlEditorHook(useAiHarness, { inMemoryEditor }) + + await waitFor(() => expect(result.current.diff.isDiffOpen).toBe(true)) + // The diff-sync effect pushes the resolved diff into the (in-memory) diff editor. + expect(inMemoryEditor.getDiffOriginal()).toBe('select 1;') + expect(inMemoryEditor.getDiffModified()).toBe('select 42;') + expect(sqlEditorDiffRequestState.pending).toBeUndefined() + // The editor's own contents are untouched until the diff is accepted. + expect(inMemoryEditor.editor.getValue()).toBe('select 1;') + }) + + it('drains a pending request exactly once across editor remounts', async () => { + sqlEditorDiffRequestState.requestDiff('select 100;', DiffType.Modification) + const inMemoryEditor = createInMemoryEditor('') + + const { rerender } = renderSqlEditorHook(useAiHarness, { + inMemoryEditor, + initialProps: { editorMountCount: 1 }, + }) + + await waitFor(() => expect(inMemoryEditor.editor.getValue()).toBe('select 100;')) + + // Simulate a fresh editor mount: the consumed request must not re-apply. + inMemoryEditor.setValue('edited by user') + rerender({ editorMountCount: 2 }) + + await new Promise((r) => setTimeout(r, 20)) + expect(inMemoryEditor.editor.getValue()).toBe('edited by user') + }) +}) + +describe('useSqlEditorAi — accept / discard diff', () => { + async function openModificationDiff() { + sqlEditorDiffRequestState.requestDiff('select 42;', DiffType.Modification) + const inMemoryEditor = createInMemoryEditor('select 1;') + const utils = renderSqlEditorHook(useAiHarness, { inMemoryEditor }) + await waitFor(() => expect(utils.result.current.diff.isDiffOpen).toBe(true)) + return { ...utils, inMemoryEditor } + } + + it('accepting a modification writes the diff result back into the editor and closes the diff', async () => { + const { result, inMemoryEditor } = await openModificationDiff() + + await act(async () => { + await result.current.ai.acceptAiHandler() + }) + + await waitFor(() => expect(result.current.diff.isDiffOpen).toBe(false)) + expect(inMemoryEditor.editor.getValue()).toBe('select 42;') + }) + + it('discarding a diff leaves the editor untouched and closes the diff', async () => { + const { result, inMemoryEditor } = await openModificationDiff() + + await act(async () => { + result.current.ai.discardAiHandler() + }) + + await waitFor(() => expect(result.current.diff.isDiffOpen).toBe(false)) + expect(inMemoryEditor.editor.getValue()).toBe('select 1;') + }) +}) + +describe('useSqlEditorAi — debug', () => { + it('onDebug opens the assistant sidebar and starts a debug chat from the failing snippet', async () => { + seedSnippet({ id: SNIPPET_ID, name: 'Broken query', sql: 'selct 1;' }) + sqlEditorSessionState.addResultError(SNIPPET_ID, { message: 'syntax error at or near "selct"' }) + + const { result, aiAssistantState } = renderSqlEditorHook(useAiHarness) + + await act(async () => { + await result.current.ai.onDebug() + }) + + // No sidebar is registered in the test tree, so opening queues a pending open. + expect(sidebarManagerState.pendingSidebarOpen).toBe(SIDEBAR_KEYS.AI_ASSISTANT) + + const activeChat = aiAssistantState.chats[aiAssistantState.activeChatId ?? ''] + expect(activeChat?.name).toBe('Debug SQL snippet') + expect(aiAssistantState.sqlSnippets).toEqual(['selct 1;']) + expect(aiAssistantState.initialInput).toContain('syntax error at or near "selct"') + }) + + it('buildDebugPrompt embeds the snippet SQL and its error message', async () => { + seedSnippet({ id: SNIPPET_ID, name: 'Broken query', sql: 'selct 1;' }) + sqlEditorSessionState.addResultError(SNIPPET_ID, { message: 'boom' }) + + const { result } = renderSqlEditorHook(useAiHarness) + + const promptText = result.current.ai.buildDebugPrompt() + expect(promptText).toContain('boom') + expect(promptText).toContain('selct 1;') + }) +}) diff --git a/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.test.tsx b/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.test.tsx new file mode 100644 index 00000000000..4c4be71b112 --- /dev/null +++ b/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.test.tsx @@ -0,0 +1,263 @@ +import { acceptUntrustedSql, untrustedSql, type SafeSqlFragment } from '@supabase/pg-meta' +import { act, waitFor } from '@testing-library/react' +import { HttpResponse } from 'msw' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { useSqlEditorExecution } from './useSqlEditorExecution' +import { useReadReplicasQuery } from '@/data/read-replicas/replicas-query' +import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' +import { createDatabaseSelectorState } from '@/state/database-selector' +import { sqlEditorSessionState } from '@/state/sql-editor/sql-editor-session-state' +import { addAPIMock } from '@/tests/lib/msw' +import { + renderSqlEditorHook, + resetSqlEditorStores, + seedSnippet, + setupSqlEditorMocks, +} from '@/tests/lib/sql-editor-test-utils' + +const SNIPPET_ID = 'execution-snippet' + +/** Promote raw text to the `SafeSqlFragment` the run pipeline expects, exactly + * as the toolbar/editor-panel promote it right at the user action. */ +const sql = (text: string): SafeSqlFragment => acceptUntrustedSql(untrustedSql(text)) + +type ExecutionArgs = Parameters[0] + +/** + * Wraps the hook with the two queries it depends on so tests can wait for the + * project + read-replicas data to load before firing a run (an unloaded project + * or connection string silently short-circuits `executeQuery`). + */ +function useExecutionHarness(args: ExecutionArgs) { + const { data: project } = useSelectedProjectQuery() + const { data: databases } = useReadReplicasQuery({ projectRef: 'default' }) + const execution = useSqlEditorExecution(args) + return { execution, isReady: !!project && !!databases } +} + +/** Registers a query-endpoint resolver that records every executed SQL body. */ +function captureExecutedQueries(rows: unknown[] = []) { + const queries: string[] = [] + addAPIMock({ + method: 'post', + path: '/platform/pg-meta/:ref/query', + response: async ({ request }) => { + const body = (await request.json()) as { query: string } + queries.push(body.query) + return HttpResponse.json(rows) + }, + }) + return queries +} + +function renderExecution( + args: Partial = {}, + { selectedDatabaseId = 'default' }: { selectedDatabaseId?: string } = {} +) { + const databaseSelectorState = createDatabaseSelectorState() + databaseSelectorState.setSelectedDatabaseId(selectedDatabaseId) + + const setAiTitle = vi.fn() + const initialProps: ExecutionArgs = { + id: SNIPPET_ID, + isDiffOpen: false, + hasSelection: false, + setAiTitle, + ...args, + } + + const utils = renderSqlEditorHook((props: ExecutionArgs) => useExecutionHarness(props), { + initialProps, + databaseSelectorState, + }) + + return { ...utils, setAiTitle } +} + +beforeEach(() => { + resetSqlEditorStores() + setupSqlEditorMocks() + seedSnippet({ id: SNIPPET_ID, name: 'My query', sql: 'select 1;' }) +}) + +afterEach(() => { + vi.restoreAllMocks() +}) + +describe('useSqlEditorExecution', () => { + it('runs a non-destructive query and writes the result to the session store', async () => { + const rows = [{ id: 1, name: 'row-1' }] + const queries = captureExecutedQueries(rows) + + const { result } = renderExecution() + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('select 1')) + }) + + await waitFor(() => expect(sqlEditorSessionState.results[SNIPPET_ID]).toBeDefined()) + expect(sqlEditorSessionState.results[SNIPPET_ID][0].rows).toEqual(rows) + expect(queries.some((q) => /select 1/i.test(q))).toBe(true) + }) + + it('appends the auto-limit to a bare SELECT and records it on the result', async () => { + const queries = captureExecutedQueries([]) + + const { result } = renderExecution() + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('select 1')) + }) + + await waitFor(() => expect(sqlEditorSessionState.results[SNIPPET_ID]).toBeDefined()) + // The session store's `limit` defaults to 100 (see resetSqlEditorStores). + expect(sqlEditorSessionState.results[SNIPPET_ID][0].autoLimit).toBe(100) + expect(queries.some((q) => /limit 100/i.test(q))).toBe(true) + }) + + it('gates a destructive query behind potentialIssues instead of running it', async () => { + const queries = captureExecutedQueries([]) + + const { result } = renderExecution() + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('drop table foo;')) + }) + + // The warning-modal gate fired: issues are surfaced and nothing was executed. + await waitFor(() => expect(result.current.execution.potentialIssues).toBeDefined()) + expect(result.current.execution.potentialIssues?.hasDestructiveOperations).toBe(true) + expect(queries.some((q) => /drop table foo/i.test(q))).toBe(false) + expect(sqlEditorSessionState.results[SNIPPET_ID]).toBeUndefined() + }) + + it('runs a destructive query when forced', async () => { + const queries = captureExecutedQueries([]) + + const { result } = renderExecution() + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('drop table foo;'), true) + }) + + await waitFor(() => expect(queries.some((q) => /drop table foo/i.test(q))).toBe(true)) + await waitFor(() => expect(sqlEditorSessionState.results[SNIPPET_ID]).toBeDefined()) + }) + + it('highlights the error line and records the error on failure', async () => { + addAPIMock({ + method: 'post', + path: '/platform/pg-meta/:ref/query', + response: async ({ request }) => { + const body = (await request.json()) as { query: string } + // The event-triggers probe shares this endpoint; only fail the real run. + if (/select 1/i.test(body.query)) { + return HttpResponse.json( + { + message: 'syntax error', + position: '8', + formattedError: + 'ERROR: syntax error at or near "slect"\nLINE 3: slect 1;\n ^', + }, + { status: 400 } + ) + } + return HttpResponse.json([]) + }, + }) + + const { result, inMemoryEditor } = renderExecution() + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('select 1')) + }) + + await waitFor(() => expect(sqlEditorSessionState.results[SNIPPET_ID]?.[0].error).toBeDefined()) + // No selection → base line 0 + parsed `LINE 3` = 3. + expect(inMemoryEditor.getHighlightedLine()).toBe(3) + expect(inMemoryEditor.getRevealedLine()).toBe(3) + }) + + it("sends the selected database's connection string as the connection header", async () => { + const connectionString = 'postgresql://postgres@replica.example:5432/postgres' + const connectionHeaders: (string | null)[] = [] + addAPIMock({ + method: 'post', + path: '/platform/pg-meta/:ref/query', + response: async ({ request }) => { + const body = (await request.json()) as { query: string } + if (/select 1/i.test(body.query)) { + connectionHeaders.push(request.headers.get('x-connection-encrypted')) + } + return HttpResponse.json([]) + }, + }) + + // Point the read-replicas list + selector at a replica with its own conn string. + addAPIMock({ + method: 'get', + path: '/platform/projects/:ref/databases', + response: [ + { + identifier: 'replica-1', + connectionString, + cloud_provider: 'AWS', + db_host: 'db.replica.supabase.co', + db_name: 'postgres', + db_port: 5432, + db_user: 'postgres', + inserted_at: '2024-01-01T00:00:00Z', + region: 'us-east-1', + restUrl: 'https://replica.supabase.co/rest/v1/', + size: 'ci_micro', + status: 'ACTIVE_HEALTHY', + }, + ], + }) + + const { result } = renderExecution({}, { selectedDatabaseId: 'replica-1' }) + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('select 1')) + }) + + await waitFor(() => expect(connectionHeaders.length).toBeGreaterThan(0)) + expect(connectionHeaders[0]).toBe(connectionString) + }) + + it('short-circuits while a diff is open', async () => { + const queries = captureExecutedQueries([]) + + const { result } = renderExecution({ isDiffOpen: true }) + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('select 1')) + }) + + await new Promise((r) => setTimeout(r, 20)) + expect(queries.some((q) => /select 1/i.test(q))).toBe(false) + expect(sqlEditorSessionState.results[SNIPPET_ID]).toBeUndefined() + }) + + it('does not auto-generate a title for an already-named snippet', async () => { + captureExecutedQueries([]) + + const { result, setAiTitle } = renderExecution() + await waitFor(() => expect(result.current.isReady).toBe(true)) + + await act(async () => { + await result.current.execution.executeQuery(sql('select 1')) + }) + + await waitFor(() => expect(sqlEditorSessionState.results[SNIPPET_ID]).toBeDefined()) + expect(setAiTitle).not.toHaveBeenCalled() + }) +})