Chore/foreign key management in column editor (#21683)

* Support adding foreign key via ColumnEditor * Support updating/deleting foreign key from ColumnEditor
2026-10-05 17:35:10 +03:00 · 2024-03-04 16:00:00 +08:00
1 parent 74563fbeba
commit bb5108e154
@@ -159,16 +159,20 @@ const GridHeaderActions = ({ table, isViewSelected, isTableSelected }: GridHeade
            <div className="flex items-center gap-1">
              {policies.length < 1 ? (
                <Tooltip.Root delayDuration={0}>
-                  <Tooltip.Trigger className="w-full">
-                    <Link passHref href={`/project/${projectRef}/auth/policies?search=${table.id}`}>
-                      <Button
-                        type="default"
-                        className="group !h-[28px] !py-0"
-                        icon={<IconPlusCircle size={12} />}
+                  <Tooltip.Trigger asChild className="w-full">
+                    <Button
+                      asChild
+                      type="default"
+                      className="group !h-[28px] !py-0"
+                      icon={<IconPlusCircle size={12} />}
+                    >
+                      <Link
+                        passHref
+                        href={`/project/${projectRef}/auth/policies?search=${table.id}`}
                      >
                        Add RLS policy
-                      </Button>
-                    </Link>
+                      </Link>
+                    </Button>
                  </Tooltip.Trigger>
                  <Tooltip.Portal>
                    <Tooltip.Content side="bottom">
@@ -190,23 +194,24 @@ const GridHeaderActions = ({ table, isViewSelected, isTableSelected }: GridHeade
                  </Tooltip.Portal>
                </Tooltip.Root>
              ) : (
-                <Link passHref href={`/project/${projectRef}/auth/policies?search=${table.id}`}>
-                  <Button
-                    type={policies.length < 1 ? 'warning' : 'default'}
-                    className="group !h-[28px] !py-0"
-                    icon={
-                      policies.length > 0 ? (
-                        <span className="text-right text-xs rounded-xl px-[6px] bg-foreground-lighter/30 text-brand-1100">
-                          {policies.length}
-                        </span>
-                      ) : (
-                        <IconPlusCircle size={12} />
-                      )
-                    }
-                  >
+                <Button
+                  asChild
+                  type={policies.length < 1 ? 'warning' : 'default'}
+                  className="group !h-[28px] !py-0"
+                  icon={
+                    policies.length > 0 ? (
+                      <span className="text-right text-xs rounded-xl px-[6px] bg-foreground-lighter/30 text-brand-1100">
+                        {policies.length}
+                      </span>
+                    ) : (
+                      <IconPlusCircle size={12} />
+                    )
+                  }
+                >
+                  <Link passHref href={`/project/${projectRef}/auth/policies?search=${table.id}`}>
                    Auth {policies.length > 1 ? 'policies' : 'policy'}
-                  </Button>
-                </Link>
+                  </Link>
+                </Button>
              )}
            </div>
          ) : (
@@ -25,7 +25,10 @@ import {
  Constraint,
  useTableConstraintsQuery,
 } from 'data/database/constraints-query'
-import { useForeignKeyConstraintsQuery } from 'data/database/foreign-key-constraints-query'
+import {
+  ForeignKeyConstraint,
+  useForeignKeyConstraintsQuery,
+} from 'data/database/foreign-key-constraints-query'
 import { useEnumeratedTypesQuery } from 'data/enumerated-types/enumerated-types-query'
 import { EXCLUDED_SCHEMAS_WITHOUT_EXTENSIONS } from 'lib/constants/schemas'
 import { Dictionary } from 'types'
@@ -45,6 +48,7 @@ import {
 import ColumnForeignKey from './ColumnForeignKey'
 import ColumnType from './ColumnType'
 import HeaderTitle from './HeaderTitle'
+import toast from 'react-hot-toast'

 export interface ColumnEditorProps {
  column?: PostgresColumn
@@ -57,6 +61,8 @@ export interface ColumnEditorProps {
    configuration: {
      columnId?: string
      primaryKey?: Constraint
+      foreignKeyRelations: ForeignKey[]
+      existingForeignKeyRelations: ForeignKeyConstraint[]
    },
    resolve: any
  ) => void
@@ -107,6 +113,11 @@ const ColumnEditor = ({
  const foreignKeys = foreignKeyMeta.filter((relation) => {
    return relation.source_id === column?.table_id && relation.source_columns.includes(column.name)
  })
+  const lockColumnType =
+    fkRelations.find(
+      (fk) =>
+        fk.columns.find((col) => col.source === columnFields?.name) !== undefined && !fk.toRemove
+    ) !== undefined

  useEffect(() => {
    if (visible) {
@@ -147,7 +158,12 @@ const ColumnEditor = ({
        const payload = isNewRecord
          ? generateCreateColumnPayload(selectedTable.id, columnFields)
          : generateUpdateColumnPayload(column!, selectedTable, columnFields)
-        const configuration = { columnId: column?.id, primaryKey }
+        const configuration = {
+          columnId: column?.id,
+          primaryKey,
+          foreignKeyRelations: fkRelations,
+          existingForeignKeyRelations: foreignKeys,
+        }
        saveChanges(payload, isNewRecord, configuration, resolve)
      } else {
        resolve()
@@ -240,11 +256,9 @@ const ColumnEditor = ({
            enumTypes={enumTypes}
            error={errors.format}
            description={
-              columnFields.foreignKey !== undefined
-                ? 'Column type cannot be changed as it has a foreign key relation'
-                : ''
+              lockColumnType ? 'Column type cannot be changed as it has a foreign key relation' : ''
            }
-            disabled={columnFields.foreignKey !== undefined}
+            disabled={lockColumnType}
            onOptionSelect={(format: string) => onUpdateField({ format, defaultValue: null })}
          />
          {columnFields.foreignKey === undefined && (
@@ -291,13 +305,13 @@ const ColumnEditor = ({
        header={<FormSectionLabel className="lg:!col-span-4">Foreign Keys</FormSectionLabel>}
      >
        <FormSectionContent loading={false} className="lg:!col-span-8">
-          <div>
-            <ColumnForeignKey
-              column={columnFields}
-              relations={fkRelations}
-              closePanel={closePanel}
-            />
-          </div>
+          <ColumnForeignKey
+            column={columnFields}
+            relations={fkRelations}
+            closePanel={closePanel}
+            onUpdateColumnType={(format: string) => onUpdateField({ format, defaultValue: null })}
+            onUpdateFkRelations={setFkRelations}
+          />
        </FormSectionContent>
      </FormSection>
      <SidePanel.Separator />
@@ -1,29 +1,45 @@
-import { noop } from 'lodash'
-import { AlertDescription_Shadcn_, AlertTitle_Shadcn_, Alert_Shadcn_, Button } from 'ui'
+import { useParams } from 'common'
+import { useState } from 'react'
+import { Button } from 'ui'

 import { useProjectContext } from 'components/layouts/ProjectLayout/ProjectContext'
 import { useForeignKeyConstraintsQuery } from 'data/database/foreign-key-constraints-query'
+import useTable from 'hooks/misc/useTable'
+import { ForeignKeySelector } from '../ForeignKeySelector/ForeignKeySelector'
 import { ForeignKey } from '../ForeignKeySelector/ForeignKeySelector.types'
 import { ColumnField } from '../SidePanelEditor.types'
 import { ForeignKeyRow } from '../TableEditor/ForeignKeysManagement/ForeignKeyRow'
 import { checkIfRelationChanged } from '../TableEditor/ForeignKeysManagement/ForeignKeysManagement.utils'
-import { useTableEditorStateSnapshot } from 'state/table-editor'

 interface ColumnForeignKeyProps {
  column: ColumnField
  relations: ForeignKey[]
  closePanel: () => void
+  onUpdateColumnType: (type: string) => void
+  onUpdateFkRelations: (fks: ForeignKey[]) => void
 }

-const ColumnForeignKey = ({ column, relations, closePanel }: ColumnForeignKeyProps) => {
+const ColumnForeignKey = ({
+  column,
+  relations,
+  closePanel,
+  onUpdateColumnType,
+  onUpdateFkRelations,
+}: ColumnForeignKeyProps) => {
+  const { id: _id } = useParams()
+  const [open, setOpen] = useState(false)
+  const [selectedFk, setSelectedFk] = useState<ForeignKey>()
+
  const { project } = useProjectContext()
-  const snap = useTableEditorStateSnapshot()
  const { data } = useForeignKeyConstraintsQuery({
    projectRef: project?.ref,
    connectionString: project?.connectionString,
    schema: column.schema,
  })

+  const id = _id ? Number(_id) : undefined
+  const { data: table } = useTable(id)
+
  const getRelationStatus = (fk: ForeignKey) => {
    const existingRelation = (data ?? []).find((x) => x.id === fk.id)
    const stateRelation = relations.find((x) => x.id === fk.id)
@@ -38,53 +54,88 @@ const ColumnForeignKey = ({ column, relations, closePanel }: ColumnForeignKeyPro
  }

  return (
-    <div className="flex flex-col gap-y-2">
-      <div>
-        {relations.map((relation) => {
-          const status = getRelationStatus(relation)
-          return (
-            <ForeignKeyRow
-              disabled
-              key={relation.id}
-              layout="vertical"
-              status={status}
-              foreignKey={relation}
-              closePanel={closePanel}
-              onSelectEdit={noop}
-              onSelectRemove={noop}
-              onSelectUndoRemove={noop}
-            />
-          )
-        })}
-      </div>
-      <Alert_Shadcn_>
-        <AlertTitle_Shadcn_>
-          {relations.length > 0 ? (
-            <>
-              These are the foreign keys which the column{' '}
-              <code className="text-xs">{column.name}</code> is included
-            </>
-          ) : column.isNewColumn ? (
-            <>Foreign keys can be added after creating the column</>
-          ) : (
-            <>
-              Foreign keys which include the column <code className="text-xs">{column.name}</code>{' '}
-              will be shown here
-            </>
-          )}
-        </AlertTitle_Shadcn_>
-        <AlertDescription_Shadcn_>
-          All foreign keys on the table can be managed by editing the{' '}
-          <code className="text-xs">
-            {column.schema}.{column.table}
-          </code>{' '}
-          table.
-        </AlertDescription_Shadcn_>
-        <Button type="default" className="mt-3" onClick={() => snap.onEditTable()}>
-          Edit table
+    <>
+      <div className="flex flex-col gap-y-2">
+        {relations.length > 0 && (
+          <div>
+            {relations.map((relation) => {
+              const status = getRelationStatus(relation)
+              return (
+                <ForeignKeyRow
+                  key={relation.id}
+                  layout="vertical"
+                  status={status}
+                  foreignKey={relation}
+                  closePanel={closePanel}
+                  onSelectEdit={() => {
+                    setOpen(true)
+                    setSelectedFk(relation)
+                  }}
+                  onSelectRemove={() => {
+                    if (status === 'ADD') {
+                      const updatedRelations = relations.filter((x) => x.id !== relation.id)
+                      onUpdateFkRelations(updatedRelations)
+                    } else {
+                      const updatedRelations = relations.map((x) => {
+                        if (x.id === relation.id) return { ...x, toRemove: true }
+                        else return x
+                      })
+                      onUpdateFkRelations(updatedRelations)
+                    }
+                  }}
+                  onSelectUndoRemove={() => {
+                    const updatedRelations = relations.map((x) => {
+                      if (x.id === relation.id) return { ...x, toRemove: false }
+                      else return x
+                    })
+                    onUpdateFkRelations(updatedRelations)
+                  }}
+                />
+              )
+            })}
+          </div>
+        )}
+
+        <Button type="default" className="w-min" onClick={() => setOpen(true)}>
+          Add foreign key
        </Button>
-      </Alert_Shadcn_>
-    </div>
+      </div>
+
+      {table !== undefined && (
+        <ForeignKeySelector
+          visible={open}
+          column={column}
+          table={{
+            id: table.id,
+            name: table.name,
+            columns:
+              column.isNewColumn && column.name
+                ? (table.columns as any[]).concat(column)
+                : (table.columns as any[]),
+          }}
+          foreignKey={selectedFk}
+          onClose={() => {
+            setOpen(false)
+            setSelectedFk(undefined)
+          }}
+          onSaveRelation={(fk) => {
+            const existingRelationIds = relations.map((x) => x.id)
+            if (fk.id !== undefined && existingRelationIds.includes(fk.id)) {
+              onUpdateFkRelations(
+                relations.map((x) => {
+                  if (x.id === fk.id) return fk
+                  return x
+                })
+              )
+            } else {
+              onUpdateFkRelations(relations.concat([fk]))
+            }
+            const targetType = fk.columns.find((col) => col.source === column.name)?.targetType
+            if (targetType) onUpdateColumnType(targetType)
+          }}
+        />
+      )}
+    </>
  )
 }

@@ -28,6 +28,7 @@ import { NUMERICAL_TYPES, TEXT_TYPES } from '../SidePanelEditor.constants'
 import { FOREIGN_KEY_CASCADE_OPTIONS } from './ForeignKeySelector.constants'
 import { ForeignKey } from './ForeignKeySelector.types'
 import { generateCascadeActionDescription } from './ForeignKeySelector.utils'
+import { ColumnField } from '../SidePanelEditor.types'

 const EMPTY_STATE: ForeignKey = {
  id: undefined,
@@ -40,7 +41,12 @@ const EMPTY_STATE: ForeignKey = {

 interface ForeignKeySelectorProps {
  visible: boolean
-  table: { id: number; name: string; columns: any[] }
+  table: {
+    id: number
+    name: string
+    columns: { id: string; name: string; format: string; isNewColumn: boolean }[]
+  }
+  column?: ColumnField // For ColumnEditor, to prefill when adding a new foreign key
  foreignKey?: ForeignKey
  onClose: () => void
  onSaveRelation: (fk: ForeignKey) => void
@@ -49,6 +55,7 @@ interface ForeignKeySelectorProps {
 export const ForeignKeySelector = ({
  visible,
  table,
+  column,
  foreignKey,
  onClose,
  onSaveRelation,
@@ -83,17 +90,8 @@ export const ForeignKeySelector = ({

  const updateSelectedTable = (tableId: number) => {
    setErrors({})
-    if (!tableId) {
-      return setFk({
-        ...EMPTY_STATE,
-        id: fk.id,
-        name: fk.name,
-        schema: fk.schema,
-        columns: [{ source: '', target: '' }],
-      })
-    }
    const table = (tables ?? []).find((x) => x.id === tableId)
-    if (table)
+    if (table) {
      setFk({
        ...EMPTY_STATE,
        id: fk.id,
@@ -101,8 +99,12 @@ export const ForeignKeySelector = ({
        tableId: table.id,
        schema: table.schema,
        table: table.name,
-        columns: [{ source: '', target: '' }],
+        columns:
+          column !== undefined
+            ? [{ source: column.name, target: '' }]
+            : [{ source: '', target: '' }],
      })
+    }
  }

  const addColumn = () => {
@@ -160,7 +162,7 @@ export const ForeignKeySelector = ({
    fk.columns.forEach((column) => {
      const { source, target, sourceType: sType, targetType: tType } = column
      const sourceColumn = table.columns.find((col) => col.name === source)
-      const sourceType = sType ?? sourceColumn?.format
+      const sourceType = sType ?? sourceColumn?.format ?? ''
      const targetType =
        tType ?? selectedTable?.columns?.find((col) => col.name === target)?.format ?? ''

@@ -246,6 +248,7 @@ export const ForeignKeySelector = ({
                  key={schema.id}
                  value={schema.name}
                  label={schema.name}
+                  className="min-w-96"
                  addOnBefore={() => <IconDatabase size={16} strokeWidth={1.5} />}
                >
                  <div className="flex items-center gap-2">
@@ -264,7 +267,7 @@ export const ForeignKeySelector = ({
            value={selectedTable?.id ?? 1}
            onChange={(value: string) => updateSelectedTable(Number(value))}
          >
-            <Listbox.Option key="empty" value={1} label="---">
+            <Listbox.Option key="empty" className="min-w-96" value={1} label="---">
              ---
            </Listbox.Option>
            {sortBy(tables, ['schema']).map((table) => {
@@ -273,6 +276,7 @@ export const ForeignKeySelector = ({
                  key={table.id}
                  value={table.id}
                  label={table.name}
+                  className="min-w-96"
                  addOnBefore={() => <Table size={16} strokeWidth={1.5} />}
                >
                  <div className="flex items-center gap-2">
@@ -312,7 +316,7 @@ export const ForeignKeySelector = ({
                    </Alert_Shadcn_>
                  )}
                  {fk.columns.map((_, idx) => (
-                    <Fragment key={`${uuidv4}`}>
+                    <Fragment key={`${uuidv4()}`}>
                      <div className="col-span-4">
                        <Listbox
                          id="column"
@@ -202,11 +202,16 @@ const SidePanelEditor = ({
  const saveColumn = async (
    payload: CreateColumnPayload | UpdateColumnPayload,
    isNewRecord: boolean,
-    configuration: { columnId?: string; primaryKey?: Constraint },
+    configuration: {
+      columnId?: string
+      primaryKey?: Constraint
+      foreignKeyRelations: ForeignKey[]
+      existingForeignKeyRelations: ForeignKeyConstraint[]
+    },
    resolve: any
  ) => {
    const selectedColumnToEdit = snap.sidePanel?.type === 'column' && snap.sidePanel.column
-    const { columnId, primaryKey } = configuration
+    const { columnId, primaryKey, foreignKeyRelations, existingForeignKeyRelations } = configuration

    const response = isNewRecord
      ? await createColumn({
@@ -215,6 +220,7 @@ const SidePanelEditor = ({
          payload: payload as CreateColumnPayload,
          selectedTable: selectedTable as PostgresTable,
          primaryKey,
+          foreignKeyRelations,
        })
      : await updateColumn({
          projectRef: project?.ref!,
@@ -223,6 +229,8 @@ const SidePanelEditor = ({
          payload: payload as UpdateColumnPayload,
          selectedTable: selectedTable as PostgresTable,
          primaryKey,
+          foreignKeyRelations,
+          existingForeignKeyRelations,
        })

    if (response?.error) {
@@ -269,12 +269,14 @@ export const createColumn = async ({
  payload,
  selectedTable,
  primaryKey,
+  foreignKeyRelations = [],
 }: {
  projectRef: string
  connectionString: string | undefined
  payload: CreateColumnPayload
  selectedTable: PostgresTable
  primaryKey?: Constraint
+  foreignKeyRelations?: ForeignKey[]
 }) => {
  const toastId = toast.loading(`Creating column "${payload.name}"...`)
  try {
@@ -311,6 +313,17 @@ export const createColumn = async ({
        primaryKeyColumns
      )
    }
+
+    // Then add the foreign key constraints here
+    if (foreignKeyRelations.length > 0) {
+      await addForeignKey({
+        projectRef,
+        connectionString,
+        table: { schema: column.schema, name: column.table },
+        foreignKeys: foreignKeyRelations,
+      })
+    }
+
    toast.success(`Successfully created column "${column.name}"`, { id: toastId })
  } catch (error: any) {
    toast.error(`An error occurred while creating the column "${payload.name}"`, { id: toastId })
@@ -325,6 +338,8 @@ export const updateColumn = async ({
  payload,
  selectedTable,
  primaryKey,
+  foreignKeyRelations = [],
+  existingForeignKeyRelations = [],
  skipPKCreation,
  skipSuccessMessage = false,
 }: {
@@ -334,14 +349,16 @@ export const updateColumn = async ({
  payload: UpdateColumnPayload
  selectedTable: PostgresTable
  primaryKey?: Constraint
+  foreignKeyRelations?: ForeignKey[]
+  existingForeignKeyRelations?: ForeignKeyConstraint[]
  skipPKCreation?: boolean
  skipSuccessMessage?: boolean
 }) => {
  try {
    const { isPrimaryKey, ...formattedPayload } = payload
    const column = await updateDatabaseColumn({
-      projectRef: projectRef,
-      connectionString: connectionString,
+      projectRef,
+      connectionString,
      id,
      payload: formattedPayload,
    })
@@ -374,6 +391,18 @@ export const updateColumn = async ({
        )
      }
    }
+
+    // Then update foreign keys
+    if (foreignKeyRelations.length > 0) {
+      await updateForeignKeys({
+        projectRef,
+        connectionString,
+        table: { schema: column.schema, name: column.table },
+        foreignKeys: foreignKeyRelations,
+        existingForeignKeyRelations,
+      })
+    }
+
    if (!skipSuccessMessage) toast.success(`Successfully updated column "${column.name}"`)
  } catch (error: any) {
    return { error }
@@ -762,43 +791,13 @@ export const updateTable = async ({
  }

  // Foreign keys will get updated here accordingly
-  const relationsToAdd = foreignKeyRelations.filter((x) => typeof x.id === 'string')
-  if (relationsToAdd.length > 0) {
-    await addForeignKey({
-      projectRef,
-      connectionString,
-      table: updatedTable,
-      foreignKeys: relationsToAdd,
-    })
-  }
-
-  const relationsToRemove = foreignKeyRelations.filter((x) => x.toRemove)
-  if (relationsToRemove.length > 0) {
-    await removeForeignKey({
-      projectRef,
-      connectionString,
-      table: updatedTable,
-      foreignKeys: relationsToRemove,
-    })
-  }
-
-  const remainingRelations = foreignKeyRelations.filter(
-    (x) => typeof x.id === 'number' && !x.toRemove
-  )
-  const relationsToUpdate = remainingRelations.filter((x) => {
-    const existingRelation = existingForeignKeyRelations.find((y) => x.id === y.id)
-    if (existingRelation !== undefined) {
-      return checkIfRelationChanged(existingRelation as unknown as ForeignKeyConstraint, x)
-    } else return false
+  await updateForeignKeys({
+    projectRef,
+    connectionString,
+    table: updatedTable,
+    foreignKeys: foreignKeyRelations,
+    existingForeignKeyRelations,
  })
-  if (relationsToUpdate.length > 0) {
-    await updateForeignKey({
-      projectRef,
-      connectionString,
-      table: updatedTable,
-      foreignKeys: relationsToUpdate,
-    })
-  }

  const queryClient = getQueryClient()

@@ -939,3 +938,54 @@ export const insertTableRows = async (
  }
  return { error: insertError }
 }
+
+const updateForeignKeys = async ({
+  projectRef,
+  connectionString,
+  table,
+  foreignKeys,
+  existingForeignKeyRelations,
+}: {
+  projectRef: string
+  connectionString?: string
+  table: { schema: string; name: string }
+  foreignKeys: ForeignKey[]
+  existingForeignKeyRelations: ForeignKeyConstraint[]
+}) => {
+  // Foreign keys will get updated here accordingly
+  const relationsToAdd = foreignKeys.filter((x) => typeof x.id === 'string')
+  if (relationsToAdd.length > 0) {
+    await addForeignKey({
+      projectRef,
+      connectionString,
+      table,
+      foreignKeys: relationsToAdd,
+    })
+  }
+
+  const relationsToRemove = foreignKeys.filter((x) => x.toRemove)
+  if (relationsToRemove.length > 0) {
+    await removeForeignKey({
+      projectRef,
+      connectionString,
+      table,
+      foreignKeys: relationsToRemove,
+    })
+  }
+
+  const remainingRelations = foreignKeys.filter((x) => typeof x.id === 'number' && !x.toRemove)
+  const relationsToUpdate = remainingRelations.filter((x) => {
+    const existingRelation = existingForeignKeyRelations.find((y) => x.id === y.id)
+    if (existingRelation !== undefined) {
+      return checkIfRelationChanged(existingRelation as unknown as ForeignKeyConstraint, x)
+    } else return false
+  })
+  if (relationsToUpdate.length > 0) {
+    await updateForeignKey({
+      projectRef,
+      connectionString,
+      table,
+      foreignKeys: relationsToUpdate,
+    })
+  }
+}