From 11fb71514905d73c006da32bdbcbcc0d3274ba31 Mon Sep 17 00:00:00 2001 From: Charis <26616127+charislam@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:30:44 -0400 Subject: [PATCH 1/2] refactor(sql-editor): extract query-issue analysis + connection string resolution (#47980) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Part 1/6 of the SQL Editor testability follow-up (extracting pure decision logic out of the SQL editor hooks so it can be exhaustively unit-tested without mocking). - Extracts `analyzeQueryIssues` and `hasBlockingIssues` out of `useSqlEditorExecution`'s inline destructive-query / warning-modal-gating logic into `SQLEditor.utils.ts`. - Extracts `resolveConnectionString`, deduping the `databases?.find(...)` lookup that was duplicated verbatim in both `useSqlEditorExecution` and `useSqlEditorExplain`. - Behavior-preserving — same runtime logic, now unit-testable in isolation. ## Test plan - [x] `pnpm --filter studio typecheck` - [x] `pnpm test:studio -- SQLEditor.utils` (159 tests passing, includes new exhaustive-permutation cases for the three extracted functions) ## Summary by CodeRabbit * **Bug Fixes** * Improved SQL safety checks before execution, including destructive queries, unsafe updates, database-altering commands, and tables missing row-level security. * Correctly recognizes tables protected by active security triggers. * Improved database connection selection for primary and read-replica databases. * Preserved the ability to run queries with force enabled when safety warnings are present. --- .../SQLEditor/SQLEditor.utils.test.ts | 110 ++++++++++++++++++ .../interfaces/SQLEditor/SQLEditor.utils.ts | 50 +++++++- .../SQLEditor/useSqlEditorExecution.ts | 40 ++----- 3 files changed, 169 insertions(+), 31 deletions(-) diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts index 389a06b9ff0f4..63f45bc3a1170 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts @@ -4,6 +4,7 @@ import { describe, expect, it, test } from 'vitest' import type { IStandaloneCodeEditor } from './SQLEditor.types' import { + analyzeQueryIssues, appendEnableRLSStatements, assembleCompletionDiff, buildDebugPromptText, @@ -15,10 +16,13 @@ import { getCreateTablesMissingRLS, getEditorSql, hasActiveEnsureRLSTrigger, + hasBlockingIssues, isUpdateWithoutWhere, + resolveConnectionString, suffixWithLimit, } from './SQLEditor.utils' import type { DatabaseEventTrigger } from '@/data/database-event-triggers/database-event-triggers-query' +import type { Database } from '@/data/read-replicas/replicas-query' const buildTrigger = (overrides: Partial = {}): DatabaseEventTrigger => ({ oid: 1, @@ -1045,3 +1049,109 @@ describe('SQLEditor.utils:buildDebugPromptText', () => { expect(result).toContain('```sql\nselect 1;\n```') }) }) + +describe('SQLEditor.utils:analyzeQueryIssues', () => { + it('flags a destructive query with no event triggers', () => { + expect(analyzeQueryIssues('drop table countries;', undefined)).toEqual({ + hasDestructiveOperations: true, + hasUpdateWithoutWhere: false, + hasAlterDatabasePreventConnection: false, + createTablesMissingRLS: [], + }) + }) + + it('filters out a CREATE TABLE covered by an active ensure_rls trigger', () => { + const result = analyzeQueryIssues('create table foo (id int);', [buildTrigger()]) + expect(result.createTablesMissingRLS).toEqual([]) + }) + + it('reports no issues for a clean select query', () => { + expect(analyzeQueryIssues('select * from countries;', undefined)).toEqual({ + hasDestructiveOperations: false, + hasUpdateWithoutWhere: false, + hasAlterDatabasePreventConnection: false, + createTablesMissingRLS: [], + }) + }) +}) + +describe('SQLEditor.utils:hasBlockingIssues', () => { + const noIssues = { + hasDestructiveOperations: false, + hasUpdateWithoutWhere: false, + hasAlterDatabasePreventConnection: false, + createTablesMissingRLS: [], + } + + it('returns false when there are no issues', () => { + expect(hasBlockingIssues(noIssues, false)).toBe(false) + }) + + it('returns true when hasDestructiveOperations is set', () => { + expect(hasBlockingIssues({ ...noIssues, hasDestructiveOperations: true }, false)).toBe(true) + }) + + it('returns true when hasUpdateWithoutWhere is set', () => { + expect(hasBlockingIssues({ ...noIssues, hasUpdateWithoutWhere: true }, false)).toBe(true) + }) + + it('returns true when hasAlterDatabasePreventConnection is set', () => { + expect(hasBlockingIssues({ ...noIssues, hasAlterDatabasePreventConnection: true }, false)).toBe( + true + ) + }) + + it('returns true when createTablesMissingRLS is non-empty', () => { + expect( + hasBlockingIssues({ ...noIssues, createTablesMissingRLS: [{ tableName: 'foo' }] }, false) + ).toBe(true) + }) + + it('returns false when force is true, even with issues present', () => { + expect(hasBlockingIssues({ ...noIssues, hasDestructiveOperations: true }, true)).toBe(false) + }) + + it('returns false when force is true and there are no issues', () => { + expect(hasBlockingIssues(noIssues, true)).toBe(false) + }) +}) + +const buildDatabase = (overrides: Partial = {}): Database => ({ + cloud_provider: 'AWS', + connectionString: 'postgres://primary', + db_host: 'db.example.com', + db_name: 'postgres', + db_port: 5432, + db_user: 'postgres', + identifier: 'primary', + inserted_at: '2024-01-01T00:00:00Z', + region: 'us-east-1', + restUrl: 'https://example.com', + size: 'micro', + status: 'ACTIVE_HEALTHY', + ...overrides, +}) + +describe('SQLEditor.utils:resolveConnectionString', () => { + it('returns the matching database connection string', () => { + const databases = [ + buildDatabase({ identifier: 'primary', connectionString: 'postgres://primary' }), + buildDatabase({ identifier: 'replica-1', connectionString: 'postgres://replica-1' }), + ] + expect(resolveConnectionString(databases, 'replica-1')).toBe('postgres://replica-1') + }) + + it('returns undefined when no database matches', () => { + const databases = [buildDatabase({ identifier: 'primary' })] + expect(resolveConnectionString(databases, 'unknown')).toBeUndefined() + }) + + it('returns undefined when databases is undefined', () => { + expect(resolveConnectionString(undefined, 'primary')).toBeUndefined() + }) + + it('normalizes a null connectionString to undefined', () => { + const databases = [buildDatabase({ identifier: 'primary', connectionString: null })] + expect(resolveConnectionString(databases, 'primary')).toBeUndefined() + }) +}) diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts index 88a89a18d3ed7..27454efa19ae1 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts @@ -8,9 +8,10 @@ import { sqlAiDisclaimerComment, updateWithoutWhereRegex, } from './SQLEditor.constants' -import { ContentDiff, type IStandaloneCodeEditor } from './SQLEditor.types' +import { ContentDiff, type IStandaloneCodeEditor, type PotentialIssues } from './SQLEditor.types' import type { SnippetWithContent } from '@/data/content/sql-folders-query' import type { DatabaseEventTrigger } from '@/data/database-event-triggers/database-event-triggers-query' +import type { Database } from '@/data/read-replicas/replicas-query' import { generateUuid } from '@/lib/api/snippets.browser' import { removeCommentsFromSql } from '@/lib/helpers' import { sqlEventParser } from '@/lib/sql-event-parser' @@ -177,6 +178,53 @@ export function checkAlterDatabaseConnection(sql: string): boolean { ) } +/** + * Runs every pre-execution safety check on a query and packages the results as + * `PotentialIssues`, used both to decide whether to show the warning modal and + * to render its content. + */ +export function analyzeQueryIssues( + sql: string, + eventTriggers: DatabaseEventTrigger[] | undefined +): PotentialIssues { + return { + hasDestructiveOperations: checkDestructiveQuery(sql), + hasUpdateWithoutWhere: isUpdateWithoutWhere(sql), + hasAlterDatabasePreventConnection: checkAlterDatabaseConnection(sql), + createTablesMissingRLS: filterTablesCoveredByEnsureRLSTrigger( + getCreateTablesMissingRLS(sql), + hasActiveEnsureRLSTrigger(eventTriggers) + ), + } +} + +/** + * Whether `issues` should block an unforced run behind the warning modal. + */ +export function hasBlockingIssues(issues: PotentialIssues, force: boolean): boolean { + return ( + !force && + (!!issues.hasDestructiveOperations || + !!issues.hasUpdateWithoutWhere || + !!issues.hasAlterDatabasePreventConnection || + (issues.createTablesMissingRLS?.length ?? 0) > 0) + ) +} + +/** + * Resolves the connection string for the currently selected database (primary + * or read replica) from the read-replicas list. Shared by the run and explain + * flows so the lookup isn't duplicated. + */ +export function resolveConnectionString( + databases: Database[] | undefined, + selectedDatabaseId: string | undefined +): string | undefined { + return ( + databases?.find((db) => db.identifier === selectedDatabaseId)?.connectionString ?? undefined + ) +} + export const generateMigrationCliCommand = (id: string, name: string, isNpx = false) => ` ${isNpx ? 'npx ' : ''}supabase snippets download ${id} | diff --git a/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts b/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts index 7b926512e0b63..85d76becfe41f 100644 --- a/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts +++ b/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts @@ -7,13 +7,10 @@ import { toast } from 'sonner' import { untitledSnippetTitle } from './SQLEditor.constants' import type { PotentialIssues } from './SQLEditor.types' import { - checkAlterDatabaseConnection, - checkDestructiveQuery, + analyzeQueryIssues, checkIfAppendLimitRequired, - filterTablesCoveredByEnsureRLSTrigger, - getCreateTablesMissingRLS, - hasActiveEnsureRLSTrigger, - isUpdateWithoutWhere, + hasBlockingIssues, + resolveConnectionString, suffixWithLimit, } from './SQLEditor.utils' import { useSQLEditorContext } from './SQLEditorContext' @@ -108,28 +105,10 @@ export function useSqlEditorExecution({ return } - const hasDestructiveOperations = checkDestructiveQuery(sql) - const hasUpdateWithoutWhere = isUpdateWithoutWhere(sql) - const hasAlterDatabasePreventConnection = checkAlterDatabaseConnection(sql) - const createTablesMissingRLS = filterTablesCoveredByEnsureRLSTrigger( - getCreateTablesMissingRLS(sql), - hasActiveEnsureRLSTrigger(eventTriggers) - ) + const issues = analyzeQueryIssues(sql, eventTriggers) - const queryHasIssues = - !force && - (hasDestructiveOperations || - hasUpdateWithoutWhere || - hasAlterDatabasePreventConnection || - createTablesMissingRLS.length > 0) - - if (queryHasIssues) { - setPotentialIssues({ - hasDestructiveOperations, - hasUpdateWithoutWhere, - hasAlterDatabasePreventConnection, - createTablesMissingRLS, - }) + if (hasBlockingIssues(issues, force)) { + setPotentialIssues(issues) return } @@ -149,9 +128,10 @@ export function useSqlEditorExecution({ clearHighlights() const impersonatedRoleState = getImpersonatedRoleState() - const connectionString = databases?.find( - (db) => db.identifier === databaseSelectorState.selectedDatabaseId - )?.connectionString + const connectionString = resolveConnectionString( + databases, + databaseSelectorState.selectedDatabaseId + ) if (!isValidConnString(connectionString)) { clearPendingRunRefocus() return toast.error('Unable to run query: Connection string is missing') From 1c827c5cbb29cacc6e9052adff2e1659e3cb05fb Mon Sep 17 00:00:00 2001 From: Charis <26616127+charislam@users.noreply.github.com> Date: Thu, 16 Jul 2026 17:08:26 -0400 Subject: [PATCH 2/2] refactor(sql-editor): extract title-gen/execute-params + merge auto-limit functions (#48013) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Part 2/6 of the SQL Editor testability follow-up, stacked on #47980 (the analyzeQueryIssues/resolveConnectionString PR). - Extracts `shouldAutoGenerateTitle` and `buildExecuteParams` out of `useSqlEditorExecution`'s inline logic into `SQLEditor.utils.ts`. - Merges `checkIfAppendLimitRequired` and `suffixWithLimit` into a single `applyAutoLimit` function — the two were only ever called together and re-parsed the same query twice at every call site. `applyAutoLimit` only accepts `SafeSqlFragment` (never a plain string) and composes the `LIMIT` suffix through `safeSql`/`literal` rather than raw template concatenation, so the only place in the file that reasserts the `SafeSqlFragment` brand on a derived string is the small, dedicated `trimTrailingSemicolons` helper — removing existing terminators can't introduce unsafe content, unlike gluing new text onto the fragment. - Updates the two other `checkIfAppendLimitRequired`/`suffixWithLimit` call sites (`EditorPanel.tsx`, `ReportBlock.tsx`) accordingly; `ReportBlock` now promotes its report SQL once and reuses the result for both its display-only auto-limit hint and its execution, instead of promoting twice. ## Test plan - [x] `pnpm --filter studio typecheck` - [x] `pnpm test:studio -- SQLEditor ReportBlock EditorPanel` (239 tests passing) --- .../Reports/ReportBlock/ReportBlock.tsx | 16 +- .../SQLEditor/SQLEditor.utils.test.ts | 219 +++++++++++++----- .../interfaces/SQLEditor/SQLEditor.utils.ts | 105 ++++++++- .../SQLEditor/useSqlEditorExecution.ts | 38 ++- .../components/ui/EditorPanel/EditorPanel.tsx | 4 +- .../sql-editor/sql-editor-session-state.ts | 2 +- 6 files changed, 284 insertions(+), 100 deletions(-) diff --git a/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx b/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx index f0e8152f3b241..ba68a57a12723 100644 --- a/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx +++ b/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx @@ -1,11 +1,11 @@ -import { acceptUntrustedSql } from '@supabase/pg-meta' +import { acceptUntrustedSql, safeSql } from '@supabase/pg-meta' import { useQuery } from '@tanstack/react-query' import { useParams } from 'common' import { X } from 'lucide-react' import { useEffect, useState } from 'react' import { toast } from 'sonner' -import { checkIfAppendLimitRequired, suffixWithLimit } from '../../SQLEditor/SQLEditor.utils' +import { applyAutoLimit } from '../../SQLEditor/SQLEditor.utils' import { BURSTABLE_IO_METRIC_KEYS, DEPRECATED_REPORTS } from '../Reports.constants' import { ChartBlock } from './ChartBlock' import { DeprecatedChartBlock } from './DeprecatedChartBlock' @@ -81,7 +81,11 @@ export const ReportBlock = ({ const autoLimit = 100 const sql = isSnippet ? (data?.content as SqlSnippets.Content)?.unchecked_sql : undefined - const { appendAutoLimit } = checkIfAppendLimitRequired(sql ?? '', autoLimit) + // acceptUntrustedSql is usually not allowed outside a user-action event + // handler, but it's explicitly fine here: adding this block to a report is + // itself the user action that approves running its SQL. + const acceptedSql = sql !== undefined ? acceptUntrustedSql(sql) : safeSql`` + const { sql: formattedSql, appendAutoLimit } = applyAutoLimit(acceptedSql, autoLimit) const chartConfig = { ...DEFAULT_CHART_CONFIG, ...(item.chartConfig ?? {}) } const isDeprecatedChart = DEPRECATED_REPORTS.includes(item.attribute) @@ -97,6 +101,7 @@ export const ReportBlock = ({ isPending: executeSqlLoading, refetch, } = useQuery({ + // eslint-disable-next-line @tanstack/query/exhaustive-deps -- formattedSql/appendAutoLimit are fully derived from sql/autoLimit, both already in the key queryKey: sqlKeys.query(projectRef, [ item.id, sql, @@ -114,14 +119,9 @@ export const ReportBlock = ({ return null } - const formattedSql = suffixWithLimit(acceptUntrustedSql(sql), autoLimit) - return executeSql({ projectRef, connectionString, - // acceptUntrustedSql is usually not allowed in an auto-run position, - // but in this case we are explicitly allowing it because adding a block - // to a report is an explicit user action. sql: formattedSql, }) }, diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts index 63f45bc3a1170..56d728c654fe2 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.test.ts @@ -2,15 +2,17 @@ import { safeSql, untrustedSql } from '@supabase/pg-meta' import { stripIndent } from 'common-tags' import { describe, expect, it, test } from 'vitest' +import { untitledSnippetTitle } from './SQLEditor.constants' import type { IStandaloneCodeEditor } from './SQLEditor.types' import { analyzeQueryIssues, appendEnableRLSStatements, + applyAutoLimit, assembleCompletionDiff, buildDebugPromptText, + buildExecuteParams, checkAlterDatabaseConnection, checkDestructiveQuery, - checkIfAppendLimitRequired, computeErrorHighlightLine, filterTablesCoveredByEnsureRLSTrigger, getCreateTablesMissingRLS, @@ -19,10 +21,12 @@ import { hasBlockingIssues, isUpdateWithoutWhere, resolveConnectionString, - suffixWithLimit, + shouldAutoGenerateTitle, + trimTrailingSemicolons, } from './SQLEditor.utils' import type { DatabaseEventTrigger } from '@/data/database-event-triggers/database-event-triggers-query' import type { Database } from '@/data/read-replicas/replicas-query' +import { createRoleImpersonationState } from '@/state/role-impersonation-state' const buildTrigger = (overrides: Partial = {}): DatabaseEventTrigger => ({ oid: 1, @@ -37,150 +41,253 @@ const buildTrigger = (overrides: Partial = {}): DatabaseEv ...overrides, }) -describe('SQLEditor.utils.ts:checkIfAppendLimitRequired', () => { +describe('SQLEditor.utils.ts:trimTrailingSemicolons', () => { + test('removes a single trailing semicolon', () => { + const sql = safeSql`select * from countries;` + expect(trimTrailingSemicolons(sql)).toBe('select * from countries') + }) + test('removes multiple trailing semicolons', () => { + const sql = safeSql`select * from countries;;;;;;;` + expect(trimTrailingSemicolons(sql)).toBe('select * from countries') + }) + test('leaves a fragment with no trailing semicolon unchanged', () => { + const sql = safeSql`select * from countries` + expect(trimTrailingSemicolons(sql)).toBe('select * from countries') + }) + test('does not touch semicolons that are not trailing', () => { + const sql = safeSql`select 1; select 2` + expect(trimTrailingSemicolons(sql)).toBe('select 1; select 2') + }) +}) + +describe('SQLEditor.utils.ts:applyAutoLimit', () => { test('Should return false if limit passed is <= 0', () => { - const sql = 'select * from countries;' + const sql = safeSql`select * from countries;` const limit = -1 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return true if limit passed is > 0', () => { - const sql = 'select * from countries;' + const sql = safeSql`select * from countries;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(true) }) test('Should return false if query already has a limit', () => { - const sql = 'select * from countries limit 10;' + const sql = safeSql`select * from countries limit 10;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query already has a limit (check for case-insensitiveness)', () => { - const sql = 'SELECT * FROM countries LIMIT 10;' + const sql = safeSql`SELECT * FROM countries LIMIT 10;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query already has a limit with whitespace before the semi colon', () => { - const sql = 'select * from countries limit 10 ;' + const sql = safeSql`select * from countries limit 10 ;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query already has a limit and offset', () => { - const sql = 'select * from countries limit 10 offset 0;' + const sql = safeSql`select * from countries limit 10 offset 0;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query already has a limit and offset with whitespace before the semi colon', () => { - const sql = 'select * from countries limit 10 offset 0 ;' + const sql = safeSql`select * from countries limit 10 offset 0 ;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query already has a limit and offset (flip order of limit and offset)', () => { - const sql = 'select * from countries offset 0 limit 1;' + const sql = safeSql`select * from countries offset 0 limit 1;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query already has a limit, even if no value provided for limit', () => { - const sql = 'select * from countries limit' + const sql = safeSql`select * from countries limit` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query uses `FETCH FIRST` instead of limit ', () => { - const sql = 'select * from countries FETCH FIRST 5 rows only' + const sql = safeSql`select * from countries FETCH FIRST 5 rows only` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query uses `fetch first` instead of limit ', () => { - const sql = 'select * from countries fetch first 5 rows only' + const sql = safeSql`select * from countries fetch first 5 rows only` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query uses `fetch first` (with random spaces) instead of limit ', () => { - const sql = 'select * from countries FETCH FIRST 5 rows only' + const sql = safeSql`select * from countries FETCH FIRST 5 rows only` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query is not a select statement', () => { - const sql = 'create table test (id int8 primary key, name varchar);' + const sql = safeSql`create table test (id int8 primary key, name varchar);` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if there are multiple queries I', () => { - const sql1 = ` -select * from countries; -select * from cities; -`.trim() + const sql1 = safeSql`select * from countries; +select * from cities;` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql1, limit) + const { appendAutoLimit } = applyAutoLimit(sql1, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if there are multiple queries II', () => { - const sql1 = ` -select * from countries; -select * from cities -`.trim() + const sql1 = safeSql`select * from countries; +select * from cities` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql1, limit) + const { appendAutoLimit } = applyAutoLimit(sql1, limit) expect(appendAutoLimit).toBe(false) }) // [Joshen] Opting to just avoid appending in this case to prevent making the logic overly complex atm test('Should return false if query has with a comment I', () => { - const sql = ` --- This is a comment -select * from cities -`.trim() + const sql = safeSql`-- This is a comment +select * from cities` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) test('Should return false if query has with a comment II', () => { - const sql = ` -select * from cities --- This is a comment -`.trim() + const sql = safeSql`select * from cities +-- This is a comment` const limit = 100 - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) + const { appendAutoLimit } = applyAutoLimit(sql, limit) expect(appendAutoLimit).toBe(false) }) -}) -// [Joshen] These will just need to test the cases when appendAutoLimit returns true then -describe('SQLEditor.utils.ts:suffixWithLimit', () => { + // [Joshen] These will just need to test the cases when appendAutoLimit returns true then test('Should add the limit param properly if query ends without a semi colon', () => { const sql = safeSql`select * from countries` const limit = 100 - const formattedSql = suffixWithLimit(sql, limit) + const { sql: formattedSql } = applyAutoLimit(sql, limit) expect(formattedSql).toBe('select * from countries limit 100;') }) test('Should add the limit param properly if query ends with a semi colon', () => { const sql = safeSql`select * from countries;` const limit = 100 - const formattedSql = suffixWithLimit(sql, limit) + const { sql: formattedSql } = applyAutoLimit(sql, limit) expect(formattedSql).toBe('select * from countries limit 100;') }) test('Should add the limit param properly if query ends with multiple semi colon', () => { const sql = safeSql`select * from countries;;;;;;;` const limit = 100 - const formattedSql = suffixWithLimit(sql, limit) + const { sql: formattedSql } = applyAutoLimit(sql, limit) expect(formattedSql).toBe('select * from countries limit 100;') }) test('Should not append a limit if query already has one with whitespace before the semi colon', () => { const sql = safeSql`select * from countries limit 10 ;` const limit = 100 - const formattedSql = suffixWithLimit(sql, limit) + const { sql: formattedSql } = applyAutoLimit(sql, limit) expect(formattedSql).toBe('select * from countries limit 10 ;') }) + test('returns the SafeSqlFragment result unchanged when no limit is appended', () => { + const sql = safeSql`select * from countries limit 10;` + const { sql: formattedSql } = applyAutoLimit(sql, 100) + expect(formattedSql).toBe(sql) + }) +}) + +describe('SQLEditor.utils.ts:shouldAutoGenerateTitle', () => { + test('returns true when AI is enabled, the name is still the placeholder, and on platform', () => { + expect( + shouldAutoGenerateTitle({ + aiOptInLevel: 'default', + snippetName: `${untitledSnippetTitle} 1`, + isPlatform: true, + }) + ).toBe(true) + }) + test('returns false when AI opt-in is disabled', () => { + expect( + shouldAutoGenerateTitle({ + aiOptInLevel: 'disabled', + snippetName: `${untitledSnippetTitle} 1`, + isPlatform: true, + }) + ).toBe(false) + }) + test('returns false when the snippet already has a custom name', () => { + expect( + shouldAutoGenerateTitle({ + aiOptInLevel: 'default', + snippetName: 'My saved query', + isPlatform: true, + }) + ).toBe(false) + }) + test('returns false when not running on the hosted platform', () => { + expect( + shouldAutoGenerateTitle({ + aiOptInLevel: 'default', + snippetName: `${untitledSnippetTitle} 1`, + isPlatform: false, + }) + ).toBe(false) + }) + test('returns false when the snippet name is undefined', () => { + expect( + shouldAutoGenerateTitle({ + aiOptInLevel: 'default', + snippetName: undefined, + isPlatform: true, + }) + ).toBe(false) + }) +}) + +describe('SQLEditor.utils.ts:buildExecuteParams', () => { + const impersonatedRoleState = createRoleImpersonationState('default', { + current: async () => ({}), + }) + + test('applies the auto-limit suffix when appropriate', () => { + const params = buildExecuteParams({ + sql: safeSql`select * from countries`, + limit: 100, + connectionString: 'postgresql://example', + projectRef: 'default', + impersonatedRoleState, + }) + expect(params.sql).toBe('select * from countries limit 100;') + expect(params.autoLimit).toBe(100) + }) + test('leaves autoLimit undefined when a limit is already present', () => { + const params = buildExecuteParams({ + sql: safeSql`select * from countries limit 10;`, + limit: 100, + connectionString: 'postgresql://example', + projectRef: 'default', + impersonatedRoleState, + }) + expect(params.sql).toBe('select * from countries limit 10;') + expect(params.autoLimit).toBeUndefined() + }) + test('reflects isRoleImpersonationEnabled from the impersonation state', () => { + const params = buildExecuteParams({ + sql: safeSql`select * from countries`, + limit: 0, + connectionString: 'postgresql://example', + projectRef: 'default', + impersonatedRoleState, + }) + expect(params.isRoleImpersonationEnabled).toBe(false) + expect(params.isStatementTimeoutDisabled).toBe(true) + expect(params.contextualInvalidation).toBe(true) + }) }) describe(`SQLEditor.utils.ts:checkDestructiveQuery`, () => { diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts index 27454efa19ae1..ab35c5a25ae3b 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts @@ -1,4 +1,10 @@ -import { untrustedSql, type SafeSqlFragment, type UntrustedSqlFragment } from '@supabase/pg-meta' +import { + literal, + safeSql, + untrustedSql, + type SafeSqlFragment, + type UntrustedSqlFragment, +} from '@supabase/pg-meta' import { TABLE_EVENT_ACTIONS } from 'common/telemetry-constants' import { @@ -6,6 +12,7 @@ import { destructiveSqlRegex, NEW_SQL_SNIPPET_SKELETON, sqlAiDisclaimerComment, + untitledSnippetTitle, updateWithoutWhereRegex, } from './SQLEditor.constants' import { ContentDiff, type IStandaloneCodeEditor, type PotentialIssues } from './SQLEditor.types' @@ -14,7 +21,12 @@ import type { DatabaseEventTrigger } from '@/data/database-event-triggers/databa import type { Database } from '@/data/read-replicas/replicas-query' import { generateUuid } from '@/lib/api/snippets.browser' import { removeCommentsFromSql } from '@/lib/helpers' +import { wrapWithRoleImpersonation } from '@/lib/role-impersonation' import { sqlEventParser } from '@/lib/sql-event-parser' +import { + isRoleImpersonationEnabled, + type RoleImpersonationState, +} from '@/state/role-impersonation-state' export type CreateTableWithoutRLS = { schema?: string @@ -225,6 +237,59 @@ export function resolveConnectionString( ) } +/** + * Whether a query run should lazily kick off AI title generation for the + * snippet: only when the org has AI enabled (not disabled/HIPAA — which would + * silently forward the query to the AI provider without consent), the + * snippet still has its placeholder name, and we're running on the hosted + * platform. + */ +export function shouldAutoGenerateTitle({ + aiOptInLevel, + snippetName, + isPlatform, +}: { + aiOptInLevel: string + snippetName: string | undefined + isPlatform: boolean +}): boolean { + return ( + aiOptInLevel !== 'disabled' && !!snippetName?.startsWith(untitledSnippetTitle) && isPlatform + ) +} + +/** + * Builds the params passed to `useExecuteSqlMutation`'s `execute`: applies the + * auto-limit suffix and role impersonation to the SQL, and derives the + * `autoLimit`/`isRoleImpersonationEnabled` flags. Callers still attach their + * own `handleError`. + */ +export function buildExecuteParams({ + sql, + limit, + connectionString, + projectRef, + impersonatedRoleState, +}: { + sql: SafeSqlFragment + limit: number + connectionString: string | undefined + projectRef: string + impersonatedRoleState: RoleImpersonationState +}) { + const { sql: formattedSql, appendAutoLimit } = applyAutoLimit(sql, limit) + + return { + projectRef, + connectionString, + sql: wrapWithRoleImpersonation(formattedSql, impersonatedRoleState), + autoLimit: appendAutoLimit ? limit : undefined, + isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRoleState.role), + isStatementTimeoutDisabled: true as const, + contextualInvalidation: true as const, + } +} + export const generateMigrationCliCommand = (id: string, name: string, isNpx = false) => ` ${isNpx ? 'npx ' : ''}supabase snippets download ${id} | @@ -270,13 +335,35 @@ export const compareAsNewSnippet = (sqlDiff: ContentDiff) => { } } +/** + * Removes trailing `;` characters from a safe SQL fragment. Only ever removes + * existing terminators — never adds text — so the result is exactly as safe + * as the input; the brand carries over intentionally. This is the one place + * in the file allowed to reassert `SafeSqlFragment` on a derived string — + * every other function composes new fragments through `safeSql`/`literal`. + */ +export function trimTrailingSemicolons(sql: SafeSqlFragment): SafeSqlFragment { + return sql.replace(/;+\s*$/, '') as SafeSqlFragment +} + // [Joshen] Just FYI as well the checks here on whether to append limit is quite restricted // This is to prevent dashboard from accidentally appending limit to the end of a query // thats not supposed to have any, since there's too many cases to cover. // We can however look into making this logic better in the future // i.e It's harder to append the limit param, than just leaving the query as it is // Otherwise we'd need a full on parser to do this properly -export const checkIfAppendLimitRequired = (sql: string, limit: number = 0) => { +// +// Only accepts `SafeSqlFragment`: this decides whether to build (and builds) +// a new SQL fragment that gets executed, so every caller — including ones +// that only want the `appendAutoLimit` flag for a display hint — must already +// hold safe SQL. Composes the ` limit N;` suffix through `safeSql`/`literal` +// rather than gluing raw template-literal text onto the fragment and casting +// the result, so the only new content this function ever stamps safe is an +// internally-generated integer literal, never arbitrary concatenated text. +export function applyAutoLimit( + sql: SafeSqlFragment, + limit: number = 0 +): { sql: SafeSqlFragment; appendAutoLimit: boolean } { // Remove lines and whitespaces to use for checking const cleanedSql = sql.trim().replaceAll('\n', ' ').replaceAll(/\s+/g, ' ') @@ -299,15 +386,13 @@ export const checkIfAppendLimitRequired = (sql: string, limit: number = 0) => { !cleanedSql.match(/limit;$/i) && !cleanedSql.match(/limit [0-9]* offset [0-9]*\s*[;]?$/i) && !cleanedSql.match(/limit [0-9]*\s*[;]?$/i) - return { cleanedSql, appendAutoLimit } -} -export const suffixWithLimit = (sql: SafeSqlFragment, limit: number = 0): SafeSqlFragment => { - const { cleanedSql, appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) - if (!appendAutoLimit) return sql - return ( - cleanedSql.endsWith(';') ? sql.replace(/[;]+$/, ` limit ${limit};`) : `${sql} limit ${limit};` - ) as SafeSqlFragment + if (!appendAutoLimit) return { sql, appendAutoLimit: false } + + const core = cleanedSql.endsWith(';') ? trimTrailingSemicolons(sql) : sql + const suffixed = safeSql`${core} limit ${literal(limit)};` + + return { sql: suffixed, appendAutoLimit: true } } /** diff --git a/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts b/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts index 85d76becfe41f..1cc661c0db7da 100644 --- a/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts +++ b/apps/studio/components/interfaces/SQLEditor/useSqlEditorExecution.ts @@ -4,14 +4,13 @@ import { IS_PLATFORM, useParams } from 'common' import { useCallback, useState } from 'react' import { toast } from 'sonner' -import { untitledSnippetTitle } from './SQLEditor.constants' import type { PotentialIssues } from './SQLEditor.types' import { analyzeQueryIssues, - checkIfAppendLimitRequired, + buildExecuteParams, hasBlockingIssues, resolveConnectionString, - suffixWithLimit, + shouldAutoGenerateTitle, } from './SQLEditor.utils' import { useSQLEditorContext } from './SQLEditorContext' import { useDatabaseEventTriggersQuery } from '@/data/database-event-triggers/database-event-triggers-query' @@ -21,13 +20,9 @@ import { useReadReplicasQuery } from '@/data/read-replicas/replicas-query' import { useExecuteSqlMutation } from '@/data/sql/execute-sql-mutation' import { useOrgAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' -import { wrapWithRoleImpersonation } from '@/lib/role-impersonation' import { useTrack } from '@/lib/telemetry/track' import { useDatabaseSelectorStateSnapshot } from '@/state/database-selector' -import { - isRoleImpersonationEnabled, - useGetImpersonatedRoleState, -} from '@/state/role-impersonation-state' +import { useGetImpersonatedRoleState } from '@/state/role-impersonation-state' import { useSqlEditorSessionSnapshot } from '@/state/sql-editor/sql-editor-session-state' import { getSqlEditorV2StateSnapshot } from '@/state/sql-editor/sql-editor-state' @@ -115,11 +110,11 @@ export function useSqlEditorExecution({ // use the latest state for the title-generation check const snippet = getSqlEditorV2StateSnapshot().snippets[id] if ( - // Don't auto-generate a title when the org has disabled AI or is a HIPAA project, - // as that would silently forward the query to the AI provider without consent - aiOptInLevel !== 'disabled' && - snippet?.snippet.name.startsWith(untitledSnippetTitle) && - IS_PLATFORM + shouldAutoGenerateTitle({ + aiOptInLevel, + snippetName: snippet?.snippet.name, + isPlatform: IS_PLATFORM, + }) ) { // Intentionally don't await title gen (lazy) setAiTitle(id, sql) @@ -137,17 +132,14 @@ export function useSqlEditorExecution({ return toast.error('Unable to run query: Connection string is missing') } - const { appendAutoLimit } = checkIfAppendLimitRequired(sql, limit) - const formattedSql = suffixWithLimit(sql, limit) - execute({ - projectRef: project.ref, - connectionString: connectionString, - sql: wrapWithRoleImpersonation(formattedSql, impersonatedRoleState), - autoLimit: appendAutoLimit ? limit : undefined, - isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRoleState.role), - isStatementTimeoutDisabled: true, - contextualInvalidation: true, + ...buildExecuteParams({ + sql, + limit, + connectionString, + projectRef: project.ref, + impersonatedRoleState, + }), handleError: (error) => { throw error }, diff --git a/apps/studio/components/ui/EditorPanel/EditorPanel.tsx b/apps/studio/components/ui/EditorPanel/EditorPanel.tsx index 8a557c2b082b3..036c73653d315 100644 --- a/apps/studio/components/ui/EditorPanel/EditorPanel.tsx +++ b/apps/studio/components/ui/EditorPanel/EditorPanel.tsx @@ -46,8 +46,8 @@ import { SaveSnippetDialog } from './SaveSnippetDialog' import { isExplainQuery } from '@/components/interfaces/ExplainVisualizer/ExplainVisualizer.utils' import { generateSnippetTitle } from '@/components/interfaces/SQLEditor/SQLEditor.constants' import { + applyAutoLimit, createSqlSnippetSkeletonV2, - suffixWithLimit, } from '@/components/interfaces/SQLEditor/SQLEditor.utils' import { useAddDefinitions } from '@/components/interfaces/SQLEditor/useAddDefinitions' import { Results } from '@/components/interfaces/SQLEditor/UtilityPanel/Results' @@ -231,7 +231,7 @@ export const EditorPanel = () => { } executeSql({ - sql: suffixWithLimit(acceptUntrustedSql(currentValue), 100), + sql: applyAutoLimit(acceptUntrustedSql(currentValue), 100).sql, projectRef: project?.ref, connectionString: project?.connectionString, isStatementTimeoutDisabled: true, diff --git a/apps/studio/state/sql-editor/sql-editor-session-state.ts b/apps/studio/state/sql-editor/sql-editor-session-state.ts index 54e82f46c1d99..338e7eb365146 100644 --- a/apps/studio/state/sql-editor/sql-editor-session-state.ts +++ b/apps/studio/state/sql-editor/sql-editor-session-state.ts @@ -23,7 +23,7 @@ export const sqlEditorSessionState = proxy({ /** * UI-imposed limit for the number of rows a query can return (a safeguard * against accidentally taking down the database with a huge SELECT). Related - * to `autoLimit` in `results`; see `checkIfAppendLimitRequired`/`suffixWithLimit`. + * to `autoLimit` in `results`; see `applyAutoLimit`. */ limit: 100,