supabase/apps/studio/components/interfaces/Database/Functions/CreateFunction/index.tsx

import { zodResolver } from '@hookform/resolvers/zod'
import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta/src/pg-format'
import { isEmpty, isNull, keyBy, mapValues, partition } from 'lodash'
import { Plus, Trash } from 'lucide-react'
import { useEffect, useMemo, useState } from 'react'
import { SubmitHandler, useFieldArray, useForm, useWatch } from 'react-hook-form'
import { toast } from 'sonner'
import {
  Button,
  cn,
  Form,
  FormControl,
  FormDescription,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
  Input,
  RadioGroupStacked,
  RadioGroupStackedItem,
  ScrollArea,
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
  Separator,
  Sheet,
  SheetContent,
  SheetFooter,
  SheetSection,
  Switch,
} from 'ui'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import z from 'zod'

import { convertArgumentTypes, convertConfigParams } from '../Functions.utils'
import { CreateFunctionConfigParamsSection } from './CreateFunctionConfigParamsSection'
import { CreateFunctionHeader } from './CreateFunctionHeader'
import { FunctionEditor } from './FunctionEditor'
import { POSTGRES_DATA_TYPES } from '@/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.constants'
import { DiscardChangesConfirmationDialog } from '@/components/ui-patterns/Dialogs/DiscardChangesConfirmationDialog'
import { SchemaSelector } from '@/components/ui/SchemaSelector'
import { useDatabaseExtensionsQuery } from '@/data/database-extensions/database-extensions-query'
import { useDatabaseFunctionCreateMutation } from '@/data/database-functions/database-functions-create-mutation'
import type { SavedDatabaseFunction } from '@/data/database-functions/database-functions-query'
import { useDatabaseFunctionUpdateMutation } from '@/data/database-functions/database-functions-update-mutation'
import { useQuerySchemaState } from '@/hooks/misc/useSchemaQueryState'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { useConfirmOnClose } from '@/hooks/ui/useConfirmOnClose'
import { useProtectedSchemas } from '@/hooks/useProtectedSchemas'

const FORM_ID = 'create-function-sidepanel'

interface CreateFunctionProps {
  func?: SavedDatabaseFunction
  isDuplicating?: boolean
  visible: boolean
  onClose: () => void
}

const FormSchema = z.object({
  name: z.string().trim().min(1),
  type: z.enum(['function', 'procedure']),
  schema: z.string().trim().min(1),
  args: z.array(z.object({ name: z.string().trim().min(1), type: z.string().trim() })),
  behavior: z.enum(['IMMUTABLE', 'STABLE', 'VOLATILE']),
  definition: z.string().trim().min(1),
  language: z.string().trim(),
  return_type: z.string().trim(),
  security_definer: z.boolean(),
  config_params: z
    .array(z.object({ name: z.string().trim().min(1), value: z.string().trim().min(1) }))
    .optional(),
})

