handle source column differences in embeddable as well

Author: flash1293

src/plugins/discover/public/application/components/discover.tsx

@@ -34,7 +34,6 @@ import { SkipBottomButton } from './skip_bottom_button';
 import { esFilters, IndexPatternField, search } from '../../../../data/public';
 import { DiscoverSidebarResponsive } from './sidebar';
 import { DiscoverProps } from './types';
-import { getDisplayedColumns } from '../helpers/columns';
 import { SortPairArr } from '../angular/doc_table/lib/get_sort';
 import { SEARCH_FIELDS_FROM_SOURCE } from '../../../common';
 import { popularizeField } from '../helpers/popularize_field';
@@ -390,7 +389,7 @@ export function Discover({
                           <div className="dscDiscoverGrid">
                             <DataGridMemoized
                               ariaLabelledBy="documentsAriaLabel"
-                              columns={getDisplayedColumns(state.columns, indexPattern)}
+                              columns={columns}
                               expandedDoc={expandedDoc}
                               indexPattern={indexPattern}
                               rows={rows}

src/plugins/discover/public/application/components/discover_grid/discover_grid.tsx

@@ -35,6 +35,7 @@ import {
 } from './discover_grid_columns';
 import { defaultPageSize, gridStyle, pageSizeArr, toolbarVisibility } from './constants';
 import { DiscoverServices } from '../../../build_services';
+import { getDisplayedColumns } from '../../helpers/columns';
 
 interface SortObj {
   id: string;
@@ -146,7 +147,8 @@ export const DiscoverGrid = ({
   showTimeCol,
   sort,
 }: DiscoverGridProps) => {
-  const defaultColumns = columns.includes('_source');
+  const displayedColumns = getDisplayedColumns(columns, indexPattern);
+  const defaultColumns = displayedColumns.includes('_source');
 
   /**
    * Pagination
@@ -207,19 +209,19 @@ export const DiscoverGrid = ({
   const randomId = useMemo(() => htmlIdGenerator()(), []);
 
   const euiGridColumns = useMemo(
-    () => getEuiGridColumns(columns, settings, indexPattern, showTimeCol, defaultColumns),
-    [columns, indexPattern, showTimeCol, settings, defaultColumns]
+    () => getEuiGridColumns(displayedColumns, settings, indexPattern, showTimeCol, defaultColumns),
+    [displayedColumns, indexPattern, showTimeCol, settings, defaultColumns]
   );
   const schemaDetectors = useMemo(() => getSchemaDetectors(), []);
   const popoverContents = useMemo(() => getPopoverContents(), []);
   const columnsVisibility = useMemo(
     () => ({
-      visibleColumns: getVisibleColumns(columns, indexPattern, showTimeCol) as string[],
+      visibleColumns: getVisibleColumns(displayedColumns, indexPattern, showTimeCol) as string[],
       setVisibleColumns: (newColumns: string[]) => {
         onSetColumns(newColumns);
       },
     }),
-    [columns, indexPattern, showTimeCol, onSetColumns]
+    [displayedColumns, indexPattern, showTimeCol, onSetColumns]
   );
   const sorting = useMemo(() => ({ columns: sortingColumns, onSort: onTableSort }), [
     sortingColumns,
@@ -316,7 +318,7 @@ export const DiscoverGrid = ({
             indexPattern={indexPattern}
             hit={expandedDoc}
             // if default columns are used, dont make them part of the URL - the context state handling will take care to restore them
-            columns={defaultColumns ? [] : columns}
+            columns={defaultColumns ? [] : displayedColumns}
             onFilter={onFilter}
             onRemoveColumn={onRemoveColumn}
             onAddColumn={onAddColumn}

src/plugins/discover/public/application/embeddable/search_embeddable.ts

@@ -47,6 +47,7 @@ import { DiscoverGridSettings } from '../components/discover_grid/types';
 import { DiscoverServices } from '../../build_services';
 import { ElasticSearchHit } from '../doc_views/doc_views_types';
 import { getDefaultSort } from '../angular/doc_table/lib/get_default_sort';
+import { handleSourceColumnState } from '../angular/helpers';
 
 interface SearchScope extends ng.IScope {
   columns?: string[];
@@ -374,7 +375,10 @@ export class SearchEmbeddable
 
     // If there is column or sort data on the panel, that means the original columns or sort settings have
     // been overridden in a dashboard.
-    searchScope.columns = this.input.columns || this.savedSearch.columns;
+    searchScope.columns = handleSourceColumnState(
+      { columns: this.input.columns || this.savedSearch.columns },
+      this.services.core.uiSettings
+    ).columns;
     const savedSearchSort =
       this.savedSearch.sort && this.savedSearch.sort.length
         ? this.savedSearch.sort