diff --git a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx index a5b13a8aa1f..0d63b7684fb 100644 --- a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx +++ b/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() diff --git a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx index bfacdc7e05e..b0f85d5ac50 100644 --- a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx +++ b/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) diff --git a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.utils.tsx b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.utils.tsx index 7587564ac8b..0ad5c98e484 100644 --- a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.utils.tsx +++ b/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) } } diff --git a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableEditor.tsx b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableEditor.tsx index 2433c8c413c..d69ee875b54 100644 --- a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableEditor.tsx +++ b/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(false) const [rlsConfirmVisible, setRlsConfirmVisible] = useState(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) diff --git a/apps/studio/data/database/constraints-query.ts b/apps/studio/data/database/constraints-query.ts new file mode 100644 index 00000000000..2c7db7134e7 --- /dev/null +++ b/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 = ( + { projectRef, connectionString, schema, table }: TableConstraintsVariables, + options: UseQueryOptions = {} +) => { + 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, + } + ) +} diff --git a/apps/studio/data/database/database-query-constants.ts b/apps/studio/data/database/database-query-constants.ts index 31b677d87d9..67934d5f686 100644 --- a/apps/studio/data/database/database-query-constants.ts +++ b/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