From 8ee14306651d77278090036c72a1fbcddd50d460 Mon Sep 17 00:00:00 2001
From: tygao <tygao@amazon.com>
Date: Thu, 17 Oct 2024 16:55:53 +0800
Subject: [PATCH 1/2] add DQS type support for association modal

Signed-off-by: tygao <tygao@amazon.com>
---
 .../data_source_association.test.tsx          | 37 +++++++++++++++++++
 .../data_source_association.tsx               | 13 +++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

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 38391b194af6..40a5ee10300c 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 19a14946d1b7..0c821923a0dd 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 '../workspace_detail/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 {

From 4b21c33fdc202858f2c1692d5acabd83e8d796c0 Mon Sep 17 00:00:00 2001
From: "opensearch-changeset-bot[bot]"
 <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
Date: Thu, 17 Oct 2024 09:02:37 +0000
Subject: [PATCH 2/2] Changeset file for PR #8620 created/updated

---
 changelogs/fragments/8620.yml | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 changelogs/fragments/8620.yml

diff --git a/changelogs/fragments/8620.yml b/changelogs/fragments/8620.yml
new file mode 100644
index 000000000000..0bddae141f95
--- /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