From 225d4613814dc828d6431ec4646525e1acbd8e55 Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Fri, 8 Nov 2024 18:03:30 +0800
Subject: [PATCH] Hook up RLS to new assistant if feature flag is enabled
 (#30357)

* Hook up RLS to new assistant if feature flag is enabled

* Reset conversation if opening a different entity

* Add prompt to use alter policy if policy definition is provided

* Remove console log
---
 .../AIPolicyEditorPanel.utils.ts              |  3 +-
 .../LockedQuerySection.tsx                    | 47 --------------
 .../Policies/AIPolicyEditorPanel/index.tsx    |  2 -
 .../Policies/PolicyTableRow/PolicyRow.tsx     | 26 +++++++-
 .../PolicyTableRow/PolicyTableRowHeader.tsx   | 61 +++++++++--------
 .../Functions/FunctionsList/FunctionList.tsx  | 28 +++++---
 .../ui/AIAssistantPanel/AIAssistant.tsx       | 65 ++++++++++++++-----
 .../ui/AIAssistantPanel/AIAssistant.utils.ts  | 34 +++++++++-
 .../ui/AIAssistantPanel/AIAssistantPanel.tsx  | 13 ++--
 .../AIAssistantPanel/EntitiesDropdownMenu.tsx | 10 +--
 .../data/database/entity-definition-query.ts  | 59 +++++++++++++++--
 apps/studio/pages/api/ai/sql/generate-v2.ts   | 36 ++++++----
 apps/studio/pages/api/ai/sql/suggest.ts       |  2 +-
 apps/studio/state/app-state.ts                |  5 +-
 packages/ai-commands/src/sql/chat.ts          |  2 +
 packages/ai-commands/src/sql/rls.ts           |  2 +
 16 files changed, 254 insertions(+), 141 deletions(-)

diff --git a/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils.ts b/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils.ts
index e68356d2d98..937eaaa4f1d 100644
--- a/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils.ts
+++ b/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils.ts
@@ -1,7 +1,8 @@
 import type { PostgresPolicy } from '@supabase/postgres-meta'
+import { isEqual } from 'lodash'
+
 import type { Message } from 'ai/react'
 import { uuidv4 } from 'lib/helpers'
-import { isEqual } from 'lodash'
 
 export type MessageWithDebug = Message & { isDebug: boolean }
 
diff --git a/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/LockedQuerySection.tsx b/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/LockedQuerySection.tsx
index 429de771a3e..5f4e045b7d1 100644
--- a/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/LockedQuerySection.tsx
+++ b/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/LockedQuerySection.tsx
@@ -1,29 +1,17 @@
 import { PostgresPolicy } from '@supabase/postgres-meta'
 import { Lock } from 'lucide-react'
-import { useRouter } from 'next/router'
-
-import { useParams } from 'common'
-import { Button } from 'ui'
-import { generateAlterPolicyQuery, generateCreatePolicyQuery } from './AIPolicyEditorPanel.utils'
 
 interface LockedCreateQuerySection {
   schema: string
   selectedPolicy?: PostgresPolicy
   formFields: { name: string; table: string; behavior: string; command: string; roles: string }
-  editorOneRef: any
-  editorTwoRef: any
 }
 
 export const LockedCreateQuerySection = ({
   schema,
   selectedPolicy,
   formFields,
-  editorOneRef,
-  editorTwoRef,
 }: LockedCreateQuerySection) => {
-  const router = useRouter()
-  const { ref } = useParams()
-
   const isEditing = selectedPolicy !== undefined
   const { name, table, behavior, command, roles } = formFields
 
@@ -38,41 +26,6 @@ export const LockedCreateQuerySection = ({
             Use options above to edit
           </p>
         </div>
-        {/* <Button
-          type="default"
-          onClick={() => {
-            const query = isEditing
-              ? generateAlterPolicyQuery({
-                  name: '',
-                  newName: name,
-                  schema,
-                  table,
-                  command,
-                  roles: roles.length === 0 ? 'public' : roles,
-                  using: (editorOneRef.current?.getValue() ?? undefined)?.trim(),
-                  check:
-                    command === 'insert'
-                      ? (editorOneRef.current?.getValue() ?? undefined)?.trim()
-                      : (editorTwoRef.current?.getValue() ?? undefined)?.trim(),
-                })
-              : generateCreatePolicyQuery({
-                  name,
-                  schema,
-                  table,
-                  behavior,
-                  command,
-                  roles: roles.length === 0 ? 'public' : roles,
-                  using: (editorOneRef.current?.getValue() ?? undefined)?.trim(),
-                  check:
-                    command === 'insert'
-                      ? (editorOneRef.current?.getValue() ?? undefined)?.trim()
-                      : (editorTwoRef.current?.getValue() ?? undefined)?.trim(),
-                })
-            router.push(`/project/${ref}/sql/new?content=${query}`)
-          }}
-        >
-          Open in SQL Editor
-        </Button> */}
       </div>
       <div className="flex items-start" style={{ fontSize: '14px' }}>
         <p className="px-6 font-mono text-sm text-foreground-light select-none">1</p>
diff --git a/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/index.tsx b/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/index.tsx
index ba6aa9ab2ae..62dfc833efb 100644
--- a/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/index.tsx
+++ b/apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/index.tsx
@@ -595,8 +595,6 @@ export const AIPolicyEditorPanel = memo(function ({
                         <LockedCreateQuerySection
                           schema={schema}
                           selectedPolicy={selectedPolicy}
-                          editorOneRef={editorOneRef}
-                          editorTwoRef={editorTwoRef}
                           formFields={{ name, table, behavior, command, roles }}
                         />
 
diff --git a/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyRow.tsx b/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyRow.tsx
index d6607237fea..a00bc79cba1 100644
--- a/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyRow.tsx
+++ b/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyRow.tsx
@@ -21,6 +21,8 @@ import {
   TooltipContent_Shadcn_,
   TooltipTrigger_Shadcn_,
 } from 'ui'
+import { useIsDatabaseFunctionsAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
+import { useAppStateSnapshot } from 'state/app-state'
 
 interface PolicyRowProps {
   policy: PostgresPolicy
@@ -33,6 +35,8 @@ const PolicyRow = ({
   onSelectEditPolicy = noop,
   onSelectDeletePolicy = noop,
 }: PolicyRowProps) => {
+  const { setAiAssistantPanel } = useAppStateSnapshot()
+  const enableAssistantV2 = useIsDatabaseFunctionsAssistantEnabled()
   const canUpdatePolicies = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_WRITE, 'policies')
 
   const { project } = useProjectContext()
@@ -88,11 +92,31 @@ const PolicyRow = ({
           <DropdownMenuTrigger asChild>
             <Button type="default" className="px-1.5" icon={<MoreVertical />} />
           </DropdownMenuTrigger>
-          <DropdownMenuContent side="bottom" align="end" className="w-40">
+          <DropdownMenuContent
+            side="bottom"
+            align="end"
+            className={cn(enableAssistantV2 ? 'w-52' : 'w-40')}
+          >
             <DropdownMenuItem className="gap-x-2" onClick={() => onSelectEditPolicy(policy)}>
               <Edit size={14} />
               <p>Edit policy</p>
             </DropdownMenuItem>
+            {enableAssistantV2 && (
+              <DropdownMenuItem
+                className="space-x-2"
+                onClick={() => {
+                  setAiAssistantPanel({
+                    open: true,
+                    editor: 'rls-policies',
+                    entity: policy,
+                    tables: [{ schema: policy.schema, name: policy.table }],
+                  })
+                }}
+              >
+                <Edit size={14} />
+                <p>Edit policy with Assistant</p>
+              </DropdownMenuItem>
+            )}
             <DropdownMenuSeparator />
             <DropdownMenuItemTooltip
               className="gap-x-2"
diff --git a/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx b/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx
index 9c4edfa0dd2..c75be788248 100644
--- a/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx
+++ b/apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx
@@ -1,20 +1,15 @@
+import { PermissionAction } from '@supabase/shared-types/out/constants'
 import { noop } from 'lodash'
 import { Lock, Unlock } from 'lucide-react'
 import { useQueryState } from 'nuqs'
 
-import { PermissionAction } from '@supabase/shared-types/out/constants'
 import { useParams } from 'common'
+import { useIsDatabaseFunctionsAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
 import { ButtonTooltip } from 'components/ui/ButtonTooltip'
 import { EditorTablePageLink } from 'data/prefetchers/project.$ref.editor.$id'
 import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
-import {
-  AiIconAnimation,
-  Badge,
-  Button,
-  TooltipContent_Shadcn_,
-  TooltipTrigger_Shadcn_,
-  Tooltip_Shadcn_,
-} from 'ui'
+import { useAppStateSnapshot } from 'state/app-state'
+import { AiIconAnimation, Badge } from 'ui'
 
 interface PolicyTableRowHeaderProps {
   table: {
@@ -40,7 +35,9 @@ const PolicyTableRowHeader = ({
   onSelectCreatePolicy,
 }: PolicyTableRowHeaderProps) => {
   const { ref } = useParams()
+  const { setAiAssistantPanel } = useAppStateSnapshot()
 
+  const enableAssistantV2 = useIsDatabaseFunctionsAssistantEnabled()
   const canToggleRLS = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_WRITE, 'tables')
 
   const isRealtimeSchema = table.schema === 'realtime'
@@ -115,25 +112,33 @@ const PolicyTableRowHeader = ({
               Create policy
             </ButtonTooltip>
 
-            <Tooltip_Shadcn_>
-              <TooltipTrigger_Shadcn_ asChild>
-                <Button
-                  type="default"
-                  className="px-1"
-                  onClick={() => {
-                    onSelectCreatePolicy()
-                    setEditView('conversation')
-                  }}
-                >
-                  <AiIconAnimation className="scale-75 [&>div>div]:border-black dark:[&>div>div]:border-white" />
-                </Button>
-              </TooltipTrigger_Shadcn_>
-              <TooltipContent_Shadcn_ side="bottom">
-                {!canToggleRLS
-                  ? 'You need additional permissions to create RLS policies'
-                  : 'Create with Supabase Assistant'}
-              </TooltipContent_Shadcn_>
-            </Tooltip_Shadcn_>
+            <ButtonTooltip
+              type="default"
+              className="px-1"
+              onClick={() => {
+                if (enableAssistantV2) {
+                  setAiAssistantPanel({
+                    open: true,
+                    editor: 'rls-policies',
+                    entity: undefined,
+                    tables: [{ schema: table.schema, name: table.name }],
+                  })
+                } else {
+                  onSelectCreatePolicy()
+                  setEditView('conversation')
+                }
+              }}
+              tooltip={{
+                content: {
+                  side: 'bottom',
+                  text: !canToggleRLS
+                    ? 'You need additional permissions to create RLS policies'
+                    : 'Create with Supabase Assistant',
+                },
+              }}
+            >
+              <AiIconAnimation className="scale-75 [&>div>div]:border-black dark:[&>div>div]:border-white" />
+            </ButtonTooltip>
           </div>
         </div>
       )}
diff --git a/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx b/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx
index b5f5bc8021d..263da4ab55b 100644
--- a/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx
+++ b/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx
@@ -11,12 +11,14 @@ import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useAppStateSnapshot } from 'state/app-state'
 import {
   Button,
+  cn,
   DropdownMenu,
   DropdownMenuContent,
   DropdownMenuItem,
   DropdownMenuSeparator,
   DropdownMenuTrigger,
 } from 'ui'
+import { useIsDatabaseFunctionsAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
 
 interface FunctionListProps {
   schema: string
@@ -35,6 +37,7 @@ const FunctionList = ({
 }: FunctionListProps) => {
   const router = useRouter()
   const { project: selectedProject } = useProjectContext()
+  const enableAssistantV2 = useIsDatabaseFunctionsAssistantEnabled()
   const { setAiAssistantPanel } = useAppStateSnapshot()
 
   const { data: functions } = useDatabaseFunctionsQuery({
@@ -108,7 +111,10 @@ const FunctionList = ({
                       <DropdownMenuTrigger asChild>
                         <Button type="default" className="px-1" icon={<MoreVertical />} />
                       </DropdownMenuTrigger>
-                      <DropdownMenuContent side="left">
+                      <DropdownMenuContent
+                        side="left"
+                        className={cn(enableAssistantV2 ? 'w-52' : 'w-40')}
+                      >
                         {isApiDocumentAvailable && (
                           <DropdownMenuItem
                             className="space-x-2"
@@ -122,15 +128,17 @@ const FunctionList = ({
                           <Edit2 size={14} />
                           <p>Edit function</p>
                         </DropdownMenuItem>
-                        <DropdownMenuItem
-                          className="space-x-2"
-                          onClick={() => {
-                            setAiAssistantPanel({ open: true, editor: 'functions', entity: x })
-                          }}
-                        >
-                          <Edit size={14} />
-                          <p>Edit function with Assistant</p>
-                        </DropdownMenuItem>
+                        {enableAssistantV2 && (
+                          <DropdownMenuItem
+                            className="space-x-2"
+                            onClick={() => {
+                              setAiAssistantPanel({ open: true, editor: 'functions', entity: x })
+                            }}
+                          >
+                            <Edit size={14} />
+                            <p>Edit function with Assistant</p>
+                          </DropdownMenuItem>
+                        )}
                         <DropdownMenuSeparator />
                         <DropdownMenuItem className="space-x-2" onClick={() => deleteFunction(x)}>
                           <Trash size={14} className="text-destructive" />
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
index 6105135bba8..08942f7d655 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
@@ -64,6 +64,7 @@ import { ContextBadge } from './ContextBadge'
 import { EntitiesDropdownMenu } from './EntitiesDropdownMenu'
 import { Message } from './Message'
 import { SchemasDropdownMenu } from './SchemasDropdownMenu'
+import { useDatabasePoliciesQuery } from 'data/database-policies/database-policies-query'
 
 const ANIMATION_DURATION = 0.3
 
@@ -91,8 +92,8 @@ export const AIAssistant = ({
   const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM
 
   const disablePrompts = useFlag('disableAssistantPrompts')
-  const { aiAssistantPanel } = useAppStateSnapshot()
-  const { editor, entity } = aiAssistantPanel
+  const { aiAssistantPanel, setAiAssistantPanel } = useAppStateSnapshot()
+  const { editor, entity, tables: selectedTables } = aiAssistantPanel
 
   const inputRef = useRef<HTMLTextAreaElement>(null)
   const bottomRef = useRef<HTMLDivElement>(null)
@@ -102,7 +103,6 @@ export const AIAssistant = ({
     SupportedAssistantEntities | ''
   >('')
   const [selectedSchemas, setSelectedSchemas] = useSchemasForAi(project?.ref!)
-  const [selectedTables, setSelectedTables] = useState<{ schema: string; name: string }[]>([])
   const [contextHistory, setContextHistory] = useState<{
     [key: string]: { entity: string; schemas: string[]; tables: string[] }
   }>({})
@@ -121,6 +121,19 @@ export const AIAssistant = ({
   const { data: check } = useCheckOpenAIKeyQuery()
   const isApiKeySet = IS_PLATFORM || !!check?.hasKey
 
+  const { data: policies } = useDatabasePoliciesQuery(
+    {
+      projectRef: project?.ref,
+      connectionString: project?.connectionString,
+    },
+    { enabled: editor === 'rls-policies' }
+  )
+  const existingPolicies = (policies ?? [])
+    .filter((policy) =>
+      selectedTables.some((x) => policy.schema === x.schema && policy.table === x.name)
+    )
+    .sort((a, b) => a.name.localeCompare(b.name))
+
   const { data: existingDefinition } = useEntityDefinitionQuery({
     id: entity?.id,
     type: editor,
@@ -165,7 +178,12 @@ export const AIAssistant = ({
   } = useChat({
     id,
     api: `${BASE_PATH}/api/ai/sql/generate-v2`,
-    body: { entityDefinitions, context: selectedDatabaseEntity, existingSql: existingDefinition },
+    body: {
+      entityDefinitions,
+      context: selectedDatabaseEntity,
+      existingSql: existingDefinition,
+      existingPolicies,
+    },
     onError: (error) => setAssistantError(JSON.parse(error.message).error),
   })
 
@@ -219,12 +237,14 @@ export const AIAssistant = ({
   const toggleEntity = ({ schema, name }: { schema: string; name: string }) => {
     const isExisting = selectedTables.find((x) => x.schema === schema && x.name === name)
     if (isExisting) {
-      setSelectedTables(selectedTables.filter((x) => !(x.schema === schema && x.name === name)))
+      setAiAssistantPanel({
+        tables: selectedTables.filter((x) => !(x.schema === schema && x.name === name)),
+      })
     } else {
       const newselectedTables = [...selectedTables, { schema, name }].sort(
         (a, b) => a.schema.localeCompare(b.schema) || a.name.localeCompare(b.name)
       )
-      setSelectedTables(newselectedTables)
+      setAiAssistantPanel({ tables: newselectedTables })
       sendTelemetryEvent(TELEMETRY_ACTIONS.TABLE_CONTEXT_ADDED)
     }
   }
@@ -435,19 +455,28 @@ export const AIAssistant = ({
                   transition={{ duration: ANIMATION_DURATION }}
                 >
                   <p className="text-center text-base text-foreground-light">
-                    How can I help you
-                    {!!entityContext ? (
+                    {entity !== undefined && !!entityContext ? (
                       <>
-                        {' '}
-                        with{' '}
-                        <span className="text-foreground">
-                          {entityContext.id === 'rls-policies'
-                            ? entityContext.label
-                            : `Database ${entityContext.label}`}
-                        </span>
+                        Need help with updating this{' '}
+                        <span className="text-foreground">{entityContext.name}</span>
                       </>
                     ) : (
-                      ' today'
+                      <>
+                        How can I help you
+                        {!!entityContext ? (
+                          <>
+                            {' '}
+                            with{' '}
+                            <span className="text-foreground">
+                              {entityContext.id === 'rls-policies'
+                                ? entityContext.label
+                                : `Database ${entityContext.label}`}
+                            </span>
+                          </>
+                        ) : (
+                          ' today'
+                        )}
+                      </>
                     )}
                     ?
                   </p>
@@ -582,7 +611,7 @@ export const AIAssistant = ({
                       value={`${selectedSchemas.slice(0, 2).join(', ')}${selectedSchemas.length > 2 ? ` and ${selectedSchemas.length - 2} other${selectedSchemas.length > 3 ? 's' : ''}` : ''}`}
                       onRemove={() => {
                         setSelectedSchemas([])
-                        setSelectedTables([])
+                        setAiAssistantPanel({ tables: [] })
                       }}
                       tooltip={
                         selectedSchemas.length > 2 ? (
@@ -610,7 +639,7 @@ export const AIAssistant = ({
                         .join(
                           ', '
                         )}${selectedTables.length > 2 ? ` and ${selectedTables.length - 2} other${selectedTables.length > 3 ? 's' : ''}` : ''}`}
-                      onRemove={() => setSelectedTables([])}
+                      onRemove={() => setAiAssistantPanel({ tables: [] })}
                       tooltip={
                         selectedTables.length > 2 ? (
                           <>
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts
index 10285b026b6..e18d6b42af8 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts
@@ -76,7 +76,8 @@ ${existingDefinition
 `
       }
     case 'rls-policies':
-      return `${PLACEHOLDER_PREFIX}
+      if (entity === undefined) {
+        return `${PLACEHOLDER_PREFIX}
 CREATE POLICY *name* ON *table_name*\n
 AS PERMISSIVE -- PERMISSIVE | RESTRICTIVE\n
 FOR ALL -- ALL | SELECT | INSERT | UPDATE | DELETE\n
@@ -84,6 +85,33 @@ TO *role_name* -- Default: public\n
 USING ( *using_expression* )\n
 WITH CHECK ( *check_expression* );
 `
+      } else {
+        let expression = ''
+        if (entity.definition !== null && entity.definition !== undefined) {
+          expression += `USING ( *${entity.definition}* )${
+            entity.check === null || entity.check === undefined ? ';' : ''
+          }\n`
+        }
+        if (entity.check !== null && entity.check !== undefined) {
+          expression += `WITH CHECK ( *${entity.check}* );\n`
+        }
+        return `${PLACEHOLDER_PREFIX}
+BEGIN;\n
+&nbsp;\n
+-- To update your policy definition\n
+ALTER POLICY "${entity.name}"\n
+ON "${entity.schema}"."${entity.table}"\n
+TO *${(entity.roles ?? []).join(', ')}*\n
+${expression}
+&nbsp;\n
+-- To rename the policy\n
+ALTER POLICY "${entity.name}"\n
+ON "${entity.schema}"."${entity.table}"\n
+RENAME TO "*New Policy Name*";\n
+&nbsp;\n
+COMMIT;
+`
+      }
     default:
       return undefined
   }
@@ -111,7 +139,7 @@ export const validateQuery = (editor: SupportedAssistantEntities | null, query:
         formattedQuery.includes('create or replace function')
       )
     case 'rls-policies':
-      return formattedQuery.includes('create policy')
+      return formattedQuery.includes('create policy') || formattedQuery.includes('alter policy')
     default:
       return true
   }
@@ -126,7 +154,7 @@ export const generatePrompt = ({
   type: SupportedAssistantQuickPromptTypes
   context: SupportedAssistantEntities
   schemas: string[]
-  tables: { schema: string; name: string }[]
+  tables: readonly { schema: string; name: string }[]
 }) => {
   if (type === 'examples') {
     return `What are some common examples of user-defined database ${context}? ${PLACEHOLDER_LIMIT}`
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx
index 937e55bdce5..702c1f0cdbe 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx
@@ -16,6 +16,7 @@ import { suffixWithLimit } from 'components/interfaces/SQLEditor/SQLEditor.utils
 import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results'
 import { useSqlDebugMutation } from 'data/ai/sql-debug-mutation'
 import { databasePoliciesKeys } from 'data/database-policies/keys'
+import { useEntityDefinitionQuery } from 'data/database/entity-definition-query'
 import { QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import { sqlKeys } from 'data/sql/keys'
 import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
@@ -48,7 +49,6 @@ import CodeEditor from '../CodeEditor/CodeEditor'
 import { AIAssistant } from './AIAssistant'
 import { generateCTA, generatePlaceholder, generateTitle, validateQuery } from './AIAssistant.utils'
 import { ASSISTANT_SUPPORT_ENTITIES } from './AiAssistant.constants'
-import { useEntityDefinitionQuery } from 'data/database/entity-definition-query'
 
 export const AiAssistantPanel = () => {
   const os = detectOS()
@@ -61,6 +61,7 @@ export const AiAssistantPanel = () => {
 
   const { open, editor, content, entity } = aiAssistantPanel
   const previousEditor = usePrevious(editor)
+  const previousEntity = usePrevious(entity)
 
   const [isAcknowledged, setIsAcknowledged] = useLocalStorage(
     LOCAL_STORAGE_KEYS.SQL_SCRATCH_PAD_BANNER_ACKNOWLEDGED,
@@ -108,7 +109,7 @@ export const AiAssistantPanel = () => {
     onSuccess: async (res) => {
       // [Joshen] If in a specific editor context mode, assume that intent was to create/update
       // a database entity - so close it once success. Otherwise it's in Quick SQL mode and we
-      // show the results. Currently though it assumes we're "creating", thinking need to support "updating" too
+      // show the results.
       if (editor !== null) {
         switch (editor) {
           case 'functions':
@@ -119,7 +120,9 @@ export const AiAssistantPanel = () => {
             break
         }
 
-        toast.success(`Successfully created ${entityContext?.name}!`)
+        toast.success(
+          `Successfully ${entity === undefined ? 'created' : 'updated'} ${entityContext?.name}!`
+        )
         setAiAssistantPanel({ open: false })
       } else {
         setShowResults(true)
@@ -257,8 +260,8 @@ export const AiAssistantPanel = () => {
 
   useEffect(() => {
     if (open) {
-      // [Joshen] Only reset the assistant if the editor changed
-      if (previousEditor !== editor) {
+      // [Joshen] Only reset the assistant if the editor changed or if the provided entity has changed
+      if (previousEditor !== editor || previousEntity !== entity) {
         setChatId(uuidv4())
         setError(undefined)
         setShowWarning(false)
diff --git a/apps/studio/components/ui/AIAssistantPanel/EntitiesDropdownMenu.tsx b/apps/studio/components/ui/AIAssistantPanel/EntitiesDropdownMenu.tsx
index d301fb5814e..f2120c27ca3 100644
--- a/apps/studio/components/ui/AIAssistantPanel/EntitiesDropdownMenu.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/EntitiesDropdownMenu.tsx
@@ -1,23 +1,23 @@
-import { useEntityTypesQuery } from 'data/entity-types/entity-types-infinite-query'
-import { useSelectedProject } from 'hooks/misc/useSelectedProject'
 import { debounce } from 'lodash'
 import { Check, Loader2, Search } from 'lucide-react'
 import { useCallback, useState } from 'react'
+
+import { useEntityTypesQuery } from 'data/entity-types/entity-types-infinite-query'
+import { useSelectedProject } from 'hooks/misc/useSelectedProject'
 import {
   Command_Shadcn_,
   CommandEmpty_Shadcn_,
   CommandGroup_Shadcn_,
   CommandItem_Shadcn_,
   CommandList_Shadcn_,
-  ScrollArea,
 } from 'ui'
+import { Admonition } from 'ui-patterns'
 import { Input } from 'ui-patterns/DataInputs/Input'
 import { DropdownScrollArea } from './DropdownScrollArea'
-import { Admonition } from 'ui-patterns'
 
 interface EntitiesDropdownMenuProps {
   selectedSchemas: string[]
-  selectedEntities: { schema: string; name: string }[]
+  selectedEntities: readonly { schema: string; name: string }[]
   onToggleEntity: (value: { schema: string; name: string }) => void
 }
 
diff --git a/apps/studio/data/database/entity-definition-query.ts b/apps/studio/data/database/entity-definition-query.ts
index ddc44aa59fd..5c8ffcb8832 100644
--- a/apps/studio/data/database/entity-definition-query.ts
+++ b/apps/studio/data/database/entity-definition-query.ts
@@ -1,7 +1,24 @@
 import { UseQueryOptions } from '@tanstack/react-query'
+
+import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types'
 import { ExecuteSqlData, ExecuteSqlError, useExecuteSqlQuery } from '../sql/execute-sql-query'
 import { databaseKeys } from './keys'
-import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types'
+
+const generatePolicyDefinition = (policy: {
+  command: string
+  name: string
+  permissive: string
+  roles: string
+  schema: string
+  table: string
+  qual: string | null
+  with_check: string | null
+}) => {
+  const roles = policy.roles.replace('{', '').replace('}', '').split(',')
+
+  return `
+CREATE POLICY "${policy.name}" on "${policy.schema}"."${policy.table}" AS ${policy.permissive} FOR ${policy.command} TO ${roles.join(', ')} ${policy.qual ? `USING (${policy.qual})` : ''} ${policy.with_check ? `WITH CHECK (${policy.with_check})` : ''};`.trim()
+}
 
 // [Joshen] Eventually should support table definition and view definition as well if possible
 export const getEntityDefinitionQuery = ({
@@ -19,9 +36,36 @@ export const getEntityDefinitionQuery = ({
       select pg_get_functiondef(${id})
     `.trim()
     case 'rls-policies':
-      // [Joshen] Eventually to-do, unless we have to piece it manually?
       return /* SQL */ `
-      select 1;
+      SELECT 
+        n.nspname AS schema,
+        c.relname AS table,
+        pol.polname AS name,
+            CASE
+                WHEN pol.polpermissive THEN 'PERMISSIVE'::text
+                ELSE 'RESTRICTIVE'::text
+            END AS permissive,
+            CASE
+                WHEN pol.polroles = '{0}'::oid[] THEN string_to_array('public'::text, ''::text)::name[]
+                ELSE ARRAY( SELECT pg_authid.rolname
+                  FROM pg_authid
+                  WHERE pg_authid.oid = ANY (pol.polroles)
+                  ORDER BY pg_authid.rolname)
+            END AS roles,
+            CASE pol.polcmd
+                WHEN 'r'::"char" THEN 'SELECT'::text
+                WHEN 'a'::"char" THEN 'INSERT'::text
+                WHEN 'w'::"char" THEN 'UPDATE'::text
+                WHEN 'd'::"char" THEN 'DELETE'::text
+                WHEN '*'::"char" THEN 'ALL'::text
+                ELSE NULL::text
+            END AS command,
+        pg_get_expr(pol.polqual, pol.polrelid) AS qual,
+        pg_get_expr(pol.polwithcheck, pol.polrelid) AS with_check
+      FROM pg_policy pol
+        JOIN pg_class c ON c.oid = pol.polrelid
+        LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
+      WHERE pol.oid = ${id};
     `.trim()
   }
 
@@ -35,7 +79,6 @@ export type EntityDefinitionVariables = {
   connectionString?: string
 }
 
-// type EntityDefinition = { id: number; sql: string }
 export type EntityDefinitionData = string
 export type EntityDefinitionError = ExecuteSqlError
 
@@ -52,7 +95,13 @@ export const useEntityDefinitionQuery = <TData extends EntityDefinitionData = En
     },
     {
       select(data) {
-        return data.result[0].pg_get_functiondef
+        if (type === 'functions') {
+          return data.result[0].pg_get_functiondef
+        } else if (type === 'rls-policies') {
+          return generatePolicyDefinition(data.result[0])
+        } else {
+          return data.result[0]
+        }
       },
       enabled: enabled && typeof id !== 'undefined' && typeof type !== 'undefined',
       ...options,
diff --git a/apps/studio/pages/api/ai/sql/generate-v2.ts b/apps/studio/pages/api/ai/sql/generate-v2.ts
index d21b384a479..c3e8e8d7fdf 100644
--- a/apps/studio/pages/api/ai/sql/generate-v2.ts
+++ b/apps/studio/pages/api/ai/sql/generate-v2.ts
@@ -1,6 +1,7 @@
 import { StreamingTextResponse } from 'ai'
-import { chatSql } from 'ai-commands/edge'
+import { chatRlsPolicy, chatSql } from 'ai-commands/edge'
 import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types'
+import { DatabasePoliciesData } from 'data/database-policies/database-policies-query'
 import { NextRequest } from 'next/server'
 import OpenAI from 'openai'
 
@@ -68,17 +69,29 @@ async function handlePost(request: NextRequest) {
   const openai = new OpenAI({ apiKey: openAiKey })
 
   const body = await (request.json() as Promise<{
+    context?: SupportedAssistantEntities
     messages: { content: string; role: 'user' | 'assistant' }[]
     existingSql?: string
     entityDefinitions: string[]
-    context?: SupportedAssistantEntities
+    existingPolicies?: DatabasePoliciesData
   }>)
 
-  const { messages, existingSql, entityDefinitions, context } = body
+  const { messages, existingSql, entityDefinitions, context, existingPolicies } = body
 
   try {
-    const stream = await chatSql(openai, messages, existingSql, entityDefinitions, context)
-    return new StreamingTextResponse(stream)
+    if (context === 'rls-policies') {
+      const stream = await chatRlsPolicy(
+        openai,
+        messages,
+        entityDefinitions,
+        existingPolicies ?? [],
+        existingSql
+      )
+      return new StreamingTextResponse(stream)
+    } else {
+      const stream = await chatSql(openai, messages, existingSql, entityDefinitions, context)
+      return new StreamingTextResponse(stream)
+    }
   } catch (error) {
     if (error instanceof Error) {
       console.error(`AI SQL generation-v2 failed: ${error.message}`)
@@ -86,14 +99,9 @@ async function handlePost(request: NextRequest) {
       console.error(`AI SQL generation-v2 failed: ${error}`)
     }
 
-    return new Response(
-      JSON.stringify({
-        error: 'There was an error processing your request',
-      }),
-      {
-        status: 500,
-        headers: { 'Content-Type': 'application/json' },
-      }
-    )
+    return new Response(JSON.stringify({ error: 'There was an error processing your request' }), {
+      status: 500,
+      headers: { 'Content-Type': 'application/json' },
+    })
   }
 }
diff --git a/apps/studio/pages/api/ai/sql/suggest.ts b/apps/studio/pages/api/ai/sql/suggest.ts
index b32c923039c..8244d6799ed 100644
--- a/apps/studio/pages/api/ai/sql/suggest.ts
+++ b/apps/studio/pages/api/ai/sql/suggest.ts
@@ -2,7 +2,7 @@ import { StreamingTextResponse } from 'ai'
 import { chatRlsPolicy } from 'ai-commands/edge'
 import { NextRequest } from 'next/server'
 import OpenAI from 'openai'
-import { DatabasePoliciesData } from '../../../../data/database-policies/database-policies-query'
+import { DatabasePoliciesData } from 'data/database-policies/database-policies-query'
 
 export const config = {
   runtime: 'edge',
diff --git a/apps/studio/state/app-state.ts b/apps/studio/state/app-state.ts
index 71f10820099..48e35d43b16 100644
--- a/apps/studio/state/app-state.ts
+++ b/apps/studio/state/app-state.ts
@@ -26,6 +26,7 @@ type AiAssistantPanelType = {
   content?: string
   // Mainly used for editing a database entity (e.g editing a function, RLS policy etc)
   entity?: CommonDatabaseEntity
+  tables: { schema: string; name: string }[]
 }
 
 export const appState = proxy({
@@ -112,9 +113,11 @@ export const appState = proxy({
     editor: null,
     content: '',
     entity: undefined,
+    tables: [],
   } as AiAssistantPanelType,
-  setAiAssistantPanel: (value: AiAssistantPanelType) => {
+  setAiAssistantPanel: (value: Partial<AiAssistantPanelType>) => {
     const hasEntityChanged = value.entity?.id !== appState.aiAssistantPanel.entity?.id
+
     appState.aiAssistantPanel = {
       ...appState.aiAssistantPanel,
       content: hasEntityChanged ? '' : appState.aiAssistantPanel.content,
diff --git a/packages/ai-commands/src/sql/chat.ts b/packages/ai-commands/src/sql/chat.ts
index 75e4b02e978..426d9f8f617 100644
--- a/packages/ai-commands/src/sql/chat.ts
+++ b/packages/ai-commands/src/sql/chat.ts
@@ -43,6 +43,8 @@ export async function chatSql(
     // is definitely not optimal at all, but just to get an idea started
     const generateInstructionsBasedOnContext = () => {
       switch (context) {
+        // [Joshen] Sorry for the mess, there's duplicate logic here between this and rls.ts - i'm wondering what should be the best practice
+        // here? Do we (1) Put ALL logic into this file, or (2) we split each context into their own files? Latter might be cleaner?
         case 'rls-policies':
           return stripIndent`
           You're a Supabase Postgres expert in writing row level security policies. Your purpose is to
diff --git a/packages/ai-commands/src/sql/rls.ts b/packages/ai-commands/src/sql/rls.ts
index 905bb276b67..dd9fbd23846 100644
--- a/packages/ai-commands/src/sql/rls.ts
+++ b/packages/ai-commands/src/sql/rls.ts
@@ -301,6 +301,8 @@ export async function chatRlsPolicy(
       content: codeBlock`
         Here is my policy definition for reference:
         ${definitionBlock}
+        
+        I'm requesting to update this policy instead so please opt to use "alter policy" instead of "create policy" where appropriate.
       `.trim(),
     })
   }