supabase/apps/studio/components/interfaces/Integrations/Wrappers/EditWrapperSheet.tsx

import { zodResolver } from '@hookform/resolvers/zod'
import { useQueryClient } from '@tanstack/react-query'
import { compact } from 'lodash'
import { Edit, 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,
  Form,
  FormControl,
  FormField,
  Input,
  Separator,
  SheetFooter,
  SheetHeader,
  SheetTitle,
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from 'ui'
import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import * as z from 'zod'

import InputField from './InputField'
import { WrapperMeta } from './Wrappers.types'
import {
  convertKVStringArrayToJson,
  FormattedWrapperTable,
  formatWrapperTables,
  getEditionFormSchema,
  NewTable,
} from './Wrappers.utils'
import WrapperTableEditor from './WrapperTableEditor'
import { DiscardChangesConfirmationDialog } from '@/components/ui-patterns/Dialogs/DiscardChangesConfirmationDialog'
import { ButtonTooltip } from '@/components/ui/ButtonTooltip'
import {
  FormSection,
  FormSectionContent,
  FormSectionLabel,
} from '@/components/ui/Forms/FormSection'
import { invalidateSchemasQuery } from '@/data/database/schemas-query'
import { useFDWUpdateMutation } from '@/data/fdw/fdw-update-mutation'
import { FDW } from '@/data/fdw/fdws-query'
import { getDecryptedValues } from '@/data/vault/vault-secret-decrypted-value-query'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { useConfirmOnClose } from '@/hooks/ui/useConfirmOnClose'
import { UUID_REGEX } from '@/lib/constants'

export interface EditWrapperSheetProps {
  wrapper: FDW
  isClosing: boolean
  wrapperMeta: WrapperMeta
  setIsClosing: (v: boolean) => void
  onClose: () => void
}

const FORM_ID = 'edit-wrapper-form'

export const EditWrapperSheet = ({
  wrapper,
  wrapperMeta,
  isClosing,
  setIsClosing,
  onClose,
}: EditWrapperSheetProps) => {
  const queryClient = useQueryClient()
  const { data: project } = useSelectedProjectQuery()

  const { mutate: updateFDW, isPending: isSaving } = useFDWUpdateMutation({
    onSuccess: () => {
      toast.success(`Successfully updated ${wrapperMeta?.label} foreign data wrapper`)
      const { tables } = getValues()
      const hasNewSchema = (tables as Record<string, any>[]).some((table) => table.is_new_schema)
      if (hasNewSchema) invalidateSchemasQuery(queryClient, project?.ref)

      onClose()
    },
  })

  const initialValues: Record<string, any> = useMemo(
    () => ({
      wrapper_name: wrapper?.name,
      server_name: wrapper?.server_name,
      ...convertKVStringArrayToJson(wrapper?.server_options ?? []),
      tables: formatWrapperTables(wrapper, wrapperMeta),
    }),
    [wrapper, wrapperMeta]
  )

  const formSchema = getEditionFormSchema(wrapperMeta)
  type FormSchema = z.infer<typeof formSchema>
  const form = useForm<FormSchema>({
    defaultValues: initialValues,
    resolver: zodResolver(formSchema),
  })

  const { getValues, resetField, setError } = form
  const { errors, isDirty, isSubmitting } = form.formState

  const {
    fields: tablesField,
    append: appendTable,
    remove: removeTable,
    update: updateTable,
  } = useFieldArray({
    control: form.control,
    name: 'tables',
  })

  const [selectedTableToEdit, setSelectedTableToEdit] = useState<FormattedWrapperTable | undefined>(
    undefined
  )
  const [isUpdateConfirmationOpen, setIsUpdateConfirmationOpen] = useState(false)

  const onUpdateTable = (values: FormattedWrapperTable) => {
    if (values.index !== undefined) {
      updateTable(values.index, values)
    } else {
      appendTable(values)
    }
    setSelectedTableToEdit(undefined)
  }

  const onSubmit: SubmitHandler<FormSchema> = async (values) => {
    const { tables } = values
    if (tables.length === 0) {
      setError('tables', {
        type: 'validate',
        message: 'Please provide at least one table.',
      })
      return
    }
    setIsUpdateConfirmationOpen(true)
  }

  const { confirmOnClose, modalProps } = useConfirmOnClose({
    checkIsDirty: () => isDirty,
    onClose,
  })

  useEffect(() => {
    if (!isClosing) return
    if (isDirty) {
      confirmOnClose()
    } else {
      onClose()
    }
    setIsClosing(false)
  }, [isDirty, confirmOnClose, isClosing, onClose, setIsClosing])

  const wrapper_name = useWatch({ name: 'wrapper_name', control: form.control })

  const [isLoadingSecrets, setIsLoadingSecrets] = useState(false)
  const [secretsReady, setSecretsReady] = useState(false)

  useEffect(() => {
    let isCurrent = true

    const encryptedOptions = wrapperMeta.server.options.filter((option) => option.encrypted)

    const encryptedIdsToFetch = compact(
      encryptedOptions.map((option) => {
        const value = initialValues[option.name]
        return value ?? null
      })
    ).filter((x) => UUID_REGEX.test(x))

    if (encryptedIdsToFetch.length === 0) {
      setSecretsReady(true)
      setIsLoadingSecrets(false)
      return
    }

    setSecretsReady(false)

    const fetchEncryptedValues = async (ids: string[]) => {
      try {
        setIsLoadingSecrets(true)
        // If the secrets haven't loaded, escape and run the effect again when they're loaded
        const decryptedValues = await getDecryptedValues({
          projectRef: project?.ref,
          connectionString: project?.connectionString,
          ids: ids,
        })
        if (!isCurrent) return

        encryptedOptions.forEach((option) => {
          const encryptedId = initialValues[option.name]

          resetField(option.name, { defaultValue: decryptedValues[encryptedId] })
        })
        setSecretsReady(true)
      } catch (error) {
        if (!isCurrent) return
        toast.error('Failed to fetch encrypted values')
      } finally {
        if (isCurrent) setIsLoadingSecrets(false)
      }
    }

    fetchEncryptedValues(encryptedIdsToFetch)

    return () => {
      isCurrent = false
    }
  }, [initialValues, wrapperMeta, resetField, project?.ref, project?.connectionString])

  return (
    <>
      <div className="flex flex-col h-full" tabIndex={-1}>
        <Form {...form}>
          <form
            id={FORM_ID}
            onSubmit={form.handleSubmit(onSubmit)}
            className="flex flex-col h-full"
          >
            <SheetHeader>
              <SheetTitle>
                Edit {wrapperMeta.label} wrapper: {wrapper.name}
              </SheetTitle>
            </SheetHeader>
            <div className="grow overflow-y-auto">
              <FormSection header={<FormSectionLabel>Wrapper Configuration</FormSectionLabel>}>
                <FormSectionContent className="flex flex-col space-y-2" loading={false}>
                  <FormField
                    control={form.control}
                    name="wrapper_name"
                    render={({ field }) => (
                      <FormItemLayout
                        layout="vertical"
                        label="Wrapper Name"
                        description={
                          wrapper_name !== initialValues.wrapper_name ? (
                            <>
                              Your wrapper's server name will be updated to{' '}
                              <code className="text-code-inline">{wrapper_name}_server</code>
                            </>
                          ) : (
                            <>
                              Your wrapper's server name is{' '}
                              <code className="text-code-inline">{wrapper_name}_server</code>
                            </>
                          )
                        }
                      >
                        <FormControl>
                          <Input {...field} />
                        </FormControl>
                      </FormItemLayout>
                    )}
                  />
                </FormSectionContent>
              </FormSection>
              <Separator />

              <FormSection
                header={<FormSectionLabel>{wrapperMeta.label} Configuration</FormSectionLabel>}
              >
                <FormSectionContent className="flex flex-col space-y-2" loading={false}>
                  {wrapperMeta.server.options
                    .filter((option) => !option.hidden)
                    .map((option) => (
                      <InputField
                        key={option.name}
                        option={option}
                        control={form.control}
                        placeholder={option.defaultValue}
                        loading={option.secureEntry ? isLoadingSecrets : undefined}
                      />
                    ))}
                </FormSectionContent>
              </FormSection>
              <Separator />
              <FormSection
                header={
                  <FormSectionLabel>
                    <p>Foreign Tables</p>
                    <p className="text-foreground-light mt-2 w-[90%]">
                      You can query your data from these foreign tables after the wrapper is created
                    </p>
                  </FormSectionLabel>
                }
              >
                <FormSectionContent className="flex flex-col space-y-2" loading={false}>
                  {tablesField.map((t, tableIndex) => {
                    // FIXME: make inference work
                    const table = t as unknown as FormattedWrapperTable
                    return (
                      <div
                        key={t.id}
                        className="flex items-center justify-between px-4 py-2 border rounded-md border-control"
                      >
                        <div>
                          <p className="text-sm">
                            {table.schema_name}.{table.table_name}
                          </p>
                          <p className="text-sm text-foreground-light">
                            Columns:{' '}
                            {(table.columns ?? []).map((column: any) => column.name).join(', ')}
                          </p>
                        </div>
                        <div className="flex items-center space-x-2">
                          <Tooltip>
                            <TooltipTrigger asChild>
                              <Button
                                className="px-1"
                                icon={<Edit />}
                                onClick={() => {
                                  setSelectedTableToEdit(table)
                                }}
                                aria-label={`Edit ${table.table_name} foreign table`}
                                // Tooltip repeats the label; screen readers would read it twice
                                aria-describedby={undefined}
                              />
                            </TooltipTrigger>
                            <TooltipContent side="bottom">{`Edit ${table.table_name} foreign table`}</TooltipContent>
                          </Tooltip>
                          <Tooltip>
                            <TooltipTrigger asChild>
                              <Button
                                className="px-1"
                                icon={<Trash />}
                                onClick={() => {
                                  removeTable(tableIndex)
                                }}
                                aria-label={`Remove ${table.table_name} foreign table`}
                                // Tooltip repeats the label; screen readers would read it twice
                                aria-describedby={undefined}
                              />
                            </TooltipTrigger>
                            <TooltipContent side="bottom">{`Remove ${table.table_name} foreign table`}</TooltipContent>
                          </Tooltip>
                        </div>
                      </div>
                    )
                  })}

                  <div className="flex justify-end">
                    <Button onClick={() => setSelectedTableToEdit(NewTable)}>
                      Add foreign table
                    </Button>
                  </div>
                  {tablesField.length === 0 && errors.tables && (
                    <p className="text-sm text-right text-red-900">
                      {errors.tables.message?.toString()}
                    </p>
                  )}
                </FormSectionContent>
              </FormSection>
            </div>
            <SheetFooter>
              <Button size="tiny" type="button" onClick={confirmOnClose} disabled={isSubmitting}>
                Cancel
              </Button>
              <ButtonTooltip
                size="tiny"
                variant="primary"
                form={FORM_ID}
                type="submit"
                disabled={isSubmitting || !isDirty || !secretsReady}
                loading={isSubmitting}
                tooltip={{
                  content: {
                    side: 'top',
                    text: !secretsReady ? 'Waiting for encrypted values to load' : undefined,
                  },
                }}
              >
                Save wrapper
              </ButtonTooltip>
            </SheetFooter>
          </form>
        </Form>
      </div>

      <ConfirmationModal
        visible={isUpdateConfirmationOpen}
        title="Save wrapper changes?"
        size="small"
        variant="warning"
        confirmLabel="Save changes"
        confirmLabelLoading="Saving changes"
        loading={isSaving}
        onCancel={() => {
          setIsUpdateConfirmationOpen(false)
          onClose()
        }}
        onConfirm={() => {
          const { tables, ...values } = getValues()
          updateFDW({
            projectRef: project?.ref,
            connectionString: project?.connectionString,
            wrapper,
            wrapperMeta,
            formState: values,
            tables,
          })
          setIsUpdateConfirmationOpen(false)
        }}
      >
        <p className="text-sm text-foreground-light">
          Removing a table or retyping a column may break views or functions that reference it.
        </p>
        <p className="text-sm text-foreground-light mt-2">Are you sure you want to continue?</p>
      </ConfirmationModal>

      <DiscardChangesConfirmationDialog {...modalProps} />

      <WrapperTableEditor
        visible={selectedTableToEdit != null}
        tables={wrapperMeta.tables}
        onCancel={() => {
          setSelectedTableToEdit(undefined)
        }}
        onSave={onUpdateTable}
        initialData={selectedTableToEdit}
      />
    </>
  )
}