diff --git a/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx b/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx index 87b9579e860cf..ba4730da723c4 100644 --- a/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx +++ b/apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx @@ -35,6 +35,13 @@ import { QueryResultRenderer } from './QueryResultRenderer' import { QuerySourceMenu } from './QuerySourceMenu' import { useQueryEditorAi } from './useQueryEditorAi' import { LegacyLogsRewriteBanner } from '@/components/interfaces/Settings/Logs/LegacyLogsRewriteBanner' +import { RunQueryWarningModal } from '@/components/interfaces/SQLEditor/RunQueryWarningModal' +import type { PotentialIssues } from '@/components/interfaces/SQLEditor/SQLEditor.types' +import { + analyzeQueryIssues, + appendEnableRLSStatements, + hasBlockingIssues, +} from '@/components/interfaces/SQLEditor/SQLEditor.utils' import { useAddDefinitions } from '@/components/interfaces/SQLEditor/useAddDefinitions' import { ResizableAIWidget } from '@/components/ui/AIEditor/ResizableAIWidget' import { getEditorSelectionParts, type EditorSelection } from '@/components/ui/AIEditor/utils' @@ -44,6 +51,7 @@ import { type DatabaseSourceParameters, type LogsSourceParameters, } from '@/data/content/notebooks/notebook-schema' +import { useDatabaseEventTriggersQuery } from '@/data/database-event-triggers/database-event-triggers-query' import { isValidConnString } from '@/data/fetchers' import { useExecuteLogsSqlMutation } from '@/data/logs/execute-logs-sql-mutation' import { @@ -169,6 +177,7 @@ export const QueryEditor = forwardRef(funct const databaseIdentifier = query._tag === 'database' ? query.database_identifier : undefined const [promptInput, setPromptInput] = useState('') + const [pendingRun, setPendingRun] = useState<{ sql: string; issues: PotentialIssues }>() const [pendingProposal, setPendingProposal] = useState(null) const pendingProposalRef = useLatest(pendingProposal) @@ -193,6 +202,16 @@ export const QueryEditor = forwardRef(funct } ) + const connectionString = + databaseIdentifier === undefined || databaseIdentifier === project?.ref + ? project?.connectionString + : databases?.find((database) => database.identifier === databaseIdentifier)?.connectionString + + const { data: eventTriggers } = useDatabaseEventTriggersQuery( + { projectRef: project?.ref, connectionString }, + { enabled: query._tag === 'database' && isValidConnString(connectionString) } + ) + const { mutateAsync: executeSql, isPending: isExecutingSql } = useExecuteSqlMutation({ onSuccess: (data) => onResultChange({ rows: data.result }), onError: (error) => onResultChange({ error }), @@ -215,11 +234,27 @@ export const QueryEditor = forwardRef(funct * decided by `query._tag`, the same discriminant that picks the execution endpoint, so * Postgres SQL cannot reach the analytics wire or vice versa. */ - const handleRunQuery = async (rawSql: string = sql) => { + const handleRunQuery = async ({ + rawSql = sql, + shouldForce = false, + }: { + rawSql?: string + shouldForce?: boolean + } = {}) => { if (!project || isBusy || pendingProposal || isRunDisabled || rawSql.trim().length === 0) return + if (query._tag === 'database') { + const issues = analyzeQueryIssues(rawSql, eventTriggers) + if (hasBlockingIssues(issues, shouldForce)) { + setPendingRun({ sql: rawSql, issues }) + return + } + } + onRun?.() - onSqlCommit?.(rawSql) + // [Joshen] This is deliberate to commit the sql, rather than the passed rawSql + // As we want to save the cell's content into the store, rather than what's getting run + onSqlCommit?.(sql) if (query._tag === 'logs') { if (!isOtelLogsEnabled) { @@ -240,11 +275,6 @@ export const QueryEditor = forwardRef(funct const safeSql = acceptUntrustedSql(untrustedSql(rawSql)) const limitedSql = applyAutoLimit(safeSql, rowLimit) - const connectionString = - databaseIdentifier === undefined || databaseIdentifier === project.ref - ? project.connectionString - : databases?.find((database) => database.identifier === databaseIdentifier) - ?.connectionString if (!isValidConnString(connectionString)) { onResultChange({ error: { message: 'Unable to run query: Connection string is missing' } }) @@ -262,6 +292,22 @@ export const QueryEditor = forwardRef(funct }).catch(() => {}) } + const handleConfirmPendingRun = () => { + if (!pendingRun) return + const runSql = pendingRun.sql + setPendingRun(undefined) + handleRunQuery({ rawSql: runSql, shouldForce: true }) + } + + const handleConfirmPendingRunWithRLS = () => { + if (!pendingRun) return + const tables = pendingRun.issues.createTablesMissingRLS ?? [] + if (tables.length === 0) return + const rewrittenSql = appendEnableRLSStatements(pendingRun.sql, tables) + setPendingRun(undefined) + handleRunQuery({ rawSql: rewrittenSql, shouldForce: true }) + } + const acceptSqlProposal = () => { if (isReadOnly || !pendingProposal) return if (sql === pendingProposal.original) { @@ -271,8 +317,6 @@ export const QueryEditor = forwardRef(funct setPendingProposal(null) } - const discardSqlProposal = () => setPendingProposal(null) - const closePrompt = () => { setPromptState(null) setPromptInput('') @@ -312,172 +356,194 @@ export const QueryEditor = forwardRef(funct }, [promptState?.isOpen]) return ( - - - - - - {title} - - {toolbarActions} - {onSourceChange && ( - + + + + + + {title} + + {toolbarActions} + {onSourceChange && ( + { + setPendingProposal(null) + onSourceChange(source) + }} + rowLimit={rowLimit} + onRowLimitChange={onRowLimitChange} + roleImpersonationState={roleImpersonationState} + /> + )} + {display && onDisplayChange && ( + + )} + : } disabled={pendingProposal !== null} - source={toQuerySourceBinding(query)} - onSourceChange={(source) => { - setPendingProposal(null) - onSourceChange(source) - }} - rowLimit={rowLimit} - onRowLimitChange={onRowLimitChange} - roleImpersonationState={roleImpersonationState} + tooltip={showQuery ? 'Hide query' : 'Show query'} + onClick={() => onShowQueryChange(!showQuery)} /> - )} - {display && onDisplayChange && ( - - )} - : } - disabled={pendingProposal !== null} - tooltip={showQuery ? 'Hide query' : 'Show query'} - onClick={() => onShowQueryChange(!showQuery)} - /> - } - tooltip="Run query" - disabled={ - isBusy || pendingProposal !== null || isRunDisabled || sql.trim().length === 0 - } - onClick={() => handleRunQuery()} - > - Run - - - - - {showQuery && ( - <> - sqlRef.current} - onProposal={({ original, modified }) => - setPendingProposal({ - original, - modified, - label: 'Review the ClickHouse SQL rewrite before accepting it', - }) - } - hidden={pendingProposal !== null} - /> - - onSqlChange(value ?? '')} - onMount={(editor, monaco) => { - editor.onDidBlurEditorWidget(() => onSqlCommitRef.current?.(sqlRef.current)) - editorInstanceRef.current = editor - - editor.addAction({ - id: 'generate-sql', - label: 'Generate SQL', - keybindings: [monaco.KeyMod.CtrlCmd | monaco.KeyMod.Shift | monaco.KeyCode.KeyK], - run: () => { - if (pendingProposalRef.current) return - const selectionParts = getEditorSelectionParts(editor) - if (selectionParts) setPromptState({ isOpen: true, ...selectionParts }) - }, + } + tooltip="Run query" + disabled={ + isBusy || pendingProposal !== null || isRunDisabled || sql.trim().length === 0 + } + onClick={() => handleRunQuery()} + > + Run + + + + + {showQuery && ( + <> + sqlRef.current} + onProposal={({ original, modified }) => + setPendingProposal({ + original, + modified, + label: 'Review the ClickHouse SQL rewrite before accepting it', }) - }} + } + hidden={pendingProposal !== null} /> - - {promptState?.isOpen && editorInstanceRef.current && !pendingProposal && ( - + onSqlChange(value ?? '')} + onMount={(editor, monaco) => { + editor.onDidBlurEditorWidget(() => onSqlCommitRef.current?.(sqlRef.current)) + editorInstanceRef.current = editor + + editor.addAction({ + id: 'generate-sql', + label: 'Generate SQL', + keybindings: [ + monaco.KeyMod.CtrlCmd | monaco.KeyMod.Shift | monaco.KeyCode.KeyK, + ], + run: () => { + if (pendingProposalRef.current) return + const selectionParts = getEditorSelectionParts(editor) + if (selectionParts) setPromptState({ isOpen: true, ...selectionParts }) + }, + }) + }} /> - )} - {pendingProposal && ( -
-
-
-

{pendingProposal.label}

- {pendingProposal.prompt && ( -

- Prompt: {pendingProposal.prompt} -

- )} + {promptState?.isOpen && editorInstanceRef.current && !pendingProposal && ( + + )} + + {pendingProposal && ( +
+
+
+

{pendingProposal.label}

+ {pendingProposal.prompt && ( +

+ Prompt: {pendingProposal.prompt} +

+ )} +
+
+ + +
-
- - +
+
-
- -
-
- )} - - - )} - - - - - - -

{(result?.rows ?? []).length.toLocaleString()} rows

- {rowLimit && ( - <> -

·

-

{rowLimit < 0 ? 'No row limit' : `Limit ${rowLimit} rows`}

+ )} + )} - - + + + + + + +

{(result?.rows ?? []).length.toLocaleString()} rows

+ {rowLimit && ( + <> +

·

+

{rowLimit < 0 ? 'No row limit' : `Limit ${rowLimit} rows`}

+ + )} + + + + {query._tag === 'database' && ( + setPendingRun(undefined)} + onConfirm={handleConfirmPendingRun} + onConfirmWithRLS={ + (pendingRun?.issues.createTablesMissingRLS?.length ?? 0) > 0 + ? handleConfirmPendingRunWithRLS + : undefined + } + /> + )} + ) }) diff --git a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.test.tsx b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.test.tsx index 9b23ad6fce495..8cde3bb59448e 100644 --- a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.test.tsx +++ b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.test.tsx @@ -127,16 +127,20 @@ describe('ExplorerNotebookTab', () => { it('runs every database and log cell, and skips markdown cells, on "Run notebook"', async () => { // `useAddDefinitions` fires its own background keywords/functions/schemas/table-columns // fetches against this same generic pg-meta query endpoint (differentiated by the `key` - // search param) as soon as a database cell's editor mounts — those are expected and - // unrelated to the actual cell run. + // search param) as soon as a database cell's editor mounts, and `QueryEditor` runs an + // event-trigger lookup for `analyzeQueryIssues` — both are expected and unrelated to the + // actual cell run. const INTELLISENSE_KEYS = ['keywords', 'database-functions', 'schemas', 'table-columns'] const dbRequests: Request[] = [] addAPIMock({ method: 'post', path: '/platform/pg-meta/:ref/query', - response: ({ request }) => { + response: async ({ request }) => { const key = new URL(request.url).searchParams.get('key') - if (!key || !INTELLISENSE_KEYS.includes(key)) dbRequests.push(request) + const { query } = (await request.clone().json()) as { query?: string } + const isIntellisenseRequest = !!key && INTELLISENSE_KEYS.includes(key) + const isEventTriggerRequest = !!query?.includes('pg_event_trigger') + if (!isIntellisenseRequest && !isEventTriggerRequest) dbRequests.push(request) return HttpResponse.json([{ result: 1 }]) }, }) diff --git a/apps/studio/components/interfaces/Explorer/__tests__/QueryTab.test.tsx b/apps/studio/components/interfaces/Explorer/__tests__/QueryTab.test.tsx index 9157409706477..1801ffd177c5a 100644 --- a/apps/studio/components/interfaces/Explorer/__tests__/QueryTab.test.tsx +++ b/apps/studio/components/interfaces/Explorer/__tests__/QueryTab.test.tsx @@ -59,13 +59,14 @@ const createDraft = ( | { _tag: 'logs' time_range: { _tag: 'relative_time_range'; amount: number; unit: 'hour' } - } + }, + sql: string = 'select 1' ) => { explorerQueryState.removeDraft({ id: 'query-test', projectRef: 'default' }) explorerQueryState.createDraft({ id: 'query-test', projectRef: 'default', - sql: 'select 1', + sql, source, }) } @@ -212,4 +213,84 @@ describe('QueryTab execution', () => { explorerQueryState.removeDraft({ id: 'query-test-2', projectRef: 'default' }) }) }) + + it('blocks a destructive query behind a confirmation modal, then runs it once confirmed', async () => { + createDraft({ _tag: 'database' }, 'delete from foo') + const executedQueries: string[] = [] + addAPIMock({ + method: 'post', + path: '/platform/pg-meta/:ref/query', + response: async ({ request }) => { + const { query } = (await request.json()) as { query: string } + if (query.trim().toLowerCase().startsWith('delete')) executedQueries.push(query) + return HttpResponse.json([]) + }, + }) + + renderQueryTab() + const runButton = await screen.findByRole('button', { name: 'Run' }) + await waitFor(() => expect(runButton).toBeEnabled()) + await userEvent.click(runButton) + + expect(await screen.findByText('Potential issue detected')).toBeInTheDocument() + expect(executedQueries).toHaveLength(0) + + await userEvent.click(screen.getByRole('button', { name: 'Run query' })) + + await waitFor(() => expect(executedQueries).toHaveLength(1)) + }) + + it('cancels a blocked query without running it', async () => { + createDraft({ _tag: 'database' }, 'update foo set bar = 1') + const executedQueries: string[] = [] + addAPIMock({ + method: 'post', + path: '/platform/pg-meta/:ref/query', + response: async ({ request }) => { + const { query } = (await request.json()) as { query: string } + if (query.trim().toLowerCase().startsWith('update')) executedQueries.push(query) + return HttpResponse.json([]) + }, + }) + + renderQueryTab() + const runButton = await screen.findByRole('button', { name: 'Run' }) + await waitFor(() => expect(runButton).toBeEnabled()) + await userEvent.click(runButton) + + expect(await screen.findByText('Potential issue detected')).toBeInTheDocument() + + await userEvent.click(screen.getByRole('button', { name: 'Cancel' })) + + expect(screen.queryByText('Potential issue detected')).not.toBeInTheDocument() + expect(executedQueries).toHaveLength(0) + }) + + it('runs a CREATE TABLE query with RLS enabled once "Run and enable RLS" is chosen', async () => { + createDraft({ _tag: 'database' }, 'create table foo (id int)') + const executedQueries: string[] = [] + addAPIMock({ + method: 'post', + path: '/platform/pg-meta/:ref/query', + response: async ({ request }) => { + const { query } = (await request.json()) as { query: string } + if (query.trim().toLowerCase().startsWith('create table')) executedQueries.push(query) + return HttpResponse.json([]) + }, + }) + + renderQueryTab() + const runButton = await screen.findByRole('button', { name: 'Run' }) + await waitFor(() => expect(runButton).toBeEnabled()) + await userEvent.click(runButton) + + expect(await screen.findByText('Potential issue detected')).toBeInTheDocument() + expect(executedQueries).toHaveLength(0) + + await userEvent.click(screen.getByRole('button', { name: 'Run and enable RLS' })) + + await waitFor(() => expect(executedQueries).toHaveLength(1)) + expect(executedQueries[0]).toContain('create table foo (id int)') + expect(executedQueries[0]).toContain('ALTER TABLE foo ENABLE ROW LEVEL SECURITY;') + }) })