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

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

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 } }
+  }) => (
+    <div data-testid="impersonated-role">
+      {roleImpersonationState?.role?.type === 'postgrest'
+        ? roleImpersonationState.role.role
+        : 'none'}
+    </div>
+  ),
+}))
 
 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(
+      <TabsStateContext.Provider value={createTabsState('default')}>
+        <QueryTab />
+      </TabsStateContext.Provider>
+    )
+
+    expect(await screen.findByTestId('impersonated-role')).toHaveTextContent('none')
+
+    await act(async () => {
+      explorerQueryState.removeDraft({ id: 'query-test-2', projectRef: 'default' })
+    })
+  })
 })

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<string>()
 
@@ -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
 

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)
 }

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)

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<string, PersistedExplorerQueryDraft>
@@ -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

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<typeof createRoleImpersonationState>
 
-/**
- * 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<typeof us
  * A role impersonation controller scoped to a single component instance instead of the
  * shared project-wide context — for surfaces (e.g. a notebook query cell) that need their
  * own independent "run as" selection rather than sharing the one global impersonation state.
+ *
+ * [Joshen] FWIW we might deprecate this method after hooking things up E2E since QueryCell's
+ * role impersonation likely needs to be persisted also
  */
 export function useLocalRoleImpersonationState(): RoleImpersonationController {
   const { data: project } = useSelectedProjectQuery()
@@ -172,6 +173,86 @@ export function useLocalRoleImpersonationState(): RoleImpersonationController {
   return { role, claims, setRole }
 }
 
+/**
+ * Like `useLocalRoleImpersonationState`, but the role selection is controlled externally
+ * (e.g. persisted per query tab) instead of held in local component state — for surfaces
+ * where the selection needs to survive beyond this component instance's lifetime. Claims
+ * stay local regardless: they're derived, time-bound tokens, so they're recomputed whenever
+ * the controlled role changes rather than persisted alongside it.
+ */
+export function useControlledRoleImpersonationState(
+  role: ImpersonationRole | undefined,
+  onRoleChange: (role: ImpersonationRole | undefined) => 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<PostgrestClaims | undefined>(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)