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)