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 {
+      ...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 {
+      ...this.dataSourceService.start(),
+    };
   }
 
-  public stop() {}
+  public stop() {
+    this.dataSourceService.stop();
+  }
 }

public/services/__tests__/data_source_service.test.ts

@@ -0,0 +1,213 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BehaviorSubject, Subject } from 'rxjs';
+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 }) => {
+  const dataSourceSelection$ = new BehaviorSubject<Map<string, DataSourceOption[]>>(new Map());
+  const uiSettings = uiSettingsServiceMock.createSetupContract();
+  const dataSourceManagement: DataSourceManagementPluginSetup = {
+    dataSourceSelection: {
+      getSelection$: () => dataSourceSelection$,
+    },
+  };
+  const dataSource = new DataSourceService();
+  const defaultDataSourceSelection$ = new Subject();
+  uiSettings.get$.mockReturnValueOnce(defaultDataSourceSelection$);
+  const setupResult = dataSource.setup({
+    uiSettings,
+    dataSourceManagement:
+      options && 'dataSourceManagement' in options
+        ? options.dataSourceManagement
+        : dataSourceManagement,
+  });
+
+  return {
+    dataSource,
+    uiSettings,
+    dataSourceSelection$,
+    defaultDataSourceSelection$,
+    setupResult,
+  };
+};
+
+describe('DataSourceService', () => {
+  it('should return data source selection provided data source id', () => {
+    const { dataSource, dataSourceSelection$ } = setup();
+
+    expect(dataSource.getDataSourceId()).toBe(null);
+
+    dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+
+    expect(dataSource.getDataSourceId()).toBe('foo');
+  });
+  describe('initDefaultDataSourceIdIfNeed', () => {
+    it('should return ui settings provided data source id', () => {
+      const { dataSource, uiSettings } = setup();
+
+      uiSettings.get.mockReturnValueOnce('foo');
+
+      expect(dataSource.getDataSourceId()).toBe(null);
+
+      dataSource.initDefaultDataSourceIdIfNeed();
+
+      expect(dataSource.getDataSourceId()).toBe('foo');
+    });
+    it('should return data source selection provided data source', () => {
+      const { dataSource, dataSourceSelection$, uiSettings } = setup();
+
+      uiSettings.get.mockReturnValueOnce('bar');
+
+      expect(dataSource.getDataSourceId()).toBe(null);
+
+      dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+
+      expect(dataSource.getDataSourceId()).toBe('foo');
+
+      dataSource.initDefaultDataSourceIdIfNeed();
+      expect(dataSource.getDataSourceId()).toBe('foo');
+    });
+  });
+  it('should update data source id after default data source id changed', () => {
+    const { dataSource, defaultDataSourceSelection$, uiSettings } = setup();
+
+    uiSettings.get.mockReturnValueOnce('foo');
+    dataSource.initDefaultDataSourceIdIfNeed();
+    expect(dataSource.getDataSourceId()).toBe('foo');
+    defaultDataSourceSelection$.next('bar');
+    expect(dataSource.getDataSourceId()).toBe('bar');
+  });
+  it('should not update data source id when data source id not from ui settings', () => {
+    const { dataSource, dataSourceSelection$, defaultDataSourceSelection$ } = setup();
+
+    expect(dataSource.getDataSourceId()).toBe(null);
+
+    dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+    defaultDataSourceSelection$.next('bar');
+    expect(dataSource.getDataSourceId()).toBe('foo');
+  });
+  it('should return null for multi data source selection', () => {
+    const { dataSource, dataSourceSelection$ } = setup();
+
+    expect(dataSource.getDataSourceId()).toBe(null);
+
+    dataSourceSelection$.next(
+      new Map([
+        [
+          'test',
+          [
+            { label: 'Foo', id: 'foo' },
+            { label: 'Bar', id: 'bar' },
+          ],
+        ],
+      ])
+    );
+    expect(dataSource.getDataSourceId()).toBe(null);
+
+    dataSourceSelection$.next(
+      new Map([
+        ['component1', [{ label: 'Foo', id: 'foo' }]],
+        ['component2', [{ label: 'Bar', id: 'bar' }]],
+      ])
+    );
+    expect(dataSource.getDataSourceId()).toBe(null);
+  });
+  it('should return null for empty data source selection', () => {
+    const { dataSource, dataSourceSelection$ } = setup();
+
+    expect(dataSource.getDataSourceId()).toBe(null);
+
+    dataSourceSelection$.next(new Map());
+    expect(dataSource.getDataSourceId()).toBe(null);
+  });
+  it('should able to subscribe data source id changes', () => {
+    const { dataSource } = setup();
+    const mockFn = jest.fn();
+    dataSource.subscribeDataSourceId({ next: mockFn });
+
+    dataSource.setDataSourceId('foo', undefined);
+    expect(mockFn).toHaveBeenCalledWith('foo');
+    expect(mockFn).toHaveBeenCalledTimes(2);
+
+    dataSource.setDataSourceId('foo', undefined);
+    expect(mockFn).toHaveBeenCalledTimes(2);
+
+    dataSource.setDataSourceId('bar', undefined);
+    expect(mockFn).toHaveBeenCalledWith('bar');
+    expect(mockFn).toHaveBeenCalledTimes(3);
+  });
+  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', () => {
+      const { dataSource } = setup({ dataSourceManagement: undefined });
+      expect(dataSource.getDataSourceQuery()).toEqual({});
+    });
+    it('should return empty object if data source id is empty', () => {
+      const { dataSource, dataSourceSelection$ } = setup();
+      dataSourceSelection$.next(new Map([['test', [{ label: '', id: '' }]]]));
+      expect(dataSource.getDataSourceQuery()).toEqual({});
+    });
+    it('should return query object with provided data source id', () => {
+      const { dataSource, dataSourceSelection$ } = setup();
+      dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+      expect(dataSource.getDataSourceQuery()).toEqual({ dataSourceId: 'foo' });
+    });
+    it('should throw error if data source id not exists', () => {
+      const { dataSource } = setup();
+      let error;
+      try {
+        dataSource.getDataSourceQuery();
+      } catch (e) {
+        error = e;
+      }
+      expect(error).toBeTruthy();
+    });
+  });
+  it('should clear data source id', () => {
+    const { dataSource, dataSourceSelection$ } = setup();
+    dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+    expect(dataSource.getDataSourceId()).toEqual('foo');
+    dataSource.clearDataSourceId();
+    expect(dataSource.getDataSourceId()).toEqual(null);
+  });
+  it('should able to change data source id from outside', () => {
+    const { dataSource, setupResult } = setup();
+    setupResult.setDataSourceId('foo');
+    expect(dataSource.getDataSourceId()).toBe('foo');
+    dataSource.start().setDataSourceId('bar');
+    expect(dataSource.getDataSourceId()).toBe('bar');
+  });
+  describe('stop', () => {
+    it('should unsubscribe data source selection', () => {
+      const { dataSource, dataSourceSelection$ } = setup();
+      dataSource.stop();
+      dataSourceSelection$.next(new Map([['test', [{ label: 'Foo', id: 'foo' }]]]));
+      expect(dataSource.getDataSourceId()).toBe(null);
+    });
+    it('should complete data source id', () => {
+      const { dataSource } = setup();
+      const mockFn = jest.fn();
+      dataSource.subscribeDataSourceId({
+        complete: mockFn,
+      });
+      dataSource.stop();
+      expect(mockFn).toHaveBeenCalled();
+    });
+  });
+});