diff --git a/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx b/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx
index df28a376973..680f7584ae9 100644
--- a/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx
+++ b/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx
@@ -1,24 +1,16 @@
 import { PermissionAction } from '@supabase/shared-types/out/constants'
 import { noop } from 'lodash'
 import { Lock, Unlock } from 'lucide-react'
+import Link from 'next/link'
+import { useRouter } from 'next/router'
 
 import { useParams } from 'common'
 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,
-  Button,
-  DropdownMenu,
-  DropdownMenuContent,
-  DropdownMenuItem,
-  DropdownMenuTrigger,
-} from 'ui'
-import { useRouter } from 'next/router'
-import Link from 'next/link'
 import { useFlag } from 'hooks/ui/useFlag'
+import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state'
+import { AiIconAnimation, Badge, Button } from 'ui'
 
 interface PolicyTableRowHeaderProps {
   table: {
diff --git a/apps/studio/components/interfaces/TableGridEditor/GridHeaderActions.tsx b/apps/studio/components/interfaces/TableGridEditor/GridHeaderActions.tsx
index ffa60b56b99..bd1d036ac96 100644
--- a/apps/studio/components/interfaces/TableGridEditor/GridHeaderActions.tsx
+++ b/apps/studio/components/interfaces/TableGridEditor/GridHeaderActions.tsx
@@ -1,7 +1,7 @@
 import { PermissionAction } from '@supabase/shared-types/out/constants'
 import { Lock, MousePointer2, PlusCircle, Unlock } from 'lucide-react'
 import Link from 'next/link'
-import { useState } from 'react'
+import { useEffect, useState } from 'react'
 import { toast } from 'sonner'
 
 import { useParams } from 'common'
@@ -24,17 +24,16 @@ import { useTableUpdateMutation } from 'data/tables/table-update-mutation'
 import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
 import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useIsFeatureEnabled } from 'hooks/misc/useIsFeatureEnabled'
-import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
+import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { PROTECTED_SCHEMAS } from 'lib/constants/schemas'
-import { useTableEditorTableStateSnapshot } from 'state/table-editor-table'
+import { useQueryState } from 'nuqs'
 import { useTableEditorStateSnapshot } from 'state/table-editor'
+import { useTableEditorTableStateSnapshot } from 'state/table-editor-table'
 import {
   Button,
-  Label_Shadcn_ as Label,
   PopoverContent_Shadcn_,
   PopoverTrigger_Shadcn_,
   Popover_Shadcn_,
-  Switch,
   Tooltip,
   TooltipContent,
   TooltipTrigger,
@@ -44,8 +43,6 @@ 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
@@ -54,15 +51,11 @@ export interface GridHeaderActionsProps {
 const GridHeaderActions = ({ table }: GridHeaderActionsProps) => {
   const { ref } = useParams()
   const { project } = useProjectContext()
-  const org = useSelectedOrganization()
-  const router = useRouter()
-  const tableEditorSnap = useTableEditorStateSnapshot()
+  const [isTestMode] = useQueryState('isTestMode')
+  const [impersonate] = useQueryState('impersonate')
 
-  useEffect(() => {
-    if (router.query.isTestMode === 'true') {
-      tableEditorSnap.setIsTestMode(true)
-    }
-  }, [router.query.isTestMode])
+  const { data: org } = useSelectedOrganizationQuery()
+  const tableEditorSnap = useTableEditorStateSnapshot()
 
   // need project lints to get security status for views
   const { data: lints = [] } = useProjectLintsQuery({ projectRef: project?.ref })
@@ -197,6 +190,13 @@ const GridHeaderActions = ({ table }: GridHeaderActionsProps) => {
     })
   }
 
+  useEffect(() => {
+    if (isTestMode === 'true') {
+      tableEditorSnap.setIsTestMode(true)
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [isTestMode])
+
   return (
     <div className="sb-grid-header__inner">
       {showHeaderActions && (
@@ -467,10 +467,8 @@ const GridHeaderActions = ({ table }: GridHeaderActionsProps) => {
           {doesHaveAutoGeneratedAPIDocs && <APIDocsButton section={['entities', table.name]} />}
           <RoleImpersonationPopover
             serviceRoleLabel="postgres"
-            initialOpen={router.query.impersonate === 'true'}
-            initialSelectedOption={
-              router.query.impersonate === 'true' ? 'authenticated' : undefined
-            }
+            initialOpen={impersonate === 'true'}
+            initialSelectedOption={impersonate === 'true' ? 'authenticated' : undefined}
           />
         </div>
       )}
diff --git a/apps/studio/components/layouts/TransactionBanner.tsx b/apps/studio/components/layouts/TransactionBanner.tsx
index abe7695f7e3..be232b28fe6 100644
--- a/apps/studio/components/layouts/TransactionBanner.tsx
+++ b/apps/studio/components/layouts/TransactionBanner.tsx
@@ -1,21 +1,20 @@
 import { useTableEditorStateSnapshot } from 'state/table-editor'
-import { Button } from 'ui'
+import { Button, cn } from 'ui'
 
 export const TransactionBanner = ({ sticky = false }: { sticky?: boolean }) => {
   const tableEditorSnap = useTableEditorStateSnapshot()
 
-  if (!tableEditorSnap.isTestMode) {
-    return null
-  }
+  if (!tableEditorSnap.isTestMode) return null
 
   return (
     <div
-      className={`flex items-center gap-x-3 bg-warning py-2 px-3 text-warning-100 ${
-        sticky ? 'sticky top-0 z-10' : ''
-      }`}
+      className={cn(
+        'flex items-center gap-x-3 bg-warning py-2 px-3 text-warning-100',
+        sticky && 'sticky top-0 z-10'
+      )}
     >
       <p className="text-xs flex-1">
-        You are in test mode. Any data changes you make will be rolled back.
+        You are in test mode. Any data changes you make will be rolled back immediately.
       </p>
       <Button
         type="text"
diff --git a/apps/studio/data/table-rows/table-row-update-mutation.ts b/apps/studio/data/table-rows/table-row-update-mutation.ts
index 5e8d4792acf..e4ca9a99d17 100644
--- a/apps/studio/data/table-rows/table-row-update-mutation.ts
+++ b/apps/studio/data/table-rows/table-row-update-mutation.ts
@@ -3,11 +3,11 @@ import { toast } from 'sonner'
 
 import { Query } from '@supabase/pg-meta/src/query'
 import { executeSql } from 'data/sql/execute-sql-query'
+import { wrapWithTransaction } from 'lib/helpers'
 import { RoleImpersonationState, wrapWithRoleImpersonation } from 'lib/role-impersonation'
 import { isRoleImpersonationEnabled } from 'state/role-impersonation-state'
 import type { ResponseError } from 'types'
 import { tableRowKeys } from './keys'
-import { wrapWithTransaction } from 'lib/helpers'
 
 export type TableRowUpdateVariables = {
   projectRef: string
@@ -58,8 +58,6 @@ export async function updateTableRow({
     sql = wrapWithTransaction(sql)
   }
 
-  console.log('sql', sql)
-
   const { result } = await executeSql({
     projectRef,
     connectionString,