From bfab3090f52841d31b028f061237811b6f93edba Mon Sep 17 00:00:00 2001 From: Joshen Lim Date: Tue, 18 Aug 2026 11:27:09 +0800 Subject: [PATCH] QueryTab: Scope role impersonation to each tab instead of global (#49139) ## Context Previous PR [here](https://github.com/supabase/supabase/pull/49101) introduced role impersonation to the Explorer -> Query Tab, but the setting was global (e.g selected role would be the same despite switching query tabs) Changes here shifts the scope of the role impersonation into the query draft so that the value is tied to each individual query tab instead ## Summary by CodeRabbit * **New Features** * Query drafts now remember selected impersonated roles when switching between drafts or returning later. * Added support for clearing saved impersonated roles. * Impersonation state remains isolated across query tabs. * **Bug Fixes** * Prevented impersonation settings from carrying over between unrelated drafts. * Invalid saved role data is safely ignored during restoration. * Logs drafts no longer persist or update impersonated roles. --- .../interfaces/Explorer/QueryTab.test.tsx | 40 +++++++- .../interfaces/Explorer/QueryTab.tsx | 15 ++- apps/studio/lib/role-impersonation.ts | 53 +++++++++++ apps/studio/state/explorer-query.test.ts | 64 +++++++++++++ apps/studio/state/explorer-query.ts | 26 +++++ .../studio/state/role-impersonation-state.tsx | 95 +++++++++++++++++-- 6 files changed, 282 insertions(+), 11 deletions(-) diff --git a/apps/studio/components/interfaces/Explorer/QueryTab.test.tsx b/apps/studio/components/interfaces/Explorer/QueryTab.test.tsx index 61c303017ad..21da1d78574 100644 --- a/apps/studio/components/interfaces/Explorer/QueryTab.test.tsx +++ b/apps/studio/components/interfaces/Explorer/QueryTab.test.tsx @@ -32,7 +32,19 @@ vi.mock('@/components/ui/CodeEditor/CodeEditor', () => ({ ), })) -vi.mock('./ExplorerQuerySourceMenu', () => ({ ExplorerQuerySourceMenu: () => null })) +vi.mock('./ExplorerQuerySourceMenu', () => ({ + ExplorerQuerySourceMenu: ({ + roleImpersonationState, + }: { + roleImpersonationState?: { role?: { type: string; role?: string } } + }) => ( +
+ {roleImpersonationState?.role?.type === 'postgrest' + ? roleImpersonationState.role.role + : 'none'} +
+ ), +})) const renderQueryTab = () => customRender( @@ -168,4 +180,30 @@ describe('QueryTab execution', () => { new Date(bodies[0].iso_timestamp_start).getTime() ).toBe(2 * 60 * 60 * 1000) }) + + it('isolates the impersonated role selection per query tab', async () => { + createDraft({ _tag: 'database' }) + explorerQueryState.removeDraft({ id: 'query-test-2', projectRef: 'default' }) + explorerQueryState.createDraft({ id: 'query-test-2', projectRef: 'default', sql: 'select 2' }) + explorerQueryState.setRole({ + id: 'query-test', + role: { type: 'postgrest', role: 'service_role' }, + }) + + const { rerender } = renderQueryTab() + expect(await screen.findByTestId('impersonated-role')).toHaveTextContent('service_role') + + testContext.params = { ref: 'default', id: 'query-test-2' } + rerender( + + + + ) + + expect(await screen.findByTestId('impersonated-role')).toHaveTextContent('none') + + await act(async () => { + explorerQueryState.removeDraft({ id: 'query-test-2', projectRef: 'default' }) + }) + }) }) diff --git a/apps/studio/components/interfaces/Explorer/QueryTab.tsx b/apps/studio/components/interfaces/Explorer/QueryTab.tsx index 00027068624..306b03b4171 100644 --- a/apps/studio/components/interfaces/Explorer/QueryTab.tsx +++ b/apps/studio/components/interfaces/Explorer/QueryTab.tsx @@ -1,14 +1,14 @@ import { useParams } from 'common' import { Loader2, SquareCode } from 'lucide-react' import { useRouter } from 'next/router' -import { useContext, useEffect, useState } from 'react' +import { useCallback, useContext, useEffect, useState } from 'react' import { Button } from 'ui' import { QueryEditor, type ExplorerQueryModel } from './QueryEditor' import { type QueryResult } from './types' import { toQuerySourceBinding } from '@/data/query-sources/query-source-registry' import { explorerQueryState, useExplorerQueryStateSnapshot } from '@/state/explorer-query' -import { useLocalRoleImpersonationState } from '@/state/role-impersonation-state' +import { useControlledRoleImpersonationState } from '@/state/role-impersonation-state' import { createTabId, TabsStateContext } from '@/state/tabs' /** Query-tab lifecycle adapter around the shared QueryEditor. */ @@ -17,7 +17,6 @@ export const QueryTab = () => { const router = useRouter() const tabs = useContext(TabsStateContext) const querySnap = useExplorerQueryStateSnapshot() - const roleImpersonationState = useLocalRoleImpersonationState() const [restoredQueryKey, setRestoredQueryKey] = useState() @@ -26,6 +25,16 @@ export const QueryTab = () => { const result = draft && id ? querySnap.results[id] : undefined const queryKey = id && ref ? `${ref}:${id}` : undefined + const roleImpersonationState = useControlledRoleImpersonationState( + draft?._tag === 'database' ? draft.role : undefined, + useCallback( + (role) => { + if (id) explorerQueryState.setRole({ id, role }) + }, + [id] + ) + ) + useEffect(() => { if (!id || !ref) return diff --git a/apps/studio/lib/role-impersonation.ts b/apps/studio/lib/role-impersonation.ts index 0f6353f12b7..96fd694deaf 100644 --- a/apps/studio/lib/role-impersonation.ts +++ b/apps/studio/lib/role-impersonation.ts @@ -1,4 +1,5 @@ import { getImpersonationSQL, type SafeSqlFragment } from '@supabase/pg-meta' +import { z } from 'zod' import { uuidv4 } from './helpers' import type { User } from '@/data/auth/users-infinite-query' @@ -40,6 +41,58 @@ type CustomImpersonationRole = { export type ImpersonationRole = PostgrestImpersonationRole | CustomImpersonationRole +/** + * The impersonated `user` is the same generated `User` shape already persisted verbatim to + * localStorage elsewhere (see `USER_IMPERSONATION_SELECTOR_PREVIOUS_SEARCHES`) — trusted + * as-is rather than re-validated field-by-field, since it only ever round-trips our own + * writes and its shape tracks a generated API type this schema shouldn't have to mirror. + */ +const impersonatedUserSchema = z + .record(z.string(), z.unknown()) + .transform((value) => value as unknown as User) + +const aalSchema = z.enum(['aal1', 'aal2']) + +const postgrestImpersonationRoleSchema = z.union([ + z.object({ type: z.literal('postgrest'), role: z.literal('anon') }).strict(), + z.object({ type: z.literal('postgrest'), role: z.literal('service_role') }).strict(), + z + .object({ + type: z.literal('postgrest'), + role: z.literal('authenticated'), + userType: z.literal('native'), + user: impersonatedUserSchema.optional(), + aal: aalSchema.optional(), + }) + .strict(), + z + .object({ + type: z.literal('postgrest'), + role: z.literal('authenticated'), + userType: z.literal('external'), + externalAuth: z + .object({ + sub: z.string(), + additionalClaims: z.record(z.string(), z.unknown()).optional(), + }) + .optional(), + aal: aalSchema.optional(), + }) + .strict(), +]) + +const customImpersonationRoleSchema = z + .object({ type: z.literal('custom'), role: z.string() }) + .strict() + +/** Parses to `ImpersonationRole` — verified at the `role` field assignment in `toDraft` + * (`state/explorer-query.ts`), since annotating the schema type directly here would also + * constrain its *input* type, which is narrower than `ImpersonationRole` pre-transform. */ +export const impersonationRoleSchema = z.union([ + postgrestImpersonationRoleSchema, + customImpersonationRoleSchema, +]) + export function getExp1HourFromNow() { return Math.floor((Date.now() + 60 * 60 * 1000) / 1000) } diff --git a/apps/studio/state/explorer-query.test.ts b/apps/studio/state/explorer-query.test.ts index 5a8c954e391..d3268a2606f 100644 --- a/apps/studio/state/explorer-query.test.ts +++ b/apps/studio/state/explorer-query.test.ts @@ -293,6 +293,70 @@ describe('explorer query drafts', () => { expect(state.drafts['query-1']).toMatchObject({ rowLimit: 100, uncheckedSql: 'select 1' }) }) + it('persists and restores a per-draft impersonated role independently of other drafts', () => { + const storage = createMemoryStorage() + const state = createExplorerQueryState(storage) + + state.createDraft({ id: 'query-1', projectRef: 'project-a' }) + state.createDraft({ id: 'query-2', projectRef: 'project-a' }) + state.setRole({ id: 'query-1', role: { type: 'postgrest', role: 'anon' } }) + + expect(state.drafts['query-1']).toMatchObject({ role: { type: 'postgrest', role: 'anon' } }) + expect(state.drafts['query-2']).toMatchObject({ role: undefined }) + + const restored = createExplorerQueryState(storage) + expect(restored.restoreDraft({ id: 'query-1', projectRef: 'project-a' })).toBe(true) + expect(restored.restoreDraft({ id: 'query-2', projectRef: 'project-a' })).toBe(true) + expect(restored.drafts['query-1']).toMatchObject({ + role: { type: 'postgrest', role: 'anon' }, + }) + expect(restored.drafts['query-2']).toMatchObject({ role: undefined }) + }) + + it('clears a persisted role when set back to undefined', () => { + const storage = createMemoryStorage() + const state = createExplorerQueryState(storage) + + state.createDraft({ id: 'query-1', projectRef: 'project-a' }) + state.setRole({ id: 'query-1', role: { type: 'postgrest', role: 'service_role' } }) + state.setRole({ id: 'query-1', role: undefined }) + + expect(state.drafts['query-1']).toMatchObject({ role: undefined }) + + const restored = createExplorerQueryState(storage) + expect(restored.restoreDraft({ id: 'query-1', projectRef: 'project-a' })).toBe(true) + expect(restored.drafts['query-1']).toMatchObject({ role: undefined }) + }) + + it('drops a malformed persisted role rather than failing to restore the draft', () => { + const storage = createMemoryStorage() + storage.setItem( + LOCAL_STORAGE_KEYS.EXPLORER_QUERY_DRAFTS('project-a'), + JSON.stringify({ + 'query-1': { + name: 'Query with bad role', + sql: 'select 1', + updatedAt: 1, + role: { type: 'postgrest', role: 'not-a-real-role' }, + }, + }) + ) + + const state = createExplorerQueryState(storage) + expect(state.restoreDraft({ id: 'query-1', projectRef: 'project-a' })).toBe(true) + expect(state.drafts['query-1']).toMatchObject({ role: undefined }) + }) + + it('ignores role updates for logs drafts', () => { + const storage = createMemoryStorage() + const state = createExplorerQueryState(storage) + + state.createDraft({ id: 'query-1', projectRef: 'project-a', source: LOGS_SOURCE }) + state.setRole({ id: 'query-1', role: { type: 'postgrest', role: 'anon' } }) + + expect(state.drafts['query-1']).not.toHaveProperty('role') + }) + it('removes the persisted draft and its session result when its tab closes', () => { const storage = createMemoryStorage() const state = createExplorerQueryState(storage) diff --git a/apps/studio/state/explorer-query.ts b/apps/studio/state/explorer-query.ts index 5e6f6cb6898..eb15cd9f3bf 100644 --- a/apps/studio/state/explorer-query.ts +++ b/apps/studio/state/explorer-query.ts @@ -17,6 +17,7 @@ import { toQuerySourceBinding, type QuerySourceBinding, } from '@/data/query-sources/query-source-registry' +import { impersonationRoleSchema, type ImpersonationRole } from '@/lib/role-impersonation' type ExplorerQueryDraftBase = { id: string @@ -36,6 +37,7 @@ export type DatabaseQueryDraft = ExplorerQueryDraftBase & _tag: 'database' uncheckedSql: UntrustedSqlFragment rowLimit: number + role?: ImpersonationRole } export type LogsQueryDraft = ExplorerQueryDraftBase & @@ -61,6 +63,7 @@ type PersistedExplorerQueryDraft = { sql: string updatedAt: number rowLimit?: number + role?: ImpersonationRole } type PersistedExplorerQueryDrafts = Record @@ -77,6 +80,7 @@ const persistedDraftSchema = z.object({ updatedAt: z.number(), source: z.unknown().optional(), rowLimit: z.unknown().optional(), + role: z.unknown().optional(), }) const VALID_ROW_LIMITS = ROWS_PER_PAGE_OPTIONS.map((option) => option.value) @@ -124,6 +128,7 @@ const toDraft = ({ database_identifier: persisted.source.database_identifier, uncheckedSql: untrustedSql(persisted.sql), rowLimit: persisted.rowLimit ?? DEFAULT_CELL_ROW_LIMIT, + role: persisted.role, } } @@ -145,6 +150,8 @@ const readPersistedDrafts = (storage: StorageLike, projectRef: string) => { ? parsedSource.data : createDefaultSourceBinding('database') + const parsedRole = impersonationRoleSchema.safeParse(draft.data.role) + const role = parsedRole.success ? parsedRole.data : undefined const rowLimit = rowLimitSchema.parse(draft.data.rowLimit) return [ @@ -155,6 +162,7 @@ const readPersistedDrafts = (storage: StorageLike, projectRef: string) => { source, sql: draft.data.sql, updatedAt: draft.data.updatedAt, + role, rowLimit, }, ], @@ -196,6 +204,7 @@ export const createExplorerQueryState = (storage: StorageLike = safeLocalStorage sql: draft.uncheckedSql, updatedAt: draft.updatedAt, rowLimit: draft._tag === 'database' ? draft.rowLimit : undefined, + role: draft._tag === 'database' ? draft.role : undefined, } writePersistedDrafts(storage, draft.projectRef, persisted) } @@ -279,6 +288,7 @@ export const createExplorerQueryState = (storage: StorageLike = safeLocalStorage if (nextSource !== undefined && nextSource._tag !== draft._tag) delete state.results[id] const currentRowLimit = draft._tag === 'database' ? draft.rowLimit : undefined + const currentRole = draft._tag === 'database' ? draft.role : undefined state.drafts[id] = toDraft({ id, @@ -289,6 +299,7 @@ export const createExplorerQueryState = (storage: StorageLike = safeLocalStorage sql: sql ?? draft.uncheckedSql, updatedAt: Date.now(), rowLimit: rowLimit ?? currentRowLimit, + role: currentRole, }, }) @@ -337,6 +348,21 @@ export const createExplorerQueryState = (storage: StorageLike = safeLocalStorage setResult: ({ id, result }: { id: string; result: ExplorerQueryResult }) => { state.results[id] = ref(result) }, + + /** + * Separate from `updateDraft` because `undefined` is a meaningful value here (clearing + * the impersonated role), whereas `updateDraft`'s optional fields all use `undefined` + * to mean "leave unchanged." Logs drafts have no impersonation concept, so this is a + * no-op for them. + */ + setRole: ({ id, role }: { id: string; role: ImpersonationRole | undefined }) => { + const draft = state.drafts[id] + if (!draft || draft._tag !== 'database') return + + const updatedDraft: DatabaseQueryDraft = { ...draft, role, updatedAt: Date.now() } + state.drafts[id] = updatedDraft + persistDraft(updatedDraft) + }, }) return state diff --git a/apps/studio/state/role-impersonation-state.tsx b/apps/studio/state/role-impersonation-state.tsx index e1ba37494ff..4142c61d696 100644 --- a/apps/studio/state/role-impersonation-state.tsx +++ b/apps/studio/state/role-impersonation-state.tsx @@ -6,11 +6,15 @@ import { useCallback, useContext, useEffect, + useRef, useState, } from 'react' import { proxy, snapshot, subscribe, useSnapshot } from 'valtio' -import { CustomAccessTokenHookDetails } from '../hooks/misc/useCustomAccessTokenHookDetails' +import { + CustomAccessTokenHookDetails, + useCustomAccessTokenHookDetails, +} from '../hooks/misc/useCustomAccessTokenHookDetails' import { executeSql } from '@/data/sql/execute-sql-mutation' import { useLatest } from '@/hooks/misc/useLatest' import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' @@ -100,12 +104,6 @@ export function createRoleImpersonationState( export type RoleImpersonationState = ReturnType -/** - * The subset of `RoleImpersonationState` a role-picker UI needs: the current selection, its - * resolved claims, and the setter. Satisfied by both the shared project-wide context (via - * `useRoleImpersonationStateSnapshot`) and `useLocalRoleImpersonationState`, so role-picking - * components can work against either without knowing which one they got. - */ export type RoleImpersonationController = Pick< RoleImpersonationState, 'role' | 'claims' | 'setRole' @@ -141,6 +139,9 @@ export function useRoleImpersonationStateSnapshot(options?: Parameters void +): RoleImpersonationController { + const { data: project } = useSelectedProjectQuery() + const customizeAccessToken = useCustomizeAccessToken(project?.ref, project?.connectionString) + + const projectRef = project?.ref ?? '' + const customizeAccessTokenRef = useLatest(customizeAccessToken) + const customAccessTokenHookDetails = useCustomAccessTokenHookDetails(project?.ref) + const customAccessTokenHookDetailsRef = useLatest(customAccessTokenHookDetails) + const onRoleChangeRef = useLatest(onRoleChange) + + const [claims, setClaims] = useState(undefined) + + // Guards against re-resolving claims for a role change that `setRole` below just resolved + // itself — without it, every selection would re-run the (possibly RPC-backed) resolution + // twice: once eagerly in `setRole`, once again here once `role` updates on the next render. + const skipNextResolveRef = useRef(false) + + useEffect(() => { + if (skipNextResolveRef.current) { + skipNextResolveRef.current = false + return + } + + let cancelled = false + + resolveRoleClaims( + projectRef, + role, + customAccessTokenHookDetailsRef.current, + customizeAccessTokenRef.current + ).then((nextClaims) => { + if (!cancelled) setClaims(nextClaims) + }) + + return () => { + cancelled = true + } + // Resolves only when the controlled role identity changes (e.g. switching tabs) + }, [customAccessTokenHookDetailsRef, customizeAccessTokenRef, projectRef, role]) + + const setRole = useCallback( + async ( + nextRole: ImpersonationRole | undefined, + customAccessTokenHookDetails?: CustomAccessTokenHookDetails + ) => { + // Captured before the await: if the controlling tab changes while this resolution is + // in flight, `onRoleChangeRef.current` will point at a different tab's callback by the + // time we get here. Comparing against the captured reference lets us detect that and + // discard the result instead of writing this role/claims into the wrong tab. + const onRoleChangeAtCallTime = onRoleChangeRef.current + + const nextClaims = await resolveRoleClaims( + projectRef, + nextRole, + customAccessTokenHookDetails ?? customAccessTokenHookDetailsRef.current, + customizeAccessTokenRef.current + ) + + if (onRoleChangeRef.current !== onRoleChangeAtCallTime) return + + skipNextResolveRef.current = true + onRoleChangeAtCallTime(nextRole) + setClaims(nextClaims) + }, + [projectRef, customizeAccessTokenRef, customAccessTokenHookDetailsRef, onRoleChangeRef] + ) + + return { role, claims, setRole } +} + export function useGetImpersonatedRoleState() { const roleImpersonationState = useContext(RoleImpersonationStateContext)