Merge remote-tracking branch 'upstream/main' into feature/loading-pro…

…gress

public/data_source_connection/utils/create_extension.tsx

@@ -4,6 +4,7 @@
  */
 
 import React from 'react';
+import { of } from 'rxjs';
 import { QueryEditorExtensionConfig } from '../../../../../src/plugins/data/public/ui/query_editor';
 import { ConfigSchema } from '../../../common/config';
 import { ConnectionsBar } from '../components';
@@ -16,8 +17,8 @@ export const createDataSourceConnectionExtension = (
   return {
     id: 'data-source-connection',
     order: 2000,
-    isEnabled: async (dependencies) => {
-      return true;
+    isEnabled$: (dependencies) => {
+      return of(true);
     },
     getComponent: (dependencies) => {
       return <ConnectionsBar dependencies={dependencies} connectionsService={connectionsService} />;

public/plugin.tsx

@@ -8,8 +8,9 @@ import { Trigger } from 'src/plugins/ui_actions/public';
 import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../src/core/public';
 import { IStorageWrapper, Storage } from '../../../src/plugins/opensearch_dashboards_utils/public';
 import { ConfigSchema } from '../common/config';
+import { ConnectionsService, createDataSourceConnectionExtension } from './data_source_connection';
 import { createQueryAssistExtension } from './query_assist';
-import { PPLSearchInterceptor, SQLSearchInterceptor, SQLAsyncSearchInterceptor } from './search';
+import { PPLSearchInterceptor, SQLAsyncSearchInterceptor, SQLSearchInterceptor } from './search';
 import { setData, setStorage } from './services';
 import {
   QueryEnhancementsPluginSetup,
@@ -20,7 +21,6 @@ import {
 import { ASYNC_TRIGGER_ID } from '../common';
 
 export type PublicConfig = Pick<ConfigSchema, 'queryAssist'>;
-import { ConnectionsService, createDataSourceConnectionExtension } from './data_source_connection';
 
 export class QueryEnhancementsPlugin
   implements
@@ -154,7 +154,11 @@ export class QueryEnhancementsPlugin
 
     data.__enhance({
       ui: {
-        queryEditorExtension: createQueryAssistExtension(core.http, this.config.queryAssist),
+        queryEditorExtension: createQueryAssistExtension(
+          core.http,
+          this.connectionsService,
+          this.config.queryAssist
+        ),
       },
     });
 

public/query_assist/components/query_assist_bar.tsx

@@ -9,16 +9,18 @@ import { IDataPluginServices, PersistedLog } from '../../../../../src/plugins/da
 import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension';
 import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
 import { QueryAssistParameters } from '../../../common/query_assist';
+import { ConnectionsService } from '../../data_source_connection';
 import { getStorage } from '../../services';
 import { useGenerateQuery } from '../hooks';
-import { getMdsDataSourceId, getPersistedLog, ProhibitedQueryError } from '../utils';
+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,11 +39,13 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
   const previousQuestionRef = useRef<string>();
 
   useEffect(() => {
-    // TODO need proper way to get dataSourceId when discover index pattern selector is removed
-    getMdsDataSourceId(services.data.indexPatterns, props.dependencies.indexPatterns?.at(0)).then(
-      (id) => (dataSourceIdRef.current = id)
-    );
-  }, [props.dependencies.indexPatterns, services.data.indexPatterns]);
+    const subscription = props.connectionsService
+      .getSelectedConnection$()
+      .subscribe((connection) => {
+        dataSourceIdRef.current = connection?.id;
+      });
+    return () => subscription.unsubscribe();
+  }, [props.connectionsService]);
 
   const onSubmit = async (e: SyntheticEvent) => {
     e.preventDefault();

public/query_assist/utils/create_extension.test.tsx

@@ -3,27 +3,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { firstValueFrom } from '@osd/std';
 import { act, render, screen } from '@testing-library/react';
 import React from 'react';
 import { coreMock } from '../../../../../src/core/public/mocks';
 import { IIndexPattern } from '../../../../../src/plugins/data/public';
 import { ConfigSchema } from '../../../common/config';
+import { ConnectionsService } from '../../data_source_connection';
+import { Connection } from '../../types';
 import { createQueryAssistExtension } from './create_extension';
 
-const coreSetupMock = coreMock.createSetup();
-const httpMock = coreSetupMock.http;
-
-jest.mock('../../services', () => ({
-  getData: jest.fn().mockReturnValue({
-    indexPatterns: {
-      get: jest.fn().mockResolvedValue({ id: 'test-pattern' }),
+const coreSetupMock = coreMock.createSetup({
+  pluginStartDeps: {
+    data: {
+      ui: {},
     },
-  }),
-}));
-
-jest.mock('.', () => ({
-  getMdsDataSourceId: jest.fn().mockResolvedValue('mock-data-source-id'),
-}));
+  },
+});
+const httpMock = coreSetupMock.http;
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
@@ -41,11 +38,18 @@ 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$({ 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, config);
-    const isEnabled = await extension.isEnabled({ language: 'PPL' });
+    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
+    const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' }));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
       query: { dataSourceId: 'mock-data-source-id' },
@@ -54,8 +58,8 @@ describe('CreateExtension', () => {
 
   it('should be disabled for unsupported language', async () => {
     httpMock.get.mockRejectedValueOnce(new Error('network failure'));
-    const extension = createQueryAssistExtension(httpMock, config);
-    const isEnabled = await extension.isEnabled({ language: 'PPL' });
+    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
+    const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' }));
     expect(isEnabled).toBeFalsy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
       query: { dataSourceId: 'mock-data-source-id' },
@@ -64,7 +68,7 @@ describe('CreateExtension', () => {
 
   it('should render the component if language is supported', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(httpMock, config);
+    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
     const component = extension.getComponent?.({
       language: 'PPL',
       indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],
@@ -81,7 +85,7 @@ describe('CreateExtension', () => {
 
   it('should render the banner if language is not supported', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(httpMock, config);
+    const extension = createQueryAssistExtension(httpMock, connectionsService, config);
     const banner = extension.getBanner?.({
       language: 'DQL',
       indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],

public/query_assist/utils/create_extension.tsx

@@ -5,73 +5,72 @@
 
 import { HttpSetup } from 'opensearch-dashboards/public';
 import React, { useEffect, useState } from 'react';
-import { getMdsDataSourceId } from '.';
+import { of } from 'rxjs';
+import { distinctUntilChanged, switchMap, map } from 'rxjs/operators';
 import { QueryEditorExtensionConfig } from '../../../../../src/plugins/data/public/ui/query_editor';
 import { QueryEditorExtensionDependencies } from '../../../../../src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension';
 import { API } from '../../../common';
 import { ConfigSchema } from '../../../common/config';
-import { getData } from '../../services';
+import { ConnectionsService } from '../../data_source_connection';
 import { QueryAssistBar } from '../components';
 import { QueryAssistBanner } from '../components/query_assist_banner';
 
 /**
- * @param dependencies - QueryEditorExtensionDependencies.
- * @param http - HttpSetup.
- * @returns list of query assist agents configured languages in the data source
- * associated with the currently selected index pattern.
+ * @returns observable list of query assist agent configured languages in the
+ * selected data source.
  */
-const getAvailableLanguages = async (
+const getAvailableLanguages$ = (
   availableLanguagesByDataSource: Map<string | undefined, string[]>,
-  dependencies: QueryEditorExtensionDependencies,
+  connectionsService: ConnectionsService,
   http: HttpSetup
-) => {
-  const dataSourceId = await getMdsDataSourceId(
-    getData().indexPatterns,
-    dependencies.indexPatterns?.at(0)
-  );
-  const cached = availableLanguagesByDataSource.get(dataSourceId);
-  if (cached !== undefined) return cached;
-
-  const languages = await http
-    .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
-      query: { dataSourceId },
+) =>
+  connectionsService.getSelectedConnection$().pipe(
+    distinctUntilChanged(),
+    switchMap(async (connection) => {
+      const dataSourceId = connection?.id;
+      const cached = availableLanguagesByDataSource.get(dataSourceId);
+      if (cached !== undefined) return cached;
+      const languages = await http
+        .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
+          query: { dataSourceId },
+        })
+        .then((response) => response.configuredLanguages)
+        .catch(() => []);
+      availableLanguagesByDataSource.set(dataSourceId, languages);
+      return languages;
     })
-    .then((response) => response.configuredLanguages)
-    .catch(() => []);
-  availableLanguagesByDataSource.set(dataSourceId, languages);
-  return languages;
-};
+  );
 
 export const createQueryAssistExtension = (
   http: HttpSetup,
+  connectionsService: ConnectionsService,
   config: ConfigSchema['queryAssist']
 ): QueryEditorExtensionConfig => {
   const availableLanguagesByDataSource: Map<string | undefined, string[]> = new Map();
 
   return {
     id: 'query-assist',
     order: 1000,
-    isEnabled: async (dependencies) => {
+    isEnabled$: (dependencies) => {
       // currently query assist tool relies on opensearch API to get index
       // mappings, non-default data source types are not supported
-      if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default') return false;
+      if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default')
+        return of(false);
 
-      const languages = await getAvailableLanguages(
-        availableLanguagesByDataSource,
-        dependencies,
-        http
+      return getAvailableLanguages$(availableLanguagesByDataSource, connectionsService, http).pipe(
+        map((languages) => languages.length > 0)
       );
-      return languages.length > 0;
     },
     getComponent: (dependencies) => {
       // only show the component if user is on a supported language.
       return (
         <QueryAssistWrapper
           availableLanguagesByDataSource={availableLanguagesByDataSource}
           dependencies={dependencies}
+          connectionsService={connectionsService}
           http={http}
         >
-          <QueryAssistBar dependencies={dependencies} />
+          <QueryAssistBar dependencies={dependencies} connectionsService={connectionsService} />
         </QueryAssistWrapper>
       );
     },
@@ -81,6 +80,7 @@ export const createQueryAssistExtension = (
         <QueryAssistWrapper
           availableLanguagesByDataSource={availableLanguagesByDataSource}
           dependencies={dependencies}
+          connectionsService={connectionsService}
           http={http}
           invert
         >
@@ -94,6 +94,7 @@ export const createQueryAssistExtension = (
 interface QueryAssistWrapperProps {
   availableLanguagesByDataSource: Map<string | undefined, string[]>;
   dependencies: QueryEditorExtensionDependencies;
+  connectionsService: ConnectionsService;
   http: HttpSetup;
   invert?: boolean;
 }
@@ -104,19 +105,18 @@ const QueryAssistWrapper: React.FC<QueryAssistWrapperProps> = (props) => {
   useEffect(() => {
     let mounted = true;
 
-    (async () => {
-      const available = (
-        await getAvailableLanguages(
-          props.availableLanguagesByDataSource,
-          props.dependencies,
-          props.http
-        )
-      ).includes(props.dependencies.language);
+    const subscription = getAvailableLanguages$(
+      props.availableLanguagesByDataSource,
+      props.connectionsService,
+      props.http
+    ).subscribe((languages) => {
+      const available = languages.includes(props.dependencies.language);
       if (mounted) setVisible(props.invert ? !available : available);
-    })();
+    });
 
     return () => {
       mounted = false;
+      subscription.unsubscribe();
     };
   }, [props]);
 

public/query_assist/utils/index.ts

@@ -5,5 +5,4 @@
 
 export * from './create_extension';
 export * from './errors';
-export * from './get_mds_id';
 export * from './get_persisted_log';

server/routes/data_source_connection/routes.ts

@@ -42,7 +42,7 @@ export function registerDataSourceConnectionsRoutes(router: IRouter) {
         'data-source',
         request.params.dataSourceId
       );
-      return response.ok({ body: { dataSource: resp } });
+      return response.ok({ body: resp });
     }
   );
 }