From 3170394ae4b4dcadc8b3f32e043e070d16ac12d6 Mon Sep 17 00:00:00 2001
From: Arjunlal B <63222211+arjunlalb@users.noreply.github.com>
Date: Tue, 24 Oct 2023 02:22:08 -0700
Subject: [PATCH] fix: use keys from across all rows to construct csv object
 (#2478)

Co-authored-by: Adithya Sreyaj <adithya.sreyaj@traceable.ai>
---
 .../service/file-download.service.test.ts     | 39 +++++++++++++++++--
 .../service/file-download.service.ts          | 28 +++++--------
 2 files changed, 45 insertions(+), 22 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..7065fca82 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
@@ -7,10 +7,10 @@ import { FileDownloadEventType, FileDownloadService } from './file-download.serv
 
 describe('File Download Service', () => {
   const mockElement = document.createElement('a');
-
+  const mockBlobResponse = new Blob(['test'], { type: 'text/csv' });
   Object.defineProperty(document, 'createElement', { value: jest.fn().mockReturnValue(mockElement) });
   Object.defineProperty(window.URL, 'createObjectURL', { value: jest.fn().mockReturnValue('') });
-
+  const blobConstructorSpy = jest.spyOn(window, 'Blob').mockReturnValue(mockBlobResponse);
   const createService = createServiceFactory({
     service: FileDownloadService,
     providers: [
@@ -49,8 +49,10 @@ describe('File Download Service', () => {
 
   test('should download as csv correctly', () => {
     const spectator = createService();
-    const csvData$ = of([{ name: 'traceable', headCount: 123 }]);
-
+    let csvData$ = of([
+      { name: 'example', headCount: 123 },
+      { name: 'hypertrace', headCount: 456, optionalValue: 1 }
+    ]);
     // With correct data
     runFakeRxjs(({ expectObservable }) => {
       expectObservable(spectator.service.downloadAsCsv({ dataSource: csvData$, fileName: 'download.csv' })).toBe(
@@ -60,5 +62,34 @@ describe('File Download Service', () => {
         }
       );
     });
+
+    // CSV conversion should work as expected
+    expect(blobConstructorSpy).toHaveBeenLastCalledWith(
+      [`Name,Head Count,Optional Value\r\n"example",123,"-"\r\n"hypertrace",456,1`],
+      { type: 'text/csv' }
+    );
+
+    //<------ Dataset-2 -  by passing headers explicitly ------>
+    csvData$ = of([
+      { name: 'example', headCount: 123, foo: 'bar' },
+      { name: 'hypertrace', headCount: 456, optionalValue: 1, foo: undefined, bar: 'baz' }
+    ]);
+
+    runFakeRxjs(({ expectObservable }) => {
+      expectObservable(
+        spectator.service.downloadAsCsv({
+          dataSource: csvData$,
+          fileName: 'download.csv',
+          header: ['name', 'foo', 'optionalValue']
+        })
+      ).toBe('(x|)', {
+        x: { type: FileDownloadEventType.Success }
+      });
+    });
+    // CSV conversion should work as expected
+    expect(blobConstructorSpy).toHaveBeenLastCalledWith(
+      [`Name,Foo,Optional Value\r\n"example","bar","-"\r\n"hypertrace","-",1`],
+      { type: 'text/csv' }
+    );
   });
 });
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 2f832b489..c45124228 100644
--- a/projects/components/src/download-file/service/file-download.service.ts
+++ b/projects/components/src/download-file/service/file-download.service.ts
@@ -1,8 +1,8 @@
 import { HttpErrorResponse } from '@angular/common/http';
 import { Injectable } from '@angular/core';
 import { Dictionary } from '@hypertrace/common';
-import { isEmpty, startCase } from 'lodash-es';
-import { combineLatest, Observable, of } from 'rxjs';
+import { isEmpty, startCase, uniq } from 'lodash-es';
+import { Observable, of } from 'rxjs';
 import { catchError, map, take } from 'rxjs/operators';
 import { NotificationService } from '../../notification/notification.service';
 
@@ -31,25 +31,17 @@ export class FileDownloadService {
    * @param config Csv download config
    */
   public downloadAsCsv(config: CsvDownloadFileConfig): Observable<FileDownloadEvent> {
-    // 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))
-        )
-      : of(config.header!);
-
+    const getAllDataKeys = (data: Dictionary<unknown>[]): string[] => uniq(data.map(row => Object.keys(row)).flat());
     // Value replacer for null and undefined values
     const replacer = (_: string, value: string) => value ?? '-';
+    const csvData$ = config.dataSource.pipe(
+      map(data => {
+        const allKeys = isEmpty(config.header) ? getAllDataKeys(data) : config.header!;
+        const header = allKeys.map(startCase);
+        const values = data.map(row => allKeys.flatMap(key => JSON.stringify(row?.[key], replacer)));
 
-    // Convert values into strings
-    const values$ = config.dataSource.pipe(
-      map(data => data.map(datum => Object.values(datum).map(value => JSON.stringify(value, replacer))))
-    );
-
-    // Create csv data as string
-    const csvData$ = combineLatest([header$, values$]).pipe(
-      map(([header, values]) => [header, ...values]), // Merge header and values
+        return [header, ...values];
+      }),
       map(data => data.map(datum => datum.join(',')).join('\r\n')) // Join data to create a string
     );