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

import { zodResolver } from '@hookform/resolvers/zod'
import { useEffect, useRef } from 'react'
import { SubmitHandler, useForm, useWatch } from 'react-hook-form'
import { toast } from 'sonner'
import {
  Button,
  Form,
  FormControl,
  FormField,
  Input,
  RadioGroupStacked,
  RadioGroupStackedItem,
  SheetFooter,
  SheetHeader,
  SheetSection,
  SheetTitle,
} from 'ui'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import * as z from 'zod'

import { CreateWrapperSheetProps } from './CreateWrapperSheet'
import InputField from './InputField'
import {
  FormSection,
  FormSectionContent,
  FormSectionLabel,
} from '@/components/ui/Forms/FormSection'
import { useSchemaCreateMutation } from '@/data/database/schema-create-mutation'
import { useSchemasQuery } from '@/data/database/schemas-query'
import { useFDWCreateMutation } from '@/data/fdw/fdw-create-mutation'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { useTrack } from '@/lib/telemetry/track'

const FORM_ID = 'create-wrapper-form'

const S3TableSchema = z.object({
  target: z.literal('S3Tables'),
  source_schema: z.string().min(1, 'Please provide a namespace name'),
  wrapper_name: z.string().min(1, 'Please provide a name for your wrapper'),
  target_schema: z.string().min(1, 'Please provide an unique target schema'),
  vault_aws_access_key_id: z.string().min(1, 'Required'),
  vault_aws_secret_access_key: z.string().min(1, 'Required'),
  region_name: z.string().min(1, 'Required'),
  vault_aws_s3table_bucket_arn: z.string().min(1, 'Required'),
})
const R2CatalogSchema = z.object({
  target: z.literal('R2Catalog'),
  source_schema: z.string().min(1, 'Please provide a namespace name'),
  wrapper_name: z.string().min(1, 'Please provide a name for your wrapper'),
  target_schema: z.string().min(1, 'Please provide an unique target schema'),
  vault_aws_access_key_id: z.string().min(1, 'Required'),
  vault_aws_secret_access_key: z.string().min(1, 'Required'),
  vault_token: z.string().min(1, 'Required'),
  warehouse: z.string().min(1, 'Required'),
  s3: z.object({ endpoint: z.string().min(1, 'Required') }),
  catalog_uri: z.string().min(1, 'Required'),
})
const IcebergRestCatalogSchema = z.object({
  target: z.literal('IcebergRestCatalog'),
  source_schema: z.string().min(1, 'Please provide a namespace name'),
  wrapper_name: z.string().min(1, 'Please provide a name for your wrapper'),
  target_schema: z.string().min(1, 'Please provide an unique target schema'),
  vault_aws_access_key_id: z.string().optional(),
  vault_aws_secret_access_key: z.string().optional(),
  region_name: z.string().optional(),
  vault_aws_s3table_bucket_arn: z.string().optional(),
  vault_token: z.string().optional(),
  warehouse: z.string().optional(),
  s3: z.object({ endpoint: z.string().min(1, 'Required') }),
  catalog_uri: z.string().optional(),
})
const formSchema = z.discriminatedUnion('target', [
  S3TableSchema,
  R2CatalogSchema,
  IcebergRestCatalogSchema,
])

type FormSchema = z.infer<typeof formSchema>

const targetFields: Record<Target, { name: string; required: boolean }[]> = {
  S3Tables: [
    { name: 'vault_aws_access_key_id', required: true },
    { name: 'vault_aws_secret_access_key', required: true },
    { name: 'region_name', required: true },
    { name: 'vault_aws_s3table_bucket_arn', required: true },
  ],
  R2Catalog: [
    { name: 'vault_aws_access_key_id', required: true },
    { name: 'vault_aws_secret_access_key', required: true },
    { name: 'vault_token', required: true },
    { name: 'warehouse', required: true },
    { name: 's3.endpoint', required: true },
    { name: 'catalog_uri', required: true },
  ],
  IcebergRestCatalog: [
    { name: 'vault_aws_access_key_id', required: false },
    { name: 'vault_aws_secret_access_key', required: false },
    { name: 'region_name', required: false },
    { name: 'vault_aws_s3table_bucket_arn', required: false },
    { name: 'vault_token', required: false },
    { name: 'warehouse', required: false },
    { name: 's3.endpoint', required: false },
    { name: 'catalog_uri', required: false },
  ],
} as const

