From fc30fc50b21a572cb21d63faddd5d07174b24560 Mon Sep 17 00:00:00 2001
From: Rohan Port <59544282+rohan-bes@users.noreply.github.com>
Date: Mon, 17 Oct 2022 15:45:21 +1100
Subject: [PATCH] [no-issue]: Fix upserting dynamic columns (#4226)

---
 .../transform/updateColumns.test.ts           | 20 +++++++++++++++++++
 .../transform/functions/insertColumns.ts      | 19 ++++++++++++------
 .../transform/functions/updateColumns.ts      |  4 +---
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/packages/report-server/src/__tests__/reportBuilder/transform/updateColumns.test.ts b/packages/report-server/src/__tests__/reportBuilder/transform/updateColumns.test.ts
index 39dbe6cf0e..798c65de54 100644
--- a/packages/report-server/src/__tests__/reportBuilder/transform/updateColumns.test.ts
+++ b/packages/report-server/src/__tests__/reportBuilder/transform/updateColumns.test.ts
@@ -186,4 +186,24 @@ describe('updateColumns', () => {
       ]),
     );
   });
+
+  it('can upsert a column dynamically', () => {
+    const transform = buildTransform([
+      {
+        transform: 'updateColumns',
+        insert: {
+          '=$name': '=$value',
+        },
+        exclude: '*',
+      },
+    ]);
+    expect(
+      transform(
+        TransformTable.fromRows([
+          { name: 'value', value: 7 },
+          { name: 'total', value: 10 },
+        ]),
+      ),
+    ).toStrictEqual(TransformTable.fromRows([{ value: 7 }, { total: 10 }]));
+  });
 });
diff --git a/packages/report-server/src/reportBuilder/transform/functions/insertColumns.ts b/packages/report-server/src/reportBuilder/transform/functions/insertColumns.ts
index f049d0c313..26504c1775 100644
--- a/packages/report-server/src/reportBuilder/transform/functions/insertColumns.ts
+++ b/packages/report-server/src/reportBuilder/transform/functions/insertColumns.ts
@@ -6,7 +6,7 @@
 import { yup } from '@tupaia/utils';
 
 import { Context } from '../../context';
-import { FieldValue } from '../../types';
+import { FieldValue, Row } from '../../types';
 import { TransformParser } from '../parser';
 import { buildWhere } from './where';
 import { mapStringToStringValidator } from './transformValidators';
@@ -25,9 +25,11 @@ export const paramsValidator = yup.object().shape({
 const insertColumns = (table: TransformTable, params: InsertColumnsParams, context: Context) => {
   const parser = new TransformParser(table, context);
   const newColumns: Record<string, FieldValue[]> = {};
-  table.getRows().forEach((_, rowIndex) => {
+  const skippedRows: Record<number, Row> = {};
+  table.getRows().forEach((row, rowIndex) => {
     const shouldEditThisRow = params.where(parser);
     if (!shouldEditThisRow) {
+      skippedRows[rowIndex] = row;
       parser.next();
       return;
     }
@@ -36,9 +38,7 @@ const insertColumns = (table: TransformTable, params: InsertColumnsParams, conte
       const columnName = parser.evaluate(key);
       const columnValue = parser.evaluate(expression);
       if (!newColumns[columnName]) {
-        newColumns[columnName] = table.hasColumn(columnName)
-          ? table.getColumnValues(columnName) // Upserting a column, so fill with current column values
-          : new Array(table.length()).fill(undefined); // Creating a new column, so fill with undefined
+        newColumns[columnName] = new Array(table.length()).fill(undefined);
       }
       newColumns[columnName].splice(rowIndex, 1, columnValue);
     });
@@ -51,7 +51,14 @@ const insertColumns = (table: TransformTable, params: InsertColumnsParams, conte
     values,
   }));
 
-  return table.upsertColumns(columnUpserts);
+  // Drop, then re-insert the original skipped rows
+  const rowsToDrop = Object.keys(skippedRows).map(rowIndexString => parseInt(rowIndexString));
+  const rowReinserts = Object.entries(skippedRows).map(([rowIndexString, row]) => ({
+    row,
+    index: parseInt(rowIndexString),
+  }));
+
+  return table.upsertColumns(columnUpserts).dropRows(rowsToDrop).insertRows(rowReinserts);
 };
 
 const buildParams = (params: unknown): InsertColumnsParams => {
diff --git a/packages/report-server/src/reportBuilder/transform/functions/updateColumns.ts b/packages/report-server/src/reportBuilder/transform/functions/updateColumns.ts
index d86ebff0a7..58751f491b 100644
--- a/packages/report-server/src/reportBuilder/transform/functions/updateColumns.ts
+++ b/packages/report-server/src/reportBuilder/transform/functions/updateColumns.ts
@@ -46,9 +46,7 @@ const updateColumns = (table: TransformTable, params: UpdateColumnsParams, conte
       const columnName = parser.evaluate(key);
       const columnValue = parser.evaluate(expression);
       if (!newColumns[columnName]) {
-        newColumns[columnName] = table.hasColumn(columnName)
-          ? table.getColumnValues(columnName) // Upserting a column, so fill with current column values
-          : new Array(table.length()).fill(undefined); // Creating a new column, so fill with undefined
+        newColumns[columnName] = new Array(table.length()).fill(undefined);
       }
       newColumns[columnName].splice(rowIndex, 1, columnValue);
     });