supabase/apps/studio/components/interfaces/Storage/ImportForeignSchemaDialog.tsx

import { zodResolver } from '@hookform/resolvers/zod'
import { useParams } from 'common'
import { uniq } from 'lodash'
import { useEffect, useState } from 'react'
import { SubmitHandler, useForm } from 'react-hook-form'
import { toast } from 'sonner'
import {
  Button,
  Dialog,
  DialogContent,
  DialogFooter,
  DialogHeader,
  DialogSection,
  DialogSectionSeparator,
  DialogTitle,
  Form,
  FormField,
  Input,
  Separator,
} from 'ui'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import z from 'zod'

import { formatWrapperTables } from '../Integrations/Wrappers/Wrappers.utils'
import { SchemaEditor } from '../TableGridEditor/SidePanelEditor/SchemaEditor'
import { getAnalyticsBucketFDWServerName } from './AnalyticsBuckets/AnalyticsBucketDetails/AnalyticsBucketDetails.utils'
import { useAnalyticsBucketAssociatedEntities } from './AnalyticsBuckets/AnalyticsBucketDetails/useAnalyticsBucketAssociatedEntities'
import { getDecryptedParameters } from './Storage.utils'
import { useSchemaCreateMutation } from '@/data/database/schema-create-mutation'
import { useSchemasQuery } from '@/data/database/schemas-query'
import { useFDWImportForeignSchemaMutation } from '@/data/fdw/fdw-import-foreign-schema-mutation'
import { useFDWUpdateMutation } from '@/data/fdw/fdw-update-mutation'
import { getFDWs } from '@/data/fdw/fdws-query'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'

export interface ImportForeignSchemaDialogProps {
  namespace: string
  circumstance?: 'fresh' | 'clash'
  visible: boolean
  onClose: () => void
}

export const ImportForeignSchemaDialog = ({
  namespace,
  visible,
  onClose,
  circumstance = 'fresh',
}: ImportForeignSchemaDialogProps) => {
  const { ref, bucketId: bucketName } = useParams()
  const { data: project } = useSelectedProjectQuery()
  const [loading, setLoading] = useState(false)
  const [createSchemaSheetOpen, setCreateSchemaSheetOpen] = useState(false)

  const { data: schemas } = useSchemasQuery({ projectRef: project?.ref! })
  const { icebergWrapperMeta: wrapperMeta } = useAnalyticsBucketAssociatedEntities({
    projectRef: ref,
    bucketId: bucketName,
  })

  const { mutateAsync: createSchema } = useSchemaCreateMutation()
  const { mutateAsync: importForeignSchema } = useFDWImportForeignSchemaMutation({})
  const { mutateAsync: updateFDW } = useFDWUpdateMutation({
    onSuccess: () => {
      toast.success(`Successfully connected “${bucketName}” to the database.`)
      onClose()
    },
  })

  const FormSchema = z.object({
    bucketName: z.string().trim(),
    sourceNamespace: z.string().trim(),
    targetSchema: z
      .string()
      .trim()
      .min(1, 'Schema name is required')
      .refine(
        (val) => {
          return !schemas?.find((s) => s.name === val)
        },
        {
          message: 'This schema already exists. Please specify a unique schema name.',
        }
      ),
  })

  const form = useForm<z.infer<typeof FormSchema>>({
    resolver: zodResolver(FormSchema),
    defaultValues: {
      bucketName,
      sourceNamespace: namespace,
      targetSchema: `fdw_analytics_${namespace}`,
    },
  })

  const onSubmit: SubmitHandler<z.infer<typeof FormSchema>> = async (values) => {
    const serverName = getAnalyticsBucketFDWServerName(values.bucketName)

    if (!ref) return console.error('Project ref is required')
    if (!wrapperMeta) return console.error('Wrapper meta is required')
    setLoading(true)

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

      await importForeignSchema({
        projectRef: ref,
        connectionString: project?.connectionString,
        serverName: serverName,
        sourceSchema: values.sourceNamespace,
        targetSchema: values.targetSchema,
      })

      const FDWs = await getFDWs({ projectRef: ref, connectionString: project?.connectionString })
      const wrapper = FDWs.find((fdw) => fdw.server_name === serverName)
      if (!wrapper) {
        throw new Error(`Foreign data wrapper with server name ${serverName} not found`)
      }

      const serverOptions = await getDecryptedParameters({
        ref: project?.ref,
        connectionString: project?.connectionString ?? undefined,
        wrapper,
        wrapperMeta,
      })

      const formValues: Record<string, string> = {
        wrapper_name: wrapper.name,
        server_name: wrapper.server_name,
        ...serverOptions,
      }

      const targetSchemas = (formValues['supabase_target_schema'] || '')
        .split(',')
        .map((s) => s.trim())

      const wrapperTables = formatWrapperTables(wrapper, wrapperMeta)

      await updateFDW({
        projectRef: project?.ref,
        connectionString: project?.connectionString,
        wrapper: wrapper,
        wrapperMeta: wrapperMeta,
        formState: {
          ...formValues,
          server_name: serverName,
          supabase_target_schema: uniq([...targetSchemas, values.targetSchema])
            .filter(Boolean)
            .join(','),
        },
        tables: wrapperTables,
      })
    } catch (error: any) {
      // error will be handled by the mutation onError callback
    } finally {
      setLoading(false)
    }
  }

  useEffect(() => {
    if (visible) {
      form.reset({
        bucketName,
        sourceNamespace: namespace,
        targetSchema: `fdw_analytics_${namespace}`,
      })
    }
  }, [visible, form, bucketName, namespace])

  return (
    <Dialog open={visible} onOpenChange={onClose}>
      <DialogContent size="medium">
        <DialogHeader>
          <DialogTitle>Create target schema</DialogTitle>
        </DialogHeader>
        <DialogSectionSeparator />
        <Form {...form}>
          <form onSubmit={form.handleSubmit(onSubmit)}>
            <DialogSection className="flex flex-col gap-y-4">
              <p className="text-sm">
                Namespace “<strong>{namespace}</strong>”{' '}
                {circumstance === 'fresh'
                  ? 'must be linked to a new schema before tables can be paired.'
                  : 'clashes with an existing database schema. Create a new schema to use as the destination for this data.'}
              </p>
              <Separator />
              <FormField
                control={form.control}
                name="targetSchema"
                render={({ field }) => (
                  <FormItemLayout
                    layout="vertical"
                    label="Target schema"
                    description="Where your analytics tables will be stored."
                  >
                    <Input {...field} placeholder="Enter schema name" />
                  </FormItemLayout>
                )}
              />
            </DialogSection>
            <DialogFooter className="flex items-center space-x-2 justify-end">
              <Button type="button" disabled={loading} onClick={onClose}>
                Cancel
              </Button>
              <Button variant="primary" type="submit" loading={loading}>
                Create
              </Button>
            </DialogFooter>
          </form>
        </Form>
        <SchemaEditor
          visible={createSchemaSheetOpen}
          closePanel={() => setCreateSchemaSheetOpen(false)}
          onSuccess={(schema) => {
            form.setValue('targetSchema', schema)
            setCreateSchemaSheetOpen(false)
          }}
        />
      </DialogContent>
    </Dialog>
  )
}