From 0ef551b853b9d8a7402772ddeda39e1c0aee1598 Mon Sep 17 00:00:00 2001
From: Pranav C <pranavxc@gmail.com>
Date: Mon, 19 Feb 2024 05:12:41 +0000
Subject: [PATCH] refactor: use composite key extraction in bulk api
 implementations

---
 packages/nocodb/src/db/BaseModelSqlv2.ts | 100 ++++++++++++++++-------
 1 file changed, 72 insertions(+), 28 deletions(-)

diff --git a/packages/nocodb/src/db/BaseModelSqlv2.ts b/packages/nocodb/src/db/BaseModelSqlv2.ts
index c8704e0244..66c0b2d5b8 100644
--- a/packages/nocodb/src/db/BaseModelSqlv2.ts
+++ b/packages/nocodb/src/db/BaseModelSqlv2.ts
@@ -2449,7 +2449,11 @@ class BaseModelSqlv2 {
       if (ag) {
         if (!response) await this.execAndParse(query);
         response = await this.readByPk(
-          insertObj[ag.column_name],
+          this.extractCompositePK({
+            rowId: insertObj[ag.column_name],
+            insertObj,
+            ag,
+          }),
           false,
           {},
           { ignoreView: true, getHiddenColumn: true },
@@ -2492,7 +2496,7 @@ class BaseModelSqlv2 {
             ).id;
           }
           response = await this.readByPk(
-            id,
+            this.extractCompositePK({ rowId: id, insertObj, ag }),
             false,
             {},
             { ignoreView: true, getHiddenColumn: true },
@@ -2505,7 +2509,7 @@ class BaseModelSqlv2 {
           ? response?.[0]?.[ai.id]
           : response?.[ai.id];
         response = await this.readByPk(
-          id,
+          this.extractCompositePK({ rowId: id, insertObj, ag }),
           false,
           {},
           { ignoreView: true, getHiddenColumn: true },
@@ -2869,22 +2873,7 @@ class BaseModelSqlv2 {
           ? response?.[0]?.[ai.id]
           : 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;
-      }
+      rowId = this.extractCompositePK({ ai, ag, rowId, insertObj });
 
       await Promise.all(postInsertOps.map((f) => f(rowId)));
 
@@ -2905,6 +2894,35 @@ class BaseModelSqlv2 {
     }
   }
 
+  private extractCompositePK({
+    ai,
+    ag,
+    rowId,
+    insertObj,
+  }: {
+    ai?: Column<any>;
+    ag?: Column<any>;
+    rowId;
+    insertObj: Record<string, any>;
+  }) {
+    // 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;
+    }
+    return rowId;
+  }
+
   private async prepareNestedLinkQb({
     nestedCols,
     data,
@@ -3198,27 +3216,43 @@ class BaseModelSqlv2 {
 
       let responses;
 
+      const aiPkCol = this.model.primaryKeys.find((pk) => pk.ai);
+      const agPkCol = this.model.primaryKeys.find((pk) => pk.meta?.ag);
+
       // insert one by one as fallback to get ids for sqlite and mysql
       if (insertOneByOneAsFallback && (this.isSqlite || this.isMySQL)) {
         // sqlite and mysql doesn't support returning, so insert one by one and return ids
         responses = [];
 
-        const aiPkCol = this.model.primaryKeys.find((pk) => pk.ai);
-
         for (const insertData of insertDatas) {
           const query = trx(this.tnPath).insert(insertData);
-          const id = (await query)[0];
-          responses.push(aiPkCol ? { [aiPkCol.title]: id } : id);
+          let id = (await query)[0];
+
+          if (agPkCol) {
+            id = insertData[agPkCol.column_name];
+          }
+
+          responses.push(
+            this.extractCompositePK({
+              rowId: id,
+              ai: aiPkCol,
+              ag: agPkCol,
+              insertObj: insertData,
+            }) || insertData,
+          );
         }
       } else {
+        const returningObj: Record<string, string> = {};
+
+        for (const col of this.model.primaryKeys) {
+          returningObj[col.title] = col.column_name;
+        }
+
         responses =
           !raw && (this.isPg || this.isMssql)
             ? await trx
                 .batchInsert(this.tnPath, insertDatas, chunkSize)
-                .returning({
-                  [this.model.primaryKey?.title]:
-                    this.model.primaryKey?.column_name,
-                })
+                .returning(returningObj)
             : await trx.batchInsert(this.tnPath, insertDatas, chunkSize);
       }
 
@@ -3232,7 +3266,17 @@ class BaseModelSqlv2 {
 
       // insert nested link data for single record insertion
       if (isSingleRecordInsertion) {
-        const rowId = responses[0][this.model.primaryKey?.title];
+        let rowId = responses[0][this.model.primaryKey?.title];
+
+        if (aiPkCol || agPkCol) {
+          rowId = this.extractCompositePK({
+            rowId,
+            ai: aiPkCol,
+            ag: agPkCol,
+            insertObj: insertDatas[0],
+          });
+        }
+
         await Promise.all(postInsertOps.map((f) => f(rowId, trx)));
       }