[discover-next][bug] add back data set navigator to control state

src/plugins/data/common/data_frames/types.ts

@@ -12,6 +12,7 @@ export * from './_df_cache';
 export enum DATA_FRAME_TYPES {
   DEFAULT = 'data_frame',
   POLLING = 'data_frame_polling',
+  ERROR = 'data_frame_error',
 }
 
 export interface DataFrameService {
@@ -46,6 +47,12 @@ export interface DataFrameBucketAgg extends DataFrameAgg {
   key: string;
 }
 
+export interface DataFrameQueryConfig {
+  dataSourceId?: string;
+  dataSourceName?: string;
+  timeFieldName?: string;
+}
+
 /**
  * This configuration is used to define how the aggregation should be performed.
  */

src/plugins/data/common/data_frames/utils.ts

@@ -13,6 +13,7 @@
   IDataFrameWithAggs,
   IDataFrameResponse,
   PartialDataFrame,
+  DataFrameQueryConfig,
 } from './types';
 import { IFieldType } from './fields';
 import { IndexPatternFieldMap, IndexPatternSpec } from '../index_patterns';
@@ -165,16 +166,6 @@
   }
   const data = body as IDataFrame;
   const hits: any[] = [];
-  for (let index = 0; index < data.size; index++) {
-    const hit: { [key: string]: any } = {};
-    data.fields.forEach((field) => {
-      hit[field.name] = field.values[index];
-    });
-    hits.push({
-      _index: data.name,
-      _source: hit,
-    });
-  }
   const searchResponse: SearchResponse<any> = {
     took: response.took,
     timed_out: false,
@@ -187,10 +178,24 @@
     hits: {
       total: 0,
       max_score: 0,
-      hits,
+      hits: [],
     },
   };
 
