add test mode

apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx

@@ -7,7 +7,17 @@ import { ButtonTooltip } from 'components/ui/ButtonTooltip'
 import { EditorTablePageLink } from 'data/prefetchers/project.$ref.editor.$id'
 import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state'
-import { AiIconAnimation, Badge } from 'ui'
+import {
+  AiIconAnimation,
+  Badge,
+  Button,
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from 'ui'
+import { useRouter } from 'next/router'
+import Link from 'next/link'
 
 interface PolicyTableRowHeaderProps {
   table: {
@@ -33,6 +43,7 @@ const PolicyTableRowHeader = ({
   onSelectCreatePolicy,
 }: PolicyTableRowHeaderProps) => {
   const { ref } = useParams()
+  const router = useRouter()
   const aiSnap = useAiAssistantStateSnapshot()
 
   const canCreatePolicies = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_WRITE, 'policies')
@@ -42,6 +53,10 @@ const PolicyTableRowHeader = ({
   const isRealtimeMessagesTable = isRealtimeSchema && table.name === 'messages'
   const isTableLocked = isRealtimeSchema ? !isRealtimeMessagesTable : isLocked
 
+  const handleTestItemClick = (operation: string) => {
+    router.push(`/project/${ref}/editor/${table.id}?impersonate=true&as=${operation}`)
+  }
+
   return (
     <div id={table.id.toString()} className="flex w-full items-center justify-between">
       <div className="flex gap-x-4 text-left">
@@ -91,6 +106,13 @@ const PolicyTableRowHeader = ({
                 {table.rls_enabled ? 'Disable RLS' : 'Enable RLS'}
               </ButtonTooltip>
             )}
+            <Button asChild type="default">
+              <Link
+                href={`/project/${ref}/editor/${table.id}?impersonate=true&hasTransaction=true`}
+              >
+                Test
+              </Link>
+            </Button>
             <ButtonTooltip
               type="default"
               disabled={!canToggleRLS || !canCreatePolicies}

apps/studio/components/interfaces/RoleImpersonationSelector/RoleImpersonationPopover.tsx

@@ -16,6 +16,8 @@ export interface RoleImpersonationPopoverProps {
   disabled?: boolean
   disabledTooltip?: ReactNode
   disallowAuthenticatedOption?: boolean
+  initialOpen?: boolean
+  initialSelectedOption?: 'authenticated' | 'anon' | 'service_role'
 }
 
 const RoleImpersonationPopover = ({
@@ -24,6 +26,8 @@ const RoleImpersonationPopover = ({
   variant = 'regular',
   align = 'end',
   disallowAuthenticatedOption = false,
+  initialOpen = false,
+  initialSelectedOption,
 
   // [Joshen] We can clean these up once the API keys fix is done
   disabled = false,
@@ -31,7 +35,7 @@ const RoleImpersonationPopover = ({
 }: RoleImpersonationPopoverProps) => {
   const state = useRoleImpersonationStateSnapshot()
 
-  const [isOpen, setIsOpen] = useState(false)
+  const [isOpen, setIsOpen] = useState(initialOpen)
 
   const currentRole = state.role?.role ?? serviceRoleLabel ?? 'service role'
 
@@ -53,7 +57,7 @@ const RoleImpersonationPopover = ({
           }}
         >
           <div className="flex items-center gap-1">
-            <span className="text-foreground-muted">Role</span>
+            <span className="text-foreground-lighter">Role</span>
             <span>{currentRole}</span>
             {state.role?.type === 'postgrest' && state.role.role === 'authenticated' && (
               <>
@@ -62,7 +66,7 @@ const RoleImpersonationPopover = ({
                 ) : state.role.userType === 'external' && state.role.externalAuth ? (
                   <ExternalAuthButtonSection sub={state.role.externalAuth.sub} />
                 ) : null}
-                <span className="text-xs text-foreground-lighter font-light">
+                <span className="text-xs text-foreground-light font-light">
                   {state.role.aal === 'aal2' ? 'AAL2' : 'AAL1'}
                 </span>
               </>
@@ -80,6 +84,7 @@ const RoleImpersonationPopover = ({
         <RoleImpersonationSelector
           serviceRoleLabel={serviceRoleLabel}
           disallowAuthenticatedOption={disallowAuthenticatedOption}
+          initialSelectedOption={initialSelectedOption}
         />
       </PopoverContent_Shadcn_>
     </Popover_Shadcn_>

apps/studio/components/interfaces/RoleImpersonationSelector/RoleImpersonationSelector.tsx

@@ -11,16 +11,21 @@ export interface RoleImpersonationSelectorProps {
   serviceRoleLabel?: string
   padded?: boolean
   disallowAuthenticatedOption?: boolean
+  initialSelectedOption?: PostgrestRole
 }
 
 const RoleImpersonationSelector = ({
   serviceRoleLabel,
   padded = true,
   disallowAuthenticatedOption = false,
+  initialSelectedOption,
 }: RoleImpersonationSelectorProps) => {
   const state = useRoleImpersonationStateSnapshot()
 
   const [selectedOption, setSelectedOption] = useState<PostgrestRole | undefined>(() => {
+    if (initialSelectedOption) {
+      return initialSelectedOption
+    }
     if (
       state.role?.type === 'postgrest' &&
       (state.role.role === 'anon' || state.role.role === 'authenticated')

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

@@ -74,7 +74,11 @@ const DeleteConfirmationDialogs = ({
       if (snap.confirmationDialog?.type === 'row') {
         snap.confirmationDialog.callback?.()
       }
-      toast.success(`Successfully deleted selected row(s)`)
+      if (getImpersonatedRoleState().hasTransaction) {
+        toast.success(`Successfully deleted selected row(s) and rolled back`)
+      } else {
+        toast.success(`Successfully deleted selected row(s)`)
+      }
     },
     onSettled: () => {
       snap.closeConfirmationDialog()
@@ -86,7 +90,11 @@ const DeleteConfirmationDialogs = ({
       if (snap.confirmationDialog?.type === 'row') {
         snap.confirmationDialog.callback?.()
       }
-      toast.success(`Successfully deleted selected rows`)
+      if (getImpersonatedRoleState().hasTransaction) {
+        toast.success(`Successfully deleted selected rows and rolled back`)
+      } else {
+        toast.success(`Successfully deleted selected rows`)
+      }
     },
     onError: (error) => {
       toast.error(`Failed to delete rows: ${error.message}`)

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

@@ -26,12 +26,15 @@ import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useIsFeatureEnabled } from 'hooks/misc/useIsFeatureEnabled'
 import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
 import { PROTECTED_SCHEMAS } from 'lib/constants/schemas'
+import { useRoleImpersonationStateSnapshot } from 'state/role-impersonation-state'
 import { useTableEditorTableStateSnapshot } from 'state/table-editor-table'
 import {
   Button,
+  Label_Shadcn_ as Label,
   PopoverContent_Shadcn_,
   PopoverTrigger_Shadcn_,
   Popover_Shadcn_,
+  Switch,
   Tooltip,
   TooltipContent,
   TooltipTrigger,
@@ -41,6 +44,8 @@ import ConfirmModal from 'ui-patterns/Dialogs/ConfirmDialog'
 import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
 import { RoleImpersonationPopover } from '../RoleImpersonationSelector'
 import ViewEntityAutofixSecurityModal from './ViewEntityAutofixSecurityModal'
+import { useEffect } from 'react'
+import { useRouter } from 'next/router'
 
 export interface GridHeaderActionsProps {
   table: Entity
@@ -50,6 +55,15 @@ const GridHeaderActions = ({ table }: GridHeaderActionsProps) => {
   const { ref } = useParams()
   const { project } = useProjectContext()
   const org = useSelectedOrganization()
+  const router = useRouter()
+  const { hasTransaction, setHasTransaction } = useRoleImpersonationStateSnapshot()
+
+  useEffect(() => {
+    const { impersonate } = router.query
+    if (impersonate === 'true') {
+      setHasTransaction(true)
+    }
+  }, [router.query.impersonate])
 
   // need project lints to get security status for views
   const { data: lints = [] } = useProjectLintsQuery({ projectRef: project?.ref })
@@ -425,7 +439,14 @@ const GridHeaderActions = ({ table }: GridHeaderActionsProps) => {
               </PopoverContent_Shadcn_>
             </Popover_Shadcn_>
           )}
-          <RoleImpersonationPopover serviceRoleLabel="postgres" />
+
+          <RoleImpersonationPopover
+            serviceRoleLabel="postgres"
+            initialOpen={router.query.impersonate === 'true'}
+            initialSelectedOption={
+              router.query.impersonate === 'true' ? 'authenticated' : undefined
+            }
+          />
           {isTable && realtimeEnabled && (
             <Button
               type="default"
@@ -441,6 +462,16 @@ const GridHeaderActions = ({ table }: GridHeaderActionsProps) => {
             </Button>
           )}
           {doesHaveAutoGeneratedAPIDocs && <APIDocsButton section={['entities', table.name]} />}
+          <div className="flex items-center gap-x-2">
+            <Label htmlFor="transaction-mode-switch" className="text-foreground-lighter text-xs">
+              Test mode
+            </Label>
+            <Switch
+              id="transaction-mode-switch"
+              checked={hasTransaction}
+              onCheckedChange={setHasTransaction}
+            />
+          </div>
         </div>
       )}
       <ConfirmationModal

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

@@ -21,6 +21,8 @@ import {
   validateFields,
 } from './RowEditor.utils'
 import { TextEditor } from './TextEditor'
+import { useRoleImpersonationStateSnapshot } from 'state/role-impersonation-state'
+import { AlertTriangle } from 'lucide-react'
 
 export interface RowEditorProps {
   row?: Dictionary<any>
@@ -43,6 +45,7 @@ const RowEditor = ({
 }: RowEditorProps) => {
   const [errors, setErrors] = useState<Dictionary<any>>({})
   const [rowFields, setRowFields] = useState<RowField[]>([])
+  const { hasTransaction } = useRoleImpersonationStateSnapshot()
 
   const [selectedValueForTextEdit, setSelectedValueForTextEdit] = useState<EditValue>()
   const [selectedValueForJsonEdit, setSelectedValueForJsonEdit] = useState<EditValue>()
@@ -167,6 +170,11 @@ const RowEditor = ({
       }`}
       onCancel={closePanel}
     >
+      {hasTransaction && (
+        <div className="flex items-center gap-x-2 bg-warning py-2 px-6 text-black sticky top-0 z-10">
+          <p className="text-xs">You are in test mode. Any changes you make will be rolled back.</p>
+        </div>
+      )}
       <form onSubmit={(e) => onSaveChanges(e)} className="h-full">
         <div className="flex h-full flex-col">
           <div className="flex flex-grow flex-col">

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

@@ -50,6 +50,7 @@ import {
 import SpreadsheetImport from './SpreadsheetImport/SpreadsheetImport'
 import TableEditor from './TableEditor/TableEditor'
 import type { ImportContent } from './TableEditor/TableEditor.types'
+import { useRoleImpersonationStateSnapshot } from 'state/role-impersonation-state'
 
 export interface SidePanelEditorProps {
   editable?: boolean
@@ -74,7 +75,7 @@ const SidePanelEditor = ({
 
   const queryClient = useQueryClient()
   const { project } = useProjectContext()
-
+  const { hasTransaction } = useRoleImpersonationStateSnapshot()
   const [isEdited, setIsEdited] = useState<boolean>(false)
   const [isClosingPanel, setIsClosingPanel] = useState<boolean>(false)
 
@@ -89,11 +90,7 @@ const SidePanelEditor = ({
       toast.success('Successfully created row')
     },
   })
-  const { mutateAsync: updateTableRow } = useTableRowUpdateMutation({
-    onSuccess() {
-      toast.success('Successfully updated row')
-    },
-  })
+  const { mutateAsync: updateTableRow } = useTableRowUpdateMutation()
   const { data: publications } = useDatabasePublicationsQuery({
     projectRef: project?.ref,
     connectionString: project?.connectionString,
@@ -158,6 +155,11 @@ const SidePanelEditor = ({
     onComplete(saveRowError)
     if (!saveRowError) {
       setIsEdited(false)
+      if (hasTransaction) {
+        toast.success('Successfully updated row and rolled back')
+      } else {
+        toast.success('Successfully updated row')
+      }
       snap.closeSidePanel()
     }
   }

apps/studio/components/layouts/editors/EditorBaseLayout.tsx

@@ -8,6 +8,8 @@ import { ProjectLayoutWithAuth } from '../ProjectLayout/ProjectLayout'
 import { CollapseButton } from '../Tabs/CollapseButton'
 import { EditorTabs } from '../Tabs/Tabs'
 import { useEditorType } from './EditorsLayout.hooks'
+import { useRoleImpersonationStateSnapshot } from 'state/role-impersonation-state'
+import { AlertTriangle } from 'lucide-react'
 
 export interface ExplorerLayoutProps extends ComponentProps<typeof ProjectLayoutWithAuth> {
   children: ReactNode
@@ -21,6 +23,7 @@ export const EditorBaseLayout = ({ children, title, product, ...props }: Explore
   const pathname = usePathname()
   const editor = useEditorType()
   const tabs = useTabsStateSnapshot()
+  const { hasTransaction } = useRoleImpersonationStateSnapshot()
 
   const hasNoOpenTabs =
     editor === 'table' ? tabs.openTabs.filter((x) => !x.startsWith('sql')).length === 0 : false
@@ -29,7 +32,19 @@ export const EditorBaseLayout = ({ children, title, product, ...props }: Explore
 
   return (
     <ProjectLayoutWithAuth resizableSidebar title={title} product={product} {...props}>
-      <div className="flex flex-col h-full">
+      <div
+        className={cn(
+          'flex flex-col h-full',
+          hasTransaction && 'border-2 border-t-0 border-warning'
+        )}
+      >
+        {hasTransaction && (
+          <div className="flex items-center gap-x-2 bg-warning py-2 px-3 text-warning-foreground">
+            <p className="text-xs text-warning-100">
+              You are in test mode. Any changes to database records will be rolled back.
+            </p>
+          </div>
+        )}
         <div
           className={cn(
             'h-10 flex items-center',

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

@@ -41,7 +41,8 @@ export async function createTableRow({
 }: TableRowCreateVariables) {
   const sql = wrapWithRoleImpersonation(
     getTableRowCreateSql({ table, payload, enumArrayColumns, returning }),
-    roleImpersonationState
+    roleImpersonationState,
+    true
   )
 
   const { result } = await executeSql({

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

@@ -44,7 +44,8 @@ export async function deleteAllTableRow({
 }: TableRowDeleteAllVariables) {
   const sql = wrapWithRoleImpersonation(
     getTableRowDeleteAllSql({ table, filters }),
-    roleImpersonationState
+    roleImpersonationState,
+    true
   )
 
   const { result } = await executeSql({

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

@@ -46,7 +46,8 @@ export async function deleteTableRow({
 }: TableRowDeleteVariables) {
   const sql = wrapWithRoleImpersonation(
     getTableRowDeleteSql({ table, rows }),
-    roleImpersonationState
+    roleImpersonationState,
+    true
   )
 
   const { result } = await executeSql({

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

@@ -48,7 +48,8 @@ export async function updateTableRow({
 }: TableRowUpdateVariables) {
   const sql = wrapWithRoleImpersonation(
     getTableRowUpdateSql({ table, configuration, payload, enumArrayColumns, returning }),
-    roleImpersonationState
+    roleImpersonationState,
+    true
   )
 
   const { result } = await executeSql({

apps/studio/lib/role-impersonation.ts

@@ -117,13 +117,26 @@ function getCustomRoleImpersonationSql(roleName: string) {
   `.trim()
 }
 
-export type RoleImpersonationState = Pick<ValtioRoleImpersonationState, 'role' | 'claims'>
+export type RoleImpersonationState = Pick<
+  ValtioRoleImpersonationState,
+  'role' | 'claims' | 'hasTransaction'
+>
 
-export function wrapWithRoleImpersonation(sql: string, state?: RoleImpersonationState) {
-  const { role, claims } = state ?? { role: undefined, claims: undefined }
+export function wrapWithRoleImpersonation(
+  sql: string,
+  state?: RoleImpersonationState,
+  isMutation: boolean = false
+) {
+  const { role, claims, hasTransaction } = state ?? {
+    role: undefined,
+    claims: undefined,
+    hasTransaction: false,
+  }
+
+  const shouldTransaction = isMutation && hasTransaction
 
   if (role === undefined) {
-    return sql
+    return shouldTransaction ? wrapWithTransaction(sql) : sql
   }
 
   const impersonationSql =
@@ -133,15 +146,20 @@ export function wrapWithRoleImpersonation(sql: string, state?: RoleImpersonation
         : ''
       : getCustomRoleImpersonationSql(role.role)
 
-  return /* SQL */ `
+  const sqlWithImpersonation = `
     ${impersonationSql}
-
     -- If the users sql returns no rows, pg-meta will
     -- fallback to returning the result of the impersonation sql.
     select 1 as "${ROLE_IMPERSONATION_NO_RESULTS}";
 
     ${sql}
   `.trim()
+
+  return shouldTransaction ? wrapWithTransaction(sqlWithImpersonation) : sqlWithImpersonation
+}
+
+function wrapWithTransaction(sql: string) {
+  return `begin;${sql}rollback;`
 }
 
 function encodeText(data: string) {

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

@@ -22,6 +22,7 @@ export function createRoleImpersonationState(
     projectRef,
     role: undefined as ImpersonationRole | undefined,
     claims: undefined as ReturnType<typeof getPostgrestClaims> | undefined,
+    hasTransaction: false,
 
     setRole: async (
       role: ImpersonationRole | undefined,
@@ -46,6 +47,9 @@ export function createRoleImpersonationState(
         roleImpersonationState.claims = claims
       }
     },
+    setHasTransaction: (hasTransaction: boolean) => {
+      roleImpersonationState.hasTransaction = hasTransaction
+    },
   })
 
   return roleImpersonationState