From 896ee963f7000445818192d4f8466e4c1eb3ffa4 Mon Sep 17 00:00:00 2001
From: Tianyu Gao <tygao@amazon.com>
Date: Fri, 18 Oct 2024 15:02:27 +0800
Subject: [PATCH] [Workspace] add data-connection type support for association
 modal (#8620)

* add DQS type support for association modal

Signed-off-by: tygao <tygao@amazon.com>

* Changeset file for PR #8620 created/updated

---------

Signed-off-by: tygao <tygao@amazon.com>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
---
 changelogs/fragments/8620.yml                 |  2 +
 .../data_source_association.test.tsx          | 37 +++++++++++++++++++
 .../data_source_association.tsx               | 13 +++++--
 3 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 changelogs/fragments/8620.yml

diff --git a/changelogs/fragments/8620.yml b/changelogs/fragments/8620.yml
new file mode 100644
index 00000000000..0bddae141f9
--- /dev/null
+++ b/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))
\ No newline at end of file
diff --git a/src/plugins/workspace/public/components/data_source_association/data_source_association.test.tsx b/src/plugins/workspace/public/components/data_source_association/data_source_association.test.tsx
index 781ab0cea88..7d8cf2cec08 100644
--- a/src/plugins/workspace/public/components/data_source_association/data_source_association.test.tsx
+++ b/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>({
diff --git a/src/plugins/workspace/public/components/data_source_association/data_source_association.tsx b/src/plugins/workspace/public/components/data_source_association/data_source_association.tsx
index 9aa54ff2e4c..d5546173357 100644
--- a/src/plugins/workspace/public/components/data_source_association/data_source_association.tsx
+++ b/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 {