diff --git a/changelogs/fragments/7943.yml b/changelogs/fragments/7943.yml new file mode 100644 index 000000000000..d7b08ffce00b --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/src/plugins/data/common/search/search_source/search_source.ts b/src/plugins/data/common/search/search_source/search_source.ts index c95e54ffd35d..c8eb99b038a7 100644 --- a/src/plugins/data/common/search/search_source/search_source.ts +++ b/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); }); diff --git a/src/plugins/data/public/index.ts b/src/plugins/data/public/index.ts index f2696e824112..0ca84c184bf4 100644 --- a/src/plugins/data/public/index.ts +++ b/src/plugins/data/public/index.ts @@ -471,6 +471,8 @@ export { LanguageConfig, LanguageService, LanguageServiceContract, + RecentQueriesTable, + QueryControls, SavedQuery, SavedQueryService, SavedQueryTimeFilter, diff --git a/src/plugins/data/public/plugin.ts b/src/plugins/data/public/plugin.ts index ede785e14b84..4744606a3141 100644 --- a/src/plugins/data/public/plugin.ts +++ b/src/plugins/data/public/plugin.ts @@ -116,6 +116,7 @@ export class DataPublicPlugin private readonly fieldFormatsService: FieldFormatsService; private readonly queryService: QueryService; private readonly storage: DataStorage; + private readonly sessionStorage: DataStorage; constructor(initializerContext: PluginInitializerContext) { this.searchService = new SearchService(initializerContext); @@ -123,7 +124,11 @@ export class DataPublicPlugin 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(), }); diff --git a/src/plugins/data/public/query/query_service.ts b/src/plugins/data/public/query/query_service.ts index dba19cca19b0..52ca446b03af 100644 --- a/src/plugins/data/public/query/query_service.ts +++ b/src/plugins/data/public/query/query_service.ts @@ -58,8 +58,9 @@ export class QueryService { state$!: ReturnType; 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, diff --git a/src/plugins/data/public/query/query_string/dataset_service/dataset_service.test.ts b/src/plugins/data/public/query/query_string/dataset_service/dataset_service.test.ts new file mode 100644 index 000000000000..f72fd0ea5b46 --- /dev/null +++ b/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['uiSettings']; + let sessionStorage: DataStorage; + let mockDataPluginServices: jest.Mocked; + + beforeEach(() => { + uiSettings = coreMock.createSetup().uiSettings; + sessionStorage = new DataStorage(window.sessionStorage, 'opensearchDashboards.'); + mockDataPluginServices = {} as jest.Mocked; + + 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); + }); +}); diff --git a/src/plugins/data/public/query/query_string/dataset_service/dataset_service.ts b/src/plugins/data/public/query/query_string/dataset_service/dataset_service.ts index faa35328075f..2f9a0884442f 100644 --- a/src/plugins/data/public/query/query_string/dataset_service/dataset_service.ts +++ b/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 = 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 { 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(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( + `${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 { diff --git a/src/plugins/data/public/query/query_string/dataset_service/lib/index_pattern_type.test.ts b/src/plugins/data/public/query/query_string/dataset_service/lib/index_pattern_type.test.ts new file mode 100644 index 000000000000..58961289b298 --- /dev/null +++ b/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', + ]); + }); +}); diff --git a/src/plugins/data/public/query/query_string/dataset_service/lib/index_type.test.ts b/src/plugins/data/public/query/query_string/dataset_service/lib/index_type.test.ts new file mode 100644 index 000000000000..d6847f8d239a --- /dev/null +++ b/src/plugins/data/public/query/query_string/dataset_service/lib/index_type.test.ts @@ -0,0 +1,80 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +// index_type.test.ts + +import { indexTypeConfig } from './index_type'; +import { SavedObjectsClientContract } from 'opensearch-dashboards/public'; +import { DATA_STRUCTURE_META_TYPES, DataStructure, Dataset } from '../../../../../common'; +import * as services from '../../../../services'; + +jest.mock('../../../../services', () => ({ + getSearchService: jest.fn(), + getIndexPatterns: jest.fn(), +})); + +describe('indexTypeConfig', () => { + const mockSavedObjectsClient = {} as SavedObjectsClientContract; + const mockServices = { + savedObjects: { client: mockSavedObjectsClient }, + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + test('toDataset converts DataStructure to Dataset', () => { + const mockPath: DataStructure[] = [ + { + id: 'datasource1', + title: 'DataSource 1', + type: 'DATA_SOURCE', + }, + { + id: 'index1', + title: 'Index 1', + type: 'INDEX', + meta: { timeFieldName: '@timestamp', type: DATA_STRUCTURE_META_TYPES.CUSTOM }, + }, + ]; + + const result = indexTypeConfig.toDataset(mockPath); + + expect(result).toEqual({ + id: 'index1', + title: 'Index 1', + type: 'INDEXES', + timeFieldName: '@timestamp', + dataSource: { + id: 'datasource1', + title: 'DataSource 1', + type: 'DATA_SOURCE', + }, + }); + }); + + test('fetchFields returns fields from index', async () => { + const mockFields = [ + { name: 'field1', type: 'string' }, + { name: 'field2', type: 'number' }, + ]; + const mockGetFieldsForWildcard = jest.fn().mockResolvedValue(mockFields); + (services.getIndexPatterns as jest.Mock).mockReturnValue({ + getFieldsForWildcard: mockGetFieldsForWildcard, + }); + + const mockDataset: Dataset = { id: 'index1', title: 'Index 1', type: 'INDEX' }; + const result = await indexTypeConfig.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: 'index1', title: 'Index 1', type: 'INDEX' }; + expect(indexTypeConfig.supportedLanguages(mockDataset)).toEqual(['SQL', 'PPL']); + }); +}); diff --git a/src/plugins/data/public/query/query_string/index.ts b/src/plugins/data/public/query/query_string/index.ts index 9e5b584a5e2e..c8f6df921ad3 100644 --- a/src/plugins/data/public/query/query_string/index.ts +++ b/src/plugins/data/public/query/query_string/index.ts @@ -35,4 +35,6 @@ export { LanguageService, LanguageConfig, EditorEnhancements, + RecentQueriesTable, + QueryControls, } from './language_service'; diff --git a/src/plugins/data/public/query/query_string/language_service/index.ts b/src/plugins/data/public/query/query_string/language_service/index.ts index 79aea071de3f..cd04fcb50724 100644 --- a/src/plugins/data/public/query/query_string/language_service/index.ts +++ b/src/plugins/data/public/query/query_string/language_service/index.ts @@ -5,3 +5,4 @@ export * from './types'; export { LanguageServiceContract, LanguageService } from './language_service'; +export { RecentQueriesTable, QueryControls } from './lib'; diff --git a/src/plugins/data/public/query/query_string/language_service/language_service.mock.ts b/src/plugins/data/public/query/query_string/language_service/language_service.mock.ts index 80b6aa478b00..936ff690353d 100644 --- a/src/plugins/data/public/query/query_string/language_service/language_service.mock.ts +++ b/src/plugins/data/public/query/query_string/language_service/language_service.mock.ts @@ -60,6 +60,7 @@ const createSetupLanguageServiceMock = (): jest.Mocked setUserQuerySessionId: jest.fn(), setUserQuerySessionIdByObj: jest.fn(), getUserQuerySessionId: jest.fn().mockReturnValue(null), + createDefaultLanguageReference: jest.fn(), }; }; diff --git a/src/plugins/data/public/query/query_string/language_service/language_service.test.ts b/src/plugins/data/public/query/query_string/language_service/language_service.test.ts new file mode 100644 index 000000000000..5220fa9bdc29 --- /dev/null +++ b/src/plugins/data/public/query/query_string/language_service/language_service.test.ts @@ -0,0 +1,154 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { LanguageService } from './language_service'; +import { ISearchInterceptor } from '../../../search'; +import { DataStorage } from '../../../../common'; +import { LanguageConfig } from './types'; + +describe('LanguageService', () => { + let service: LanguageService; + let mockSearchInterceptor: jest.Mocked; + let mockStorage: jest.Mocked; + + beforeEach(() => { + mockSearchInterceptor = {} as jest.Mocked; + mockStorage = ({ + get: jest.fn(), + set: jest.fn(), + remove: jest.fn(), + } as unknown) as jest.Mocked; + + service = new LanguageService(mockSearchInterceptor, mockStorage); + }); + + test('registerLanguage and getLanguage', () => { + const mockLanguage: LanguageConfig = { + id: 'test-language', + title: 'Test Language', + search: {} as any, + getQueryString: jest.fn(), + editor: {} as any, + fields: {}, + showDocLinks: true, + editorSupportedAppNames: ['test-app'], + }; + + service.registerLanguage(mockLanguage); + expect(service.getLanguage('test-language')).toBe(mockLanguage); + }); + + test('getLanguages returns all registered languages', () => { + const languages = service.getLanguages(); + expect(languages).toHaveLength(2); // DQL and Lucene are registered by default + expect(languages[0].id).toBe('kuery'); + expect(languages[1].id).toBe('lucene'); + }); + + test('getDefaultLanguage returns DQL by default', () => { + const defaultLanguage = service.getDefaultLanguage(); + expect(defaultLanguage.id).toBe('kuery'); + }); + + test('getUserQueryLanguageBlocklist', () => { + mockStorage.get.mockReturnValue(['sql', 'ppl']); + expect(service.getUserQueryLanguageBlocklist()).toEqual(['sql', 'ppl']); + }); + + test('setUserQueryLanguageBlocklist', () => { + service.setUserQueryLanguageBlocklist(['sql', 'ppl']); + expect(mockStorage.set).toHaveBeenCalledWith('userQueryLanguageBlocklist', ['sql', 'ppl']); + }); + + test('getUserQueryLanguage', () => { + mockStorage.get.mockReturnValue('sql'); + expect(service.getUserQueryLanguage()).toBe('sql'); + }); + + test('setUserQueryLanguage', () => { + service.setUserQueryLanguage('sql'); + expect(mockStorage.set).toHaveBeenCalledWith('userQueryLanguage', 'sql'); + }); + + test('getUserQueryString', () => { + mockStorage.get.mockReturnValue('SELECT * FROM table'); + expect(service.getUserQueryString()).toBe('SELECT * FROM table'); + }); + + test('setUserQueryString', () => { + service.setUserQueryString('SELECT * FROM table'); + expect(mockStorage.set).toHaveBeenCalledWith('userQueryString', 'SELECT * FROM table'); + }); + + test('getUiOverrides', () => { + const mockOverrides = { fields: { filterable: true } }; + mockStorage.get.mockReturnValue(mockOverrides); + expect(service.getUiOverrides()).toEqual(mockOverrides); + }); + + test('setUiOverrides', () => { + const mockOverrides = { fields: { filterable: true } }; + service.setUiOverrides(mockOverrides); + expect(mockStorage.set).toHaveBeenCalledWith('uiOverrides', mockOverrides); + }); + + test('setUiOverrides with undefined clears overrides', () => { + service.setUiOverrides(undefined); + expect(mockStorage.remove).toHaveBeenCalledWith('uiOverrides'); + }); + + test('setUiOverridesByUserQueryLanguage', () => { + const mockLanguage: LanguageConfig = { + id: 'test-language', + title: 'Test Language', + search: {} as any, + getQueryString: jest.fn(), + editor: {} as any, + fields: { filterable: true }, + showDocLinks: true, + editorSupportedAppNames: ['test-app'], + }; + service.registerLanguage(mockLanguage); + + service.setUiOverridesByUserQueryLanguage('test-language'); + expect(mockStorage.set).toHaveBeenCalledWith('uiOverrides', { fields: { filterable: true } }); + }); + + test('setUserQuerySessionId', () => { + const mockSetItem = jest.spyOn(Storage.prototype, 'setItem'); + service.setUserQuerySessionId('test-source', 'test-session-id'); + expect(mockSetItem).toHaveBeenCalledWith( + 'async-query-session-id_test-source', + 'test-session-id' + ); + }); + + test('getUserQuerySessionId', () => { + const mockGetItem = jest.spyOn(Storage.prototype, 'getItem').mockReturnValue('test-session-id'); + expect(service.getUserQuerySessionId('test-source')).toBe('test-session-id'); + expect(mockGetItem).toHaveBeenCalledWith('async-query-session-id_test-source'); + }); + + test('setUserQuerySessionIdByObj', () => { + const mockSetItem = jest.spyOn(Storage.prototype, 'setItem'); + service.setUserQuerySessionIdByObj('test-source', { sessionId: 'test-session-id' }); + expect(mockSetItem).toHaveBeenCalledWith( + 'async-query-session-id_test-source', + 'test-session-id' + ); + }); + + test('resetUserQuery', () => { + service.resetUserQuery(); + expect(mockStorage.set).toHaveBeenCalledWith('userQueryLanguage', 'kuery'); + expect(mockStorage.set).toHaveBeenCalledWith('userQueryString', ''); + }); + + test('__enhance adds query editor extensions', () => { + const mockExtension = { id: 'test-extension' }; + service.__enhance({ queryEditorExtension: mockExtension }); + expect(service.getQueryEditorExtensionMap()).toEqual({ 'test-extension': mockExtension }); + }); +}); diff --git a/src/plugins/data/public/query/query_string/language_service/language_service.ts b/src/plugins/data/public/query/query_string/language_service/language_service.ts index 899915248e52..6fb250d8d359 100644 --- a/src/plugins/data/public/query/query_string/language_service/language_service.ts +++ b/src/plugins/data/public/query/query_string/language_service/language_service.ts @@ -14,7 +14,7 @@ import { UiEnhancements, } from '../../../ui'; import { DataStorage, setOverrides as setFieldOverrides } from '../../../../common'; -import { createDefaultLanguageReference } from './default_language_reference'; +import { createDefaultLanguageReference } from './lib/default_language_reference'; export class LanguageService { private languages: Map = new Map(); diff --git a/src/plugins/data/public/query/query_string/language_service/lib/_index.scss b/src/plugins/data/public/query/query_string/language_service/lib/_index.scss new file mode 100644 index 000000000000..3ebab52467b2 --- /dev/null +++ b/src/plugins/data/public/query/query_string/language_service/lib/_index.scss @@ -0,0 +1 @@ +@import "./recent_query"; diff --git a/src/plugins/data/public/query/query_string/language_service/_recent_query.scss b/src/plugins/data/public/query/query_string/language_service/lib/_recent_query.scss similarity index 100% rename from src/plugins/data/public/query/query_string/language_service/_recent_query.scss rename to src/plugins/data/public/query/query_string/language_service/lib/_recent_query.scss diff --git a/src/plugins/data/public/query/query_string/language_service/default_language_reference.tsx b/src/plugins/data/public/query/query_string/language_service/lib/default_language_reference.tsx similarity index 93% rename from src/plugins/data/public/query/query_string/language_service/default_language_reference.tsx rename to src/plugins/data/public/query/query_string/language_service/lib/default_language_reference.tsx index 56dda2550426..89f188363775 100644 --- a/src/plugins/data/public/query/query_string/language_service/default_language_reference.tsx +++ b/src/plugins/data/public/query/query_string/language_service/lib/default_language_reference.tsx @@ -9,8 +9,8 @@ import { EuiButtonIcon, EuiLink, EuiPopover, EuiPopoverTitle, EuiText } from '@e import React from 'react'; import { FormattedMessage } from 'react-intl'; -import { IDataPluginServices } from '../../../types'; -import { useOpenSearchDashboards } from '../../../../../opensearch_dashboards_react/public'; +import { IDataPluginServices } from '../../../../types'; +import { useOpenSearchDashboards } from '../../../../../../opensearch_dashboards_react/public'; export const DefaultLanguageReference = () => { const opensearchDashboards = useOpenSearchDashboards(); diff --git a/src/plugins/data/public/query/query_string/language_service/get_query_control_links.tsx b/src/plugins/data/public/query/query_string/language_service/lib/get_query_control_links.tsx similarity index 100% rename from src/plugins/data/public/query/query_string/language_service/get_query_control_links.tsx rename to src/plugins/data/public/query/query_string/language_service/lib/get_query_control_links.tsx diff --git a/src/plugins/data/public/query/query_string/language_service/lib/index.ts b/src/plugins/data/public/query/query_string/language_service/lib/index.ts index 42b8f7f2fb18..ca93870dfd4d 100644 --- a/src/plugins/data/public/query/query_string/language_service/lib/index.ts +++ b/src/plugins/data/public/query/query_string/language_service/lib/index.ts @@ -5,3 +5,7 @@ export * from './dql_language'; export * from './lucene_language'; + +export * from './default_language_reference'; +export * from './get_query_control_links'; +export * from './recent_query'; diff --git a/src/plugins/data/public/query/query_string/language_service/recent_query.tsx b/src/plugins/data/public/query/query_string/language_service/lib/recent_query.tsx similarity index 84% rename from src/plugins/data/public/query/query_string/language_service/recent_query.tsx rename to src/plugins/data/public/query/query_string/language_service/lib/recent_query.tsx index 0db291af8d31..2e3fdb7f55aa 100644 --- a/src/plugins/data/public/query/query_string/language_service/recent_query.tsx +++ b/src/plugins/data/public/query/query_string/language_service/lib/recent_query.tsx @@ -3,31 +3,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -import './_recent_query.scss'; +import './_index.scss'; import React, { useEffect, useState } from 'react'; import { EuiBasicTable, EuiBasicTableColumn, EuiButtonIcon, EuiCopy } from '@elastic/eui'; import moment from 'moment'; -import { Query, TimeRange } from 'src/plugins/data/common'; -import { QueryStringContract } from '../query_string_manager'; - -interface RecentQueryItem { - query: Query; - time: number; - timeRange?: TimeRange; -} - -interface RecentQueryTableItem { - id: number; - query: Query['query']; - time: string; -} - -interface RecentQueriesTableProps { - queryString: QueryStringContract; - onClickRecentQuery: (query: Query, timeRange?: TimeRange) => void; - isVisible: boolean; -} +import { RecentQueriesTableProps, RecentQueryItem, RecentQueryTableItem } from '../types'; export const MAX_RECENT_QUERY_SIZE = 10; diff --git a/src/plugins/data/public/query/query_string/language_service/types.ts b/src/plugins/data/public/query/query_string/language_service/types.ts index 5e556e188068..c351cc597917 100644 --- a/src/plugins/data/public/query/query_string/language_service/types.ts +++ b/src/plugins/data/public/query/query_string/language_service/types.ts @@ -4,9 +4,32 @@ */ import { ISearchInterceptor } from '../../../search'; -import { Query, QueryEditorExtensionConfig } from '../../../../public'; +import { + Query, + QueryEditorExtensionConfig, + QueryStringContract, + TimeRange, +} from '../../../../public'; import { EditorInstance } from '../../../ui/query_editor/editors'; +export interface RecentQueryItem { + query: Query; + time: number; + timeRange?: TimeRange; +} + +export interface RecentQueryTableItem { + id: number; + query: Query['query']; + time: string; +} + +export interface RecentQueriesTableProps { + queryString: QueryStringContract; + onClickRecentQuery: (query: Query, timeRange?: TimeRange) => void; + isVisible: boolean; +} + export interface EditorEnhancements { queryEditorExtension?: QueryEditorExtensionConfig; } diff --git a/src/plugins/data/public/query/query_string/query_string_manager.test.ts b/src/plugins/data/public/query/query_string/query_string_manager.test.ts index 1d7a8d9cf00b..4db9b55b622d 100644 --- a/src/plugins/data/public/query/query_string/query_string_manager.test.ts +++ b/src/plugins/data/public/query/query_string/query_string_manager.test.ts @@ -32,19 +32,22 @@ import { QueryStringManager } from './query_string_manager'; import { coreMock } from '../../../../../core/public/mocks'; import { Query } from '../../../common/query'; import { ISearchInterceptor } from '../../search'; -import { DataStorage } from 'src/plugins/data/common'; +import { DataStorage, DEFAULT_DATA } from 'src/plugins/data/common'; describe('QueryStringManager', () => { let service: QueryStringManager; let storage: DataStorage; + let sessionStorage: DataStorage; let mockSearchInterceptor: jest.Mocked; beforeEach(() => { - storage = new DataStorage(window.localStorage, 'opensearch_dashboards.'); + storage = new DataStorage(window.localStorage, 'opensearchDashboards.'); + sessionStorage = new DataStorage(window.sessionStorage, 'opensearchDashboards.'); mockSearchInterceptor = {} as jest.Mocked; service = new QueryStringManager( storage, + sessionStorage, coreMock.createSetup().uiSettings, mockSearchInterceptor ); @@ -66,4 +69,74 @@ describe('QueryStringManager', () => { service.setQuery({ ...newQuery }); expect(emittedValues).toHaveLength(1); }); + + test('getQuery returns the current query', () => { + const initialQuery = service.getQuery(); + expect(initialQuery).toHaveProperty('query'); + expect(initialQuery).toHaveProperty('language'); + + const newQuery = { query: 'test query', language: 'sql' }; + service.setQuery(newQuery); + expect(service.getQuery()).toEqual(newQuery); + }); + + test('clearQuery resets to default query', () => { + const newQuery = { query: 'test query', language: 'sql' }; + service.setQuery(newQuery); + expect(service.getQuery()).toEqual(newQuery); + + service.clearQuery(); + const defaultQuery = service.getQuery(); + expect(defaultQuery).not.toEqual(newQuery); + expect(defaultQuery.query).toBe(''); + }); + + test('formatQuery handles different input types', () => { + const stringQuery = 'test query'; + const formattedStringQuery = service.formatQuery(stringQuery); + expect(formattedStringQuery).toHaveProperty('query', stringQuery); + expect(formattedStringQuery).toHaveProperty('language'); + + const objectQuery = { query: 'object query', language: 'sql' }; + const formattedObjectQuery = service.formatQuery(objectQuery); + expect(formattedObjectQuery).toEqual(objectQuery); + + const formattedUndefinedQuery = service.formatQuery(undefined); + expect(formattedUndefinedQuery).toEqual(service.getDefaultQuery()); + }); + + test('clearQueryHistory clears the query history', () => { + service.addToQueryHistory({ query: 'test query 1', language: 'sql' }); + service.addToQueryHistory({ query: 'test query 2', language: 'sql' }); + expect(service.getQueryHistory()).toHaveLength(2); + + service.clearQueryHistory(); + expect(service.getQueryHistory()).toHaveLength(0); + }); + + test('addToQueryHistory adds query to history', () => { + const query: Query = { query: 'test query', language: 'sql' }; + service.addToQueryHistory(query); + const history = service.getQueryHistory(); + expect(history).toHaveLength(1); + expect(history[0]).toHaveProperty('query', query); + }); + + test('getInitialQueryByLanguage returns correct query for language', () => { + const sqlQuery = service.getInitialQueryByLanguage('sql'); + expect(sqlQuery).toHaveProperty('language', 'sql'); + + const pplQuery = service.getInitialQueryByLanguage('ppl'); + expect(pplQuery).toHaveProperty('language', 'ppl'); + }); + + test('getInitialQueryByDataset returns correct query for dataset', () => { + const dataset = { + id: 'test-dataset', + title: 'Test Dataset', + type: DEFAULT_DATA.SET_TYPES.INDEX_PATTERN, + }; + const query = service.getInitialQueryByDataset(dataset); + expect(query).toHaveProperty('dataset', dataset); + }); }); diff --git a/src/plugins/data/public/query/query_string/query_string_manager.ts b/src/plugins/data/public/query/query_string/query_string_manager.ts index 723461224453..e26f3a042c03 100644 --- a/src/plugins/data/public/query/query_string/query_string_manager.ts +++ b/src/plugins/data/public/query/query_string/query_string_manager.ts @@ -46,12 +46,13 @@ export class QueryStringManager { constructor( private readonly storage: DataStorage, + private readonly sessionStorage: DataStorage, private readonly uiSettings: CoreStart['uiSettings'], private readonly defaultSearchInterceptor: ISearchInterceptor ) { this.query$ = new BehaviorSubject(this.getDefaultQuery()); this.queryHistory = createHistory({ storage }); - this.datasetService = new DatasetService(uiSettings); + this.datasetService = new DatasetService(uiSettings, this.sessionStorage); this.languageService = new LanguageService(this.defaultSearchInterceptor, this.storage); } diff --git a/src/plugins/data/public/query/types.ts b/src/plugins/data/public/query/types.ts index ee359f0939a2..1b4a58697a3d 100644 --- a/src/plugins/data/public/query/types.ts +++ b/src/plugins/data/public/query/types.ts @@ -34,8 +34,9 @@ export interface IQueryStart { /** @internal */ export interface QueryServiceSetupDependencies { - storage: DataStorage; uiSettings: IUiSettingsClient; + storage: DataStorage; + sessionStorage: DataStorage; defaultSearchInterceptor: ISearchInterceptor; } diff --git a/src/plugins/data/public/ui/dataset_selector/configurator.tsx b/src/plugins/data/public/ui/dataset_selector/configurator.tsx index 527d5f7e66c8..f90695ab6b84 100644 --- a/src/plugins/data/public/ui/dataset_selector/configurator.tsx +++ b/src/plugins/data/public/ui/dataset_selector/configurator.tsx @@ -151,10 +151,16 @@ export const Configurator = ({ - onConfirm({ ...baseDataset, language, timeFieldName })} fill> + { + queryString.getDatasetService().cacheDataset({ ...dataset, language, timeFieldName }); + onConfirm({ ...dataset, language, timeFieldName }); + }} + fill + > 100' + }); + } +} +``` + +### Services +#### Query String Manager +The Query String Manager provides methods to manage and interact with queries: + +- `setQuery(query: Query)`: Set the current query. +- `getQuery()`: Query: Get the current query. +- `getLanguageService()`: Access language-specific services. + +#### Dataset Service +The Dataset Service allows interaction with different data sources: + +- `registerType(typeConfig: DatasetTypeConfig)`: Register a new dataset type. +- `getType(type: string)`: DatasetTypeConfig: Get a registered dataset type. + +#### Query Languages +##### PPL (Piped Processing Language) +PPL is a query language that uses a series of commands separated by pipes (|) to process and transform data. +Example PPL query: +``` +source = my_index | where count > 100 | stats sum(price) by category +``` + +##### SQL +SQL support allows users to query data using standard SQL syntax. +Example SQL query: +```sql +SELECT category, SUM(price) FROM my_index WHERE count > 100 GROUP BY category +``` \ No newline at end of file diff --git a/src/plugins/query_enhancements/common/types.ts b/src/plugins/query_enhancements/common/types.ts index 7efb0d9f3d13..3cfaab27128e 100644 --- a/src/plugins/query_enhancements/common/types.ts +++ b/src/plugins/query_enhancements/common/types.ts @@ -4,7 +4,6 @@ */ import { CoreSetup } from 'opensearch-dashboards/public'; -import { Observable } from 'rxjs'; export interface QueryAggConfig { [key: string]: { @@ -17,10 +16,21 @@ export interface QueryAggConfig { }; } +export interface QueryStatusConfig { + sessionId: string; + queryId?: string; +} + export interface EnhancedFetchContext { http: CoreSetup['http']; path: string; signal?: AbortSignal; } -export type FetchFunction = (params?: P) => Observable; +export interface QueryStatusOptions { + fetchStatus: () => Promise; + interval?: number; + isServer?: boolean; +} + +export type FetchFunction = (params?: P) => Promise; diff --git a/src/plugins/query_enhancements/common/utils.ts b/src/plugins/query_enhancements/common/utils.ts index bdb8740e4e48..597bf3124496 100644 --- a/src/plugins/query_enhancements/common/utils.ts +++ b/src/plugins/query_enhancements/common/utils.ts @@ -3,10 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { IDataFrame, Query } from 'src/plugins/data/common'; -import { Observable, Subscription, from, throwError, timer } from 'rxjs'; -import { catchError, concatMap, last, takeWhile, tap } from 'rxjs/operators'; -import { EnhancedFetchContext, FetchFunction, QueryAggConfig } from './types'; +import { Query } from 'src/plugins/data/common'; +import { from, throwError, timer } from 'rxjs'; +import { filter, mergeMap, take, takeWhile, tap } from 'rxjs/operators'; +import { + EnhancedFetchContext, + QueryAggConfig, + QueryStatusConfig, + QueryStatusOptions, +} from './types'; export const formatDate = (dateString: string) => { const date = new Date(dateString); @@ -36,100 +41,16 @@ export const removeKeyword = (queryString: string | undefined) => { return queryString?.replace(new RegExp('.keyword'), '') ?? ''; }; -export class DataFramePolling { - public data: T | null = null; - public error: Error | null = null; - public loading: boolean = true; - private shouldPoll: boolean = false; - private intervalRef?: NodeJS.Timeout; - private subscription?: Subscription; - - constructor( - private fetchFunction: FetchFunction, - private interval: number = 5000, - private onPollingSuccess: (data: T) => boolean, - private onPollingError: (error: Error) => boolean - ) {} - - fetch(): Observable { - return timer(0, this.interval).pipe( - concatMap(() => this.fetchFunction()), - takeWhile((resp) => this.onPollingSuccess(resp), true), - tap((resp: T) => { - this.data = resp; - }), - last(), - catchError((error: Error) => { - this.onPollingError(error); - return throwError(error); - }) - ); - } - - fetchData(params?: P) { - this.loading = true; - this.subscription = this.fetchFunction(params).subscribe({ - next: (result: any) => { - this.data = result; - this.loading = false; - - if (this.onPollingSuccess && this.onPollingSuccess(result)) { - this.stopPolling(); - } - }, - error: (err: any) => { - this.error = err as Error; - this.loading = false; - - if (this.onPollingError && this.onPollingError(this.error)) { - this.stopPolling(); - } - }, - }); - } - - startPolling(params?: P) { - this.shouldPoll = true; - if (!this.intervalRef) { - this.intervalRef = setInterval(() => { - if (this.shouldPoll) { - this.fetchData(params); - } - }, this.interval); - } - } - - stopPolling() { - this.shouldPoll = false; - if (this.intervalRef) { - clearInterval(this.intervalRef); - this.intervalRef = undefined; - } - if (this.subscription) { - this.subscription.unsubscribe(); - this.subscription = undefined; - } - } - - waitForPolling(): Promise { - return new Promise((resolve) => { - const checkLoading = () => { - if (!this.loading) { - resolve(this.data); - } else { - setTimeout(checkLoading, this.interval); - } - }; - checkLoading(); - }); - } -} +export const handleFacetError = (response: any) => { + const error = new Error(response.data); + error.name = response.status; + return throwError(error); +}; -export const handleDataFrameError = (response: any) => { - const df = response.body; - if (df.error) { - const jsError = new Error(df.error.response); - return throwError(jsError); +export const handleFetchError = (response: any) => { + if (response.body.error) { + const error = new Error(response.body.error.response); + return throwError(error); } }; @@ -143,31 +64,38 @@ export const fetch = (context: EnhancedFetchContext, query: Query, aggConfig?: Q body, signal, }) - ).pipe(tap(handleDataFrameError)); + ).pipe(tap(handleFetchError)); }; -export const fetchDataFrame = (context: EnhancedFetchContext, query: Query, df: IDataFrame) => { - const { http, path, signal } = context; - const body = JSON.stringify({ query: { ...query, format: 'jdbc' }, df }); - return from( - http.fetch({ - method: 'POST', - path, - body, - signal, - }) - ).pipe(tap(handleDataFrameError)); +export const handleQueryStatus = (options: QueryStatusOptions): Promise => { + const { fetchStatus, interval = 5000, isServer = false } = options; + + return timer(0, interval) + .pipe( + mergeMap(() => fetchStatus()), + takeWhile((response) => { + const status = isServer + ? (response as any)?.data?.status?.toUpperCase() + : (response as any)?.status?.toUpperCase(); + return status !== 'SUCCESS' && status !== 'FAILED'; + }, true), + filter((response) => { + const status = isServer + ? (response as any)?.data?.status?.toUpperCase() + : (response as any)?.status?.toUpperCase(); + if (status === 'FAILED') { + throw new Error('Job failed'); + } + return status === 'SUCCESS'; + }), + take(1) + ) + .toPromise(); }; -export const fetchDataFramePolling = (context: EnhancedFetchContext, df: IDataFrame) => { - const { http, path, signal } = context; - const queryId = df.meta?.queryId; - const dataSourceId = df.meta?.queryConfig?.dataSourceId; - return from( - http.fetch({ - method: 'GET', - path: `${path}/${queryId}${dataSourceId ? `/${dataSourceId}` : ''}`, - signal, - }) - ); +export const buildQueryStatusConfig = (response: any) => { + return { + queryId: response.data.queryId, + sessionId: response.data.sessionId, + } as QueryStatusConfig; }; diff --git a/src/plugins/query_enhancements/public/datasets/index.ts b/src/plugins/query_enhancements/public/datasets/index.ts index ddd7bf74995c..49b39d052b6d 100644 --- a/src/plugins/query_enhancements/public/datasets/index.ts +++ b/src/plugins/query_enhancements/public/datasets/index.ts @@ -3,4 +3,4 @@ * SPDX-License-Identifier: Apache-2.0 */ -export * from './s3_handler'; +export * from './s3_type'; diff --git a/src/plugins/query_enhancements/public/datasets/s3_type.test.ts b/src/plugins/query_enhancements/public/datasets/s3_type.test.ts new file mode 100644 index 000000000000..8688fb0ce891 --- /dev/null +++ b/src/plugins/query_enhancements/public/datasets/s3_type.test.ts @@ -0,0 +1,74 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +// s3_type.test.ts + +import { s3TypeConfig } from './s3_type'; +import { HttpSetup, SavedObjectsClientContract } from 'opensearch-dashboards/public'; +import { DataStructure, Dataset } from '../../../data/common'; + +describe('s3TypeConfig', () => { + const mockHttp = {} as HttpSetup; + const mockSavedObjectsClient = {} as SavedObjectsClientContract; + const mockServices = { + http: mockHttp, + savedObjects: { client: mockSavedObjectsClient }, + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + test('toDataset converts DataStructure path to Dataset', () => { + const mockPath: DataStructure[] = [ + { id: 'ds1', title: 'DataSource 1', type: 'DATA_SOURCE' }, + { id: 'conn1', title: 'Connection 1', type: 'CONNECTION' }, + { id: 'db1', title: 'Database 1', type: 'DATABASE' }, + { id: 'table1', title: 'Table 1', type: 'TABLE' }, + ]; + + const result = s3TypeConfig.toDataset(mockPath); + + expect(result).toEqual({ + id: 'table1', + title: 'Connection 1.Database 1.Table 1', + type: 'S3', + dataSource: { + id: 'ds1', + title: 'DataSource 1', + type: 'DATA_SOURCE', + }, + }); + }); + + test.skip('fetch returns correct structure based on input', async () => { + const mockFetch = jest.fn().mockResolvedValue({ datarows: [['db1'], ['db2']] }); + mockHttp.fetch = mockFetch; + + const mockFind = jest.fn().mockResolvedValue({ + savedObjects: [{ id: 'ds1', attributes: { title: 'DataSource 1' } }], + }); + mockSavedObjectsClient.find = mockFind; + + const result = await s3TypeConfig.fetch(mockServices as any, [ + { id: 'ds1', title: 'DS 1', type: 'DATA_SOURCE' }, + ]); + + expect(result.children).toBeDefined(); + expect(result.hasNext).toBe(true); + }); + + test('fetchFields returns empty array', async () => { + const mockDataset: Dataset = { id: 'table1', title: 'Table 1', type: 'S3' }; + const result = await s3TypeConfig.fetchFields(mockDataset); + + expect(result).toEqual([]); + }); + + test('supportedLanguages returns SQL', () => { + const mockDataset: Dataset = { id: 'table1', title: 'Table 1', type: 'S3' }; + expect(s3TypeConfig.supportedLanguages(mockDataset)).toEqual(['SQL']); + }); +}); diff --git a/src/plugins/query_enhancements/public/datasets/s3_handler.ts b/src/plugins/query_enhancements/public/datasets/s3_type.ts similarity index 74% rename from src/plugins/query_enhancements/public/datasets/s3_handler.ts rename to src/plugins/query_enhancements/public/datasets/s3_type.ts index f54435df293f..d2bca88696b3 100644 --- a/src/plugins/query_enhancements/public/datasets/s3_handler.ts +++ b/src/plugins/query_enhancements/public/datasets/s3_type.ts @@ -4,8 +4,6 @@ */ import { HttpSetup, SavedObjectsClientContract } from 'opensearch-dashboards/public'; -import { timer } from 'rxjs'; -import { filter, map, mergeMap, takeWhile } from 'rxjs/operators'; import { DATA_STRUCTURE_META_TYPES, DEFAULT_DATA, @@ -16,7 +14,7 @@ import { DatasetField, } from '../../../data/common'; import { DatasetTypeConfig, IDataPluginServices } from '../../../data/public'; -import { DATASET } from '../../common'; +import { DATASET, handleQueryStatus } from '../../common'; const S3_ICON = 'visTable'; @@ -105,66 +103,46 @@ export const s3TypeConfig: DatasetTypeConfig = { }, }; -const fetch = ( +const fetch = async ( http: HttpSetup, path: DataStructure[], type: 'DATABASE' | 'TABLE' ): Promise => { - return new Promise((resolve, reject) => { - const dataSource = path.find((ds) => ds.type === 'DATA_SOURCE'); - const parent = path[path.length - 1]; - const meta = parent.meta as DataStructureCustomMeta; - - timer(0, 5000) - .pipe( - mergeMap(() => - http.fetch('../../api/enhancements/datasource/jobs', { - query: { - id: dataSource?.id, - queryId: meta.query.id, - }, - }) - ), - takeWhile( - (response) => response.status !== 'SUCCESS' && response.status !== 'FAILED', - true - ), - filter((response) => response.status === 'SUCCESS'), - map((response) => { - if (response.status === 'FAILED') { - throw new Error('Job failed'); - } - return response.datarows.map((item: string[]) => ({ - id: `${parent.id}.${item[type === 'DATABASE' ? 0 : 1]}`, - title: item[type === 'DATABASE' ? 0 : 1], - type, - meta: { - type: DATA_STRUCTURE_META_TYPES.CUSTOM, - query: meta.query, - session: meta.session, - } as DataStructureCustomMeta, - })); - }) - ) - .subscribe({ - next: (dataStructures) => { - resolve(dataStructures); - }, - error: (error) => { - reject(error); - }, - complete: () => { - reject(new Error('No response')); - }, - }); - }); + const dataSource = path.find((ds) => ds.type === 'DATA_SOURCE'); + const parent = path[path.length - 1]; + const meta = parent.meta as DataStructureCustomMeta; + + try { + const response = await handleQueryStatus({ + fetchStatus: () => + http.fetch('../../api/enhancements/datasource/jobs', { + query: { + id: dataSource?.id, + queryId: meta.queryId, + }, + }), + }); + + return response.datarows.map((item: string[]) => ({ + id: `${parent.id}.${item[type === 'DATABASE' ? 0 : 1]}`, + title: item[type === 'DATABASE' ? 0 : 1], + type, + meta: { + type: DATA_STRUCTURE_META_TYPES.CUSTOM, + query: meta.query, + session: meta.session, + } as DataStructureCustomMeta, + })); + } catch (error) { + throw error; + } }; const setMeta = (dataStructure: DataStructure, response: any) => { return { ...dataStructure.meta, - query: { id: response.queryId }, - session: { id: response.sessionId }, + queryId: response.queryId, + sessionId: response.sessionId, } as DataStructureCustomMeta; }; @@ -214,14 +192,17 @@ const fetchConnections = async ( const fetchDatabases = async (http: HttpSetup, path: DataStructure[]): Promise => { const dataSource = path.find((ds) => ds.type === 'DATA_SOURCE'); const connection = path[path.length - 1]; - const query = (connection.meta as DataStructureCustomMeta).query; + const meta = connection.meta as DataStructureCustomMeta; const response = await http.post(`../../api/enhancements/datasource/jobs`, { body: JSON.stringify({ lang: 'sql', query: `SHOW DATABASES in ${connection.title}`, datasource: dataSource?.title, + ...(meta.sessionId && { sessionId: meta.sessionId }), }), - query, + query: { + id: dataSource?.id, + }, }); connection.meta = setMeta(connection, response); @@ -231,12 +212,15 @@ const fetchDatabases = async (http: HttpSetup, path: DataStructure[]): Promise => { const dataSource = path.find((ds) => ds.type === 'DATA_SOURCE'); + const sessionId = (path.find((ds) => ds.type === 'CONNECTION')?.meta as DataStructureCustomMeta) + .sessionId; const database = path[path.length - 1]; const response = await http.post(`../../api/enhancements/datasource/jobs`, { body: JSON.stringify({ lang: 'sql', query: `SHOW TABLES in ${database.title}`, datasource: dataSource?.title, + ...(sessionId && { sessionId }), }), query: { id: dataSource?.id, diff --git a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts index 69e4e14af562..07f6ce7fce4a 100644 --- a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts +++ b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts @@ -5,8 +5,8 @@ import { trimEnd } from 'lodash'; import { Observable, throwError } from 'rxjs'; -import { catchError } from 'rxjs/operators'; -import { Query } from '../../../data/common'; +import { catchError, tap } from 'rxjs/operators'; +import { CoreStart } from 'opensearch-dashboards/public'; import { DataPublicPluginStart, IOpenSearchDashboardsSearchRequest, @@ -20,12 +20,14 @@ import { QueryEnhancementsPluginStartDependencies } from '../types'; export class SQLSearchInterceptor extends SearchInterceptor { protected queryService!: DataPublicPluginStart['query']; + protected notifications!: CoreStart['notifications']; constructor(deps: SearchInterceptorDeps) { super(deps); deps.startServices.then(([coreStart, depsStart]) => { this.queryService = (depsStart as QueryEnhancementsPluginStartDependencies).data.query; + this.notifications = coreStart.notifications; }); } @@ -34,16 +36,16 @@ export class SQLSearchInterceptor extends SearchInterceptor { signal?: AbortSignal, strategy?: string ): Observable { - const { id, ...searchRequest } = request; + const isAsync = strategy === SEARCH_STRATEGY.SQL_ASYNC; const context: EnhancedFetchContext = { http: this.deps.http, - path: trimEnd(API.SQL_SEARCH), + path: trimEnd(isAsync ? API.SQL_ASYNC_SEARCH : API.SQL_SEARCH), signal, }; - const query = this.buildQuery(strategy); - - return fetch(context, query).pipe( + if (isAsync) this.notifications.toasts.add('Fetching data...'); + return fetch(context, this.queryService.queryString.getQuery()).pipe( + tap(() => isAsync && this.notifications.toasts.addSuccess('Fetch complete...')), catchError((error) => { return throwError(error); }) @@ -55,27 +57,4 @@ export class SQLSearchInterceptor extends SearchInterceptor { const strategy = dataset?.type === DATASET.S3 ? SEARCH_STRATEGY.SQL_ASYNC : SEARCH_STRATEGY.SQL; return this.runSearch(request, options.abortSignal, strategy); } - - private buildQuery(strategy?: string): Query { - const query: Query = this.queryService.queryString.getQuery(); - // TODO: MQL keeping here for S3 - // const dataset = query.dataset; - - // if (strategy === SEARCH_STRATEGY.SQL_ASYNC && dataset?.dataSource) { - // const sessionId = this.queryService.queryString - // .getLanguageService() - // .getUserQuerySessionId(dataset.dataSource.title); - // if (sessionId) { - // return { - // ...query, - // meta: { - // ...query.meta, - // sessionId, - // }, - // }; - // } - // } - - return query; - } } diff --git a/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts b/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts index 856bf056f748..65eb9da67ac6 100644 --- a/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts +++ b/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts @@ -8,14 +8,13 @@ import { Observable } from 'rxjs'; import { ISearchStrategy, SearchUsage } from '../../../data/server'; import { DATA_FRAME_TYPES, - IDataFrameError, IDataFrameResponse, IDataFrameWithAggs, IOpenSearchDashboardsSearchRequest, Query, createDataFrame, } from '../../../data/common'; -import { getFields } from '../../common/utils'; +import { getFields, handleFacetError } from '../../common/utils'; import { Facet } from '../utils'; import { QueryAggConfig } from '../../common'; @@ -40,13 +39,7 @@ export const pplSearchStrategyProvider = ( const aggConfig: QueryAggConfig | undefined = request.body.aggConfig; const rawResponse: any = await pplFacet.describeQuery(context, request); - if (!rawResponse.success) { - return { - type: DATA_FRAME_TYPES.ERROR, - body: { error: rawResponse.data }, - took: rawResponse.took, - } as IDataFrameError; - } + if (!rawResponse.success) handleFacetError(rawResponse); const dataFrame = createDataFrame({ name: query.dataset?.id, diff --git a/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts index ea292c1a3218..6cf067d5614f 100644 --- a/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts +++ b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts @@ -8,14 +8,19 @@ import { Observable } from 'rxjs'; import { ISearchStrategy, SearchUsage } from '../../../data/server'; import { DATA_FRAME_TYPES, - IDataFrameError, IDataFrameResponse, IOpenSearchDashboardsSearchRequest, - PartialDataFrame, Query, createDataFrame, } from '../../../data/common'; import { Facet } from '../utils'; +import { + buildQueryStatusConfig, + getFields, + handleFacetError, + handleQueryStatus, + SEARCH_STRATEGY, +} from '../../common'; export const sqlAsyncSearchStrategyProvider = ( config$: Observable, @@ -34,78 +39,45 @@ export const sqlAsyncSearchStrategyProvider = ( return { search: async (context, request: any, options) => { try { - const query: Query = request?.body?.query; - // Create job: this should return a queryId and sessionId - if (query) { - const df = request.body?.df; - request.body = { - query: query.query, - datasource: query.dataset?.dataSource?.title, - lang: SEARCH_STRATEGY.SQL, - sessionId: df?.meta?.sessionId, - }; - const rawResponse: any = await sqlAsyncFacet.describeQuery(context, request); - // handles failure - if (!rawResponse.success) { - return { - type: DATA_FRAME_TYPES.POLLING, - body: { error: rawResponse.data }, - took: rawResponse.took, - } as IDataFrameError; - } - const queryId = rawResponse.data?.queryId; - const sessionId = rawResponse.data?.sessionId; + const query: Query = request.body.query; + const startTime = Date.now(); + request.body = { ...request.body, lang: SEARCH_STRATEGY.SQL }; + const rawResponse: any = await sqlAsyncFacet.describeQuery(context, request); + + if (!rawResponse.success) handleFacetError(rawResponse); + + const statusConfig = buildQueryStatusConfig(rawResponse); + request.params = { queryId: statusConfig.queryId }; + + const response = await handleQueryStatus({ + fetchStatus: async () => { + const status: any = await sqlAsyncJobsFacet.describeQuery(context, request); + logger.info( + `sqlAsyncSearchStrategy: JOB: ${statusConfig.queryId} - STATUS: ${status.data?.status}` + ); + return status; + }, + isServer: true, + }); - const partial: PartialDataFrame = { - ...df, - fields: rawResponse?.data?.schema || [], - }; - const dataFrame = createDataFrame(partial); - dataFrame.meta = { - ...dataFrame?.meta, - query: query.query, - queryId, - sessionId, - }; - dataFrame.name = request.body?.datasource; - return { - type: DATA_FRAME_TYPES.POLLING, - body: dataFrame, - took: rawResponse.took, - } as IDataFrameResponse; - } else { - const queryId = request.params.queryId; - request.params = { queryId }; - const asyncResponse: any = await sqlAsyncJobsFacet.describeQuery(context, request); - const status = asyncResponse.data.status; - const partial: PartialDataFrame = { - name: '', - fields: asyncResponse?.data?.schema || [], - }; - const dataFrame = createDataFrame(partial); - dataFrame.fields?.forEach((field, index) => { - field.values = asyncResponse?.data.datarows.map((row: any) => row[index]); - }); + const dataFrame = createDataFrame({ + name: query.dataset?.id, + schema: response.data.schema, + meta: statusConfig, + fields: getFields(response), + }); - dataFrame.size = asyncResponse?.data?.datarows?.length || 0; + const elapsedMs = Date.now() - startTime; - dataFrame.meta = { - ...dataFrame?.meta, - status, - queryId, - error: status === 'FAILED' && asyncResponse.data?.error, - }; - dataFrame.name = request.body?.datasource; + dataFrame.size = response.data.datarows.length; - // TODO: MQL should this be the time for polling or the time for job creation? - if (usage) usage.trackSuccess(asyncResponse.took); + if (usage) usage.trackSuccess(elapsedMs); - return { - type: DATA_FRAME_TYPES.POLLING, - body: dataFrame, - took: asyncResponse.took, - } as IDataFrameResponse; - } + return { + type: DATA_FRAME_TYPES.POLLING, + body: dataFrame, + took: elapsedMs, + } as IDataFrameResponse; } catch (e) { logger.error(`sqlAsyncSearchStrategy: ${e.message}`); if (usage) usage.trackError(); diff --git a/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts b/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts index 822534879040..2df810f7d4b5 100644 --- a/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts +++ b/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts @@ -15,9 +15,7 @@ import { SearchUsage } from '../../../data/server'; import { DATA_FRAME_TYPES, IDataFrameError, - IDataFrameResponse, IOpenSearchDashboardsSearchRequest, - Query, } from '../../../data/common'; import * as facet from '../utils/facet'; import * as utils from '../../common/utils'; diff --git a/src/plugins/query_enhancements/server/utils/facet.ts b/src/plugins/query_enhancements/server/utils/facet.ts index ca8f6940e356..587a0f933444 100644 --- a/src/plugins/query_enhancements/server/utils/facet.ts +++ b/src/plugins/query_enhancements/server/utils/facet.ts @@ -38,12 +38,20 @@ export class Facet { ): Promise => { try { const query: Query = request.body.query; - const { format, df } = request.body; + const { dataSource } = query.dataset!; + const { format, lang } = request.body; const params = { - body: { query: query.query }, + body: { + query: query.query, + ...(dataSource && { datasource: dataSource.title }), + ...(dataSource?.meta?.sessionId && { + sessionId: dataSource?.meta?.sessionId, + }), + ...(lang && { lang }), + }, ...(format !== 'jdbc' && { format }), }; - const clientId = query.dataset?.dataSource?.id ?? df?.meta?.queryConfig?.dataSourceId; + const clientId = dataSource?.id; const client = clientId ? context.dataSource.opensearch.legacy.getClient(clientId).callAPI : this.defaultClient.asScoped(request).callAsCurrentUser; @@ -67,11 +75,11 @@ export class Facet { endpoint: string ): Promise => { try { + const query: Query = request.body.query; const params = request.params; - const { df } = request.body; - const dataSourceId = df?.meta?.queryConfig?.dataSourceId; - const client = dataSourceId - ? context.dataSource.opensearch.legacy.getClient(dataSourceId).callAPI + const clientId = query.dataset?.dataSource?.id; + const client = clientId + ? context.dataSource.opensearch.legacy.getClient(clientId).callAPI : this.defaultClient.asScoped(request).callAsCurrentUser; const queryRes = await client(endpoint, params); return {