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

import { useQueryClient } from '@tanstack/react-query'
import { compact, isEmpty, mapValues } from 'lodash'
import { Edit, Trash } from 'lucide-react'
import { useCallback, useEffect, useRef, useState } from 'react'
import { toast } from 'sonner'

import { UUID_REGEX } from '@/lib/constants'
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 { Button, Form, Input, SheetFooter, SheetHeader, SheetTitle } from 'ui'
import { DiscardChangesConfirmationDialog } from 'components/ui-patterns/Dialogs/DiscardChangesConfirmationDialog'
import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
import InputField from './InputField'
import { WrapperMeta } from './Wrappers.types'
import {
  convertKVStringArrayToJson,
  FormattedWrapperTable,
  formatWrapperTables,
  makeValidateRequired,
} from './Wrappers.utils'
import WrapperTableEditor from './WrapperTableEditor'

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`)
      setWrapperTables([])

      const hasNewSchema = wrapperTables.some((table) => table.is_new_schema)
      if (hasNewSchema) invalidateSchemasQuery(queryClient, project?.ref)
    },
  })

  const [wrapperTables, setWrapperTables] = useState(() =>
    formatWrapperTables(wrapper, wrapperMeta)
  )
  const [isEditingTable, setIsEditingTable] = useState(false)
  const [selectedTableToEdit, setSelectedTableToEdit] = useState<FormattedWrapperTable | undefined>(
    undefined
  )
  const [formErrors, setFormErrors] = useState<{ [k: string]: string }>({})
  const [isUpdateConfirmationOpen, setIsUpdateConfirmationOpen] = useState(false)
  const [pendingFormState, setPendingFormState] = useState<Record<string, string> | null>(null)
  const hasChangesRef = useRef(false)

  const initialValues = {
    wrapper_name: wrapper?.name,
    server_name: wrapper?.server_name,
    ...convertKVStringArrayToJson(wrapper?.server_options ?? []),
  }

  const onUpdateTable = (values: FormattedWrapperTable) => {
    setWrapperTables((prev) => {
      // if the new values have tableIndex, we are editing an existing table
      if (values.tableIndex !== undefined) {
        const tableIndex = values.tableIndex
        const wrapperTables = [...prev]
        delete values.tableIndex
        wrapperTables[tableIndex] = values
        return wrapperTables
      }
      return [...prev, values]
    })
    setIsEditingTable(false)
    setSelectedTableToEdit(undefined)
  }

  const onSubmit = async (values: Record<string, string>) => {
    const validate = makeValidateRequired(wrapperMeta.server.options)
    const errors = validate(values)

    const { wrapper_name } = values
    if (wrapper_name.length === 0) errors.name = 'Please provide a name for your wrapper'
    if (!wrapperMeta.canTargetSchema && wrapperTables.length === 0)
      errors.tables = 'Please add at least one table'
    if (!isEmpty(errors)) {
      setFormErrors(errors)
      return
    }

    setFormErrors({})
    setPendingFormState({ ...values, server_name: `${wrapper_name}_server` })
    setIsUpdateConfirmationOpen(true)
  }

  const checkIsDirty = useCallback(() => hasChangesRef.current, [])

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

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

  return (
    <>
      <div className="flex flex-col h-full" tabIndex={-1}>
        <Form
          id={FORM_ID}
          initialValues={initialValues}
          onSubmit={onSubmit}
          className="h-full flex flex-col"
        >
          {({
            values,
            initialValues,
            resetForm,
          }: {
            values: Record<string, string>
            initialValues: Record<string, string>
            resetForm: (value: Record<string, Record<string, string>>) => void
          }) => {
            // [Alaister] although this "technically" is breaking the rules of React hooks
            // it won't error because the hooks are always rendered in the same order
            // eslint-disable-next-line react-hooks/rules-of-hooks
            const [loadingSecrets, setLoadingSecrets] = useState(false)

            const initialTables = formatWrapperTables({
              handler: wrapper.handler,
              tables: wrapper?.tables ?? [],
            })
            const hasFormChanges = JSON.stringify(values) !== JSON.stringify(initialValues)
            const hasTableChanges = JSON.stringify(initialTables) !== JSON.stringify(wrapperTables)
            const hasChanges = hasFormChanges || hasTableChanges
            hasChangesRef.current = hasChanges

            // [Alaister] although this "technically" is breaking the rules of React hooks
            // it won't error because the hooks are always rendered in the same order
            // eslint-disable-next-line react-hooks/rules-of-hooks
            useEffect(() => {
              const fetchEncryptedValues = async (ids: string[]) => {
                try {
                  setLoadingSecrets(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,
                  })

                  // replace all values which are in the decryptedValues object with the decrypted value
                  const transformValues = (values: Record<string, string>) => {
                    return mapValues(values, (value) => {
                      return decryptedValues[value] ?? value
                    })
                  }

                  resetForm({
                    values: transformValues(values),
                    initialValues: transformValues(initialValues),
                  })
                } catch (error) {
                  toast.error('Failed to fetch encrypted values')
                } finally {
                  setLoadingSecrets(false)
                }
              }

              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))
              // [Joshen] ^ Validate UUID to filter out already decrypted values

              if (encryptedIdsToFetch.length > 0) {
                fetchEncryptedValues(encryptedIdsToFetch)
              }
              /**
               * [Joshen] We're deliberately not adding values and initialValues to the dependency array here
               * as we only want to fetch the encrypted values once on load + values and initialValues will be updated
               * as a result of that
               */
              // eslint-disable-next-line react-hooks/exhaustive-deps
            }, [project?.ref, project?.connectionString])

            return (
              <>
                <SheetHeader>
                  <SheetTitle>
                    Edit {wrapperMeta.label} wrapper: {wrapper.name}
                  </SheetTitle>
                </SheetHeader>
                <div className="flex-grow overflow-y-auto">
                  <FormSection header={<FormSectionLabel>Wrapper Configuration</FormSectionLabel>}>
                    <FormSectionContent loading={false}>
                      <Input
                        id="wrapper_name"
                        label="Wrapper Name"
                        error={formErrors.wrapper_name}
                        descriptionText={
                          values.wrapper_name !== initialValues.wrapper_name ? (
                            <>
                              Your wrapper's server name will be updated to{' '}
                              <code className="text-code-inline">{values.wrapper_name}_server</code>
                            </>
                          ) : (
                            <>
                              Your wrapper's server name is{' '}
                              <code className="text-code-inline">{values.wrapper_name}_server</code>
                            </>
                          )
                        }
                      />
                    </FormSectionContent>
                  </FormSection>
                  <FormSection
                    header={<FormSectionLabel>{wrapperMeta.label} Configuration</FormSectionLabel>}
                  >
                    <FormSectionContent loading={false}>
                      {wrapperMeta.server.options
                        .filter((option) => !option.hidden)
                        .map((option) => (
                          <InputField
                            key={option.name}
                            option={option}
                            loading={option.encrypted ? loadingSecrets : false}
                            error={formErrors[option.name]}
                          />
                        ))}
                    </FormSectionContent>
                  </FormSection>
                  <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 loading={false}>
                      {wrapperTables.length === 0 ? (
                        <div className="flex justify-end translate-y-4">
                          <Button type="default" onClick={() => setIsEditingTable(true)}>
                            Add foreign table
                          </Button>
                        </div>
                      ) : (
                        <div className="space-y-2">
                          {wrapperTables.map((table, i) => {
                            const target = table?.table ?? table.object

                            return (
                              <div
                                key={`${table.schema_name}.${table.table_name}`}
                                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 mt-1">
                                    Target: {target}
                                  </p>
                                  <p className="text-sm text-foreground-light">
                                    Columns: {table.columns.map((column) => column.name).join(', ')}
                                  </p>
                                </div>
                                <div className="flex items-center space-x-2">
                                  {/* Wrappers which import foreign schema don't have tables and their tables can't be edited */}
                                  {wrapperMeta.tables.length !== 0 && (
                                    <Button
                                      type="default"
                                      className="px-1"
                                      icon={<Edit />}
                                      onClick={() => {
                                        setIsEditingTable(true)
                                        setSelectedTableToEdit({ ...table, tableIndex: i })
                                      }}
                                    />
                                  )}
                                  <Button
                                    type="default"
                                    className="px-1"
                                    icon={<Trash />}
                                    onClick={() => {
                                      setWrapperTables((prev) => prev.filter((_, j) => j !== i))
                                    }}
                                  />
                                </div>
                              </div>
                            )
                          })}
                        </div>
                      )}
                      {wrapperTables.length > 0 && (
                        <div className="flex justify-end">
                          <Button type="default" onClick={() => setIsEditingTable(true)}>
                            Add foreign table
                          </Button>
                        </div>
                      )}
                      {wrapperTables.length === 0 && formErrors.tables && (
                        <p className="text-sm text-right text-red-900">{formErrors.tables}</p>
                      )}
                    </FormSectionContent>
                  </FormSection>
                </div>
                <SheetFooter>
                  <Button
                    size="tiny"
                    type="default"
                    htmlType="button"
                    onClick={confirmOnClose}
                    disabled={isSaving}
                  >
                    Cancel
                  </Button>
                  <Button
                    size="tiny"
                    type="primary"
                    form={FORM_ID}
                    htmlType="submit"
                    disabled={isSaving}
                    loading={isSaving}
                  >
                    Save wrapper
                  </Button>
                </SheetFooter>
              </>
            )
          }}
        </Form>
      </div>

      <ConfirmationModal
        visible={isUpdateConfirmationOpen}
        title="Recreate wrapper?"
        size="medium"
        variant="warning"
        confirmLabel="Recreate wrapper"
        confirmLabelLoading="Recreating wrapper"
        loading={isSaving}
        onCancel={() => {
          setIsUpdateConfirmationOpen(false)
          setPendingFormState(null)
          onClose()
        }}
        onConfirm={() => {
          if (pendingFormState === null) return
          updateFDW({
            projectRef: project?.ref,
            connectionString: project?.connectionString,
            wrapper,
            wrapperMeta,
            formState: pendingFormState,
            tables: wrapperTables,
          })
          setIsUpdateConfirmationOpen(false)
          setPendingFormState(null)
        }}
      >
        <p className="text-sm text-foreground-light">
          Saving changes will drop the existing wrapper and recreate it. Foreign servers and tables
          will be recreated, and dependent objects like functions or views that reference those
          tables may need to be updated manually afterwards.
        </p>
        <p className="text-sm text-foreground-light mt-2">Are you sure you want to continue?</p>
      </ConfirmationModal>

      <DiscardChangesConfirmationDialog {...modalProps} />

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