chore: move role impersonation state into a context (#19510)

* chore: move role impersonation state into a context

* clean up ProjectContext.tsx imports

* move isRoleImpersonationEnabled to only relevant queries

apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx

@@ -29,6 +29,7 @@ import { useSqlGenerateMutation } from 'data/ai/sql-generate-mutation'
 import { useSqlTitleGenerateMutation } from 'data/ai/sql-title-mutation'
 import { SqlSnippet } from 'data/content/sql-snippets-query'
 import { useEntityDefinitionsQuery } from 'data/database/entity-definitions-query'
+import { useReadReplicasQuery } from 'data/read-replicas/replicas-query'
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import { useFormatQueryMutation } from 'data/sql/format-sql-query'
 import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
@@ -46,8 +47,9 @@ import { uuidv4 } from 'lib/helpers'
 import { useProfile } from 'lib/profile'
 import { wrapWithRoleImpersonation } from 'lib/role-impersonation'
 import Telemetry from 'lib/telemetry'
+import toast from 'react-hot-toast'
 import { useAppStateSnapshot } from 'state/app-state'
-import { getImpersonatedRole } from 'state/role-impersonation-state'
+import { isRoleImpersonationEnabled, useGetImpersonatedRole } from 'state/role-impersonation-state'
 import { getSqlEditorStateSnapshot, useSqlEditorStateSnapshot } from 'state/sql-editor'
 import { subscriptionHasHipaaAddon } from '../Billing/Subscription/Subscription.utils'
 import AISchemaSuggestionPopover from './AISchemaSuggestionPopover'
@@ -66,8 +68,6 @@ import {
   getDiffTypeDropdownLabel,
 } from './SQLEditor.utils'
 import UtilityPanel from './UtilityPanel/UtilityPanel'
-import { useReadReplicasQuery } from 'data/read-replicas/replicas-query'
-import toast from 'react-hot-toast'
 
 // Load the monaco editor client-side only (does not behave well server-side)
 const MonacoEditor = dynamic(() => import('./MonacoEditor'), { ssr: false })
@@ -289,6 +289,8 @@ const SQLEditor = () => {
     }
   }, [formatQuery, id, isDiffOpen, project, snap])
 
+  const getImpersonatedRole = useGetImpersonatedRole()
+
   const executeQuery = useCallback(
     async (force: boolean = false) => {
       if (isDiffOpen) return
@@ -323,6 +325,7 @@ const SQLEditor = () => {
           setLineHighlights([])
         }
 
+        const impersonatedRole = getImpersonatedRole()
         const connectionString = !readReplicasEnabled
           ? project.connectionString
           : databases?.find((db) => db.identifier === snap.selectedDatabaseId)?.connectionString
@@ -335,12 +338,23 @@ const SQLEditor = () => {
           connectionString: connectionString,
           sql: wrapWithRoleImpersonation(sql, {
             projectRef: project.ref,
-            role: getImpersonatedRole(),
+            role: impersonatedRole,
           }),
+          isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRole),
         })
       }
     },
