From cc76e34b260eede79dc23700b80d8819d0ffcc7d Mon Sep 17 00:00:00 2001
From: Arjunlal B <arjun@traceable.ai>
Date: Sat, 21 Oct 2023 13:41:00 +0530
Subject: [PATCH] fix: use keys from across all rows to construct csv object

---
 .../service/file-download.service.test.ts        |  6 ++++--
 .../service/file-download.service.ts             | 16 +++++++++++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/projects/components/src/download-file/service/file-download.service.test.ts b/projects/components/src/download-file/service/file-download.service.test.ts
index 828f7bf30..b9810d0ca 100644
--- a/projects/components/src/download-file/service/file-download.service.test.ts
+++ b/projects/components/src/download-file/service/file-download.service.test.ts
@@ -49,8 +49,10 @@ describe('File Download Service', () => {
 
   test('should download as csv correctly', () => {
     const spectator = createService();
-    const csvData$ = of([{ name: 'traceable', headCount: 123 }]);
-
+    const csvData$ = of([
+      { name: 'traceable', headCount: 123 },
+      { name: 'hypertrace', headCount: 456, optionalValue: 1 }
+    ]);
     // With correct data
     runFakeRxjs(({ expectObservable }) => {
       expectObservable(spectator.service.downloadAsCsv({ dataSource: csvData$, fileName: 'download.csv' })).toBe(
diff --git a/projects/components/src/download-file/service/file-download.service.ts b/projects/components/src/download-file/service/file-download.service.ts
index 0b5a33411..8bc6a1d1b 100644
--- a/projects/components/src/download-file/service/file-download.service.ts
+++ b/projects/components/src/download-file/service/file-download.service.ts
@@ -1,7 +1,7 @@
 import { HttpErrorResponse } from '@angular/common/http';
 import { Injectable } from '@angular/core';
 import { Dictionary } from '@hypertrace/common';
-import { isEmpty, startCase } from 'lodash-es';
+import { isEmpty, startCase, uniq } from 'lodash-es';
 import { combineLatest, Observable, of } from 'rxjs';
 import { catchError, map, take } from 'rxjs/operators';
 import { NotificationService } from '../../notification/notification.service';
@@ -31,11 +31,14 @@ export class FileDownloadService {
    * @param config Csv download config
    */
   public downloadAsCsv(config: CsvDownloadFileConfig): Observable<FileDownloadEvent> {
+    const getAllDataKeys = (data: Dictionary<unknown>[]): string[] => uniq(data.map(row => Object.keys(row)).flat());
     // If given header is empty, then create header from the data keys
     const header$ = isEmpty(config.header)
       ? config.dataSource.pipe(
-          map(data => Object.keys(data[0] ?? [])),
-          map(keys => keys.map(startCase))
+          map(
+            data => getAllDataKeys(data),
+            map((keys: string[]) => keys.map(startCase))
+          )
         )
       : of(config.header!);
 
@@ -44,6 +47,13 @@ export class FileDownloadService {
 
     // Convert values into strings
     const values$ = config.dataSource.pipe(
+      // All rows are not guaranteed to have all keys, so add all keys to all rows
+      map(data => {
+        const allKeys = getAllDataKeys(data);
+        const allKeysDict = allKeys.reduce((acc, key) => ({ ...acc, [key]: undefined }), {});
+
+        return data.map(row => ({ ...allKeysDict, ...row }));
+      }),
       map(data => data.map(datum => Object.values(datum).map(value => JSON.stringify(value, replacer))))
     );