From 0417aa8f16e62a0a2cce1dedf67f5a90e71fc4c4 Mon Sep 17 00:00:00 2001
From: Pranav C <pranavxc@gmail.com>
Date: Mon, 19 Feb 2024 05:12:40 +0000
Subject: [PATCH] fix: support composite key in v2 update

---
 packages/nocodb/src/db/BaseModelSqlv2.ts | 38 ++++++++++++++++++++----
 1 file changed, 32 insertions(+), 6 deletions(-)

diff --git a/packages/nocodb/src/db/BaseModelSqlv2.ts b/packages/nocodb/src/db/BaseModelSqlv2.ts
index b6a7c16034..c8704e0244 100644
--- a/packages/nocodb/src/db/BaseModelSqlv2.ts
+++ b/packages/nocodb/src/db/BaseModelSqlv2.ts
@@ -2870,6 +2870,22 @@ class BaseModelSqlv2 {
           : response?.[ai.id];
       }
 
+      // handle if composite primary key is used along with ai or ag
+      if ((ai || ag) && this.model.primaryKeys?.length > 1) {
+        // generate object with ai column and rest of the primary keys
+        const pkObj = {};
+        for (const pk of this.model.primaryKeys) {
+          if (ai && pk.id === ai.id) {
+            pkObj[pk.id] = rowId;
+          } else if (ag && pk.id === ag.id) {
+            pkObj[pk.id] = rowId;
+          } else {
+            pkObj[pk.id] = insertObj[pk.column_name] ?? null;
+          }
+        }
+        rowId = pkObj;
+      }
+
       await Promise.all(postInsertOps.map((f) => f(rowId)));
 
       if (rowId !== null && rowId !== undefined) {
@@ -3892,12 +3908,22 @@ class BaseModelSqlv2 {
   // todo: handle composite primary key
   protected _extractPksValues(data: any) {
     // data can be still inserted without PK
-    // TODO: return a meaningful value
-    if (!this.model.primaryKey) return 'N/A';
-    return (
-      data[this.model.primaryKey.title] ||
-      data[this.model.primaryKey.column_name]
-    );
+
+    // if composite primary key return an object with all the primary keys
+    if (this.model.primaryKeys.length > 1) {
+      const pkValues = {};
+      for (const pk of this.model.primaryKeys) {
+        pkValues[pk.title] = data[pk.title] || data[pk.column_name];
+      }
+      return pkValues;
+    } else if (this.model.primaryKey) {
+      return (
+        data[this.model.primaryKey.title] ||
+        data[this.model.primaryKey.column_name]
+      );
+    } else {
+      return 'N/A';
+    }
   }
 
   protected async errorDelete(_e, _id, _trx, _cookie) {}