Fix table editor unable to update PK after renaming table (#21285)

* Fix table editor unable to update PK after renaming table

* Lint

apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx

@@ -40,6 +40,11 @@ import {
 import ColumnForeignKey from './ColumnForeignKey'
 import ColumnType from './ColumnType'
 import HeaderTitle from './HeaderTitle'
+import {
+  CONSTRAINT_TYPE,
+  Constraint,
+  useTableConstraintsQuery,
+} from 'data/database/constraints-query'
 
 export interface ColumnEditorProps {
   column?: PostgresColumn
@@ -50,7 +55,8 @@ export interface ColumnEditorProps {
     payload: CreateColumnPayload | UpdateColumnPayload,
     isNewRecord: boolean,
     configuration: {
-      columnId: string | undefined
+      columnId?: string
+      primaryKey?: Constraint
     },
     resolve: any
   ) => void
@@ -80,6 +86,16 @@ const ColumnEditor = ({
     (type) => !EXCLUDED_SCHEMAS_WITHOUT_EXTENSIONS.includes(type.schema)
   )
 
+  const { data: constraints } = useTableConstraintsQuery({
+    projectRef: project?.ref,
+    connectionString: project?.connectionString,
+    schema: selectedTable?.schema,
+    table: selectedTable?.name,
+  })
+  const primaryKey = (constraints ?? []).find(
+    (constraint) => constraint.type === CONSTRAINT_TYPE.PRIMARY_KEY_CONSTRAINT
+  )
+
   const { data } = useForeignKeyConstraintsQuery({
     projectRef: project?.ref,
     connectionString: project?.connectionString,
@@ -131,7 +147,7 @@ const ColumnEditor = ({
         const payload = isNewRecord
           ? generateCreateColumnPayload(selectedTable.id, columnFields)
           : generateUpdateColumnPayload(column!, selectedTable, columnFields)
-        const configuration = { columnId: column?.id }
+        const configuration = { columnId: column?.id, primaryKey }
         saveChanges(payload, isNewRecord, configuration, resolve)
       } else {
         resolve()

apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx

@@ -38,6 +38,7 @@ import {
   updateTable,
 } from './SidePanelEditor.utils'
 import { ImportContent } from './TableEditor/TableEditor.types'
+import { Constraint } from 'data/database/constraints-query'
 
 export interface SidePanelEditorProps {
   editable?: boolean
@@ -201,11 +202,11 @@ const SidePanelEditor = ({
   const saveColumn = async (
     payload: CreateColumnPayload | UpdateColumnPayload,
     isNewRecord: boolean,
-    configuration: { columnId?: string },
+    configuration: { columnId?: string; primaryKey?: Constraint },
     resolve: any
   ) => {
     const selectedColumnToEdit = snap.sidePanel?.type === 'column' && snap.sidePanel.column
-    const { columnId } = configuration
+    const { columnId, primaryKey } = configuration
 
     const response = isNewRecord
       ? await createColumn({
@@ -213,6 +214,7 @@ const SidePanelEditor = ({
           connectionString: project?.connectionString,
           payload: payload as CreateColumnPayload,
           selectedTable: selectedTable as PostgresTable,
+          primaryKey,
         })
       : await updateColumn({
           projectRef: project?.ref!,
@@ -220,6 +222,7 @@ const SidePanelEditor = ({
           id: columnId as string,
           payload: payload as UpdateColumnPayload,
           selectedTable: selectedTable as PostgresTable,
+          primaryKey,
         })
 
     if (response?.error) {
@@ -361,6 +364,7 @@ const SidePanelEditor = ({
       isRealtimeEnabled: boolean
       isDuplicateRows: boolean
       existingForeignKeyRelations: ForeignKeyConstraint[]
+      primaryKey?: Constraint
     },
     resolve: any
   ) => {
@@ -372,6 +376,7 @@ const SidePanelEditor = ({
       isRealtimeEnabled,
       isDuplicateRows,
       existingForeignKeyRelations,
+      primaryKey,
     } = configuration
 
     try {
@@ -436,6 +441,7 @@ const SidePanelEditor = ({
           columns,
           foreignKeyRelations,
           existingForeignKeyRelations,
+          primaryKey,
         })
 
         await updateTableRealtime(table, isRealtimeEnabled)

apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.utils.tsx

@@ -33,6 +33,7 @@ import { ForeignKey } from './ForeignKeySelector/ForeignKeySelector.types'
 import { ColumnField, CreateColumnPayload, UpdateColumnPayload } from './SidePanelEditor.types'
 import { checkIfRelationChanged } from './TableEditor/ForeignKeysManagement/ForeignKeysManagement.utils'
 import { ImportContent } from './TableEditor/TableEditor.types'
+import { Constraint } from 'data/database/constraints-query'
 
 const BATCH_SIZE = 1000
 const CHUNK_SIZE = 1024 * 1024 * 0.1 // 0.1MB
@@ -112,18 +113,19 @@ export const addPrimaryKey = async (
   })
 }
 
-export const removePrimaryKey = async (
+export const dropConstraint = async (
   projectRef: string,
   connectionString: string | undefined,
   schema: string,
-  table: string
+  table: string,
+  name: string
 ) => {
-  const query = `ALTER TABLE "${schema}"."${table}" DROP CONSTRAINT "${table}_pkey"`
+  const query = `ALTER TABLE "${schema}"."${table}" DROP CONSTRAINT "${name}"`
   return await executeSql({
     projectRef: projectRef,
     connectionString: connectionString,
     sql: query,
-    queryKey: ['primary-keys'],
+    queryKey: ['drop-constraint'],
   })
 }
 
@@ -266,11 +268,13 @@ export const createColumn = async ({
   connectionString,
   payload,
   selectedTable,
+  primaryKey,
 }: {
   projectRef: string
   connectionString: string | undefined
   payload: CreateColumnPayload
   selectedTable: PostgresTable
+  primaryKey?: Constraint
 }) => {
   const toastId = toast.loading(`Creating column "${payload.name}"...`)
   try {
@@ -288,8 +292,14 @@ export const createColumn = async ({
       // Same logic in createTable: Remove any primary key constraints first (we'll add it back later)
       const existingPrimaryKeys = selectedTable.primary_keys.map((x) => x.name)
 
-      if (existingPrimaryKeys.length > 0) {
-        await removePrimaryKey(projectRef, connectionString, column.schema, column.table)
+      if (existingPrimaryKeys.length > 0 && primaryKey !== undefined) {
+        await dropConstraint(
+          projectRef,
+          connectionString,
+          column.schema,
+          column.table,
+          primaryKey.name
+        )
       }
 
       const primaryKeyColumns = existingPrimaryKeys.concat([column.name])
@@ -314,6 +324,7 @@ export const updateColumn = async ({
   id,
   payload,
   selectedTable,
+  primaryKey,
   skipPKCreation,
   skipSuccessMessage = false,
 }: {
@@ -322,6 +333,7 @@ export const updateColumn = async ({
   id: string
   payload: UpdateColumnPayload
   selectedTable: PostgresTable
+  primaryKey?: Constraint
   skipPKCreation?: boolean
   skipSuccessMessage?: boolean
 }) => {
@@ -338,8 +350,14 @@ export const updateColumn = async ({
       const existingPrimaryKeys = selectedTable.primary_keys.map((x) => x.name)
 
       // Primary key is getting updated for the column
-      if (existingPrimaryKeys.length > 0) {
-        await removePrimaryKey(projectRef, connectionString, column.schema, column.table)
+      if (existingPrimaryKeys.length > 0 && primaryKey !== undefined) {
+        await dropConstraint(
+          projectRef,
+          connectionString,
+          column.schema,
+          column.table,
+          primaryKey.name
+        )
       }
 
       const primaryKeyColumns = isPrimaryKey
@@ -632,6 +650,7 @@ export const updateTable = async ({
   columns,
   foreignKeyRelations,
   existingForeignKeyRelations,
+  primaryKey,
 }: {
   projectRef: string
   connectionString: string | undefined
@@ -641,6 +660,7 @@ export const updateTable = async ({
   columns: ColumnField[]
   foreignKeyRelations: ForeignKey[]
   existingForeignKeyRelations: ForeignKeyConstraint[]
+  primaryKey?: Constraint
 }) => {
   // Prepare a check to see if primary keys to the tables were updated or not
   const primaryKeyColumns = columns
@@ -655,8 +675,8 @@ export const updateTable = async ({
     // If we do it later, and if the user deleted a PK column, we'd need to do
     // an additional check when removing PK if the column in the PK was removed
     // So doing this one step earlier, lets us skip that additional check.
-    if (table.primary_keys.length > 0) {
-      await removePrimaryKey(projectRef, connectionString, table.schema, table.name)
+    if (primaryKey !== undefined) {
+      await dropConstraint(projectRef, connectionString, table.schema, table.name, primaryKey.name)
     }
   }
 

apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableEditor.tsx

@@ -32,6 +32,11 @@ import {
   generateTableFieldFromPostgresTable,
   validateFields,
 } from './TableEditor.utils'
+import {
+  CONSTRAINT_TYPE,
+  Constraint,
+  useTableConstraintsQuery,
+} from 'data/database/constraints-query'
 
 export interface TableEditorProps {
   table?: PostgresTable
@@ -54,6 +59,7 @@ export interface TableEditorProps {
       isRealtimeEnabled: boolean
       isDuplicateRows: boolean
       existingForeignKeyRelations: ForeignKeyConstraint[]
+      primaryKey?: Constraint
     },
     resolve: any
   ) => void
@@ -103,6 +109,16 @@ const TableEditor = ({
   const [isImportingSpreadsheet, setIsImportingSpreadsheet] = useState<boolean>(false)
   const [rlsConfirmVisible, setRlsConfirmVisible] = useState<boolean>(false)
 
+  const { data: constraints } = useTableConstraintsQuery({
+    projectRef: project?.ref,
+    connectionString: project?.connectionString,
+    schema: table?.schema,
+    table: table?.name,
+  })
+  const primaryKey = (constraints ?? []).find(
+    (constraint) => constraint.type === CONSTRAINT_TYPE.PRIMARY_KEY_CONSTRAINT
+  )
+
   const { data: foreignKeyMeta } = useForeignKeyConstraintsQuery({
     projectRef: project?.ref,
     connectionString: project?.connectionString,
@@ -173,6 +189,7 @@ const TableEditor = ({
           isRealtimeEnabled: tableFields.isRealtimeEnabled,
           isDuplicateRows: isDuplicateRows,
           existingForeignKeyRelations: foreignKeys,
+          primaryKey,
         }
 
         saveChanges(payload, tableFields.columns, fkRelations, isNewRecord, configuration, resolve)

apps/studio/data/database/constraints-query.ts

@@ -0,0 +1,69 @@
+import { UseQueryOptions } from '@tanstack/react-query'
+import { ExecuteSqlData, useExecuteSqlQuery } from '../sql/execute-sql-query'
+
+type GetTableConstraintsVariables = {
+  schema?: string
+  table?: string
+}
+
+export type Constraint = {
+  id: number
+  name: string
+  type: string
+}
+
+export enum CONSTRAINT_TYPE {
+  CHECK_CONSTRAINT = 'c',
+  FOREIGN_KEY_CONSTRAINT = 'f',
+  PRIMARY_KEY_CONSTRAINT = 'p',
+  UNIQUE_CONSTRAINT = 'u',
+  CONSTRAINT_TRIGGER = 't',
+  EXCLUSION_CONSTRAINT = 'x',
+}
+
+export const getTableConstraints = ({ schema, table }: GetTableConstraintsVariables) => {
+  const sql = /* SQL */ `
+  SELECT 
+    con.oid as id,
+    con.conname as name,
+    con.contype as type
+  FROM pg_catalog.pg_constraint con
+  INNER JOIN pg_catalog.pg_class rel
+            ON rel.oid = con.conrelid
+  INNER JOIN pg_catalog.pg_namespace nsp
+            ON nsp.oid = connamespace
+  WHERE nsp.nspname = '${schema}'
+        AND rel.relname = '${table}';
+`.trim()
+
+  return sql
+}
+
+export type TableConstraintsVariables = GetTableConstraintsVariables & {
+  projectRef?: string
+  connectionString?: string
+}
+
+export type TableConstraintsData = Constraint[]
+export type TableConstraintsError = unknown
+
+export const useTableConstraintsQuery = <TData extends TableConstraintsData = TableConstraintsData>(
+  { projectRef, connectionString, schema, table }: TableConstraintsVariables,
+  options: UseQueryOptions<ExecuteSqlData, TableConstraintsError, TData> = {}
+) => {
+  return useExecuteSqlQuery(
+    {
+      projectRef,
+      connectionString,
+      sql: getTableConstraints({ schema, table }),
+      queryKey: ['table-constraints'],
+    },
+    {
+      enabled: typeof schema !== 'undefined' && typeof table !== 'undefined',
+      select(data) {
+        return (data as any)?.result ?? []
+      },
+      ...options,
+    }
+  )
+}

apps/studio/data/database/database-query-constants.ts

@@ -10,15 +10,6 @@ export enum FOREIGN_KEY_CASCADE_ACTION {
   SET_DEFAULT = 'd',
 }
 
-export enum CONSTRAINT_TYPE {
-  CHECK_CONSTRAINT = 'c',
-  FOREIGN_KEY_CONSTRAINT = 'f',
-  PRIMARY_KEY_CONSTRAINT = 'p',
-  UNIQUE_CONSTRAINT = 'u',
-  CONSTRAINT_TRIGGER = 't',
-  EXCLUSION_CONSTRAINT = 'x',
-}
-
 // Derived from https://github.com/MichaelDBA/pg_get_tabledef
 // NOTE: when updating, \n must be replaced with \\n in the SQL below