feat: enhance primary key and unique field handling logic (#817)

src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-field/table-field-modal/table-field-modal.tsx

@@ -17,7 +17,7 @@ import { useTranslation } from 'react-i18next';
 import { Textarea } from '@/components/textarea/textarea';
 import { useDebounce } from '@/hooks/use-debounce';
 import equal from 'fast-deep-equal';
-import type { DatabaseType } from '@/lib/domain';
+import type { DatabaseType, DBTable } from '@/lib/domain';
 
 import {
     Select,
@@ -29,6 +29,7 @@ import {
 
 export interface TableFieldPopoverProps {
     field: DBField;
+    table: DBTable;
     databaseType: DatabaseType;
     updateField: (attrs: Partial<DBField>) => void;
     removeField: () => void;
@@ -36,6 +37,7 @@ export interface TableFieldPopoverProps {
 
 export const TableFieldPopover: React.FC<TableFieldPopoverProps> = ({
     field,
+    table,
     databaseType,
     updateField,
     removeField,
@@ -44,6 +46,19 @@ export const TableFieldPopover: React.FC<TableFieldPopoverProps> = ({
     const [localField, setLocalField] = React.useState<DBField>(field);
     const [isOpen, setIsOpen] = React.useState(false);
 
+    // Check if this field is the only primary key in the table
+    const isOnlyPrimaryKey = React.useMemo(() => {
+        if (!field.primaryKey) return false;
+
+        // Early exit if we find another primary key
+        for (const f of table.fields) {
+            if (f.id !== field.id && f.primaryKey) {
+                return false;
+            }
+        }
+        return true;
+    }, [table.fields, field.primaryKey, field.id]);
+
     useEffect(() => {
         setLocalField(field);
     }, [field]);
@@ -113,7 +128,7 @@ export const TableFieldPopover: React.FC<TableFieldPopoverProps> = ({
                             </Label>
                             <Checkbox
                                 checked={localField.unique}
-                                disabled={field.primaryKey}
+                                disabled={isOnlyPrimaryKey}
                                 onCheckedChange={(value) =>
                                     setLocalField((current) => ({
                                         ...current,

src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-field/table-field.tsx

@@ -23,8 +23,10 @@ import type {
 } from '@/components/select-box/select-box';
 import { SelectBox } from '@/components/select-box/select-box';
 import { TableFieldPopover } from './table-field-modal/table-field-modal';
+import type { DBTable } from '@/lib/domain';
 
 export interface TableFieldProps {
+    table: DBTable;
     field: DBField;
     updateField: (attrs: Partial<DBField>) => void;
     removeField: () => void;
@@ -76,6 +78,7 @@ const generateFieldRegexPatterns = (
 };
 
 export const TableField: React.FC<TableFieldProps> = ({
+    table,
     field,
     updateField,
     removeField,
@@ -83,6 +86,13 @@ export const TableField: React.FC<TableFieldProps> = ({
     const { databaseType, customTypes } = useChartDB();
     const { t } = useTranslation();
 
+    // Only calculate primary key fields, not just count
+    const primaryKeyFields = useMemo(() => {
+        return table.fields.filter((f) => f.primaryKey);
+    }, [table.fields]);
+
+    const primaryKeyCount = primaryKeyFields.length;
+
     const { attributes, listeners, setNodeRef, transform, transition } =
         useSortable({ id: field.id });
 
@@ -191,6 +201,42 @@ export const TableField: React.FC<TableFieldProps> = ({
         transition,
     };
 
+    const handlePrimaryKeyToggle = useCallback(
+        (value: boolean) => {
+            if (value) {
+                // When setting as primary key
+                const updates: Partial<DBField> = {
+                    primaryKey: true,
+                };
+                // Only auto-set unique if this will be the only primary key
+                if (primaryKeyCount === 0) {
+                    updates.unique = true;
+                }
+                updateField(updates);
+            } else {
+                // When removing primary key
+                updateField({
+                    primaryKey: false,
+                });
+            }
+        },
+        [primaryKeyCount, updateField]
+    );
+
+    const handleNullableToggle = useCallback(
+        (value: boolean) => {
+            updateField({ nullable: value });
+        },
+        [updateField]
+    );
+
+    const handleNameChange = useCallback(
+        (e: React.ChangeEvent<HTMLInputElement>) => {
+            updateField({ name: e.target.value });
+        },
+        [updateField]
+    );
+
     return (
         <div
             className="flex flex-1 touch-none flex-row justify-between gap-2 p-1"
@@ -215,11 +261,7 @@ export const TableField: React.FC<TableFieldProps> = ({
                                     'side_panel.tables_section.table.field_name'
                                 )}
                                 value={field.name}
-                                onChange={(e) =>
-                                    updateField({
-                                        name: e.target.value,
-                                    })
-                                }
+                                onChange={handleNameChange}
                             />
                         </span>
                     </TooltipTrigger>
@@ -265,11 +307,7 @@ export const TableField: React.FC<TableFieldProps> = ({
                         <span>
                             <TableFieldToggle
                                 pressed={field.nullable}
-                                onPressedChange={(value) =>
-                                    updateField({
-                                        nullable: value,
-                                    })
-                                }
+                                onPressedChange={handleNullableToggle}
                             >
                                 N
                             </TableFieldToggle>
@@ -284,12 +322,7 @@ export const TableField: React.FC<TableFieldProps> = ({
                         <span>
                             <TableFieldToggle
                                 pressed={field.primaryKey}
-                                onPressedChange={(value) =>
-                                    updateField({
-                                        unique: value,
-                                        primaryKey: value,
-                                    })
-                                }
+                                onPressedChange={handlePrimaryKeyToggle}
                             >
                                 <KeyRound className="h-3.5" />
                             </TableFieldToggle>
@@ -301,6 +334,7 @@ export const TableField: React.FC<TableFieldProps> = ({
                 </Tooltip>
                 <TableFieldPopover
                     field={field}
+                    table={table}
                     updateField={updateField}
                     removeField={removeField}
                     databaseType={databaseType}

src/pages/editor-page/side-panel/tables-section/table-list/table-list-item/table-list-item-content/table-list-item-content.tsx

@@ -56,6 +56,32 @@ export const TableListItemContent: React.FC<TableListItemContentProps> = ({
     >(['fields']);
     const sensors = useSensors(useSensor(PointerSensor));
 
+    // Create a memoized version of the field updater that handles primary key logic
+    const handleFieldUpdate = useCallback(
+        (fieldId: string, attrs: Partial<DBField>) => {
+            updateField(table.id, fieldId, attrs);
+
+            // Handle the case when removing a primary key and only one remains
+            if (attrs.primaryKey === false) {
+                const remainingPrimaryKeys = table.fields.filter(
+                    (f) => f.id !== fieldId && f.primaryKey
+                );
+                if (remainingPrimaryKeys.length === 1) {
+                    // Set the remaining primary key field as unique
+                    updateField(
+                        table.id,
+                        remainingPrimaryKeys[0].id,
+                        {
+                            unique: true,
+                        },
+                        { updateHistory: false }
+                    );
+                }
+            }
+        },
+        [table.id, table.fields, updateField]
+    );
+
     const handleDragEnd = (event: DragEndEvent) => {
         const { active, over } = event;
 
@@ -147,14 +173,9 @@ export const TableListItemContent: React.FC<TableListItemContentProps> = ({
                                     <TableField
                                         key={field.id}
                                         field={field}
-                                        updateField={(
-                                            attrs: Partial<DBField>
-                                        ) =>
-                                            updateField(
-                                                table.id,
-                                                field.id,
-                                                attrs
-                                            )
+                                        table={table}
+                                        updateField={(attrs) =>
+                                            handleFieldUpdate(field.id, attrs)
                                         }
                                         removeField={() =>
                                             removeField(table.id, field.id)