From 2999fdf85b34fe57386095cb99630ba967abc5fc Mon Sep 17 00:00:00 2001
From: Jialiang Liang <jiallian@amazon.com>
Date: Tue, 10 Sep 2024 12:02:40 -0700
Subject: [PATCH] [MDS]Add MDS support for Integration (#8008)

* Add mds support for Integration

Signed-off-by: Ryan Liang <jiallian@amazon.com>

* Add mds support for Integration flyout only for creation

Signed-off-by: Ryan Liang <jiallian@amazon.com>

* Fix the installed integration table content

Signed-off-by: Ryan Liang <jiallian@amazon.com>

* Update snapshot

Signed-off-by: Ryan Liang <jiallian@amazon.com>

* Update to fix the local cluster integration fetch and add more test case

Signed-off-by: Ryan Liang <jiallian@amazon.com>

* Update release notes

Signed-off-by: Ryan Liang <jiallian@amazon.com>

* Changeset file for PR #8008 created/updated

---------

Signed-off-by: Ryan Liang <jiallian@amazon.com>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
---
 changelogs/fragments/8008.yml                 |   2 +
 ...nsearch-dashboards.release-notes-2.17.0.md |   1 +
 .../direct_query_connection_detail.test.tsx   | 111 +++++++-
 .../direct_query_connection_detail.tsx        |  88 ++++--
 .../setup_integration.test.tsx.snap           | 262 +++++++++++++++++-
 .../create_integration_helpers.ts             |  45 ++-
 .../installed_integrations_table.tsx          |  12 +
 .../integrations/setup_integration.tsx        |  70 +++--
 .../mount_management_section.tsx              |   1 +
 9 files changed, 533 insertions(+), 59 deletions(-)
 create mode 100644 changelogs/fragments/8008.yml

diff --git a/changelogs/fragments/8008.yml b/changelogs/fragments/8008.yml
new file mode 100644
index 000000000000..a6360be04d31
--- /dev/null
+++ b/changelogs/fragments/8008.yml
@@ -0,0 +1,2 @@
+feat:
+- [MDS]Add MDS support for Integration ([#8008](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8008))
\ No newline at end of file
diff --git a/release-notes/opensearch-dashboards.release-notes-2.17.0.md b/release-notes/opensearch-dashboards.release-notes-2.17.0.md
index e671af3d2777..7be47ec89aca 100644
--- a/release-notes/opensearch-dashboards.release-notes-2.17.0.md
+++ b/release-notes/opensearch-dashboards.release-notes-2.17.0.md
@@ -70,6 +70,7 @@
  - Support DQCs in create page ([#7961](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7961))
  - [Workspace] Hide home breadcrumbs when in a workspace ([#7992](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7992))
  - [Workspace]Deny get or bulkGet for global data source ([#8043](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8043))
+ - [MDS]Add MDS support for Integration #8008 ([#8008](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8008))
 
 ### 🐛 Bug Fixes
 
diff --git a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.test.tsx b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.test.tsx
index ea13ffd2568f..334fad5f810f 100644
--- a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.test.tsx
+++ b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.test.tsx
@@ -8,8 +8,13 @@ import { render, screen, waitFor } from '@testing-library/react';
 import '@testing-library/jest-dom';
 import { MemoryRouter, Route } from 'react-router-dom';
 import { DirectQueryDataConnectionDetail } from './direct_query_connection_detail';
-import { ApplicationStart, HttpStart, NotificationsStart } from 'opensearch-dashboards/public';
-import { isPluginInstalled } from '../../utils';
+import {
+  ApplicationStart,
+  HttpStart,
+  NotificationsStart,
+  SavedObjectsStart,
+} from 'opensearch-dashboards/public';
+import * as utils from '../../utils';
 
 jest.mock('../../../constants', () => ({
   DATACONNECTIONS_BASE: '/api/dataconnections',
@@ -64,17 +69,27 @@ jest.mock('../associated_object_management/utils/associated_objects_tab_utils',
 
 jest.mock('../../utils', () => ({
   isPluginInstalled: jest.fn(),
+  getDataSourcesWithFields: jest.fn(),
 }));
 
 const renderComponent = ({
   featureFlagStatus = false,
   http = {},
-  notifications = {},
+  notifications = {
+    toasts: {
+      addDanger: jest.fn(),
+    },
+  },
   application = {},
   setBreadcrumbs = jest.fn(),
+  savedObjects = {
+    client: {
+      find: jest.fn().mockResolvedValue({ saved_objects: [] }),
+    },
+  },
 }) => {
   return render(
-    <MemoryRouter initialEntries={['/dataconnections/test']}>
+    <MemoryRouter initialEntries={['/dataconnections/test?dataSourceMDSId=test-mdsid']}>
       <Route path="/dataconnections/:dataSourceName">
         <DirectQueryDataConnectionDetail
           featureFlagStatus={featureFlagStatus}
@@ -82,6 +97,7 @@ const renderComponent = ({
           notifications={notifications as NotificationsStart}
           application={application as ApplicationStart}
           setBreadcrumbs={setBreadcrumbs}
+          savedObjects={savedObjects as SavedObjectsStart}
         />
       </Route>
     </MemoryRouter>
@@ -91,7 +107,7 @@ const renderComponent = ({
 describe('DirectQueryDataConnectionDetail', () => {
   beforeEach(() => {
     jest.clearAllMocks();
-    (isPluginInstalled as jest.Mock).mockResolvedValue(true);
+    (utils.isPluginInstalled as jest.Mock).mockResolvedValue(true);
   });
 
   test('renders without crashing', async () => {
@@ -288,4 +304,89 @@ describe('DirectQueryDataConnectionDetail', () => {
     expect(screen.getByText('Configure Integrations')).toBeInTheDocument();
     expect(screen.getByText('Installed Integrations')).toBeInTheDocument();
   });
+
+  test('filters integrations by references when featureFlagStatus is true and dataSourceMDSId exists', async () => {
+    const mockHttp = {
+      get: jest.fn().mockImplementation((url) => {
+        if (url === '/api/integrations/store') {
+          return Promise.resolve({
+            data: {
+              hits: [
+                {
+                  dataSource: 'flint_test_default',
+                  references: [{ id: 'test-mdsid', name: 'Test Integration', type: 'data-source' }],
+                },
+                {
+                  dataSource: 'flint_test_default',
+                  references: [
+                    { id: 'other-mdsid', name: 'Other Integration', type: 'data-source' },
+                  ],
+                },
+              ],
+            },
+          });
+        } else {
+          return Promise.resolve({
+            allowedRoles: ['role1'],
+            description: 'Test description',
+            name: 'Test datasource',
+            connector: 'S3GLUE',
+            properties: {},
+            status: 'ACTIVE',
+          });
+        }
+      }),
+    };
+
+    const mockNotifications = {
+      toasts: {
+        addDanger: jest.fn(),
+      },
+    };
+
+    const mockSavedObjects = {
+      client: {
+        find: jest.fn().mockResolvedValue({
+          saved_objects: [
+            {
+              id: 'test-mdsid',
+              attributes: {
+                title: 'Test Data Source',
+              },
+            },
+          ],
+        }),
+      },
+    };
+
+    (utils.getDataSourcesWithFields as jest.Mock).mockResolvedValue([
+      {
+        id: 'test-mdsid',
+        attributes: {
+          title: 'Test Data Source',
+        },
+      },
+    ]);
+
+    renderComponent({
+      featureFlagStatus: true,
+      http: mockHttp,
+      notifications: mockNotifications,
+      savedObjects: mockSavedObjects,
+    });
+
+    await waitFor(() => expect(mockHttp.get).toHaveBeenCalledWith('/api/integrations/store'), {
+      timeout: 1000,
+    });
+
+    await waitFor(
+      () => {
+        const filteredIntegration = screen.queryByText('Configure Integrations');
+        expect(filteredIntegration).toBeInTheDocument();
+      },
+      { timeout: 1000 }
+    );
+
+    expect(mockNotifications.toasts.addDanger).not.toHaveBeenCalled();
+  });
 });
diff --git a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.tsx b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.tsx
index aad3db2260c9..9593f683ff92 100644
--- a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.tsx
+++ b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/connection_detail/direct_query_connection_detail.tsx
@@ -19,7 +19,12 @@ import {
   EuiCard,
   EuiAccordion,
 } from '@elastic/eui';
-import { ApplicationStart, HttpStart, NotificationsStart } from 'opensearch-dashboards/public';
+import {
+  ApplicationStart,
+  HttpStart,
+  NotificationsStart,
+  SavedObjectsStart,
+} from 'opensearch-dashboards/public';
 import { useLocation, useParams } from 'react-router-dom';
 import { escapeRegExp } from 'lodash';
 import { DATACONNECTIONS_BASE } from '../../../constants';
@@ -46,7 +51,7 @@ import {
   IntegrationInstancesSearchResult,
 } from '../../../../framework/types';
 import { INTEGRATIONS_BASE } from '../../../../framework/utils/shared';
-import { isPluginInstalled } from '../../utils';
+import { isPluginInstalled, getDataSourcesWithFields } from '../../utils';
 
 interface DirectQueryDataConnectionDetailProps {
   featureFlagStatus: boolean;
@@ -55,6 +60,7 @@ interface DirectQueryDataConnectionDetailProps {
   application: ApplicationStart;
   setBreadcrumbs: (breadcrumbs: any) => void;
   useNewUX: boolean;
+  savedObjects: SavedObjectsStart;
 }
 
 export const DirectQueryDataConnectionDetail: React.FC<DirectQueryDataConnectionDetailProps> = ({
@@ -64,6 +70,7 @@ export const DirectQueryDataConnectionDetail: React.FC<DirectQueryDataConnection
   application,
   setBreadcrumbs,
   useNewUX,
+  savedObjects,
 }) => {
   const [observabilityDashboardsExists, setObservabilityDashboardsExists] = useState(false);
   const { dataSourceName } = useParams<{ dataSourceName: string }>();
@@ -117,24 +124,60 @@ export const DirectQueryDataConnectionDetail: React.FC<DirectQueryDataConnection
   const [refreshIntegrationsFlag, setRefreshIntegrationsFlag] = useState(false);
   const refreshInstances = () => setRefreshIntegrationsFlag((prev) => !prev);
 
+  const [clusterTitle, setDataSourceTitle] = useState<string | undefined>();
+  const fetchDataSources = async () => {
+    try {
+      const dataSources = await getDataSourcesWithFields(savedObjects.client, ['id', 'title']);
+
+      // Find the data source title based on the dataSourceMDSId
+      const foundDataSource = dataSources.find((ds: any) => ds.id === dataSourceMDSId);
+      if (foundDataSource) {
+        setDataSourceTitle(foundDataSource.attributes.title);
+      }
+    } catch (error) {
+      notifications.toasts.addDanger({
+        title: 'Failed to fetch data sources',
+        text: error.message,
+      });
+    }
+  };
+
+  useEffect(() => {
+    if (featureFlagStatus) {
+      fetchDataSources();
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [featureFlagStatus, savedObjects, notifications, dataSourceMDSId]);
+
   useEffect(() => {
     const searchDataSourcePattern = new RegExp(
       `flint_${escapeRegExp(datasourceDetails.name)}_default_.*`
     );
+
     const findIntegrations = async () => {
       const result: { data: IntegrationInstancesSearchResult } = await http!.get(
         INTEGRATIONS_BASE + `/store`
       );
+
       if (result.data?.hits) {
-        setDataSourceIntegrations(
-          result.data.hits.filter((res) => res.dataSource.match(searchDataSourcePattern))
+        let filteredIntegrations = result.data.hits.filter((res) =>
+          res.dataSource.match(searchDataSourcePattern)
         );
+
+        if (featureFlagStatus && dataSourceMDSId !== null) {
+          filteredIntegrations = filteredIntegrations.filter((res) => {
+            return res.references && res.references.some((ref) => ref.id === dataSourceMDSId);
+          });
+        }
+
+        setDataSourceIntegrations(filteredIntegrations);
       } else {
         setDataSourceIntegrations([]);
       }
     };
+
     findIntegrations();
-  }, [http, datasourceDetails.name, refreshIntegrationsFlag]);
+  }, [http, datasourceDetails.name, refreshIntegrationsFlag, featureFlagStatus, dataSourceMDSId]);
 
   const [showIntegrationsFlyout, setShowIntegrationsFlyout] = useState(false);
   const onclickIntegrationsCard = () => {
@@ -146,6 +189,8 @@ export const DirectQueryDataConnectionDetail: React.FC<DirectQueryDataConnection
       datasourceType={datasourceDetails.connector}
       datasourceName={datasourceDetails.name}
       http={http}
+      selectedDataSourceId={dataSourceMDSId || ''}
+      selectedClusterName={clusterTitle}
     />
   ) : null;
 
@@ -189,7 +234,7 @@ export const DirectQueryDataConnectionDetail: React.FC<DirectQueryDataConnection
   const DefaultDatasourceCards = () => {
     return (
       <EuiFlexGroup>
-        {!featureFlagStatus && observabilityDashboardsExists && (
+        {observabilityDashboardsExists && (
           <EuiFlexItem>
             <EuiCard
               icon={<EuiIcon size="xxl" type="integrationGeneral" />}
@@ -402,21 +447,22 @@ export const DirectQueryDataConnectionDetail: React.FC<DirectQueryDataConnection
               />
             ),
           },
-          !featureFlagStatus &&
-            observabilityDashboardsExists && {
-              id: 'installed_integrations',
-              name: 'Installed Integrations',
-              disabled: false,
-              content: (
-                <InstalledIntegrationsTable
-                  integrations={dataSourceIntegrations}
-                  datasourceType={datasourceDetails.connector}
-                  datasourceName={datasourceDetails.name}
-                  refreshInstances={refreshInstances}
-                  http={http}
-                />
-              ),
-            },
+          observabilityDashboardsExists && {
+            id: 'installed_integrations',
+            name: 'Installed Integrations',
+            disabled: false,
+            content: (
+              <InstalledIntegrationsTable
+                integrations={dataSourceIntegrations}
+                datasourceType={datasourceDetails.connector}
+                datasourceName={datasourceDetails.name}
+                refreshInstances={refreshInstances}
+                http={http}
+                selectedDataSourceId={featureFlagStatus ? dataSourceMDSId ?? undefined : undefined}
+                selectedClusterName={clusterTitle}
+              />
+            ),
+          },
         ].filter(Boolean)
       : [];
 
diff --git a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/__snapshots__/setup_integration.test.tsx.snap b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/__snapshots__/setup_integration.test.tsx.snap
index eaefe0d108f1..a0df1a8465d1 100644
--- a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/__snapshots__/setup_integration.test.tsx.snap
+++ b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/__snapshots__/setup_integration.test.tsx.snap
@@ -16,6 +16,12 @@ exports[`SetupIntegrationForm tests renders SetupIntegrationForm 1`] = `
             "enabledWorkflows": Array [],
           }
         }
+        http={
+          Object {
+            "get": [MockFunction],
+            "post": [MockFunction],
+          }
+        }
         integration={
           Object {
             "assets": Array [],
@@ -104,7 +110,54 @@ exports[`SetupIntegrationForm tests renders SetupIntegrationForm and matches sna
           },
         ],
       },
-      "post": [MockFunction],
+      "post": [MockFunction] {
+        "calls": Array [
+          Array [
+            "/api/console/proxy",
+            Object {
+              "body": "{}",
+              "query": Object {
+                "method": "GET",
+                "path": "_data_stream/ss4o_*",
+              },
+            },
+          ],
+          Array [
+            "/api/console/proxy",
+            Object {
+              "body": "{}",
+              "query": Object {
+                "method": "GET",
+                "path": "_data_stream/ss4o_*",
+              },
+            },
+          ],
+          Array [
+            "/api/console/proxy",
+            Object {
+              "body": "{}",
+              "query": Object {
+                "method": "GET",
+                "path": "_data_stream/ss4o_*",
+              },
+            },
+          ],
+        ],
+        "results": Array [
+          Object {
+            "type": "return",
+            "value": undefined,
+          },
+          Object {
+            "type": "return",
+            "value": undefined,
+          },
+          Object {
+            "type": "return",
+            "value": undefined,
+          },
+        ],
+      },
     }
   }
   integration="test_integration"
@@ -136,6 +189,85 @@ exports[`SetupIntegrationForm tests renders SetupIntegrationForm and matches sna
                   "enabledWorkflows": Array [],
                 }
               }
+              http={
+                Object {
+                  "get": [MockFunction] {
+                    "calls": Array [
+                      Array [
+                        "/api/integrations/repository/test_integration",
+                      ],
+                      Array [
+                        "/api/integrations/repository/test_integration",
+                      ],
+                      Array [
+                        "/api/integrations/repository/test_integration",
+                      ],
+                    ],
+                    "results": Array [
+                      Object {
+                        "type": "return",
+                        "value": Promise {},
+                      },
+                      Object {
+                        "type": "return",
+                        "value": Promise {},
+                      },
+                      Object {
+                        "type": "return",
+                        "value": Promise {},
+                      },
+                    ],
+                  },
+                  "post": [MockFunction] {
+                    "calls": Array [
+                      Array [
+                        "/api/console/proxy",
+                        Object {
+                          "body": "{}",
+                          "query": Object {
+                            "method": "GET",
+                            "path": "_data_stream/ss4o_*",
+                          },
+                        },
+                      ],
+                      Array [
+                        "/api/console/proxy",
+                        Object {
+                          "body": "{}",
+                          "query": Object {
+                            "method": "GET",
+                            "path": "_data_stream/ss4o_*",
+                          },
+                        },
+                      ],
+                      Array [
+                        "/api/console/proxy",
+                        Object {
+                          "body": "{}",
+                          "query": Object {
+                            "method": "GET",
+                            "path": "_data_stream/ss4o_*",
+                          },
+                        },
+                      ],
+                    ],
+                    "results": Array [
+                      Object {
+                        "type": "return",
+                        "value": undefined,
+                      },
+                      Object {
+                        "type": "return",
+                        "value": undefined,
+                      },
+                      Object {
+                        "type": "return",
+                        "value": undefined,
+                      },
+                    ],
+                  },
+                }
+              }
               integration={
                 Object {
                   "assets": Array [],
@@ -330,6 +462,85 @@ exports[`SetupIntegrationForm tests renders SetupIntegrationForm and matches sna
                         "enabledWorkflows": Array [],
                       }
                     }
+                    http={
+                      Object {
+                        "get": [MockFunction] {
+                          "calls": Array [
+                            Array [
+                              "/api/integrations/repository/test_integration",
+                            ],
+                            Array [
+                              "/api/integrations/repository/test_integration",
+                            ],
+                            Array [
+                              "/api/integrations/repository/test_integration",
+                            ],
+                          ],
+                          "results": Array [
+                            Object {
+                              "type": "return",
+                              "value": Promise {},
+                            },
+                            Object {
+                              "type": "return",
+                              "value": Promise {},
+                            },
+                            Object {
+                              "type": "return",
+                              "value": Promise {},
+                            },
+                          ],
+                        },
+                        "post": [MockFunction] {
+                          "calls": Array [
+                            Array [
+                              "/api/console/proxy",
+                              Object {
+                                "body": "{}",
+                                "query": Object {
+                                  "method": "GET",
+                                  "path": "_data_stream/ss4o_*",
+                                },
+                              },
+                            ],
+                            Array [
+                              "/api/console/proxy",
+                              Object {
+                                "body": "{}",
+                                "query": Object {
+                                  "method": "GET",
+                                  "path": "_data_stream/ss4o_*",
+                                },
+                              },
+                            ],
+                            Array [
+                              "/api/console/proxy",
+                              Object {
+                                "body": "{}",
+                                "query": Object {
+                                  "method": "GET",
+                                  "path": "_data_stream/ss4o_*",
+                                },
+                              },
+                            ],
+                          ],
+                          "results": Array [
+                            Object {
+                              "type": "return",
+                              "value": undefined,
+                            },
+                            Object {
+                              "type": "return",
+                              "value": undefined,
+                            },
+                            Object {
+                              "type": "return",
+                              "value": undefined,
+                            },
+                          ],
+                        },
+                      }
+                    }
                     integration={
                       Object {
                         "assets": Array [],
@@ -956,7 +1167,54 @@ exports[`SetupIntegrationForm tests renders SetupIntegrationForm and matches sna
                       },
                     ],
                   },
-                  "post": [MockFunction],
+                  "post": [MockFunction] {
+                    "calls": Array [
+                      Array [
+                        "/api/console/proxy",
+                        Object {
+                          "body": "{}",
+                          "query": Object {
+                            "method": "GET",
+                            "path": "_data_stream/ss4o_*",
+                          },
+                        },
+                      ],
+                      Array [
+                        "/api/console/proxy",
+                        Object {
+                          "body": "{}",
+                          "query": Object {
+                            "method": "GET",
+                            "path": "_data_stream/ss4o_*",
+                          },
+                        },
+                      ],
+                      Array [
+                        "/api/console/proxy",
+                        Object {
+                          "body": "{}",
+                          "query": Object {
+                            "method": "GET",
+                            "path": "_data_stream/ss4o_*",
+                          },
+                        },
+                      ],
+                    ],
+                    "results": Array [
+                      Object {
+                        "type": "return",
+                        "value": undefined,
+                      },
+                      Object {
+                        "type": "return",
+                        "value": undefined,
+                      },
+                      Object {
+                        "type": "return",
+                        "value": undefined,
+                      },
+                    ],
+                  },
                 }
               }
               integration={
diff --git a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/create_integration_helpers.ts b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/create_integration_helpers.ts
index e56764756eaa..1f34317c15ec 100644
--- a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/create_integration_helpers.ts
+++ b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/create_integration_helpers.ts
@@ -25,6 +25,8 @@ interface AddIntegrationRequestParams {
   skipRedirect?: boolean;
   dataSourceInfo?: { dataSource: string; tableName: string };
   http: HttpStart;
+  dataSourceMDSId?: string;
+  dataSourceMDSLabel?: string;
 }
 
 interface ComponentMappingPayload {
@@ -133,13 +135,13 @@ export const checkDataSourceName = (
 
 export const fetchDataSourceMappings = async (
   targetDataSource: string,
-  http: HttpStart
+  http: HttpStart,
+  dataSourceMDSId?: string
 ): Promise<{ [key: string]: { properties: Properties } } | null> => {
   return http
-    .post(CONSOLE_PROXY, {
+    .post(`/api/dsl/integrations/mapping`, {
       query: {
-        path: `${targetDataSource}/_mapping`,
-        method: 'GET',
+        dataSourceMDSId,
       },
     })
     .then((response) => {
@@ -209,7 +211,8 @@ export const doExistingDataSourceValidation = async (
 const createComponentMapping = async (
   componentName: string,
   payload: ComponentMappingPayload,
-  http: HttpStart
+  http: HttpStart,
+  dataSourceMDSId?: string
 ): Promise<{ [key: string]: { properties: Properties } } | null> => {
   const version = payload.template.mappings._meta.version;
   return http.post(CONSOLE_PROXY, {
@@ -217,6 +220,7 @@ const createComponentMapping = async (
     query: {
       path: `_component_template/ss4o_${componentName}-${version}-template`,
       method: 'POST',
+      dataSourceId: dataSourceMDSId,
     },
   });
 };
@@ -226,7 +230,8 @@ const createIndexMapping = async (
   payload: ComponentMappingPayload,
   dataSourceName: string,
   integration: IntegrationConfig,
-  http: HttpStart
+  http: HttpStart,
+  dataSourceMDSId?: string
 ): Promise<{ [key: string]: { properties: Properties } } | null> => {
   const version = payload.template.mappings._meta.version;
   payload.index_patterns = [dataSourceName];
@@ -235,6 +240,7 @@ const createIndexMapping = async (
     query: {
       path: `_index_template/ss4o_${componentName}-${integration.name}-${version}-sample`,
       method: 'POST',
+      dataSourceId: dataSourceMDSId,
     },
   });
 };
@@ -244,7 +250,8 @@ const createIndexPatternMappings = async (
   integrationTemplateId: string,
   integration: IntegrationConfig,
   setToast: (title: string, color?: Color, text?: string | undefined) => void,
-  http: HttpStart
+  http: HttpStart,
+  dataSourceMDSId?: string
 ): Promise<void> => {
   // TODO the nested methods still need the dataSource -> indexPattern rename applied, sub-methods
   // here still have old naming convention
@@ -266,14 +273,18 @@ const createIndexPatternMappings = async (
         if (key === integration.type) {
           return Promise.resolve();
         }
-        return createComponentMapping(key, mapping as ComponentMappingPayload, http);
+        return createComponentMapping(
+          key,
+          mapping as ComponentMappingPayload,
+          http,
+          dataSourceMDSId
+        );
       })
     );
     // In order to see our changes, we need to manually provoke a refresh
-    await http.post(CONSOLE_PROXY, {
+    await http.post(`/api/dsl/integrations/refresh`, {
       query: {
-        path: '_refresh',
-        method: 'GET',
+        dataSourceMDSId,
       },
     });
     await createIndexMapping(
@@ -281,7 +292,8 @@ const createIndexPatternMappings = async (
       mappings[integration.type],
       targetDataSource,
       integration,
-      http
+      http,
+      dataSourceMDSId
     );
   } catch (err) {
     error = err.message;
@@ -305,6 +317,8 @@ export async function addIntegrationRequest({
   skipRedirect,
   dataSourceInfo,
   http,
+  dataSourceMDSId,
+  dataSourceMDSLabel,
 }: AddIntegrationRequestParams): Promise<boolean> {
   if (addSample) {
     createIndexPatternMappings(
@@ -312,19 +326,24 @@ export async function addIntegrationRequest({
       templateName,
       integration,
       setToast,
-      http
+      http,
+      dataSourceMDSId
     );
     name = `${templateName}-sample`;
     indexPattern = `ss4o_${integration.type}-${templateName}-sample-sample`;
   }
 
   const createReqBody: {
+    dataSourceMDSId?: string;
+    dataSourceMDSLabel?: string;
     name?: string;
     indexPattern?: string;
     workflows?: string[];
     dataSource?: string;
     tableName?: string;
   } = {
+    dataSourceMDSId,
+    dataSourceMDSLabel,
     name,
     indexPattern,
     workflows,
diff --git a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/installed_integrations_table.tsx b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/installed_integrations_table.tsx
index f561179507b0..69070372e1cc 100644
--- a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/installed_integrations_table.tsx
+++ b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/installed_integrations_table.tsx
@@ -107,6 +107,8 @@ export const InstallIntegrationFlyout = ({
   datasourceType,
   datasourceName,
   refreshInstances,
+  selectedDataSourceId,
+  selectedClusterName,
   http,
 }: {
   closeFlyout: () => void;
@@ -114,6 +116,8 @@ export const InstallIntegrationFlyout = ({
   datasourceName: string;
   refreshInstances: () => void;
   http: HttpStart;
+  selectedDataSourceId?: string;
+  selectedClusterName?: string;
 }) => {
   const [availableIntegrations, setAvailableIntegrations] = useState({
     hits: [],
@@ -155,6 +159,8 @@ export const InstallIntegrationFlyout = ({
         />
       ) : (
         <SetupIntegrationForm
+          selectedClusterName={selectedClusterName}
+          selectedDataSourceId={selectedDataSourceId}
           integration={installingIntegration}
           unsetIntegration={() => setInstallingIntegration(null)}
           renderType="flyout"
@@ -186,12 +192,16 @@ export const InstalledIntegrationsTable = ({
   datasourceName,
   refreshInstances,
   http,
+  selectedDataSourceId,
+  selectedClusterName,
 }: {
   integrations: IntegrationInstanceResult[];
   datasourceType: DatasourceType;
   datasourceName: string;
   refreshInstances: () => void;
   http: HttpStart;
+  selectedDataSourceId?: string;
+  selectedClusterName?: string;
 }) => {
   const basePathLink = (link: string): string => {
     if (http.basePath) {
@@ -276,6 +286,8 @@ export const InstalledIntegrationsTable = ({
           datasourceName={datasourceName}
           refreshInstances={refreshInstances}
           http={http}
+          selectedDataSourceId={selectedDataSourceId}
+          selectedClusterName={selectedClusterName}
         />
       ) : null}
     </>
diff --git a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/setup_integration.tsx b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/setup_integration.tsx
index c0b47afa561e..3db98f0602fe 100644
--- a/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/setup_integration.tsx
+++ b/src/plugins/data_source_management/public/components/direct_query_data_sources_components/integrations/setup_integration.tsx
@@ -25,6 +25,7 @@ import { addIntegrationRequest } from './create_integration_helpers';
 import { SetupIntegrationFormInputs } from './setup_integration_inputs';
 import { CONSOLE_PROXY, INTEGRATIONS_BASE } from '../../../../framework/utils/shared';
 import { IntegrationConfig, ParsedIntegrationAsset, Result } from '../../../../framework/types';
+import { SQLService } from '../../../../framework/requests/sql';
 
 export interface IntegrationSetupInputs {
   displayName: string;
@@ -50,32 +51,33 @@ type SetupCallout = { show: true; title: string; color?: Color; text?: string }
 const runQuery = async (
   query: string,
   datasource: string,
-  sessionId: string | null,
-  http: HttpStart
+  sessionId: string | undefined,
+  http: HttpStart,
+  dataSourceMDSId?: string
 ): Promise<Result<{ poll: object; sessionId: string }>> => {
   // Used for polling
   const sleep = (ms: number) => {
     return new Promise((resolve) => setTimeout(resolve, ms));
   };
 
+  const sqlService = new SQLService(http);
+
   try {
-    const queryResponse: { queryId: string; sessionId: string } = await http.post(CONSOLE_PROXY, {
-      body: JSON.stringify({ query, datasource, lang: 'sql', sessionId }),
-      query: {
-        path: '_plugins/_async_query',
-        method: 'POST',
+    const queryResponse: { queryId: string; sessionId: string } = await sqlService.fetch(
+      {
+        query,
+        datasource,
+        lang: 'sql',
+        sessionId,
       },
-    });
+      dataSourceMDSId
+    );
+
     let poll: { status: string; error?: string } = { status: 'undefined' };
-    const [queryId, newSessionId] = [queryResponse.queryId, queryResponse.sessionId];
+    const { queryId, sessionId: newSessionId } = queryResponse;
     while (!poll.error) {
-      poll = await http.post(CONSOLE_PROXY, {
-        body: '{}',
-        query: {
-          path: '_plugins/_async_query/' + queryId,
-          method: 'GET',
-        },
-      });
+      poll = await sqlService.fetchWithJobId({ queryId }, dataSourceMDSId);
+
       if (poll.status.toLowerCase() === 'success') {
         return {
           ok: true,
@@ -131,6 +133,8 @@ const addIntegration = async ({
   setCalloutLikeToast,
   setIsInstalling,
   http,
+  dataSourceMDSId,
+  dataSourceMDSLabel,
 }: {
   config: IntegrationSetupInputs;
   integration: IntegrationConfig;
@@ -138,9 +142,11 @@ const addIntegration = async ({
   setCalloutLikeToast: (title: string, color?: Color, text?: string) => void;
   setIsInstalling?: (isInstalling: boolean, success?: boolean) => void;
   http: HttpStart;
+  dataSourceMDSId?: string;
+  dataSourceMDSLabel?: string;
 }) => {
   setLoading(true);
-  let sessionId: string | null = null;
+  let sessionId: string | undefined;
 
   if (config.connectionType === 'index') {
     const res = await addIntegrationRequest({
@@ -148,6 +154,8 @@ const addIntegration = async ({
       templateName: integration.name,
       integration,
       setToast: setCalloutLikeToast,
+      dataSourceMDSId,
+      dataSourceMDSLabel,
       name: config.displayName,
       indexPattern: config.connectionDataSource,
       skipRedirect: setIsInstalling ? true : false,
@@ -174,7 +182,13 @@ const addIntegration = async ({
       }
 
       const queryStr = prepareQuery(query.query, config);
-      const result = await runQuery(queryStr, config.connectionDataSource, sessionId, http);
+      const result = await runQuery(
+        queryStr,
+        config.connectionDataSource,
+        sessionId,
+        http,
+        dataSourceMDSId
+      );
       if (!result.ok) {
         setLoading(false);
         setCalloutLikeToast('Failed to add integration', 'danger', result.error.message);
@@ -188,6 +202,8 @@ const addIntegration = async ({
       templateName: integration.name,
       integration,
       setToast: setCalloutLikeToast,
+      dataSourceMDSId,
+      dataSourceMDSLabel,
       name: config.displayName,
       indexPattern: `flint_${config.connectionDataSource}_default_${config.connectionTableName}__*`,
       workflows: config.enabledWorkflows,
@@ -230,6 +246,8 @@ export function SetupBottomBar({
   unsetIntegration,
   setIsInstalling,
   http,
+  dataSourceMDSId,
+  dataSourceMDSLabel,
 }: {
   config: IntegrationSetupInputs;
   integration: IntegrationConfig;
@@ -239,6 +257,8 @@ export function SetupBottomBar({
   unsetIntegration?: () => void;
   setIsInstalling?: (isInstalling: boolean, success?: boolean) => void;
   http: HttpStart;
+  dataSourceMDSId?: string;
+  dataSourceMDSLabel?: string;
 }) {
   // Drop-in replacement for setToast
   const setCalloutLikeToast = (title: string, color?: Color, text?: string) =>
@@ -290,6 +310,8 @@ export function SetupBottomBar({
                   setIsInstalling(newLoading);
                 },
                 setCalloutLikeToast,
+                dataSourceMDSId,
+                dataSourceMDSLabel,
                 setIsInstalling,
                 http,
               });
@@ -299,6 +321,8 @@ export function SetupBottomBar({
                 config,
                 setLoading,
                 setCalloutLikeToast,
+                dataSourceMDSId,
+                dataSourceMDSLabel,
                 setIsInstalling,
                 http,
               });
@@ -332,6 +356,8 @@ export function SetupIntegrationForm({
   forceConnection,
   setIsInstalling,
   http,
+  selectedDataSourceId,
+  selectedClusterName,
 }: {
   integration: string;
   renderType: 'page' | 'flyout';
@@ -342,6 +368,8 @@ export function SetupIntegrationForm({
   };
   setIsInstalling?: (isInstalling: boolean, success?: boolean) => void;
   http: HttpStart;
+  selectedDataSourceId?: string | undefined;
+  selectedClusterName?: string | undefined;
 }) {
   const [integConfig, setConfig] = useState({
     displayName: `${integration} Integration`,
@@ -391,6 +419,7 @@ export function SetupIntegrationForm({
                 integration={template}
                 setupCallout={setupCallout}
                 lockConnectionType={forceConnection !== undefined}
+                http={http}
               />
             )}
           </EuiPageContentBody>
@@ -404,6 +433,8 @@ export function SetupIntegrationForm({
             setSetupCallout={setSetupCallout}
             unsetIntegration={unsetIntegration}
             setIsInstalling={setIsInstalling}
+            dataSourceMDSId={selectedDataSourceId}
+            dataSourceMDSLabel={selectedClusterName}
             http={http}
           />
         </EuiBottomBar>
@@ -422,6 +453,7 @@ export function SetupIntegrationForm({
               integration={template}
               setupCallout={setupCallout}
               lockConnectionType={forceConnection !== undefined}
+              http={http}
             />
           )}
         </EuiFlyoutBody>
@@ -434,6 +466,8 @@ export function SetupIntegrationForm({
             setSetupCallout={setSetupCallout}
             unsetIntegration={unsetIntegration}
             setIsInstalling={setIsInstalling}
+            dataSourceMDSId={selectedDataSourceId}
+            dataSourceMDSLabel={selectedClusterName}
             http={http}
           />
         </EuiFlyoutFooter>
diff --git a/src/plugins/data_source_management/public/management_app/mount_management_section.tsx b/src/plugins/data_source_management/public/management_app/mount_management_section.tsx
index e31f8e2eef17..01727d5bf2dc 100644
--- a/src/plugins/data_source_management/public/management_app/mount_management_section.tsx
+++ b/src/plugins/data_source_management/public/management_app/mount_management_section.tsx
@@ -79,6 +79,7 @@ export async function mountManagementSection(
             setBreadcrumbs={params.setBreadcrumbs}
             application={application}
             useNewUX={useNewUX}
+            savedObjects={savedObjects}
           />
         </Route>
         {canManageDataSource && (