type Target = 'S3Tables' | 'R2Catalog' | 'IcebergRestCatalog'

const INITIAL_VALUES = {
  wrapper_name: '',
  source_schema: '',
  target_schema: '',
  target: 'S3Tables',
  vault_aws_access_key_id: '',
  vault_aws_s3table_bucket_arn: '',
  vault_aws_secret_access_key: '',
  region_name: '',
} satisfies FormSchema

export const CreateIcebergWrapperSheet = ({
  wrapperMeta,
  onDirty,
  onClose,
  onCloseWithConfirmation,
}: CreateWrapperSheetProps) => {
  const { data: project } = useSelectedProjectQuery()
  const track = useTrack()

  const { mutateAsync: createFDW, isPending: isCreatingWrapper } = useFDWCreateMutation({
    onSuccess: () => {
      toast.success(`Successfully created ${wrapperMeta?.label} foreign data wrapper`)
      onClose()
    },
  })

  const { data: schemas } = useSchemasQuery({
    projectRef: project?.ref!,
    connectionString: project?.connectionString,
  })

  const { mutateAsync: createSchema } = useSchemaCreateMutation()

  const form = useForm<FormSchema>({
    resolver: zodResolver(formSchema),
    defaultValues: INITIAL_VALUES,
  })
  const { resetField, formState, setError } = form
  const { isDirty, isSubmitting } = formState

  useEffect(() => {
    onDirty(isDirty)
  }, [onDirty, isDirty])

  const currentTarget = useRef<FormSchema['target']>(INITIAL_VALUES.target)
  useEffect(() => {
    return form.subscribe({
      name: 'target',
      formState: { values: true },
      callback: ({ values }) => {
        if (!values.target || values.target === currentTarget.current) return
        currentTarget.current = values.target

        const fields = targetFields[values.target]
        if (!fields) return

        wrapperMeta.server.options.forEach((option) => {
          // @ts-expect-error Can't reconcile with form schema
          resetField(option.name, { defaultValue: option.defaultValue ?? '' })
        })
      },
    })
  }, [form, resetField, wrapperMeta])

  const onSubmit: SubmitHandler<FormSchema> = async (values) => {
    const foundSchema = schemas?.find((s) => s.name === values.target_schema)
    if (foundSchema) {
      setError('target_schema', {
        type: 'validate',
        message: 'This schema already exists. Please specify a unique schema name.',
      })
      return
    }

    let formValues: Record<string, string> = {}
    if (values.target === 'R2Catalog' || values.target === 'IcebergRestCatalog') {
      const { s3, ...otherFormValues } = values
      formValues = otherFormValues
      formValues['s3.endpoint'] = s3.endpoint
    } else {
      formValues = values
    }

    try {
      await createSchema({
        projectRef: project?.ref,
        connectionString: project?.connectionString,
        name: values.target_schema,
      })

      await createFDW({
        projectRef: project?.ref,
        connectionString: project?.connectionString,
        wrapperMeta,
        formState: {
          ...formValues,
          server_name: `${values.wrapper_name}_server`,
          supabase_target_schema: values.target_schema,
        },
        mode: 'schema',
        tables: [],
        sourceSchema: values.source_schema,
        targetSchema: values.target_schema,
      })

      track('foreign_data_wrapper_created', { wrapperType: wrapperMeta.label })
    } catch (error) {
      console.error(error)
      // The error will be handled by the mutation onError callback (toast.error)
    }
  }

  const isLoading = isCreatingWrapper || isSubmitting
  const wrapperName = useWatch({ name: 'wrapper_name', control: form.control })
  const target = useWatch({ name: 'target', control: form.control })

  const targetOptions = wrapperMeta.server.options
    .filter((option) => targetFields[target].find((field) => field.name === option.name))
    .map((option) => {
      return {
        ...option,
        required: !!targetFields[target].find((field) => field.name === option.name)?.required,
      }
    })
  return (
    <>
      <div className="h-full" tabIndex={-1}>
        <Form {...form}>
          <form
            id={FORM_ID}
            onSubmit={form.handleSubmit(onSubmit)}
            className="flex flex-col h-full"
          >
            <SheetHeader>
              <SheetTitle>Create a {wrapperMeta.label} wrapper</SheetTitle>
            </SheetHeader>
            <SheetSection 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="horizontal"
                        label="Wrapper Name"
                        description={
                          wrapperName.length > 0 ? (
                            <>
                              Your wrapper's server name will be{' '}
                              <code className="text-code-inline">{wrapperName}_server</code>
                            </>
                          ) : (
                            ''
                          )
                        }
                      >
                        <FormControl>
                          <Input {...field} />
                        </FormControl>
                      </FormItemLayout>
                    )}
                  />
                </FormSectionContent>
              </FormSection>
              <FormSection header={<FormSectionLabel>Data target</FormSectionLabel>}>
                <FormSectionContent className="flex flex-col space-y-2" loading={false}>
                  <FormField
                    control={form.control}
                    name="target"
                    render={({ field }) => (
                      <FormItemLayout layout="vertical">
                        <div>
                          <RadioGroupStacked value={field.value} onValueChange={field.onChange}>
                            <RadioGroupStackedItem
                              key="S3Tables"
                              value="S3Tables"
                              label="AWS S3 Tables"
                              showIndicator={false}
                            >
                              <div className="flex gap-x-5">
                                <div className="flex flex-col">
                                  <p className="text-foreground-light text-left">
                                    AWS S3 storage that's optimized for analytics workloads.
                                  </p>
                                </div>
                              </div>
                            </RadioGroupStackedItem>
                            <RadioGroupStackedItem
                              key="R2Catalog"
                              value="R2Catalog"
                              label="Cloudflare R2 Catalog"
                              showIndicator={false}
                            >
                              <div className="flex gap-x-5">
                                <div className="flex flex-col">
                                  <p className="text-foreground-light text-left">
                                    Managed Apache Iceberg built directly into your R2 bucket.
                                  </p>
                                </div>
                              </div>
                            </RadioGroupStackedItem>
                            <RadioGroupStackedItem
                              key="IcebergRestCatalog"
                              value="IcebergRestCatalog"
                              label="Iceberg REST Catalog"
                              showIndicator={false}
                            >
                              <div className="flex gap-x-5">
                                <div className="flex flex-col">
                                  <p className="text-foreground-light text-left">
                                    Can be used with any S3-compatible storage.
                                  </p>
                                </div>
                              </div>
                            </RadioGroupStackedItem>
                          </RadioGroupStacked>
                        </div>
                      </FormItemLayout>
                    )}
                  />
                </FormSectionContent>
              </FormSection>

              <FormSection
                header={<FormSectionLabel>{wrapperMeta.label} Configuration</FormSectionLabel>}
              >
                <FormSectionContent className="flex flex-col space-y-2" loading={false}>
                  {targetOptions.map((option) =>
                    option.hidden ? (
                      <input
                        key={`${option.name}-${option.required}-${option.hidden}`}
                        type="hidden"
                        // @ts-expect-error Can't reconcile with form schema
                        {...form.register(option.name)}
                      />
                    ) : (
                      <InputField
                        key={`${option.name}-${option.required}-${option.hidden}`}
                        control={form.control}
                        option={option}
                      />
                    )
                  )}
                </FormSectionContent>
              </FormSection>
              <FormSection
                header={
                  <FormSectionLabel>
                    <p>Foreign Schema</p>
                    <p className="text-foreground-light mt-2 w-[90%]">
                      You can query your data from the foreign tables in the specified schema after
                      the wrapper is created.
                    </p>
                  </FormSectionLabel>
                }
              >
                <FormSectionContent className="flex flex-col space-y-2" loading={false}>
                  {wrapperMeta.sourceSchemaOption && (
                    <InputField control={form.control} option={wrapperMeta.sourceSchemaOption} />
                  )}
                  <InputField
                    control={form.control}
                    option={{
                      name: 'target_schema',
                      label: 'Specify a new schema to create all wrapper tables in',
                      description:
                        'A new schema will be created. For security purposes, the wrapper tables from the foreign schema cannot be created within an existing schema.',
                      required: true,
                      encrypted: false,
                      secureEntry: false,
                    }}
                  />
                </FormSectionContent>
              </FormSection>
            </SheetSection>

            <SheetFooter>
              <Button
                size="tiny"
                variant="default"
                type="button"
                onClick={onCloseWithConfirmation}
                disabled={isLoading}
              >
                Cancel
              </Button>
              <Button
                size="tiny"
                variant="primary"
                form={FORM_ID}
                type="submit"
                loading={isLoading}
                disabled={isLoading || !isDirty}
              >
                Create wrapper
              </Button>
            </SheetFooter>
          </form>
        </Form>
      </div>
    </>
  )
}