[Workspace] add data-connection type support for association modal (o…

…pensearch-project#8620)

* add DQS type support for association modal

Signed-off-by: tygao <[email protected]>

* Changeset file for PR opensearch-project#8620 created/updated

---------

Signed-off-by: tygao <[email protected]>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>

changelogs/fragments/8620.yml

@@ -0,0 +1,2 @@
+feat:
+- Add data-connection type support for association modal ([#8620](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8620))

src/plugins/workspace/public/components/data_source_association/data_source_association.test.tsx

@@ -103,6 +103,43 @@ describe('<DataSourceAssociation />', () => {
     });
   });
 
+  it('should associate data connections successfully', async () => {
+    const associateMock = jest
+      .fn()
+      .mockResolvedValue({ success: true, result: [{ id: 'id1' }, { id: 'id2' }] });
+    servicesMock.workspaces.client$ = new BehaviorSubject<IWorkspaceClient | null>({
+      associate: associateMock,
+      copy: jest.fn(),
+      dissociate: jest.fn(),
+      ui: jest.fn(),
+    });
+    servicesMock.workspaces.currentWorkspaceId$ = new BehaviorSubject<string>('workspace_test');
+
+    (AssociationDataSourceModalContent as jest.Mock).mockImplementation((props: any) => (
+      <button
+        onClick={() =>
+          props.handleAssignDataSourceConnections([
+            { id: 'id1', connectionType: DataSourceConnectionType.DataConnection },
+            { id: 'id2', connectionType: DataSourceConnectionType.DataConnection },
+          ])
+        }
+      >
+        Mocked association button
+      </button>
+    ));
+
+    render(<DataSourceAssociation excludedDataSourceIds={[]} />);
+    fireEvent.click(screen.getByTestId('workspaceAssociateDataSourceButton'));
+    fireEvent.click(screen.getByText('Direct query data sources'));
+    fireEvent.click(screen.getByText('Mocked association button'));
+    await waitFor(() => {
+      expect(associateMock).toHaveBeenCalled();
+      expect(servicesMock.notifications.toasts.addSuccess).toHaveBeenCalledWith(
+        expect.objectContaining({ title: '2 data sources been associated to the workspace' })
+      );
+    });
+  });
+
   it('should display error toast when associate data source failed', async () => {
     const associateMock = jest.fn().mockRejectedValue(new Error());
     servicesMock.workspaces.client$ = new BehaviorSubject<IWorkspaceClient | null>({

src/plugins/workspace/public/components/data_source_association/data_source_association.tsx

@@ -22,6 +22,10 @@ import {
 import { AssociationDataSourceModalContent } from './association_data_source_modal';
 import { AssociationDataSourceModalMode } from '../../../common/constants';
 import { DataSourceConnection, DataSourceConnectionType } from '../../../common/types';
+import {
+  DATA_CONNECTION_SAVED_OBJECT_TYPE,
+  DATA_SOURCE_SAVED_OBJECT_TYPE,
+} from '../../../../data_source/common';
 
 interface Props {
   excludedDataSourceIds: string[];
@@ -54,10 +58,13 @@ export const DataSourceAssociation = ({ excludedDataSourceIds, onComplete, onErr
 
   const onAssociateDataSource = useCallback(
     async (ds: DataSourceConnection[]) => {
-      const objects = ds
+      const dataSourceObjects = ds
         .filter((d) => d.connectionType === DataSourceConnectionType.OpenSearchConnection)
-        .map((d) => ({ id: d.id, type: 'data-source' }));
-
+        .map((d) => ({ id: d.id, type: DATA_SOURCE_SAVED_OBJECT_TYPE }));
+      const dataConnectionObjects = ds
+        .filter((d) => d.connectionType === DataSourceConnectionType.DataConnection)
+        .map((d) => ({ id: d.id, type: DATA_CONNECTION_SAVED_OBJECT_TYPE }));
+      const objects = [...dataSourceObjects, ...dataConnectionObjects];
       if (workspaceClient && currentWorkspaceId) {
         let failedCount = 0;
         try {