import { zodResolver } from '@hookform/resolvers/zod' import { literal, safeSql } from '@supabase/pg-meta' import { PermissionAction } from '@supabase/shared-types/out/constants' import { Plus, Trash2 } from 'lucide-react' import { parseAsBoolean, useQueryState } from 'nuqs' import { useEffect } from 'react' import { SubmitHandler, useFieldArray, useForm } from 'react-hook-form' import { toast } from 'sonner' import { Button, Form, FormControl, FormField, Input, RadioGroupStacked, RadioGroupStackedItem, Separator, Sheet, SheetContent, SheetFooter, SheetHeader, SheetSection, SheetTitle, SheetTrigger, } from 'ui' import { Admonition } from 'ui-patterns' import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout' import z from 'zod' import { inverseValidBucketNameRegex } from '../CreateBucketModal.utils' import { useS3VectorsWrapperInstance } from './useS3VectorsWrapperInstance' import { ButtonTooltip } from '@/components/ui/ButtonTooltip' import { DocsButton } from '@/components/ui/DocsButton' import { useFDWImportForeignSchemaMutation } from '@/data/fdw/fdw-import-foreign-schema-mutation' import { useVectorBucketIndexCreateMutation } from '@/data/storage/vector-bucket-index-create-mutation' import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions' import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' import { DOCS_URL } from '@/lib/constants' const isStagingLocal = process.env.NEXT_PUBLIC_ENVIRONMENT !== 'prod' const BUCKET_INDEX_NAME_REGEX = /^[a-z0-9](?:[a-z0-9.-]{1,61})?[a-z0-9]$/ const DISTANCE_METRICS = [ { value: 'cosine', label: 'Cosine', description: 'Measures similarity between two vectors, based on directions, not magnitude.', }, { value: 'euclidean', label: 'Euclidean', description: 'Measures straight-line distance between two vectors, using both directions and magnitudes.', }, ] as const const FormSchema = z.object({ name: z .string() .trim() .min(3, 'Name must be at least 3 characters') .max(63, 'Name must be below 63 characters') .refine( (value) => value !== 'public', '"public" is a reserved name. Please choose another name' ) .superRefine((name, ctx) => { if (!BUCKET_INDEX_NAME_REGEX.test(name)) { const [match] = name.match(inverseValidBucketNameRegex) ?? [] ctx.addIssue({ path: [], code: z.ZodIssueCode.custom, message: !!match ? `Bucket name cannot contain the "${match}" character` : 'Bucket name contains an invalid special character', }) } }), dimension: z .number() .int('Dimension must be an integer') .min(1, 'Dimension must be at least 1') .max(4096, 'Dimension must be at most 4096'), distanceMetric: z.enum(['cosine', 'euclidean'], { required_error: 'Please select a distance metric', }), metadataKeys: z .array( z.object({ value: z.string().min(1, 'The metadata key needs to be at least 1 character long'), }) ) .default([]), }) const formId = 'create-vector-table-form' export type CreateVectorTableForm = z.infer interface CreateVectorTableSheetProps { bucketName?: string } export const CreateVectorTableSheet = ({ bucketName }: CreateVectorTableSheetProps) => { const { data: project } = useSelectedProjectQuery() const [visible, setVisible] = useQueryState( 'newTable', parseAsBoolean.withDefault(false).withOptions({ history: 'push', clearOnDefault: true }) ) const { can: canCreateBuckets } = useAsyncCheckPermissions(PermissionAction.STORAGE_WRITE, '*') const { data: wrapperInstance } = useS3VectorsWrapperInstance({ bucketId: bucketName }) const schema = (wrapperInstance?.server_options ?? []) .find((x) => x.startsWith('briven_target_schema')) ?.split('briven_target_schema=')[1] // [Joshen] Can remove this once this restriction is removed const showIndexCreationNotice = isStagingLocal && !!project && project?.region !== 'us-east-1' const defaultValues = { name: '', dimension: undefined, distanceMetric: 'cosine' as 'cosine' | 'euclidean', metadataKeys: [], } const form = useForm({ resolver: zodResolver(FormSchema as any), defaultValues, values: defaultValues as any, }) const { fields, append, remove } = useFieldArray({ control: form.control, name: 'metadataKeys', }) const { mutateAsync: createVectorBucketTable, isPending: isCreatingVectorBucketTable } = useVectorBucketIndexCreateMutation() const { mutateAsync: importForeignSchema, isPending: isImportingForeignSchema } = useFDWImportForeignSchemaMutation({ onError: () => {}, }) const isCreating = isCreatingVectorBucketTable || isImportingForeignSchema const onSubmit: SubmitHandler = async (values) => { if (!project?.ref) return console.error('Project ref is required') if (!bucketName) return console.error('Bucket name is required') try { await createVectorBucketTable({ projectRef: project.ref, bucketName: bucketName, indexName: values.name, dataType: 'float32', dimension: values.dimension!, distanceMetric: values.distanceMetric, metadataKeys: values.metadataKeys.map((key) => key.value), }) } catch (error: any) { toast.error(`Failed to create vector table: ${error.message}`) return } try { if (wrapperInstance && !!schema) { await importForeignSchema({ projectRef: project.ref, connectionString: project?.connectionString, serverName: wrapperInstance.server_name, sourceSchema: schema, targetSchema: schema, schemaOptions: [safeSql`bucket_name ${literal(bucketName)}`], }) } } catch (error: any) { toast.warning(`Failed to connect vector table to the database: ${error.message}`) } toast.success(`Successfully created vector table “${values.name}”`) form.reset() setVisible(false) } useEffect(() => { if (!visible) { form.reset() } }, [visible]) return ( } disabled={!canCreateBuckets} onClick={() => setVisible(true)} tooltip={{ content: { side: 'bottom', text: !canCreateBuckets ? 'You need additional permissions to create buckets' : undefined, }, }} > Create table Create vector table {showIndexCreationNotice && ( )}
( )} /> ( { const value = e.target.value field.onChange(value === '' ? undefined : Number(value)) }} value={field.value ?? ''} /> )} /> ( {DISTANCE_METRICS.map((metric) => ( ))} )} />
{fields.map((field, index) => (
( )} />
))}
) }