Add data source service

CHANGELOG.md

@@ -13,3 +13,4 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Add incontext insight component ([#53](https://github.com/opensearch-project/dashboards-assistant/pull/53))
 - Fetch root agent id before executing the agent ([#165](https://github.com/opensearch-project/dashboards-assistant/pull/165))
 - Integrate chatbot with sidecar service ([#164](https://github.com/opensearch-project/dashboards-assistant/pull/164))
+- Add data source service ([#191](https://github.com/opensearch-project/dashboards-assistant/pull/191))

opensearch_dashboards.json

@@ -11,7 +11,9 @@
     "opensearchDashboardsUtils"
   ],
   "optionalPlugins": [
-    "securityDashboards"
+    "securityDashboards",
+    "dataSource",
+    "dataSourceManagement"
   ],
   "configPath": [
     "assistant"

public/contexts/__mocks__/core_context.tsx

@@ -24,6 +24,7 @@ export const useCore = jest.fn(() => {
         load: jest.fn(),
       },
       conversationLoad: {},
+      dataSource: {},
     },
   };
   useCoreMock.services.http.delete.mockReturnValue(Promise.resolve());

public/contexts/core_context.tsx

@@ -8,13 +8,14 @@ import {
   useOpenSearchDashboards,
 } from '../../../../src/plugins/opensearch_dashboards_react/public';
 import { AssistantPluginStartDependencies, AssistantPluginSetupDependencies } from '../types';
-import { ConversationLoadService, ConversationsService } from '../services';
+import { ConversationLoadService, ConversationsService, DataSourceService } from '../services';
 
 export interface AssistantServices extends Required<OpenSearchDashboardsServices> {
   setupDeps: AssistantPluginSetupDependencies;
   startDeps: AssistantPluginStartDependencies;
   conversationLoad: ConversationLoadService;
   conversations: ConversationsService;
+  dataSource: DataSourceService;
 }
 
 export const useCore: () => OpenSearchDashboardsReactContextValue<

public/plugin.tsx

@@ -31,6 +31,7 @@ import {
   setIncontextInsightRegistry,
 } from './services';
 import { ConfigSchema } from '../common/types/config';
+import { DataSourceService } from './services/data_source_service';
 
 export const [getCoreStart, setCoreStart] = createGetterSetter<CoreStart>('CoreStart');
 
@@ -58,9 +59,11 @@ export class AssistantPlugin
     > {
   private config: ConfigSchema;
   incontextInsightRegistry: IncontextInsightRegistry | undefined;
+  private dataSourceService: DataSourceService;
 
   constructor(initializerContext: PluginInitializerContext) {
     this.config = initializerContext.config.get<ConfigSchema>();
+    this.dataSourceService = new DataSourceService();
   }
 
   public setup(
@@ -95,6 +98,11 @@ export class AssistantPlugin
     const checkAccess = (account: Awaited<ReturnType<typeof getAccount>>) =>
       account.data.roles.some((role) => ['all_access', 'assistant_user'].includes(role));
 
+    const dataSourceSetupResult = this.dataSourceService.setup({
+      uiSettings: core.uiSettings,
+      dataSourceManagement: setupDeps.dataSourceManagement,
+    });
+
     if (this.config.chat.enabled) {
       const setupChat = async () => {
         const [coreStart, startDeps] = await core.getStartServices();
@@ -105,6 +113,7 @@ export class AssistantPlugin
           startDeps,
           conversationLoad: new ConversationLoadService(coreStart.http),
           conversations: new ConversationsService(coreStart.http),
+          dataSource: this.dataSourceService,
         });
         const account = await getAccount();
         const username = account.data.user_name;
@@ -131,6 +140,7 @@ export class AssistantPlugin
     }
 
     return {
+      dataSource: dataSourceSetupResult,
       registerMessageRenderer: (contentType, render) => {
         if (contentType in messageRenderers)
           console.warn(`Content renderer type ${contentType} is already registered.`);
@@ -160,8 +170,12 @@ export class AssistantPlugin
     setChrome(core.chrome);
     setNotifications(core.notifications);
 
-    return {};
+    return {
+      dataSource: this.dataSourceService.start(),
+    };
   }
 
-  public stop() {}
+  public stop() {
+    this.dataSourceService.stop();
+  }
 }

public/services/__tests__/data_source_service.test.ts

@@ -0,0 +1,212 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BehaviorSubject } from 'rxjs';
+import { first } from 'rxjs/operators';
+
+import { uiSettingsServiceMock } from '../../../../../src/core/public/mocks';
+import { DataSourceOption } from '../../../../../src/plugins/data_source_management/public/components/data_source_menu/types';
+import { DataSourceManagementPluginSetup } from '../../types';
+import { DataSourceService } from '../data_source_service';
+
+const setup = (options?: {
+  dataSourceManagement?: DataSourceManagementPluginSetup;
+  defaultDataSourceId?: string | null;
+  dataSourceSelection?: Map<string, DataSourceOption[]>;
+}) => {
+  const dataSourceSelection$ = new BehaviorSubject<Map<string, DataSourceOption[]>>(
+    options?.dataSourceSelection ?? new Map()
+  );
+  const uiSettings = uiSettingsServiceMock.createSetupContract();
+  const dataSourceManagement: DataSourceManagementPluginSetup = {
+    dataSourceSelection: {
+      getSelection$: () => dataSourceSelection$,
+    },
+  };
+  const dataSource = new DataSourceService();
+  const defaultDataSourceSelection$ = new BehaviorSubject(options?.defaultDataSourceId ?? null);
+  uiSettings.get$.mockReturnValue(defaultDataSourceSelection$);
+  const setupResult = dataSource.setup({
+    uiSettings,
+    dataSourceManagement:
+      options && 'dataSourceManagement' in options
+        ? options.dataSourceManagement
+        : dataSourceManagement,
+  });
+
+  return {
+    dataSource,
+    dataSourceSelection$,
+    defaultDataSourceSelection$,
+    setupResult,
+  };
+};
+
+describe('DataSourceService', () => {
+  describe('getDataSourceId$', () => {
+    it('should return data source selection provided value', async () => {
+      const { dataSource } = setup({
+        defaultDataSourceId: 'foo',
+        dataSourceSelection: new Map([['test', [{ label: 'Bar', id: 'bar' }]]]),
+      });
+
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('bar');
+    });
+    it('should return data source selection provided value even default data source changed', async () => {
+      const { dataSource, defaultDataSourceSelection$ } = setup({
+        defaultDataSourceId: 'foo',
+        dataSourceSelection: new Map([['test', [{ label: 'Bar', id: 'bar' }]]]),
+      });
+
+      defaultDataSourceSelection$.next('baz');
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('bar');
+    });
+    it('should return default data source id if no data source selection', async () => {
+      const { dataSource } = setup({ defaultDataSourceId: 'foo' });
+
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('foo');
+    });
+    it('should return default data source id if data source selection become empty', () => {
+      const { dataSource, dataSourceSelection$ } = setup({
+        defaultDataSourceId: 'foo',
+        dataSourceSelection: new Map([['test', [{ label: 'Bar', id: 'bar' }]]]),
+      });
+      const observerFn = jest.fn();
+      dataSource.getDataSourceId$().subscribe(observerFn);
+      expect(observerFn).toHaveBeenLastCalledWith('bar');
+
+      dataSourceSelection$.next(new Map());
+      expect(observerFn).toHaveBeenLastCalledWith('foo');
+    });
+    it('should return default data source for multi data source selection', async () => {
+      const { dataSource, dataSourceSelection$ } = setup({
+        defaultDataSourceId: 'baz',
+        dataSourceSelection: new Map([
+          [
+            'test',
+            [
+              { label: 'Foo', id: 'foo' },
+              { label: 'Bar', id: 'bar' },
+            ],
+          ],
+        ]),
+      });
+
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('baz');
+
+      dataSourceSelection$.next(
+        new Map([
+          ['component1', [{ label: 'Foo', id: 'foo' }]],
+          ['component2', [{ label: 'Bar', id: 'bar' }]],
+        ])
+      );
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('baz');
+    });
+    it('should return default data source for empty data source selection', async () => {
+      const { dataSource } = setup({
+        defaultDataSourceId: 'foo',
+        dataSourceSelection: new Map(),
+      });
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('foo');
+    });
+  });
+
+  describe('isMDSEnabled', () => {
+    it('should return true if multi data source provided', () => {
+      const { dataSource } = setup();
+      expect(dataSource.isMDSEnabled()).toBe(true);
+    });
+    it('should return false if multi data source not provided', () => {
+      const { dataSource } = setup({ dataSourceManagement: undefined });
+      expect(dataSource.isMDSEnabled()).toBe(false);
+    });
+  });
+  describe('getDataSourceQuery', () => {
+    it('should return empty object if MDS not enabled', async () => {
+      const { dataSource } = setup({ dataSourceManagement: undefined });
+      expect(await dataSource.getDataSourceQuery()).toEqual({});
+    });
+    it('should return empty object if data source id is empty', async () => {
+      const { dataSource } = setup({
+        dataSourceSelection: new Map([['test', [{ label: '', id: '' }]]]),
+      });
+      expect(await dataSource.getDataSourceQuery()).toEqual({});
+    });
+    it('should return query object with provided data source id', async () => {
+      const { dataSource } = setup({ defaultDataSourceId: 'foo' });
+      expect(await dataSource.getDataSourceQuery()).toEqual({ dataSourceId: 'foo' });
+    });
+    it('should throw error if data source id not exists', async () => {
+      const { dataSource } = setup();
+      let error;
+      try {
+        await dataSource.getDataSourceQuery();
+      } catch (e) {
+        error = e;
+      }
+      expect(error).toBeTruthy();
+    });
+  });
+  describe('stop', () => {
+    it('should not emit after data source selection unsubscribe', async () => {
+      const { dataSource, dataSourceSelection$ } = setup();
+      const observerFn = jest.fn();
+      dataSource.getDataSourceId$().subscribe(observerFn);
+      expect(observerFn).toHaveBeenCalledTimes(1);
+      dataSource.stop();
+      dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+      expect(observerFn).toHaveBeenCalledTimes(1);
+    });
+    it('should not emit after data source id subject complete', () => {
+      const { dataSource } = setup();
+      const observerFn = jest.fn();
+      dataSource.getDataSourceId$().subscribe(observerFn);
+      expect(observerFn).toHaveBeenCalledTimes(1);
+      dataSource.stop();
+      dataSource.setDataSourceId('foo');
+      expect(observerFn).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('setup', () => {
+    it('should able to change data source id from setup result', async () => {
+      const { dataSource, setupResult } = setup();
+      setupResult.setDataSourceId('foo');
+      expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('foo');
+    });
+
+    it('should update data source id after data source selection changed', () => {
+      const { dataSource, dataSourceSelection$ } = setup();
+      const observerFn = jest.fn();
+      dataSource.getDataSourceId$().subscribe(observerFn);
+
+      dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+      expect(observerFn).toHaveBeenLastCalledWith('foo');
+
+      dataSourceSelection$.next(new Map([['test', [{ label: 'Bar', id: 'bar' }]]]));
+      expect(observerFn).toHaveBeenLastCalledWith('bar');
+    });
+  });
+
+  it('should able to change data source id from start result', async () => {
+    const { dataSource } = setup();
+    dataSource.start().setDataSourceId('bar');
+    expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('bar');
+  });
+
+  it('should not fire change when call setDataSourceId with same data source id', async () => {
+    const { dataSource } = setup();
+    const observerFn = jest.fn();
+    dataSource.getDataSourceId$().subscribe(observerFn);
+    dataSource.setDataSourceId('foo');
+    expect(observerFn).toHaveBeenCalledTimes(2);
+
+    dataSource.setDataSourceId('foo');
+    expect(observerFn).toHaveBeenCalledTimes(2);
+
+    dataSource.setDataSourceId('bar');
+    expect(observerFn).toHaveBeenCalledTimes(3);
+  });
+});