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
- {/* */}1
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 ({Edit policy
Edit policy with Assistant
+Edit function
-Edit function with Assistant
-Edit function with Assistant
+- How can I help you - {!!entityContext ? ( + {entity !== undefined && !!entityContext ? ( <> - {' '} - with{' '} - - {entityContext.id === 'rls-policies' - ? entityContext.label - : `Database ${entityContext.label}`} - + Need help with updating this{' '} + {entityContext.name} > ) : ( - ' today' + <> + How can I help you + {!!entityContext ? ( + <> + {' '} + with{' '} + + {entityContext.id === 'rls-policies' + ? entityContext.label + : `Database ${entityContext.label}`} + + > + ) : ( + ' today' + )} + > )} ?
@@ -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 + \n +-- To update your policy definition\n +ALTER POLICY "${entity.name}"\n +ON "${entity.schema}"."${entity.table}"\n +TO *${(entity.roles ?? []).join(', ')}*\n +${expression} + \n +-- To rename the policy\n +ALTER POLICY "${entity.name}"\n +ON "${entity.schema}"."${entity.table}"\n +RENAME TO "*New Policy Name*";\n + \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 =