[no-issue]: Fix upserting dynamic columns (#4226)

beyondessential · Oct 17, 2022 · fc30fc5 · fc30fc5
1 parent b51b419
commit fc30fc5
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 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
@@ -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
@@ -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
@@ -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);
     });