Reporting: Fix _index and _id columns in CSV export (#96097)

* Reporting: Fix _index and _id columns in CSV export * optimize - cache _columns and run getColumns once per execution * Update x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/generate_csv.ts Co-authored-by: Michael Dokolin <dokmic@gmail.com> * feedback * fix typescripts * fix plugin list test * fix plugin list * take away the export interface to test CI build stats Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com> Co-authored-by: Michael Dokolin <dokmic@gmail.com>
2025-04-24 17:59:23 -04:00 · 2021-04-13 08:03:09 -07:00 · 2021-04-13 08:03:09 -07:00 · f67f0e80e7
commit f67f0e80e7
parent 417776d9b6
14 changed files with 343 additions and 119 deletions
--- a/src/plugins/discover/public/application/helpers/get_sharing_data.test.ts
+++ b/src/plugins/discover/public/application/helpers/get_sharing_data.test.ts
@ -6,13 +6,12 @@
 * Side Public License, v 1.
 */

-import { Capabilities } from 'kibana/public';
-import { getSharingData, showPublicUrlSwitch } from './get_sharing_data';
-import { IUiSettingsClient } from 'kibana/public';
-import { createSearchSourceMock } from '../../../../data/common/search/search_source/mocks';
-import { indexPatternMock } from '../../__mocks__/index_pattern';
-import { DOC_HIDE_TIME_COLUMN_SETTING, SORT_DEFAULT_ORDER_SETTING } from '../../../common';
+import { Capabilities, IUiSettingsClient } from 'kibana/public';
 import { IndexPattern } from 'src/plugins/data/public';
+import { createSearchSourceMock } from '../../../../data/common/search/search_source/mocks';
+import { DOC_HIDE_TIME_COLUMN_SETTING, SORT_DEFAULT_ORDER_SETTING } from '../../../common';
+import { indexPatternMock } from '../../__mocks__/index_pattern';
+import { getSharingData, showPublicUrlSwitch } from './get_sharing_data';

 describe('getSharingData', () => {
  let mockConfig: IUiSettingsClient;
@ -36,6 +35,32 @@ describe('getSharingData', () => {
    const result = await getSharingData(searchSourceMock, { columns: [] }, mockConfig);
    expect(result).toMatchInlineSnapshot(`
      Object {
+        "columns": Array [],
+        "searchSource": Object {
+          "index": "the-index-pattern-id",
+          "sort": Array [
+            Object {
+              "_score": "desc",
+            },
+          ],
+        },
+      }
+    `);
+  });
+
+  test('returns valid data for sharing when columns are selected', async () => {
+    const searchSourceMock = createSearchSourceMock({ index: indexPatternMock });
+    const result = await getSharingData(
+      searchSourceMock,
+      { columns: ['column_a', 'column_b'] },
+      mockConfig
+    );
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "columns": Array [
+          "column_a",
+          "column_b",
+        ],
        "searchSource": Object {
          "index": "the-index-pattern-id",
          "sort": Array [
@ -69,16 +94,16 @@ describe('getSharingData', () => {
    );
    expect(result).toMatchInlineSnapshot(`
      Object {
+        "columns": Array [
+          "cool-timefield",
+          "cool-field-1",
+          "cool-field-2",
+          "cool-field-3",
+          "cool-field-4",
+          "cool-field-5",
+          "cool-field-6",
+        ],
        "searchSource": Object {
-          "fields": Array [
-            "cool-timefield",
-            "cool-field-1",
-            "cool-field-2",
-            "cool-field-3",
-            "cool-field-4",
-            "cool-field-5",
-            "cool-field-6",
-          ],
          "index": "the-index-pattern-id",
          "sort": Array [
            Object {
@ -120,15 +145,15 @@ describe('getSharingData', () => {
    );
    expect(result).toMatchInlineSnapshot(`
      Object {
+        "columns": Array [
+          "cool-field-1",
+          "cool-field-2",
+          "cool-field-3",
+          "cool-field-4",
+          "cool-field-5",
+          "cool-field-6",
+        ],
        "searchSource": Object {
-          "fields": Array [
-            "cool-field-1",
-            "cool-field-2",
-            "cool-field-3",
-            "cool-field-4",
-            "cool-field-5",
-            "cool-field-6",
-          ],
          "index": "the-index-pattern-id",
          "sort": Array [
            Object {
--- a/src/plugins/discover/public/application/helpers/get_sharing_data.ts
+++ b/src/plugins/discover/public/application/helpers/get_sharing_data.ts
@ -7,11 +7,11 @@
 */

 import type { Capabilities, IUiSettingsClient } from 'kibana/public';
-import { DOC_HIDE_TIME_COLUMN_SETTING, SORT_DEFAULT_ORDER_SETTING } from '../../../common';
-import { getSortForSearchSource } from '../angular/doc_table';
 import { ISearchSource } from '../../../../data/common';
-import { AppState } from '../angular/discover_state';
+import { DOC_HIDE_TIME_COLUMN_SETTING, SORT_DEFAULT_ORDER_SETTING } from '../../../common';
 import type { SavedSearch, SortOrder } from '../../saved_searches/types';
+import { AppState } from '../angular/discover_state';
+import { getSortForSearchSource } from '../angular/doc_table';

 /**
 * Preparing data to share the current state as link or CSV/Report
@ -23,10 +23,6 @@ export async function getSharingData(
 ) {
  const searchSource = currentSearchSource.createCopy();
  const index = searchSource.getField('index')!;
-  const fields = {
-    fields: searchSource.getField('fields'),
-    fieldsFromSource: searchSource.getField('fieldsFromSource'),
-  };

  searchSource.setField(
    'sort',
@ -37,7 +33,7 @@ export async function getSharingData(
  searchSource.removeField('aggs');
  searchSource.removeField('size');

-  // fields get re-set to match the saved search columns
+  // Columns that the user has selected in the saved search
  let columns = state.columns || [];

  if (columns && columns.length > 0) {
@ -50,14 +46,11 @@ export async function getSharingData(
    if (timeFieldName && !columns.includes(timeFieldName)) {
      columns = [timeFieldName, ...columns];
    }
-
-    // if columns were selected in the saved search, use them for the searchSource's fields
-    const fieldsKey = fields.fieldsFromSource ? 'fieldsFromSource' : 'fields';
-    searchSource.setField(fieldsKey, columns);
  }

  return {
    searchSource: searchSource.getSerializedFields(true),
+    columns,
  };
 }

--- a/x-pack/plugins/reporting/public/panel_actions/get_csv_panel_action.test.ts
+++ b/x-pack/plugins/reporting/public/panel_actions/get_csv_panel_action.test.ts
@ -16,6 +16,7 @@ describe('GetCsvReportPanelAction', () => {
  let core: any;
  let context: any;
  let mockLicense$: any;
+  let mockSearchSource: any;

  beforeAll(() => {
    if (typeof window.URL.revokeObjectURL === 'undefined') {
@ -49,22 +50,19 @@ describe('GetCsvReportPanelAction', () => {
      },
    } as any;

+    mockSearchSource = {
+      createCopy: () => mockSearchSource,
+      removeField: jest.fn(),
+      setField: jest.fn(),
+      getField: jest.fn(),
+      getSerializedFields: jest.fn().mockImplementation(() => ({})),
+    };
+
    context = {
      embeddable: {
        type: 'search',
        getSavedSearch: () => {
-          const searchSource = {
-            createCopy: () => searchSource,
-            removeField: jest.fn(),
-            setField: jest.fn(),
-            getField: jest.fn().mockImplementation((key: string) => {
-              if (key === 'index') {
-                return 'my-test-index-*';
-              }
-            }),
-            getSerializedFields: jest.fn().mockImplementation(() => ({})),
-          };
-          return { searchSource };
+          return { searchSource: mockSearchSource };
        },
        getTitle: () => `The Dude`,
        getInspectorAdapters: () => null,
@ -79,6 +77,49 @@ describe('GetCsvReportPanelAction', () => {
    } as any;
  });

+  it('translates empty embeddable context into job params', async () => {
+    const panel = new GetCsvReportPanelAction(core, mockLicense$());
+
+    await panel.execute(context);
+
+    expect(core.http.post).toHaveBeenCalledWith(
+      '/api/reporting/v1/generate/immediate/csv_searchsource',
+      {
+        body: '{"searchSource":{},"columns":[],"browserTimezone":"America/New_York"}',
+      }
+    );
+  });
+
+  it('translates embeddable context into job params', async () => {
+    // setup
+    mockSearchSource = {
+      createCopy: () => mockSearchSource,
+      removeField: jest.fn(),
+      setField: jest.fn(),
+      getField: jest.fn(),
+      getSerializedFields: jest.fn().mockImplementation(() => ({ testData: 'testDataValue' })),
+    };
+    context.embeddable.getSavedSearch = () => {
+      return {
+        searchSource: mockSearchSource,
+        columns: ['column_a', 'column_b'],
+      };
+    };
+
+    const panel = new GetCsvReportPanelAction(core, mockLicense$());
+
+    // test
+    await panel.execute(context);
+
+    expect(core.http.post).toHaveBeenCalledWith(
+      '/api/reporting/v1/generate/immediate/csv_searchsource',
+      {
+        body:
+          '{"searchSource":{"testData":"testDataValue"},"columns":["column_a","column_b"],"browserTimezone":"America/New_York"}',
+      }
+    );
+  });
+
  it('allows downloading for valid licenses', async () => {
    const panel = new GetCsvReportPanelAction(core, mockLicense$());

--- a/x-pack/plugins/reporting/public/panel_actions/get_csv_panel_action.tsx
+++ b/x-pack/plugins/reporting/public/panel_actions/get_csv_panel_action.tsx
@ -7,21 +7,19 @@

 import { i18n } from '@kbn/i18n';
 import moment from 'moment-timezone';
-import { CoreSetup } from 'src/core/public';
+import type { CoreSetup } from 'src/core/public';
+import type { ISearchEmbeddable, SavedSearch } from '../../../../../src/plugins/discover/public';
 import {
  loadSharingDataHelpers,
-  ISearchEmbeddable,
-  SavedSearch,
  SEARCH_EMBEDDABLE_TYPE,
 } from '../../../../../src/plugins/discover/public';
-import { IEmbeddable, ViewMode } from '../../../../../src/plugins/embeddable/public';
-import {
-  IncompatibleActionError,
-  UiActionsActionDefinition as ActionDefinition,
-} from '../../../../../src/plugins/ui_actions/public';
-import { LicensingPluginSetup } from '../../../licensing/public';
+import type { IEmbeddable } from '../../../../../src/plugins/embeddable/public';
+import { ViewMode } from '../../../../../src/plugins/embeddable/public';
+import type { UiActionsActionDefinition as ActionDefinition } from '../../../../../src/plugins/ui_actions/public';
+import { IncompatibleActionError } from '../../../../../src/plugins/ui_actions/public';
+import type { LicensingPluginSetup } from '../../../licensing/public';
 import { API_GENERATE_IMMEDIATE, CSV_REPORTING_ACTION } from '../../common/constants';
-import { JobParamsDownloadCSV } from '../../server/export_types/csv_searchsource_immediate/types';
+import type { JobParamsDownloadCSV } from '../../server/export_types/csv_searchsource_immediate/types';
 import { checkLicense } from '../lib/license_check';

 function isSavedSearchEmbeddable(
@ -64,14 +62,11 @@ export class GetCsvReportPanelAction implements ActionDefinition<ActionContext>

  public async getSearchSource(savedSearch: SavedSearch, embeddable: ISearchEmbeddable) {
    const { getSharingData } = await loadSharingDataHelpers();
-    const searchSource = savedSearch.searchSource.createCopy();
-    const { searchSource: serializedSearchSource } = await getSharingData(
-      searchSource,
+    return await getSharingData(
+      savedSearch.searchSource,
      savedSearch, // TODO: get unsaved state (using embeddale.searchScope): https://github.com/elastic/kibana/issues/43977
      this.core.uiSettings
    );
-
-    return serializedSearchSource;
  }

  public isCompatible = async (context: ActionContext) => {
@ -96,12 +91,13 @@ export class GetCsvReportPanelAction implements ActionDefinition<ActionContext>
    }

    const savedSearch = embeddable.getSavedSearch();
-    const searchSource = await this.getSearchSource(savedSearch, embeddable);
+    const { columns, searchSource } = await this.getSearchSource(savedSearch, embeddable);

    const kibanaTimezone = this.core.uiSettings.get('dateFormat:tz');
    const browserTimezone = kibanaTimezone === 'Browser' ? moment.tz.guess() : kibanaTimezone;
    const immediateJobParams: JobParamsDownloadCSV = {
      searchSource,
+      columns,
      browserTimezone,
      title: savedSearch.title,
    };
--- a/x-pack/plugins/reporting/public/share_context_menu/register_csv_reporting.tsx
+++ b/x-pack/plugins/reporting/public/share_context_menu/register_csv_reporting.tsx
@ -8,14 +8,15 @@
 import { i18n } from '@kbn/i18n';
 import moment from 'moment-timezone';
 import React from 'react';
-import { IUiSettingsClient, ToastsSetup } from 'src/core/public';
-import { ShareContext } from '../../../../../src/plugins/share/public';
-import { LicensingPluginSetup } from '../../../licensing/public';
+import type { IUiSettingsClient, ToastsSetup } from 'src/core/public';
+import type { SearchSourceFields } from 'src/plugins/data/common';
+import type { ShareContext } from '../../../../../src/plugins/share/public';
+import type { LicensingPluginSetup } from '../../../licensing/public';
 import { CSV_JOB_TYPE } from '../../common/constants';
-import { JobParamsCSV } from '../../server/export_types/csv_searchsource/types';
+import type { JobParamsCSV } from '../../server/export_types/csv_searchsource/types';
 import { ReportingPanelContent } from '../components/reporting_panel_content_lazy';
 import { checkLicense } from '../lib/license_check';
-import { ReportingAPIClient } from '../lib/reporting_api_client';
+import type { ReportingAPIClient } from '../lib/reporting_api_client';

 interface ReportingProvider {
  apiClient: ReportingAPIClient;
@ -65,7 +66,8 @@ export const csvReportingProvider = ({
      browserTimezone,
      title: sharingData.title as string,
      objectType,
-      searchSource: sharingData.searchSource,
+      searchSource: sharingData.searchSource as SearchSourceFields,
+      columns: sharingData.columns as string[] | undefined,
    };

    const getJobParams = () => jobParams;
--- a/x-pack/plugins/reporting/public/share_context_menu/register_pdf_png_reporting.tsx
+++ b/x-pack/plugins/reporting/public/share_context_menu/register_pdf_png_reporting.tsx
@ -8,15 +8,15 @@
 import { i18n } from '@kbn/i18n';
 import moment from 'moment-timezone';
 import React from 'react';
-import { IUiSettingsClient, ToastsSetup } from 'src/core/public';
-import { ShareContext } from '../../../../../src/plugins/share/public';
-import { LicensingPluginSetup } from '../../../licensing/public';
-import { LayoutParams } from '../../common/types';
-import { JobParamsPNG } from '../../server/export_types/png/types';
-import { JobParamsPDF } from '../../server/export_types/printable_pdf/types';
+import type { IUiSettingsClient, ToastsSetup } from 'src/core/public';
+import type { ShareContext } from '../../../../../src/plugins/share/public';
+import type { LicensingPluginSetup } from '../../../licensing/public';
+import type { LayoutParams } from '../../common/types';
+import type { JobParamsPNG } from '../../server/export_types/png/types';
+import type { JobParamsPDF } from '../../server/export_types/printable_pdf/types';
 import { ScreenCapturePanelContent } from '../components/screen_capture_panel_content_lazy';
 import { checkLicense } from '../lib/license_check';
-import { ReportingAPIClient } from '../lib/reporting_api_client';
+import type { ReportingAPIClient } from '../lib/reporting_api_client';

 interface ReportingPDFPNGProvider {
  apiClient: ReportingAPIClient;
--- a/x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/snapshots/generate_csv.test.ts.snap
+++ b/x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/snapshots/generate_csv.test.ts.snap
@ -1,18 +1,36 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP

-exports[`fields cells can be multi-value 1`] = `
+exports[`fields from job.columns (7.13+ generated) cells can be multi-value 1`] = `
+"product,category
+coconut,\\"cool, rad\\"
+"
+`;
+
+exports[`fields from job.columns (7.13+ generated) columns can be top-level fields such as _id and _index 1`] = `
+"\\"_id\\",\\"_index\\",product,category
+\\"my-cool-id\\",\\"my-cool-index\\",coconut,\\"cool, rad\\"
+"
+`;
+
+exports[`fields from job.columns (7.13+ generated) empty columns defaults to using searchSource.getFields() 1`] = `
+"product
+coconut
+"
+`;
+
+exports[`fields from job.searchSource.getFields() (7.12 generated) cells can be multi-value 1`] = `
 "\\"_id\\",sku
 \\"my-cool-id\\",\\"This is a cool SKU., This is also a cool SKU.\\"
 "
 `;

-exports[`fields provides top-level underscored fields as columns 1`] = `
+exports[`fields from job.searchSource.getFields() (7.12 generated) provides top-level underscored fields as columns 1`] = `
 "\\"_id\\",\\"_index\\",date,message
 \\"my-cool-id\\",\\"my-cool-index\\",\\"2020-12-31T00:14:28.000Z\\",\\"it's nice to see you\\"
 "
 `;

-exports[`fields sorts the fields when they are to be used as table column names 1`] = `
+exports[`fields from job.searchSource.getFields() (7.12 generated) sorts the fields when they are to be used as table column names 1`] = `
 "\\"_id\\",\\"_index\\",\\"_score\\",\\"_type\\",date,\\"message_t\\",\\"message_u\\",\\"message_v\\",\\"message_w\\",\\"message_x\\",\\"message_y\\",\\"message_z\\"
 \\"my-cool-id\\",\\"my-cool-index\\",\\"'-\\",\\"'-\\",\\"2020-12-31T00:14:28.000Z\\",\\"test field T\\",\\"test field U\\",\\"test field V\\",\\"test field W\\",\\"test field X\\",\\"test field Y\\",\\"test field Z\\"
 "
--- a/x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/generate_csv.test.ts
+++ b/x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/generate_csv.test.ts
@ -326,7 +326,7 @@ it('uses the scrollId to page all the data', async () => {
  expect(csvResult.content).toMatchSnapshot();
 });

-describe('fields', () => {
+describe('fields from job.searchSource.getFields() (7.12 generated)', () => {
  it('cells can be multi-value', async () => {
    searchSourceMock.getField = jest.fn().mockImplementation((key: string) => {
      if (key === 'fields') {
@ -497,6 +497,140 @@ describe('fields', () => {
  });
 });

+describe('fields from job.columns (7.13+ generated)', () => {
+  it('cells can be multi-value', async () => {
+    mockDataClient.search = jest.fn().mockImplementation(() =>
+      Rx.of({
+        rawResponse: {
+          hits: {
+            hits: [
+              {
+                _id: 'my-cool-id',
+                _index: 'my-cool-index',
+                _version: 4,
+                fields: {
+                  product: 'coconut',
+                  category: [`cool`, `rad`],
+                },
+              },
+            ],
+            total: 1,
+          },
+        },
+      })
+    );
+
+    const generateCsv = new CsvGenerator(
+      createMockJob({ searchSource: {}, columns: ['product', 'category'] }),
+      mockConfig,
+      {
+        es: mockEsClient,
+        data: mockDataClient,
+        uiSettings: uiSettingsClient,
+      },
+      {
+        searchSourceStart: mockSearchSourceService,
+        fieldFormatsRegistry: mockFieldFormatsRegistry,
+      },
+      new CancellationToken(),
+      logger
+    );
+    const csvResult = await generateCsv.generateData();
+
+    expect(csvResult.content).toMatchSnapshot();
+  });
+
+  it('columns can be top-level fields such as _id and _index', async () => {
+    mockDataClient.search = jest.fn().mockImplementation(() =>
+      Rx.of({
+        rawResponse: {
+          hits: {
+            hits: [
+              {
+                _id: 'my-cool-id',
+                _index: 'my-cool-index',
+                _version: 4,
+                fields: {
+                  product: 'coconut',
+                  category: [`cool`, `rad`],
+                },
+              },
+            ],
+            total: 1,
+          },
+        },
+      })
+    );
+
+    const generateCsv = new CsvGenerator(
+      createMockJob({ searchSource: {}, columns: ['_id', '_index', 'product', 'category'] }),
+      mockConfig,
+      {
+        es: mockEsClient,
+        data: mockDataClient,
+        uiSettings: uiSettingsClient,
+      },
+      {
+        searchSourceStart: mockSearchSourceService,
+        fieldFormatsRegistry: mockFieldFormatsRegistry,
+      },
+      new CancellationToken(),
+      logger
+    );
+    const csvResult = await generateCsv.generateData();
+
+    expect(csvResult.content).toMatchSnapshot();
+  });
+
+  it('empty columns defaults to using searchSource.getFields()', async () => {
+    searchSourceMock.getField = jest.fn().mockImplementation((key: string) => {
+      if (key === 'fields') {
+        return ['product'];
+      }
+      return mockSearchSourceGetFieldDefault(key);
+    });
+    mockDataClient.search = jest.fn().mockImplementation(() =>
+      Rx.of({
+        rawResponse: {
+          hits: {
+            hits: [
+              {
+                _id: 'my-cool-id',
+                _index: 'my-cool-index',
+                _version: 4,
+                fields: {
+                  product: 'coconut',
+                  category: [`cool`, `rad`],
+                },
+              },
+            ],
+            total: 1,
+          },
+        },
+      })
+    );
+
+    const generateCsv = new CsvGenerator(
+      createMockJob({ searchSource: {}, columns: [] }),
+      mockConfig,
+      {
+        es: mockEsClient,
+        data: mockDataClient,
+        uiSettings: uiSettingsClient,
+      },
+      {
+        searchSourceStart: mockSearchSourceService,
+        fieldFormatsRegistry: mockFieldFormatsRegistry,
+      },
+      new CancellationToken(),
+      logger
+    );
+    const csvResult = await generateCsv.generateData();
+
+    expect(csvResult.content).toMatchSnapshot();
+  });
+});
+
 describe('formulas', () => {
  const TEST_FORMULA = '=SUM(A1:A2)';

--- a/x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/generate_csv.ts
+++ b/x-pack/plugins/reporting/server/export_types/csv_searchsource/generate_csv/generate_csv.ts
@ -20,6 +20,7 @@ import {
  ISearchSource,
  ISearchStartSearchSource,
  SearchFieldValue,
+  SearchSourceFields,
  tabifyDocs,
 } from '../../../../../../../src/plugins/data/common';
 import { KbnServerError } from '../../../../../../../src/plugins/kibana_utils/server';
@ -60,7 +61,8 @@ function isPlainStringArray(
 }

 export class CsvGenerator {
-  private _formatters: Record<string, FieldFormat> | null = null;
+  private _columns?: string[];
+  private _formatters?: Record<string, FieldFormat>;
  private csvContainsFormulas = false;
  private maxSizeReached = false;
  private csvRowCount = 0;
@ -135,27 +137,36 @@ export class CsvGenerator {
    };
  }

-  // use fields/fieldsFromSource from the searchSource to get the ordering of columns
-  // otherwise use the table columns as they are
-  private getFields(searchSource: ISearchSource, table: Datatable): string[] {
-    const fieldValues: Record<string, string | boolean | SearchFieldValue[] | undefined> = {
-      fields: searchSource.getField('fields'),
-      fieldsFromSource: searchSource.getField('fieldsFromSource'),
-    };
-    const fieldSource = fieldValues.fieldsFromSource ? 'fieldsFromSource' : 'fields';
-    this.logger.debug(`Getting search source fields from: '${fieldSource}'`);
-
-    const fields = fieldValues[fieldSource];
-    // Check if field name values are string[] and if the fields are user-defined
-    if (isPlainStringArray(fields)) {
-      return fields;
+  private getColumns(searchSource: ISearchSource, table: Datatable) {
+    if (this._columns != null) {
+      return this._columns;
    }

-    // Default to using the table column IDs as the fields
-    const columnIds = table.columns.map((c) => c.id);
-    // Fields in the API response don't come sorted - they need to be sorted client-side
-    columnIds.sort();
-    return columnIds;
+    // if columns is not provided in job params,
+    // default to use fields/fieldsFromSource from the searchSource to get the ordering of columns
+    const getFromSearchSource = (): string[] => {
+      const fieldValues: Pick<SearchSourceFields, 'fields' | 'fieldsFromSource'> = {
+        fields: searchSource.getField('fields'),
+        fieldsFromSource: searchSource.getField('fieldsFromSource'),
+      };
+      const fieldSource = fieldValues.fieldsFromSource ? 'fieldsFromSource' : 'fields';
+      this.logger.debug(`Getting columns from '${fieldSource}' in search source.`);
+
+      const fields = fieldValues[fieldSource];
+      // Check if field name values are string[] and if the fields are user-defined
+      if (isPlainStringArray(fields)) {
+        return fields;
+      }
+
+      // Default to using the table column IDs as the fields
+      const columnIds = table.columns.map((c) => c.id);
+      // Fields in the API response don't come sorted - they need to be sorted client-side
+      columnIds.sort();
+      return columnIds;
+    };
+    this._columns = this.job.columns?.length ? this.job.columns : getFromSearchSource();
+
+    return this._columns;
  }

  private formatCellValues(formatters: Record<string, FieldFormat>) {
@ -202,16 +213,16 @@ export class CsvGenerator {
  }

  /*
-   * Use the list of fields to generate the header row
+   * Use the list of columns to generate the header row
   */
  private generateHeader(
-    fields: string[],
+    columns: string[],
    table: Datatable,
    builder: MaxSizeStringBuilder,
    settings: CsvExportSettings
  ) {
    this.logger.debug(`Building CSV header row...`);
-    const header = fields.map(this.escapeValues(settings)).join(settings.separator) + '\n';
+    const header = columns.map(this.escapeValues(settings)).join(settings.separator) + '\n';

    if (!builder.tryAppend(header)) {
      return {
@ -227,7 +238,7 @@ export class CsvGenerator {
   * Format a Datatable into rows of CSV content
   */
  private generateRows(
-    fields: string[],
+    columns: string[],
    table: Datatable,
    builder: MaxSizeStringBuilder,
    formatters: Record<string, FieldFormat>,
@ -240,7 +251,7 @@ export class CsvGenerator {
      }

      const row =
-        fields
+        columns
          .map((f) => ({ column: f, data: dataTableRow[f] }))
          .map(this.formatCellValues(formatters))
          .map(this.escapeValues(settings))
@ -338,11 +349,13 @@ export class CsvGenerator {
          break;
        }

-        const fields = this.getFields(searchSource, table);
+        // If columns exists in the job params, use it to order the CSV columns
+        // otherwise, get the ordering from the searchSource's fields / fieldsFromSource
+        const columns = this.getColumns(searchSource, table);

        if (first) {
          first = false;
-          this.generateHeader(fields, table, builder, settings);
+          this.generateHeader(columns, table, builder, settings);
        }

        if (table.rows.length < 1) {
@ -350,7 +363,7 @@ export class CsvGenerator {
        }

        const formatters = this.getFormatters(table);
-        this.generateRows(fields, table, builder, formatters, settings);
+        this.generateRows(columns, table, builder, formatters, settings);

        // update iterator
        currentRecord += table.rows.length;
--- a/x-pack/plugins/reporting/server/export_types/csv_searchsource/types.d.ts
+++ b/x-pack/plugins/reporting/server/export_types/csv_searchsource/types.d.ts
@ -5,13 +5,15 @@
 * 2.0.
 */

-import { BaseParams, BasePayload } from '../../types';
+import type { SearchSourceFields } from 'src/plugins/data/common';
+import type { BaseParams, BasePayload } from '../../types';

 export type RawValue = string | object | null | undefined;

 interface BaseParamsCSV {
  browserTimezone: string;
-  searchSource: any;
+  searchSource: SearchSourceFields;
+  columns?: string[];
 }

 export type JobParamsCSV = BaseParamsCSV & BaseParams;
--- a/x-pack/plugins/reporting/server/export_types/csv_searchsource_immediate/types.d.ts
+++ b/x-pack/plugins/reporting/server/export_types/csv_searchsource_immediate/types.d.ts
@ -5,7 +5,7 @@
 * 2.0.
 */

-import { TimeRangeParams } from '../common';
+import type { SearchSourceFields } from 'src/plugins/data/common';

 export interface FakeRequest {
  headers: Record<string, string>;
@ -14,7 +14,8 @@ export interface FakeRequest {
 export interface JobParamsDownloadCSV {
  browserTimezone: string;
  title: string;
-  searchSource: any;
+  searchSource: SearchSourceFields;
+  columns?: string[];
 }

 export interface SavedObjectServiceError {
--- a/x-pack/plugins/reporting/server/routes/csv_searchsource_immediate.ts
+++ b/x-pack/plugins/reporting/server/routes/csv_searchsource_immediate.ts
@ -44,6 +44,7 @@ export function registerGenerateCsvFromSavedObjectImmediate(
      path: `${API_BASE_GENERATE_V1}/immediate/csv_searchsource`,
      validate: {
        body: schema.object({
+          columns: schema.maybe(schema.arrayOf(schema.string())),
          searchSource: schema.object({}, { unknowns: 'allow' }),
          browserTimezone: schema.string({ defaultValue: 'UTC' }),
          title: schema.string(),
--- a/x-pack/test/functional/apps/dashboard/reporting/download_csv.ts
+++ b/x-pack/test/functional/apps/dashboard/reporting/download_csv.ts
@ -50,8 +50,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
    await testSubjects.existOrFail('csvDownloadStarted'); // validate toast panel
  };

-  // FAILING ES PROMOTION: https://github.com/elastic/kibana/issues/96000
-  describe.skip('Download CSV', () => {
+  describe('Download CSV', () => {
    before('initialize tests', async () => {
      log.debug('ReportingPage:initTests');
      await browser.setWindowSize(1600, 850);
--- a/x-pack/test/reporting_api_integration/reporting_and_security/csv_searchsource_immediate.ts
+++ b/x-pack/test/reporting_api_integration/reporting_and_security/csv_searchsource_immediate.ts
@ -10,9 +10,9 @@ import supertest from 'supertest';
 import { JobParamsDownloadCSV } from '../../../plugins/reporting/server/export_types/csv_searchsource_immediate/types';
 import { FtrProviderContext } from '../ftr_provider_context';

-const getMockJobParams = (obj: Partial<JobParamsDownloadCSV>): JobParamsDownloadCSV => ({
+const getMockJobParams = (obj: any): JobParamsDownloadCSV => ({
  title: `Mock CSV Title`,
-  ...(obj as any),
+  ...obj,
 });

 // eslint-disable-next-line import/no-default-export
@ -31,8 +31,7 @@ export default function ({ getService }: FtrProviderContext) {
    },
  };

-  // FAILING ES PROMOTION: https://github.com/elastic/kibana/issues/96000
-  describe.skip('CSV Generation from SearchSource', () => {
+  describe('CSV Generation from SearchSource', () => {
    before(async () => {
      await kibanaServer.uiSettings.update({
        'csv:quoteValues': false,
@ -387,9 +386,9 @@ export default function ({ getService }: FtrProviderContext) {
              version: true,
              index: '907bc200-a294-11e9-a900-ef10e0ac769e',
              sort: [{ date: 'desc' }],
-              fields: ['date', 'message', '_id', '_index'],
              filter: [],
            },
+            columns: ['date', 'message', '_id', '_index'],
          })
        );
        const { status: resStatus, text: resText, type: resType } = res;