[discover] async query and caching (opensearch-project#7943)

Async query feature for S3 type on Discover. Due to the time to click and verify I end up implementing the cache to avoid waiting too long per every re-build.

What this PR does:

    Poll for query
    Cache data structures and refetches them from session storage
    Poll based on the data type

What this PR does NOT do yet:

    SessionId for search strategy, working fine for selector
    isCacheable field property
    Abort signal on server

---------

Signed-off-by: Kawika Avilla <[email protected]>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
(cherry picked from commit 8c9abe2)

changelogs/fragments/7943.yml

@@ -0,0 +1,2 @@
+feat:
+- Async query search and caching, also adding tests to related components ([#7943](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7943))

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

@@ -440,11 +440,15 @@ export class SearchSource {
           await this.setDataFrame(dataFrameResponse.body as IDataFrame);
           return onResponse(searchRequest, convertResult(response as IDataFrameResponse));
         }
+        if ((response as IDataFrameResponse).type === DATA_FRAME_TYPES.POLLING) {
+          const dataFrameResponse = response as IDataFrameResponse;
+          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/index.ts

@@ -471,6 +471,8 @@ export {
   LanguageConfig,
   LanguageService,
   LanguageServiceContract,
+  RecentQueriesTable,
+  QueryControls,
   SavedQuery,
   SavedQueryService,
   SavedQueryTimeFilter,

src/plugins/data/public/plugin.ts

@@ -116,14 +116,19 @@ export class DataPublicPlugin
   private readonly fieldFormatsService: FieldFormatsService;
   private readonly queryService: QueryService;
   private readonly storage: DataStorage;
+  private readonly sessionStorage: DataStorage;
 
   constructor(initializerContext: PluginInitializerContext<ConfigSchema>) {
     this.searchService = new SearchService(initializerContext);
     this.uiService = new UiService(initializerContext);
     this.queryService = new QueryService();
     this.fieldFormatsService = new FieldFormatsService();
     this.autocomplete = new AutocompleteService(initializerContext);
-    this.storage = createStorage({ engine: window.localStorage, prefix: 'opensearch_dashboards.' });
+    this.storage = createStorage({ engine: window.localStorage, prefix: 'opensearchDashboards.' });
+    this.sessionStorage = createStorage({
+      engine: window.sessionStorage,
+      prefix: 'opensearchDashboards.',
+    });
   }
 
   public setup(
@@ -143,6 +148,7 @@ export class DataPublicPlugin
     const queryService = this.queryService.setup({
       uiSettings: core.uiSettings,
       storage: this.storage,
+      sessionStorage: this.sessionStorage,
       defaultSearchInterceptor: searchService.getDefaultSearchInterceptor(),
     });
 

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

@@ -58,8 +58,9 @@ export class QueryService {
   state$!: ReturnType<typeof createQueryStateObservable>;
 
   public setup({
-    storage,
     uiSettings,
+    storage,
+    sessionStorage,
     defaultSearchInterceptor,
   }: QueryServiceSetupDependencies): IQuerySetup {
     this.filterManager = new FilterManager(uiSettings);
@@ -70,7 +71,12 @@ export class QueryService {
       storage,
     });
 
-    this.queryStringManager = new QueryStringManager(storage, uiSettings, defaultSearchInterceptor);
+    this.queryStringManager = new QueryStringManager(
+      storage,
+      sessionStorage,
+      uiSettings,
+      defaultSearchInterceptor
+    );
 
     this.state$ = createQueryStateObservable({
       filterManager: this.filterManager,

src/plugins/data/public/query/query_string/dataset_service/dataset_service.test.ts

@@ -0,0 +1,86 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { DatasetService } from './dataset_service';
+import { coreMock } from '../../../../../../core/public/mocks';
+import { DataStorage } from 'src/plugins/data/common';
+import { DataStructure } from '../../../../common';
+import { IDataPluginServices } from '../../../types';
+
+describe('DatasetService', () => {
+  let service: DatasetService;
+  let uiSettings: ReturnType<typeof coreMock.createSetup>['uiSettings'];
+  let sessionStorage: DataStorage;
+  let mockDataPluginServices: jest.Mocked<IDataPluginServices>;
+
+  beforeEach(() => {
+    uiSettings = coreMock.createSetup().uiSettings;
+    sessionStorage = new DataStorage(window.sessionStorage, 'opensearchDashboards.');
+    mockDataPluginServices = {} as jest.Mocked<IDataPluginServices>;
+
+    service = new DatasetService(uiSettings, sessionStorage);
+  });
+
+  test('registerType and getType', () => {
+    const mockType = {
+      id: 'test-type',
+      title: 'Test Type',
+      meta: { icon: { type: 'test' } },
+      toDataset: jest.fn(),
+      fetch: jest.fn(),
+      fetchFields: jest.fn(),
+      supportedLanguages: jest.fn(),
+    };
+
+    service.registerType(mockType);
+    expect(service.getType('test-type')).toBe(mockType);
+  });
+
+  test('getTypes returns all registered types', () => {
+    const mockType1 = { id: 'type1', title: 'Type 1', meta: { icon: { type: 'test1' } } };
+    const mockType2 = { id: 'type2', title: 'Type 2', meta: { icon: { type: 'test2' } } };
+
+    service.registerType(mockType1 as any);
+    service.registerType(mockType2 as any);
+
+    const types = service.getTypes();
+    expect(types).toHaveLength(2);
+    expect(types).toContainEqual(mockType1);
+    expect(types).toContainEqual(mockType2);
+  });
+
+  test('fetchOptions caches and returns data structures', async () => {
+    const mockType = {
+      id: 'test-type',
+      title: 'Test Type',
+      meta: { icon: { type: 'test' } },
+      toDataset: jest.fn(),
+      fetch: jest.fn().mockResolvedValue({
+        id: 'test-structure',
+        title: 'Test Structure',
+        type: 'test-type',
+        children: [{ id: 'child1', title: 'Child 1', type: 'test-type' }],
+      }),
+      fetchFields: jest.fn(),
+      supportedLanguages: jest.fn(),
+    };
+
+    service.registerType(mockType);
+
+    const path: DataStructure[] = [{ id: 'root', title: 'Root', type: 'root' }];
+    const result = await service.fetchOptions(mockDataPluginServices, path, 'test-type');
+
+    expect(result).toEqual({
+      id: 'test-structure',
+      title: 'Test Structure',
+      type: 'test-type',
+      children: [{ id: 'child1', title: 'Child 1', type: 'test-type' }],
+    });
+
+    const cachedResult = await service.fetchOptions(mockDataPluginServices, path, 'test-type');
+    expect(cachedResult).toEqual(result);
+    expect(mockType.fetch).toHaveBeenCalledTimes(2);
+  });
+});

src/plugins/data/public/query/query_string/dataset_service/dataset_service.ts

@@ -11,6 +11,8 @@ import {
   DEFAULT_DATA,
   IFieldType,
   UI_SETTINGS,
+  DataStorage,
+  CachedDataStructure,
 } from '../../../../common';
 import { DatasetTypeConfig } from './types';
 import { indexPatternTypeConfig, indexTypeConfig } from './lib';
@@ -22,7 +24,10 @@ export class DatasetService {
   private defaultDataset?: Dataset;
   private typesRegistry: Map<string, DatasetTypeConfig> = new Map();
 
-  constructor(private readonly uiSettings: CoreStart['uiSettings']) {
+  constructor(
+    private readonly uiSettings: CoreStart['uiSettings'],
+    private readonly sessionStorage: DataStorage
+  ) {
     if (this.uiSettings.get(UI_SETTINGS.QUERY_ENHANCEMENTS_ENABLED)) {
       this.registerDefaultTypes();
     }
@@ -76,23 +81,87 @@ export class DatasetService {
             }
           : undefined,
       } as IndexPatternSpec;
-      const temporaryIndexPattern = await this.indexPatterns?.create(spec);
+      const temporaryIndexPattern = await this.indexPatterns?.create(spec, true);
       if (temporaryIndexPattern) {
         this.indexPatterns?.saveToCache(dataset.id, temporaryIndexPattern);
       }
     }
   }
 
-  public fetchOptions(
+  public async fetchOptions(
     services: IDataPluginServices,
     path: DataStructure[],
     dataType: string
   ): Promise<DataStructure> {
     const type = this.typesRegistry.get(dataType);
     if (!type) {
-      throw new Error(`No handler found for type: ${path[0]}`);
+      throw new Error(`No handler found for type: ${dataType}`);
     }
-    return type.fetch(services, path);
+
+    const lastPathItem = path[path.length - 1];
+    const cacheKey = `${dataType}.${lastPathItem.id}`;
+
+    const cachedDataStructure = this.sessionStorage.get<CachedDataStructure>(cacheKey);
+    if (cachedDataStructure?.children?.length > 0) {
+      return this.cacheToDataStructure(dataType, cachedDataStructure);
+    }
+
+    const fetchedDataStructure = await type.fetch(services, path);
+    this.cacheDataStructure(dataType, fetchedDataStructure);
+    return fetchedDataStructure;
+  }
+
+  private cacheToDataStructure(
+    dataType: string,
+    cachedDataStructure: CachedDataStructure
+  ): DataStructure {
+    const reconstructed: DataStructure = {
+      ...cachedDataStructure,
+      parent: undefined,
+      children: cachedDataStructure.children
+        .map((childId) => {
+          const cachedChild = this.sessionStorage.get<CachedDataStructure>(
+            `${dataType}.${childId}`
+          );
+          if (!cachedChild) return;
+          return {
+            id: cachedChild.id,
+            title: cachedChild.title,
+            type: cachedChild.type,
+            meta: cachedChild.meta,
+          } as DataStructure;
+        })
+        .filter((child): child is DataStructure => !!child),
+    };
+
+    return reconstructed;
+  }
+
+  private cacheDataStructure(dataType: string, dataStructure: DataStructure) {
+    const cachedDataStructure: CachedDataStructure = {
+      id: dataStructure.id,
+      title: dataStructure.title,
+      type: dataStructure.type,
+      parent: dataStructure.parent?.id || '',
+      children: dataStructure.children?.map((child) => child.id) || [],
+      hasNext: dataStructure.hasNext,
+      columnHeader: dataStructure.columnHeader,
+      meta: dataStructure.meta,
+    };
+
+    this.sessionStorage.set(`${dataType}.${dataStructure.id}`, cachedDataStructure);
+
+    dataStructure.children?.forEach((child) => {
+      const cachedChild: CachedDataStructure = {
+        id: child.id,
+        title: child.title,
+        type: child.type,
+        parent: dataStructure.id,
+        children: [],
+        meta: child.meta,
+      };
+      this.sessionStorage.set(`${dataType}.${child.id}`, cachedChild);
+    });
   }
 
   private async fetchDefaultDataset(): Promise<Dataset | undefined> {

src/plugins/data/public/query/query_string/dataset_service/lib/index_pattern_type.test.ts

@@ -0,0 +1,92 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// index_pattern_type.test.ts
+
+import { indexPatternTypeConfig } from './index_pattern_type';
+import { SavedObjectsClientContract } from 'opensearch-dashboards/public';
+import { DATA_STRUCTURE_META_TYPES, DataStructure, Dataset } from '../../../../../common';
+import * as services from '../../../../services';
+
+jest.mock('../../../../services', () => ({
+  getIndexPatterns: jest.fn(),
+}));
+
+jest.mock('./utils', () => ({
+  injectMetaToDataStructures: jest.fn(),
+}));
+
+describe('indexPatternTypeConfig', () => {
+  const mockSavedObjectsClient = {} as SavedObjectsClientContract;
+  const mockServices = {
+    savedObjects: { client: mockSavedObjectsClient },
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  test('toDataset converts DataStructure to Dataset', () => {
+    const mockPath: DataStructure[] = [
+      {
+        id: 'test-pattern',
+        title: 'Test Pattern',
+        type: 'INDEX_PATTERN',
+        meta: { timeFieldName: '@timestamp', type: DATA_STRUCTURE_META_TYPES.CUSTOM },
+      },
+    ];
+
+    const result = indexPatternTypeConfig.toDataset(mockPath);
+
+    expect(result).toEqual({
+      id: 'test-pattern',
+      title: 'Test Pattern',
+      type: 'INDEX_PATTERN',
+      timeFieldName: '@timestamp',
+      dataSource: undefined,
+    });
+  });
+
+  test('fetchFields returns fields from index pattern', async () => {
+    const mockIndexPattern = {
+      fields: [
+        { name: 'field1', type: 'string' },
+        { name: 'field2', type: 'number' },
+      ],
+    };
+    const mockGet = jest.fn().mockResolvedValue(mockIndexPattern);
+    (services.getIndexPatterns as jest.Mock).mockReturnValue({ get: mockGet });
+
+    const mockDataset: Dataset = { id: 'test-pattern', title: 'Test', type: 'INDEX_PATTERN' };
+    const result = await indexPatternTypeConfig.fetchFields(mockDataset);
+
+    expect(result).toHaveLength(2);
+    expect(result[0]).toEqual({ name: 'field1', type: 'string' });
+    expect(result[1]).toEqual({ name: 'field2', type: 'number' });
+  });
+
+  test('supportedLanguages returns correct languages', () => {
+    const mockDataset: Dataset = {
+      id: 'test-pattern',
+      title: 'Test',
+      type: 'INDEX_PATTERN',
+      dataSource: { id: 'dataSourceId', title: 'Cluster 1', type: 'OpenSearch' },
+    };
+    expect(indexPatternTypeConfig.supportedLanguages(mockDataset)).toEqual([
+      'DQL',
+      'Lucene',
+      'PPL',
+      'SQL',
+    ]);
+
+    mockDataset.dataSource = { ...mockDataset.dataSource!, type: 'other' };
+    expect(indexPatternTypeConfig.supportedLanguages(mockDataset)).toEqual([
+      'DQL',
+      'Lucene',
+      'PPL',
+      'SQL',
+    ]);
+  });
+});