-    [isDiffOpen, id, isExecuting, project, execute, setAiTitle, hasHipaaAddon, supabaseAIEnabled]
+    [
+      isDiffOpen,
+      id,
+      isExecuting,
+      project,
+      supabaseAIEnabled,
+      hasHipaaAddon,
+      execute,
+      getImpersonatedRole,
+      setAiTitle,
+    ]
   )
 
   const handleNewQuery = useCallback(

apps/studio/components/interfaces/TableGridEditor/DeleteConfirmationDialogs.tsx

@@ -17,7 +17,7 @@ import { useGetTables } from 'data/tables/tables-query'
 import { useStore, useUrlState } from 'hooks'
 import { TableLike } from 'hooks/misc/useTable'
 import { noop } from 'lib/void'
-import { getImpersonatedRole } from 'state/role-impersonation-state'
+import { useGetImpersonatedRole } from 'state/role-impersonation-state'
 import { useTableEditorStateSnapshot } from 'state/table-editor'
 
 export type DeleteConfirmationDialogsProps = {
@@ -196,6 +196,8 @@ const DeleteConfirmationDialogs = ({
     }
   }
 
+  const getImpersonatedRole = useGetImpersonatedRole()
+
   const onConfirmDeleteRow = async () => {
     if (!project) return console.error('Project ref is required')
     if (!selectedTable) return console.error('Selected table required')

apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx

@@ -14,7 +14,7 @@ import { useTableRowCreateMutation } from 'data/table-rows/table-row-create-muta
 import { useTableRowUpdateMutation } from 'data/table-rows/table-row-update-mutation'
 import { tableKeys } from 'data/tables/keys'
 import { useStore, useUrlState } from 'hooks'
-import { getImpersonatedRole } from 'state/role-impersonation-state'
+import { useGetImpersonatedRole } from 'state/role-impersonation-state'
 import { useTableEditorStateSnapshot } from 'state/table-editor'
 import { ColumnEditor, RowEditor, SpreadsheetImport, TableEditor } from '.'
 import ForeignRowSelector from './RowEditor/ForeignRowSelector/ForeignRowSelector'
@@ -68,6 +68,8 @@ const SidePanelEditor = ({
     },
   })
 
+  const getImpersonatedRole = useGetImpersonatedRole()
+
   const saveRow = async (
     payload: any,
     isNewRecord: boolean,

apps/studio/components/interfaces/TableGridEditor/TableGridEditor.tsx

@@ -31,8 +31,7 @@ import useEntityType from 'hooks/misc/useEntityType'
 import { TableLike } from 'hooks/misc/useTable'
 import { EXCLUDED_SCHEMAS } from 'lib/constants/schemas'
 import { EMPTY_ARR } from 'lib/void'
-import { useAppStateSnapshot } from 'state/app-state'
-import { getImpersonatedRole } from 'state/role-impersonation-state'
+import { useGetImpersonatedRole } from 'state/role-impersonation-state'
 import { useTableEditorStateSnapshot } from 'state/table-editor'
 import { SchemaView } from 'types'
 import GridHeaderActions from './GridHeaderActions'
@@ -58,10 +57,11 @@ const TableGridEditor = ({
   const { ref: projectRef, id } = useParams()
 
   const { project } = useProjectContext()
-  const appSnap = useAppStateSnapshot()
   const snap = useTableEditorStateSnapshot()
   const gridRef = useRef<SupabaseGridRef>(null)
 
+  const getImpersonatedRole = useGetImpersonatedRole()
+
   const [encryptedColumns, setEncryptedColumns] = useState([])
 
   const [{ view: selectedView = 'data' }, setUrlState] = useUrlState()

apps/studio/components/layouts/ProjectLayout/ProjectContext.tsx

@@ -3,6 +3,7 @@ import { createContext, PropsWithChildren, useContext, useMemo } from 'react'
 import { useParams } from 'common/hooks'
 import { useProjectDetailQuery } from 'data/projects/project-detail-query'
 import { PROJECT_STATUS } from 'lib/constants'
+import { RoleImpersonationStateContextProvider } from 'state/role-impersonation-state'
 import { TableEditorStateContextProvider } from 'state/table-editor'
 import { Project } from 'types'
 
@@ -40,7 +41,9 @@ export const ProjectContextProvider = ({
   return (
     <ProjectContext.Provider value={value}>
       <TableEditorStateContextProvider key={`table-editor-state-${projectRef}`}>
-        {children}
+        <RoleImpersonationStateContextProvider key={`role-impersonation-state-${projectRef}`}>
+          {children}
+        </RoleImpersonationStateContextProvider>
       </TableEditorStateContextProvider>
     </ProjectContext.Provider>
   )

apps/studio/data/sql/execute-sql-query.ts

@@ -13,7 +13,6 @@ import {
   ROLE_IMPERSONATION_NO_RESULTS,
   ROLE_IMPERSONATION_SQL_LINE_COUNT,
 } from 'lib/role-impersonation'
-import { getRoleImpersonationStateSnapshot } from 'state/role-impersonation-state'
 import { sqlKeys } from './keys'
 
 export type Error = { code: number; message: string; requestId: string }
@@ -24,6 +23,7 @@ export type ExecuteSqlVariables = {
   sql: string
   queryKey?: QueryKey
   handleError?: (error: { code: number; message: string; requestId: string }) => any
+  isRoleImpersonationEnabled?: boolean
 }
 
 export async function executeSql(
@@ -33,9 +33,15 @@ export async function executeSql(
     sql,
     queryKey,
     handleError,
+    isRoleImpersonationEnabled = false,
   }: Pick<
     ExecuteSqlVariables,
-    'projectRef' | 'connectionString' | 'sql' | 'queryKey' | 'handleError'
+    | 'projectRef'
+    | 'connectionString'
+    | 'sql'
+    | 'queryKey'
+    | 'handleError'
+    | 'isRoleImpersonationEnabled'
   >,
   signal?: AbortSignal
 ) {
@@ -44,8 +50,6 @@ export async function executeSql(
   let headers = new Headers()
   if (connectionString) headers.set('x-connection-encrypted', connectionString)
 
-  const isRoleImpersonationEnabled = getRoleImpersonationStateSnapshot().role?.type === 'postgrest'
-
   let { data, error } = await post('/platform/pg-meta/{ref}/query', {
     signal,
     params: {
@@ -107,13 +111,23 @@ export type ExecuteSqlData = Awaited<ReturnType<typeof executeSql>>
 export type ExecuteSqlError = unknown
 
 export const useExecuteSqlQuery = <TData = ExecuteSqlData>(
-  { projectRef, connectionString, sql, queryKey, handleError }: ExecuteSqlVariables,
+  {
+    projectRef,
+    connectionString,
+    sql,
+    queryKey,
+    handleError,
+    isRoleImpersonationEnabled,
+  }: ExecuteSqlVariables,
   { enabled = true, ...options }: UseQueryOptions<ExecuteSqlData, ExecuteSqlError, TData> = {}
 ) =>
   useQuery<ExecuteSqlData, ExecuteSqlError, TData>(
     sqlKeys.query(projectRef, queryKey ?? [md5(sql)]),
     ({ signal }) =>
-      executeSql({ projectRef, connectionString, sql, queryKey, handleError }, signal),
+      executeSql(
+        { projectRef, connectionString, sql, queryKey, handleError, isRoleImpersonationEnabled },
+        signal
+      ),
     { enabled: enabled && typeof projectRef !== 'undefined', ...options }
   )
 

apps/studio/data/table-rows/table-row-create-mutation.ts

@@ -5,6 +5,7 @@ import { Query, SupaTable } from 'components/grid'
 import { executeSql } from 'data/sql/execute-sql-query'
 import { sqlKeys } from 'data/sql/keys'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { isRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import { ResponseError } from 'types'
 
 export type TableRowCreateVariables = {
@@ -43,7 +44,12 @@ export async function createTableRow({
     }
   )
 
-  const { result } = await executeSql({ projectRef, connectionString, sql })
+  const { result } = await executeSql({
+    projectRef,
+    connectionString,
+    sql,
+    isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRole),
+  })
 
   return result
 }

apps/studio/data/table-rows/table-row-delete-all-mutation.ts

@@ -5,6 +5,7 @@ import { Filter, Query, SupaTable } from 'components/grid'
 import { executeSql } from 'data/sql/execute-sql-query'
 import { sqlKeys } from 'data/sql/keys'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { isRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import { ResponseError } from 'types'
 import { formatFilterValue } from './utils'
 
@@ -44,7 +45,12 @@ export async function deleteAllTableRow({
     role: impersonatedRole,
   })
 
-  const { result } = await executeSql({ projectRef, connectionString, sql })
+  const { result } = await executeSql({
+    projectRef,
+    connectionString,
+    sql,
+    isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRole),
+  })
 
   return result
 }

apps/studio/data/table-rows/table-row-delete-mutation.ts

@@ -6,6 +6,7 @@ import { executeSql } from 'data/sql/execute-sql-query'
 import { sqlKeys } from 'data/sql/keys'
 import { Table } from 'data/tables/table-query'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { isRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import { ResponseError } from 'types'
 import { getPrimaryKeys } from './utils'
 
@@ -45,7 +46,12 @@ export async function deleteTableRow({
     role: impersonatedRole,
   })
 
-  const { result } = await executeSql({ projectRef, connectionString, sql })
+  const { result } = await executeSql({
+    projectRef,
+    connectionString,
+    sql,
+    isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRole),
+  })
 
   return result
 }

apps/studio/data/table-rows/table-row-truncate-mutation.ts

@@ -5,6 +5,7 @@ import { Query, SupaTable } from 'components/grid'
 import { executeSql } from 'data/sql/execute-sql-query'
 import { sqlKeys } from 'data/sql/keys'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { isRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import { ResponseError } from 'types'
 
 export type TableRowTruncateVariables = {
@@ -31,7 +32,12 @@ export async function truncateTableRow({
     role: impersonatedRole,
   })
 
-  const { result } = await executeSql({ projectRef, connectionString, sql })
+  const { result } = await executeSql({
+    projectRef,
+    connectionString,
+    sql,
+    isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRole),
+  })
 
   return result
 }

apps/studio/data/table-rows/table-row-update-mutation.ts

@@ -5,6 +5,7 @@ import { Query, SupaTable } from 'components/grid'
 import { executeSql } from 'data/sql/execute-sql-query'
 import { sqlKeys } from 'data/sql/keys'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { isRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import { ResponseError } from 'types'
 
 export type TableRowUpdateVariables = {
@@ -47,7 +48,12 @@ export async function updateTableRow({
     }
   )
 
-  const { result } = await executeSql({ projectRef, connectionString, sql })
+  const { result } = await executeSql({
+    projectRef,
+    connectionString,
+    sql,
+    isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRole),
+  })
 
   return result
 }

apps/studio/data/table-rows/table-rows-count-query.ts

@@ -3,6 +3,7 @@ import { useCallback } from 'react'
 
 import { Filter, Query, SupaTable } from 'components/grid'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { useIsRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import { ExecuteSqlData, useExecuteSqlPrefetch, useExecuteSqlQuery } from '../sql/execute-sql-query'
 import { formatFilterValue } from './utils'
 
@@ -55,8 +56,10 @@ export const useTableRowsCountQuery = <TData extends TableRowsCountData = TableR
     ...args
   }: TableRowsCountVariables,
   options: UseQueryOptions<ExecuteSqlData, TableRowsCountError, TData> = {}
-) =>
-  useExecuteSqlQuery(
+) => {
+  const isRoleImpersonationEnabled = useIsRoleImpersonationEnabled()
+
+  return useExecuteSqlQuery(
     {
       projectRef,
       connectionString,
@@ -72,6 +75,7 @@ export const useTableRowsCountQuery = <TData extends TableRowsCountData = TableR
           ...args,
         },
       ],
+      isRoleImpersonationEnabled,
     },
     {
       select(data) {
@@ -83,6 +87,7 @@ export const useTableRowsCountQuery = <TData extends TableRowsCountData = TableR
       ...options,
     }
   )
+}
 
 /**
  * useTableRowsCountPrefetch is used for prefetching the table rows count. For example, starting a query loading before a page is navigated to.

apps/studio/data/table-rows/table-rows-query.ts

@@ -2,6 +2,7 @@ import { QueryKey, UseQueryOptions } from '@tanstack/react-query'
 import { Filter, Query, Sort, SupaRow, SupaTable } from 'components/grid'
 import { ImpersonationRole, wrapWithRoleImpersonation } from 'lib/role-impersonation'
 import { useCallback } from 'react'
+import { useIsRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import {
   ExecuteSqlData,
   executeSql,
@@ -140,8 +141,10 @@ export type TableRowsError = unknown
 export const useTableRowsQuery = <TData extends TableRowsData = TableRowsData>(
   { projectRef, connectionString, queryKey, table, impersonatedRole, ...args }: TableRowsVariables,
   options: UseQueryOptions<ExecuteSqlData, TableRowsError, TData> = {}
-) =>
-  useExecuteSqlQuery(
+) => {
+  const isRoleImpersonationEnabled = useIsRoleImpersonationEnabled()
+
+  return useExecuteSqlQuery(
     {
       projectRef,
       connectionString,
@@ -157,6 +160,7 @@ export const useTableRowsQuery = <TData extends TableRowsData = TableRowsData>(
           ...args,
         },
       ],
+      isRoleImpersonationEnabled,
     },
     {
       select(data) {
@@ -172,6 +176,7 @@ export const useTableRowsQuery = <TData extends TableRowsData = TableRowsData>(
       ...options,
     }
   )
+}
 
 /**
  * useTableRowsPrefetch is used for prefetching table rows. For example, starting a query loading before a page is navigated to.

apps/studio/pages/project/[ref]/api/graphiql.tsx

@@ -15,7 +15,7 @@ import { useProjectPostgrestConfigQuery } from 'data/config/project-postgrest-co
 import { useStore } from 'hooks'
 import { API_URL, IS_PLATFORM } from 'lib/constants'
 import { getRoleImpersonationJWT } from 'lib/role-impersonation'
-import { getImpersonatedRole } from 'state/role-impersonation-state'
+import { useGetImpersonatedRole } from 'state/role-impersonation-state'
 import { NextPageWithLayout } from 'types'
 
 const GraphiQLPage: NextPageWithLayout = () => {
@@ -45,6 +45,8 @@ const GraphiQLPage: NextPageWithLayout = () => {
     }
   }, [ui.selectedProjectRef])
 
+  const getImpersonatedRole = useGetImpersonatedRole()
+
   const fetcher = useMemo(() => {
     const fetcherFn = createGraphiQLFetcher({
       url: `${API_URL}/projects/${projectRef}/api/graphql`,
@@ -78,7 +80,7 @@ const GraphiQLPage: NextPageWithLayout = () => {
     }
 
     return customFetcher
-  }, [projectRef, jwtSecret, accessToken, serviceRoleKey])
+  }, [projectRef, getImpersonatedRole, jwtSecret, accessToken, serviceRoleKey])
 
   if ((IS_PLATFORM && !accessToken) || !isFetched || (isExtensionsLoading && !pgGraphqlExtension)) {
     return <Connecting />

apps/studio/state/role-impersonation-state.ts

@@ -1,18 +0,0 @@
-import { proxy, snapshot, useSnapshot } from 'valtio'
-import { ImpersonationRole } from 'lib/role-impersonation'
-
-export const roleImpersonationState = proxy({
-  role: undefined as ImpersonationRole | undefined,
-  setRole: (role: ImpersonationRole | undefined) => {
-    roleImpersonationState.role = role
-  },
-})
-
-export const getRoleImpersonationStateSnapshot = () => snapshot(roleImpersonationState)
-
-export const useRoleImpersonationStateSnapshot = (options?: Parameters<typeof useSnapshot>[1]) =>
-  useSnapshot(roleImpersonationState, options)
-
-export function getImpersonatedRole() {
-  return getRoleImpersonationStateSnapshot().role
-}

apps/studio/state/role-impersonation-state.tsx

@@ -0,0 +1,52 @@
+import { PropsWithChildren, createContext, useCallback, useContext } from 'react'
+import { proxy, snapshot, useSnapshot } from 'valtio'
+
+import { useConstant } from 'common'
+import { ImpersonationRole } from 'lib/role-impersonation'
+
+export function createRoleImpersonationState() {
+  const roleImpersonationState = proxy({
+    role: undefined as ImpersonationRole | undefined,
+    setRole: (role: ImpersonationRole | undefined) => {
+      roleImpersonationState.role = role
+    },
+  })
+
+  return roleImpersonationState
+}
+
+export type RoleImpersonationState = ReturnType<typeof createRoleImpersonationState>
+
+export const RoleImpersonationStateContext = createContext<RoleImpersonationState>(
+  createRoleImpersonationState()
+)
+
+export const RoleImpersonationStateContextProvider = ({ children }: PropsWithChildren) => {
+  const state = useConstant(createRoleImpersonationState)
+
+  return (
+    <RoleImpersonationStateContext.Provider value={state}>
+      {children}
+    </RoleImpersonationStateContext.Provider>
+  )
+}
+
+export function useRoleImpersonationStateSnapshot(options?: Parameters<typeof useSnapshot>[1]) {
+  const roleImpersonationState = useContext(RoleImpersonationStateContext)
+
+  return useSnapshot(roleImpersonationState, options)
+}
+
+export function useGetImpersonatedRole() {
+  const roleImpersonationState = useContext(RoleImpersonationStateContext)
+
+  return useCallback(() => snapshot(roleImpersonationState).role, [roleImpersonationState])
+}
+
+export function isRoleImpersonationEnabled(impersonationRole?: ImpersonationRole) {
+  return impersonationRole?.type === 'postgrest'
+}
+
+export function useIsRoleImpersonationEnabled() {
+  return isRoleImpersonationEnabled(useRoleImpersonationStateSnapshot().role)
+}

packages/common/hooks/index.ts

@@ -1,4 +1,5 @@
 export * from './useBreakpoint'
+export * from './useConstant'
 export * from './useDebounce'
 export * from './useParams'
 export * from './useTelemetryProps'

packages/common/hooks/useConstant.ts

@@ -0,0 +1,18 @@
+// based on https://github.com/Andarist/use-constant
+
+import { useRef } from 'react'
+
+type ResultBox<T> = { v: T }
+
+/**
+ * React hook for creating a value exactly once
+ */
+export function useConstant<T>(fn: () => T): T {
+  const ref = useRef<ResultBox<T>>()
+
+  if (!ref.current) {
+    ref.current = { v: fn() }
+  }
+
+  return ref.current.v
+}