+  if (data && data.fields && data.fields.length > 0) {
+    for (let index = 0; index < data.size; index++) {
+      const hit: { [key: string]: any } = {};
+      data.fields.forEach((field) => {
+        hit[field.name] = field.values[index];
+      });
+      hits.push({
+        _index: data.name,
+        _source: hit,
+      });
+    }
+  }
+  searchResponse.hits.hits = hits;
+
   if (data.hasOwnProperty('aggs')) {
     const dataWithAggs = data as IDataFrameWithAggs;
     if (!dataWithAggs.aggs) {
@@ -282,9 +287,19 @@
  */
 export const getTimeField = (
   data: IDataFrame,
+  queryConfig?: DataFrameQueryConfig,
   aggConfig?: DataFrameAggConfig
 ): Partial<IFieldType> | undefined => {
+  if (queryConfig?.timeFieldName) {
+    return {
+      name: queryConfig.timeFieldName,
+      type: 'date',
+    };
+  }
   const fields = data.schema || data.fields;
+  if (!fields) {
+    throw Error('Invalid time field');
+  }
   return aggConfig && aggConfig.date_histogram && aggConfig.date_histogram.field
     ? fields.find((field) => field.name === aggConfig?.date_histogram?.field)
     : fields.find((field) => field.type === 'date');
@@ -468,7 +483,12 @@
   return {
     id: id ?? DATA_FRAME_TYPES.DEFAULT,
     title: dataFrame.name,
-    timeFieldName: getTimeField(dataFrame)?.name,
+    timeFieldName: getTimeField(dataFrame, dataFrame.meta?.queryConfig)?.name,
+    dataSourceRef: {
+      id: dataFrame.meta?.queryConfig?.dataSourceId,
+      name: dataFrame.meta?.queryConfig?.dataSourceName,
+      type: dataFrame.meta?.queryConfig?.dataSourceType,
+    },
     type: !id ? DATA_FRAME_TYPES.DEFAULT : undefined,
     fields: fields.reduce(flattenFields, {} as IndexPatternFieldMap),
   };

src/plugins/data/common/search/opensearch_search/types.ts

@@ -48,10 +48,6 @@ export interface ISearchOptions {
    * Use this option to enable support for long numerals.
    */
   withLongNumeralsSupport?: boolean;
-  /**
-   * Use this option to enable support for async.
-   */
-  isAsync?: boolean;
 }
 
 export type ISearchRequestParams<T = Record<string, any>> = {

src/plugins/data/common/search/search_source/search_source.ts

@@ -90,6 +90,7 @@
 import {
   DATA_FRAME_TYPES,
   IDataFrame,
+  IDataFrameError,
   IDataFrameResponse,
   convertResult,
   createDataFrame,
@@ -430,7 +431,7 @@
   private async fetchExternalSearch(searchRequest: SearchRequest, options: ISearchOptions) {
     const { search, getConfig, onResponse } = this.dependencies;
 
     const currentDataframe = this.getDataFrame();
     if (!currentDataframe || currentDataframe.name !== searchRequest.index?.id) {
       await this.createDataFrame(searchRequest);
     }
@@ -447,6 +448,10 @@
           await this.setDataFrame(dataFrameResponse.body as IDataFrame);
           return onResponse(searchRequest, convertResult(response as IDataFrameResponse));
         }
+        if ((response as IDataFrameResponse).type === DATA_FRAME_TYPES.ERROR) {
+          const dataFrameError = response as IDataFrameError;
+          throw new RequestFailure(null, dataFrameError);
+        }
         // TODO: MQL else if data_frame_polling then poll for the data frame updating the df fields only
       }
       return onResponse(searchRequest, response.rawResponse);

src/plugins/data/public/antlr/opensearch_sql/code_completion.ts

@@ -21,9 +21,8 @@
 import { SqlErrorListener } from './sql_error_listerner';
 import { findCursorTokenIndex } from '../shared/cursor';
 import { openSearchSqlAutocompleteData } from './opensearch_sql_autocomplete';
-import { getUiSettings } from '../../services';
 import { SQL_SYMBOLS } from './constants';
-import { QuerySuggestionGetFnArgs } from '../../autocomplete';
+import { QuerySuggestion, QuerySuggestionGetFnArgs } from '../../autocomplete';
 import { fetchColumnValues, fetchTableSchemas } from '../shared/utils';
 
 export interface SuggestionParams {
@@ -46,8 +45,8 @@
   query,
   services,
 }: QuerySuggestionGetFnArgs): Promise<QuerySuggestion[]> => {
   const { api } = services.uiSettings;
   const dataSetManager = services.data.query.dataSet;
   const suggestions = getOpenSearchSqlAutoCompleteSuggestions(query, {
     line: position?.lineNumber || selectionStart,
     column: position?.column || selectionEnd,
@@ -59,12 +58,12 @@
     // Fetch columns and values
     if ('suggestColumns' in suggestions && (suggestions.suggestColumns?.tables?.length ?? 0) > 0) {
       const tableNames = suggestions.suggestColumns?.tables?.map((table) => table.name) ?? [];
-      const schemas = await fetchTableSchemas(tableNames, api, connectionService);
+      const schemas = await fetchTableSchemas(tableNames, api, services);
 
       schemas.forEach((schema) => {
         if (schema.body?.fields?.length > 0) {
           const columns = schema.body.fields.find((col: any) => col.name === 'COLUMN_NAME');
           const fieldTypes = schema.body.fields.find((col: any) => col.name === 'DATA_TYPE');
           if (columns && fieldTypes) {
             finalSuggestions.push(
               ...columns.values.map((col: string, index: number) => ({
@@ -86,7 +85,7 @@
           tableNames,
           suggestions.suggestValuesForColumn as string,
           api,
-          connectionService
+          services
         );
         values.forEach((value) => {
           if (value.body?.fields?.length > 0) {

src/plugins/data/public/autocomplete/providers/query_suggestion_provider.ts

@@ -53,7 +53,7 @@ export interface QuerySuggestionGetFnArgs {
   signal?: AbortSignal;
   boolFilter?: any;
   position?: monaco.Position;
-  connectionService?: any; // will need to add type when ConnectionService is properly exposed from queryEnhancements
+  services?: any; // will need to add type when ConnectionService is properly exposed from queryEnhancements
 }
 
 /** @public **/

src/plugins/data/public/index.ts

@@ -446,6 +446,9 @@ export {
   // for BWC, keeping the old name
   IUiStart as DataPublicPluginStartUi,
   DataSetNavigator,
+  setAsyncSessionId,
+  getAsyncSessionId,
+  setAsyncSessionIdByObj,
 } from './ui';
 
 /**
@@ -462,6 +465,8 @@ export {
   QueryState,
   getDefaultQuery,
   FilterManager,
+  DataSetManager,
+  DataSetContract,
   SavedQuery,
   SavedQueryService,
   SavedQueryTimeFilter,

src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts

@@ -6,12 +6,13 @@
 import { DataSetManager } from './dataset_manager';
 import { coreMock } from '../../../../../core/public/mocks';
 import { SimpleDataSet } from '../../../common';
-
 describe('DataSetManager', () => {
   let service: DataSetManager;
 
   beforeEach(() => {
-    service = new DataSetManager(coreMock.createSetup().uiSettings);
+    const uiSettingsMock = coreMock.createSetup().uiSettings;
+    uiSettingsMock.get.mockReturnValue(true);
+    service = new DataSetManager(uiSettingsMock);
   });
 
   test('getUpdates$ emits initially and after data set changes', () => {
@@ -20,15 +21,15 @@ describe('DataSetManager', () => {
     obs$.subscribe((v) => {
       emittedValues.push(v);
     });
-    expect(emittedValues).toHaveLength(1);
-    expect(emittedValues[1]).toEqual(undefined);
+    expect(emittedValues).toHaveLength(0);
+    expect(emittedValues[0]).toEqual(undefined);
 
     const newDataSet: SimpleDataSet = { id: 'test_dataset', title: 'Test Dataset' };
     service.setDataSet(newDataSet);
-    expect(emittedValues).toHaveLength(2);
-    expect(emittedValues[1]).toEqual(newDataSet);
+    expect(emittedValues).toHaveLength(1);
+    expect(emittedValues[0]).toEqual(newDataSet);
 
     service.setDataSet({ ...newDataSet });
-    expect(emittedValues).toHaveLength(3);
+    expect(emittedValues).toHaveLength(2);
   });
 });

src/plugins/data/public/query/dataset_manager/dataset_manager.ts

@@ -5,29 +5,32 @@
 
 import { BehaviorSubject } from 'rxjs';
 import { CoreStart } from 'opensearch-dashboards/public';
+import { skip } from 'rxjs/operators';
 import {
-  IndexPatternsService,
   SIMPLE_DATA_SET_TYPES,
   SimpleDataSet,
   SimpleDataSource,
+  UI_SETTINGS,
 } from '../../../common';
+import { IndexPatternsContract } from '../../index_patterns';
 
 export class DataSetManager {
   private dataSet$: BehaviorSubject<SimpleDataSet | undefined>;
-  private indexPatterns?: IndexPatternsService;
+  private indexPatterns?: IndexPatternsContract;
   private defaultDataSet?: SimpleDataSet;
 
   constructor(private readonly uiSettings: CoreStart['uiSettings']) {
     this.dataSet$ = new BehaviorSubject<SimpleDataSet | undefined>(undefined);
   }
 
-  public init = async (indexPatterns: IndexPatternsService) => {
+  public init = async (indexPatterns: IndexPatternsContract) => {
     this.indexPatterns = indexPatterns;
     this.defaultDataSet = await this.fetchDefaultDataSet();
+    return this.defaultDataSet;
   };
 
   public getUpdates$ = () => {
-    return this.dataSet$.asObservable();
+    return this.dataSet$.asObservable().pipe(skip(1));
   };
 
   public getDataSet = () => {
@@ -39,17 +42,18 @@
    * @param {Query} query
    */
   public setDataSet = (dataSet: SimpleDataSet | undefined) => {
+    if (!this.uiSettings.get(UI_SETTINGS.QUERY_ENHANCEMENTS_ENABLED)) return;
     this.dataSet$.next(dataSet);
   };
 
   public getDefaultDataSet = () => {
     return this.defaultDataSet;
   };
 
-  private fetchDefaultDataSet = async (): Promise<SimpleDataSet | undefined> => {
+  public fetchDefaultDataSet = async (): Promise<SimpleDataSet | undefined> => {
     const defaultIndexPatternId = this.uiSettings.get('defaultIndex');
     if (!defaultIndexPatternId) {
       return undefined;
     }
 
     const indexPattern = await this.indexPatterns?.get(defaultIndexPatternId);
@@ -58,10 +62,10 @@
     }
 
     if (!indexPattern.id) {
       return undefined;
     }
 
     return {
       id: indexPattern.id,
       title: indexPattern.title,
       type: SIMPLE_DATA_SET_TYPES.INDEX_PATTERN,

src/plugins/data/public/query/query_service.ts

@@ -37,7 +37,7 @@ import { TimefilterService, TimefilterSetup } from './timefilter';
 import { createSavedQueryService } from './saved_query/saved_query_service';
 import { createQueryStateObservable } from './state_sync/create_global_query_observable';
 import { QueryStringManager, QueryStringContract } from './query_string';
-import { DataSetManager, DataSetContract } from './dataset_manager';
+import { DataSetManager } from './dataset_manager';
 import { buildOpenSearchQuery, getOpenSearchQueryConfig, IndexPatternsService } from '../../common';
 import { getUiSettings } from '../services';
 import { IndexPattern } from '..';
@@ -63,7 +63,7 @@ export class QueryService {
   filterManager!: FilterManager;
   timefilter!: TimefilterSetup;
   queryStringManager!: QueryStringContract;
-  dataSetManager!: DataSetContract;
+  dataSetManager!: DataSetManager;
 
   state$!: ReturnType<typeof createQueryStateObservable>;