fix(query assist): update reading data source id from dataset manager (…

…opensearch-project#7464) * revert to read datasource id from index pattern Signed-off-by: Joshua Li <[email protected]> * add dataset mock to query mock Signed-off-by: Joshua Li <[email protected]> * update query assist to use dataset manager Signed-off-by: Joshua Li <[email protected]> * use selected dataset state instead of relying on rerender Signed-off-by: Joshua Li <[email protected]> * remove skip 1 in dataset observable Signed-off-by: Joshua Li <[email protected]> * update dataset_manager tests Signed-off-by: Joshua Li <[email protected]> --------- Signed-off-by: Joshua Li <[email protected]>
Qxisylolo · Jul 25, 2024 · a526102 · a526102
1 parent 8a5faff
commit a526102
Show file tree

Hide file tree

Showing 9 changed files with 57 additions and 207 deletions.
diff --git a/src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts b/src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts
@@ -14,20 +14,21 @@ describe('DataSetManager', () => {
     service = new DataSetManager(coreMock.createSetup().uiSettings);
   });
 
-  test('getUpdates$ is a cold emits only after query changes', () => {
+  test('getUpdates$ emits initially and after data set changes', () => {
     const obs$ = service.getUpdates$();
-    const emittedValues: SimpleDataSet[] = [];
+    const emittedValues: Array<SimpleDataSet | undefined> = [];
     obs$.subscribe((v) => {
-      emittedValues.push(v!);
+      emittedValues.push(v);
     });
-    expect(emittedValues).toHaveLength(0);
+    expect(emittedValues).toHaveLength(1);
+    expect(emittedValues[1]).toEqual(undefined);
 
     const newDataSet: SimpleDataSet = { id: 'test_dataset', title: 'Test Dataset' };
     service.setDataSet(newDataSet);
-    expect(emittedValues).toHaveLength(1);
-    expect(emittedValues[0]).toEqual(newDataSet);
+    expect(emittedValues).toHaveLength(2);
+    expect(emittedValues[1]).toEqual(newDataSet);
 
     service.setDataSet({ ...newDataSet });
-    expect(emittedValues).toHaveLength(2);
+    expect(emittedValues).toHaveLength(3);
   });
 });
diff --git a/src/plugins/data/public/query/dataset_manager/dataset_manager.ts b/src/plugins/data/public/query/dataset_manager/dataset_manager.ts
@@ -26,7 +26,7 @@ export class DataSetManager {
   };
 
   public getUpdates$ = () => {
-    return this.dataSet$.asObservable().pipe(skip(1));
+    return this.dataSet$.asObservable();
   };
 
   public getDataSet = () => {

diff --git a/src/plugins/data/public/query/mocks.ts b/src/plugins/data/public/query/mocks.ts
@@ -33,6 +33,7 @@ import { QueryService, QuerySetup, QueryStart } from '.';
 import { timefilterServiceMock } from './timefilter/timefilter_service.mock';
 import { createFilterManagerMock } from './filter_manager/filter_manager.mock';
 import { queryStringManagerMock } from './query_string/query_string_manager.mock';
+import { dataSetManagerMock } from './dataset_manager/dataset_manager.mock';
 
 type QueryServiceClientContract = PublicMethodsOf<QueryService>;
 
@@ -41,6 +42,7 @@ const createSetupContractMock = () => {
     filterManager: createFilterManagerMock(),
     timefilter: timefilterServiceMock.createSetupContract(),
     queryString: queryStringManagerMock.createSetupContract(),
+    dataSet: dataSetManagerMock.createSetupContract(),
     state$: new Observable(),
   };
 
@@ -55,6 +57,7 @@ const createStartContractMock = () => {
     savedQueries: jest.fn() as any,
     state$: new Observable(),
     timefilter: timefilterServiceMock.createStartContract(),
+    dataSet: dataSetManagerMock.createStartContract(),
     getOpenSearchQuery: jest.fn(),
   };
 

diff --git a/src/plugins/query_enhancements/public/plugin.tsx b/src/plugins/query_enhancements/public/plugin.tsx
@@ -109,11 +109,7 @@ export class QueryEnhancementsPlugin
 
     data.__enhance({
       ui: {
-        queryEditorExtension: createQueryAssistExtension(
-          core.http,
-          this.connectionsService,
-          this.config.queryAssist
-        ),
+        queryEditorExtension: createQueryAssistExtension(core.http, data, this.config.queryAssist),
       },
     });
 

diff --git a/src/plugins/query_enhancements/public/query_assist/components/index_selector.tsx b/src/plugins/query_enhancements/public/query_assist/components/index_selector.tsx
diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx
@@ -5,25 +5,23 @@
 
 import { EuiFlexGroup, EuiFlexItem, EuiForm, EuiFormRow } from '@elastic/eui';
 import React, { SyntheticEvent, useEffect, useMemo, useRef, useState } from 'react';
+import { SimpleDataSet } from '../../../../data/common';
 import {
   IDataPluginServices,
   PersistedLog,
   QueryEditorExtensionDependencies,
 } from '../../../../data/public';
 import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
 import { QueryAssistParameters } from '../../../common/query_assist';
-import { ConnectionsService } from '../../services';
 import { getStorage } from '../../services';
 import { useGenerateQuery } from '../hooks';
 import { getPersistedLog, ProhibitedQueryError } from '../utils';
 import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs';
-import { IndexSelector } from './index_selector';
 import { QueryAssistInput } from './query_assist_input';
 import { QueryAssistSubmitButton } from './submit_button';
 
 interface QueryAssistInputProps {
   dependencies: QueryEditorExtensionDependencies;
-  connectionsService: ConnectionsService;
 }
 
 export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
@@ -37,18 +35,16 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   const { generateQuery, loading } = useGenerateQuery();
   const [callOutType, setCallOutType] = useState<QueryAssistCallOutType>();
   const dismissCallout = () => setCallOutType(undefined);
-  const [selectedIndex, setSelectedIndex] = useState<string>('');
-  const dataSourceIdRef = useRef<string>();
+  const [selectedDataSet, setSelectedDataSet] = useState<SimpleDataSet>();
+  const selectedIndex = selectedDataSet?.title;
   const previousQuestionRef = useRef<string>();
 
   useEffect(() => {
-    const subscription = props.connectionsService
-      .getSelectedConnection$()
-      .subscribe((connection) => {
-        dataSourceIdRef.current = connection?.dataSource?.id;
-      });
+    const subscription = services.data.query.dataSet.getUpdates$().subscribe((dataSet) => {
+      setSelectedDataSet(dataSet);
+    });
     return () => subscription.unsubscribe();
-  }, [props.connectionsService]);
+  }, [services.data.query.dataSet]);
 
   const onSubmit = async (e: SyntheticEvent) => {
     e.preventDefault();
@@ -67,7 +63,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
       question: inputRef.current.value,
       index: selectedIndex,
       language: props.dependencies.language,
-      dataSourceId: dataSourceIdRef.current,
+      dataSourceId: selectedDataSet?.dataSourceRef?.id,
     };
     const { response, error } = await generateQuery(params);
     if (error) {
@@ -90,13 +86,6 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
     <EuiForm component="form" onSubmit={onSubmit}>
       <EuiFormRow fullWidth>
         <EuiFlexGroup gutterSize="s" responsive={false} alignItems="center">
-          <EuiFlexItem grow={false}>
-            <IndexSelector
-              selectedIndex={selectedIndex}
-              setSelectedIndex={setSelectedIndex}
-              dataSourceId={dataSourceIdRef.current}
-            />
-          </EuiFlexItem>
           <EuiFlexItem>
             <QueryAssistInput
               inputRef={inputRef}

diff --git a/src/plugins/query_enhancements/public/query_assist/hooks/use_indices.ts b/src/plugins/query_enhancements/public/query_assist/hooks/use_indices.ts
diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx
@@ -6,11 +6,13 @@
 import { firstValueFrom } from '@osd/std';
 import { act, render, screen } from '@testing-library/react';
 import React from 'react';
+import { of } from 'rxjs';
 import { coreMock } from '../../../../../core/public/mocks';
+import { SimpleDataSet } from '../../../../data/common';
 import { IIndexPattern } from '../../../../data/public';
+import { dataPluginMock } from '../../../../data/public/mocks';
+import { DataSetContract } from '../../../../data/public/query';
 import { ConfigSchema } from '../../../common/config';
-import { ConnectionsService } from '../../data_source_connection';
-import { Connection } from '../../types';
 import { createQueryAssistExtension } from './create_extension';
 
 const coreSetupMock = coreMock.createSetup({
@@ -21,6 +23,18 @@ const coreSetupMock = coreMock.createSetup({
   },
 });
 const httpMock = coreSetupMock.http;
+const dataMock = dataPluginMock.createSetupContract();
+const dataSetMock = dataMock.query.dataSet as jest.Mocked<DataSetContract>;
+
+const mockSimpleDataSet = {
+  id: 'mock-data-set-id',
+  title: 'mock-title',
+  dataSourceRef: {
+    id: 'mock-data-source-id',
+  },
+} as SimpleDataSet;
+
+dataSetMock.getUpdates$.mockReturnValue(of(mockSimpleDataSet));
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
@@ -35,19 +49,10 @@ describe('CreateExtension', () => {
   const config: ConfigSchema['queryAssist'] = {
     supportedLanguages: [{ language: 'PPL', agentConfig: 'os_query_assist_ppl' }],
   };
-  const connectionsService = new ConnectionsService({
-    startServices: coreSetupMock.getStartServices(),
-    http: httpMock,
-  });
-
-  // for these tests we only need id field in the connection
-  connectionsService.setSelectedConnection$({
-    dataSource: { id: 'mock-data-source-id' },
-  } as Connection);
 
   it('should be enabled if at least one language is configured', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' }));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
@@ -57,7 +62,7 @@ describe('CreateExtension', () => {
 
   it('should be disabled for unsupported language', async () => {
     httpMock.get.mockRejectedValueOnce(new Error('network failure'));
-    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' }));
     expect(isEnabled).toBeFalsy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
@@ -67,7 +72,7 @@ describe('CreateExtension', () => {
 
   it('should render the component if language is supported', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const component = extension.getComponent?.({
       language: 'PPL',
       indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],
@@ -84,7 +89,7 @@ describe('CreateExtension', () => {
 
   it('should render the banner if language is not supported', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const banner = extension.getBanner?.({
       language: 'DQL',
       indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],