export const CreateFunction = ({
  func,
  visible,
  isDuplicating = false,
  onClose,
}: CreateFunctionProps) => {
  const { selectedSchema } = useQuerySchemaState()
  const { data: project } = useSelectedProjectQuery()

  const [advancedSettingsShown, setAdvancedSettingsShown] = useState(false)
  const [focusedEditor, setFocusedEditor] = useState(false)

  const isEditing = !isDuplicating && !!func?.id

  const form = useForm<z.infer<typeof FormSchema>>({
    resolver: zodResolver(FormSchema),
  })
  const [type, language] = useWatch({ control: form.control, name: ['type', 'language'] })

  const { confirmOnClose, handleOpenChange, modalProps } = useConfirmOnClose({
    checkIsDirty: () => form.formState.isDirty,
    onClose,
  })

  const { mutate: createDatabaseFunction, isPending: isCreating } =
    useDatabaseFunctionCreateMutation()
  const { mutate: updateDatabaseFunction, isPending: isUpdating } =
    useDatabaseFunctionUpdateMutation()

  const onSubmit: SubmitHandler<z.infer<typeof FormSchema>> = async (data) => {
    if (!project) return console.error('Project is required')
    // Submit click is the explicit user gesture that promotes form-entered SQL fragments
    // (`args` items, `return_type`, and each `config_params` value) to executable.
    const payload = {
      ...data,
      args: data.args.map((x) => acceptUntrustedSql(untrustedSql(`${x.name} ${x.type}`))),
      return_type: acceptUntrustedSql(untrustedSql(data.return_type)),
      config_params: mapValues(keyBy(data.config_params, 'name'), (item) =>
        acceptUntrustedSql(untrustedSql(item.value))
      ),
    }

    if (isEditing) {
      updateDatabaseFunction(
        {
          func,
          projectRef: project.ref,
          connectionString: project.connectionString,
          payload,
        },
        {
          onSuccess: () => {
            toast.success(`Successfully updated function ${data.name}`)
            onClose()
          },
        }
      )
    } else {
      createDatabaseFunction(
        {
          projectRef: project.ref,
          connectionString: project.connectionString,
          payload,
        },
        {
          onSuccess: () => {
            toast.success(`Successfully created function ${data.name}`)
            onClose()
          },
        }
      )
    }
  }

  useEffect(() => {
    if (visible) {
      setFocusedEditor(false)
      const type = func?.type ?? 'function'
      form.reset({
        type,
        name: func?.name ?? '',
        schema: func?.schema ?? selectedSchema,
        args: convertArgumentTypes({ type, value: func?.argument_types || '' }).value,
        behavior: func?.behavior ?? 'VOLATILE',
        definition: func?.definition ?? '',
        language: func?.language ?? 'plpgsql',
        return_type: func?.return_type ?? 'void',
        security_definer: func?.security_definer ?? false,
        config_params: convertConfigParams(func?.config_params).value,
      })
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [visible, func?.id])

  const { data: protectedSchemas } = useProtectedSchemas()

  return (
    <Sheet open={visible} onOpenChange={handleOpenChange}>
      <SheetContent
        size="default"
        className="p-0 flex flex-row gap-0 min-w-screen! lg:min-w-[600px]!"
      >
        <div className="flex flex-col grow w-full">
          <CreateFunctionHeader selectedFunction={func?.name} isDuplicating={isDuplicating} />

          <Separator />

          <Form {...form}>
            <form
              id={FORM_ID}
              className="grow overflow-auto"
              onSubmit={form.handleSubmit(onSubmit)}
            >
              <SheetSection className={cn(focusedEditor ? 'hidden' : '', 'flex flex-col gap-y-4')}>
                <FormField
                  control={form.control}
                  name="schema"
                  render={({ field }) => (
                    <FormItemLayout
                      label="Schema"
                      description={
                        <>
                          Tables made in the table editor will be in{' '}
                          <code className="text-code-inline">public</code>
                        </>
                      }
                      layout="horizontal"
                    >
                      <FormControl>
                        <SchemaSelector
                          selectedSchemaName={field.value}
                          excludedSchemas={protectedSchemas?.map((s) => s.name)}
                          size="small"
                          onSelectSchema={(name) => field.onChange(name)}
                        />
                      </FormControl>
                    </FormItemLayout>
                  )}
                />
                <FormField
                  control={form.control}
                  name="name"
                  render={({ field }) => (
                    <FormItemLayout
                      label="Name of function"
                      description="Name will also be used for the function name in postgres"
                      layout="horizontal"
                    >
                      <FormControl>
                        <Input {...field} placeholder="Name of function" />
                      </FormControl>
                    </FormItemLayout>
                  )}
                />
              </SheetSection>

              <Separator className={focusedEditor ? 'hidden' : ''} />

              <SheetSection className={focusedEditor ? 'hidden' : 'flex flex-col gap-y-4'}>
                <FormField
                  control={form.control}
                  name="type"
                  render={({ field }) => (
                    <FormItemLayout label="Type" layout="horizontal">
                      <Select
                        disabled={isEditing}
                        onValueChange={field.onChange}
                        defaultValue={field.value}
                      >
                        <SelectTrigger className="col-span-8">
                          {field.value === 'function' ? 'Function' : 'Stored procedure'}
                        </SelectTrigger>
                        <SelectContent align="end">
                          <SelectItem value="function">
                            <p>Function</p>
                            <p className="text-foreground-lighter">
                              For query logic, triggers, and RPC calls
                            </p>
                          </SelectItem>
                          <SelectItem value="procedure">
                            <p>Stored procedure</p>
                            <p className="text-foreground-lighter">
                              For batch or multi-step tasks that manages transactions
                            </p>
                          </SelectItem>
                        </SelectContent>
                      </Select>
                    </FormItemLayout>
                  )}
                />
                {type === 'function' && (
                  <FormField
                    control={form.control}
                    name="return_type"
                    render={({ field }) => (
                      <FormItemLayout label="Return type" layout="horizontal">
                        {/* Form selects don't need form controls, otherwise the CSS gets weird */}
                        <Select
                          disabled={isEditing}
                          onValueChange={field.onChange}
                          defaultValue={field.value}
                        >
                          <SelectTrigger className="col-span-8">
                            <SelectValue />
                          </SelectTrigger>
                          <SelectContent>
                            <ScrollArea className="h-52">
                              {['void', 'record', 'trigger', 'integer', ...POSTGRES_DATA_TYPES].map(
                                (option) => (
                                  <SelectItem value={option} key={option}>
                                    {option}
                                  </SelectItem>
                                )
                              )}
                            </ScrollArea>
                          </SelectContent>
                        </Select>
                      </FormItemLayout>
                    )}
                  />
                )}
              </SheetSection>

              <Separator className={focusedEditor ? 'hidden' : ''} />

              <SheetSection className={focusedEditor ? 'hidden' : ''}>
                <FormFieldArgs readonly={isEditing} />
              </SheetSection>

              <Separator className={focusedEditor ? 'hidden' : ''} />

              <SheetSection className={`${focusedEditor ? 'h-full' : ''} px-0!`}>
                <FormField
                  control={form.control}
                  name="definition"
                  render={({ field }) => (
                    <FormItem className="space-y-4 flex flex-col h-full">
                      <div className="px-content">
                        <FormLabel className="text-base text-foreground">Definition</FormLabel>
                        <FormDescription className="text-sm text-foreground-light">
                          <p>
                            The language below should be written in{' '}
                            <code className="text-code-inline">{language}</code>.
                          </p>
                          {!isEditing && <p>Change the language in the Advanced Settings below.</p>}
                        </FormDescription>
                      </div>
                      <div
                        className={cn(
                          'border border-default flex',
                          focusedEditor ? 'grow ' : 'h-72'
                        )}
                      >
                        <FunctionEditor
                          field={field}
                          language={language}
                          focused={focusedEditor}
                          setFocused={setFocusedEditor}
                        />
                      </div>

                      <FormMessage className="px-content" />
                    </FormItem>
                  )}
                />
              </SheetSection>
              <Separator className={focusedEditor ? 'hidden' : ''} />
              {isEditing ? null : (
                <>
                  <SheetSection className={focusedEditor ? 'hidden' : ''}>
                    <div className="space-y-8 rounded-sm bg-studio py-4 px-6 border border-overlay">
                      <FormItem className="flex flex-row items-center justify-between">
                        <div className="space-y-0.5">
                          <FormLabel className="text-sm">Show advanced settings</FormLabel>
                          <FormDescription className="text-foreground-lighter">
                            These are settings that might be familiar for Postgres developers
                          </FormDescription>
                        </div>
                        <FormControl>
                          <Switch
                            checked={advancedSettingsShown}
                            onCheckedChange={(checked) => setAdvancedSettingsShown(checked)}
                          />
                        </FormControl>
                      </FormItem>
                    </div>
                  </SheetSection>
                  {advancedSettingsShown && (
                    <>
                      <SheetSection className={focusedEditor ? 'hidden' : 'space-y-2 pt-0'}>
                        <FormFieldLanguage />
                        {type === 'function' && (
                          <FormField
                            control={form.control}
                            name="behavior"
                            render={({ field }) => (
                              <FormItemLayout label="Behavior" layout="horizontal">
                                {/* Form selects don't need form controls, otherwise the CSS gets weird */}
                                <Select defaultValue={field.value} onValueChange={field.onChange}>
                                  <SelectTrigger className="col-span-8">
                                    <SelectValue />
                                  </SelectTrigger>
                                  <SelectContent>
                                    <SelectItem value="IMMUTABLE" key="IMMUTABLE">
                                      immutable
                                    </SelectItem>
                                    <SelectItem value="STABLE" key="STABLE">
                                      stable
                                    </SelectItem>
                                    <SelectItem value="VOLATILE" key="VOLATILE">
                                      volatile
                                    </SelectItem>
                                  </SelectContent>
                                </Select>
                              </FormItemLayout>
                            )}
                          />
                        )}
                      </SheetSection>
                      <Separator className={focusedEditor ? 'hidden' : ''} />
                      <SheetSection className={focusedEditor ? 'hidden' : ''}>
                        <CreateFunctionConfigParamsSection />
                      </SheetSection>
                      <Separator className={focusedEditor ? 'hidden' : ''} />
                      <SheetSection className={focusedEditor ? 'hidden' : ''}>
                        <h5 className="text-base text-foreground mb-4">Type of Security</h5>
                        <FormField
                          control={form.control}
                          name="security_definer"
                          render={({ field }) => (
                            <FormItem>
                              <FormControl className="col-span-8">
                                <RadioGroupStacked
                                  onValueChange={(value) =>
                                    field.onChange(value == 'SECURITY_DEFINER')
                                  }
                                  value={field.value ? 'SECURITY_DEFINER' : 'SECURITY_INVOKER'}
                                >
                                  <RadioGroupStackedItem
                                    value="SECURITY_INVOKER"
                                    id="SECURITY_INVOKER"
                                    label="SECURITY INVOKER"
                                    description={
                                      <>
                                        Function is to be executed with the privileges of the user
                                        that <span className="text-foreground">calls it</span>.
                                      </>
                                    }
                                  />
                                  <RadioGroupStackedItem
                                    value="SECURITY_DEFINER"
                                    id="SECURITY_DEFINER"
                                    label="SECURITY DEFINER"
                                    description={
                                      <>
                                        Function is to be executed with the privileges of the user
                                        that <span className="text-foreground">created it</span>.
                                      </>
                                    }
                                  />
                                </RadioGroupStacked>
                              </FormControl>
                              <FormMessage />
                            </FormItem>
                          )}
                        />
                      </SheetSection>
                    </>
                  )}
                </>
              )}
            </form>
          </Form>
          <SheetFooter>
            <Button disabled={isCreating || isUpdating} variant="default" onClick={confirmOnClose}>
              Cancel
            </Button>
            <Button
              form={FORM_ID}
              type="submit"
              disabled={isCreating || isUpdating}
              loading={isCreating || isUpdating}
            >
              {isEditing ? 'Save' : 'Create'} function
            </Button>
          </SheetFooter>
        </div>
        <DiscardChangesConfirmationDialog {...modalProps} />
      </SheetContent>
    </Sheet>
  )
}

interface FormFieldConfigParamsProps {
  readonly?: boolean
}

const FormFieldArgs = ({ readonly }: FormFieldConfigParamsProps) => {
  const { fields, append, remove } = useFieldArray<z.infer<typeof FormSchema>>({
    name: 'args',
  })

  return (
    <div className="flex flex-col gap-y-4">
      <div>
        <h5 className="text-base text-foreground">Arguments</h5>
        <p className="text-sm text-foreground-light">
          Arguments can be referenced in the function body using either names or numbers.
        </p>
      </div>
      <div className="space-y-2">
        {readonly && isEmpty(fields) && (
          <span className="text-sm text-foreground-lighter">No argument for this function</span>
        )}
        {fields.map((field, index) => {
          return (
            <div className="flex flex-row space-x-1" key={field.id}>
              <FormField
                name={`args.${index}.name`}
                render={({ field }) => (
                  <FormItem className="flex-1">
                    <FormControl>
                      <Input {...field} disabled={readonly} placeholder="argument_name" />
                    </FormControl>
                    <FormMessage />
                  </FormItem>
                )}
              />
              <FormField
                name={`args.${index}.type`}
                render={({ field }) => (
                  <FormItem className="flex-1">
                    <FormControl>
                      {readonly ? (
                        <Input value={field.value} disabled readOnly className="h-auto" />
                      ) : (
                        <Select
                          disabled={readonly}
                          onValueChange={field.onChange}
                          defaultValue={field.value}
                        >
                          <SelectTrigger className="h-[34px]">
                            <SelectValue />
                          </SelectTrigger>
                          <SelectContent>
                            <ScrollArea className="h-52">
                              {['integer', ...POSTGRES_DATA_TYPES].map((option) => (
                                <SelectItem value={option} key={option}>
                                  {option}
                                </SelectItem>
                              ))}
                            </ScrollArea>
                          </SelectContent>
                        </Select>
                      )}
                    </FormControl>
                    <FormMessage />
                  </FormItem>
                )}
              />

              {!readonly && (
                <Button
                  variant="default"
                  icon={<Trash size={12} />}
                  onClick={() => remove(index)}
                  className="h-[34px] w-[34px]"
                />
              )}
            </div>
          )
        })}

        {!readonly && (
          <Button
            variant="default"
            icon={<Plus size={12} />}
            onClick={() => append({ name: '', type: 'integer' })}
            disabled={readonly}
          >
            Add a new argument
          </Button>
        )}
      </div>
    </div>
  )
}

const ALL_ALLOWED_LANGUAGES = ['plpgsql', 'sql', 'plcoffee', 'plv8', 'plls']

const FormFieldLanguage = () => {
  const { data: project } = useSelectedProjectQuery()

  const { data: enabledExtensions } = useDatabaseExtensionsQuery(
    {
      projectRef: project?.ref,
      connectionString: project?.connectionString,
    },
    {
      select(data) {
        return partition(data, (ext) => !isNull(ext.installed_version))[0]
      },
    }
  )

  const allowedLanguages = useMemo(() => {
    return ALL_ALLOWED_LANGUAGES.filter((lang) => {
      if (lang.startsWith('pl')) {
        return enabledExtensions?.find((ex) => ex.name === lang) !== undefined
      }

      return true
    })
  }, [enabledExtensions])

  return (
    <FormField
      name="language"
      render={({ field }) => (
        <FormItemLayout label="Language" layout="horizontal">
          {/* Form selects don't need form controls, otherwise the CSS gets weird */}
          <Select onValueChange={field.onChange} defaultValue={field.value}>
            <SelectTrigger className="col-span-8">
              <SelectValue />
            </SelectTrigger>
            <SelectContent>
              {allowedLanguages.map((option) => (
                <SelectItem value={option} key={option}>
                  {option}
                </SelectItem>
              ))}
            </SelectContent>
          </Select>
        </FormItemLayout>
      )}
    />
  )
}