From 6e866907186c1a47296f992f016f5ac8d0d0a24c Mon Sep 17 00:00:00 2001
From: Sandeep Kumawat <skumwt@amazon.com>
Date: Thu, 11 Jul 2024 15:48:39 +0530
Subject: [PATCH] Add MDS support to snapshot pages

Signed-off-by: Sandeep Kumawat <skumwt@amazon.com>
---
 .../CreateSnapshotPolicy.tsx                  |  74 ++++++++-
 public/pages/Main/Main.tsx                    |  23 +++
 .../CreateRepositoryFlyout.tsx                |  35 +++-
 .../containers/Repositories/Repositories.tsx  |  33 +++-
 .../SnapshotPolicies/SnapshotPolicies.tsx     |  37 ++++-
 .../SnapshotPolicyDetails.tsx                 |  18 ++-
 .../CreateSnapshotFlyout.tsx                  |  33 +++-
 .../RestoreSnapshotFlyout.tsx                 | 151 ++++++++++--------
 .../SnapshotFlyout/SnapshotFlyout.tsx         |  23 ++-
 .../containers/Snapshots/Snapshots.tsx        |  32 +++-
 public/services/SnapshotManagementService.ts  |  77 ++++++---
 server/plugin.ts                              |   4 +-
 server/routes/snapshotManagement.ts           |  53 +++++-
 server/services/SnapshotManagementService.ts  |  95 ++++++-----
 14 files changed, 508 insertions(+), 180 deletions(-)

diff --git a/public/pages/CreateSnapshotPolicy/containers/CreateSnapshotPolicy/CreateSnapshotPolicy.tsx b/public/pages/CreateSnapshotPolicy/containers/CreateSnapshotPolicy/CreateSnapshotPolicy.tsx
index aa9674ae8..02c73fac8 100644
--- a/public/pages/CreateSnapshotPolicy/containers/CreateSnapshotPolicy/CreateSnapshotPolicy.tsx
+++ b/public/pages/CreateSnapshotPolicy/containers/CreateSnapshotPolicy/CreateSnapshotPolicy.tsx
@@ -27,7 +27,7 @@ import {
   EuiButtonIcon,
   EuiLink,
 } from "@elastic/eui";
-import React, { ChangeEvent, Component } from "react";
+import React, { ChangeEvent, Component, useContext } from "react";
 import { RouteComponentProps } from "react-router-dom";
 import { CoreServicesContext } from "../../../../components/core_services";
 import { CatRepository, CreateRepositoryBody, CreateRepositorySettings, FeatureChannelList } from "../../../../../server/models/interfaces";
@@ -52,15 +52,18 @@ import SnapshotIndicesRepoInput from "../../components/SnapshotIndicesRepoInput"
 import CronSchedule from "../../components/CronSchedule";
 import SnapshotAdvancedSettings from "../../components/SnapshotAdvancedSettings";
 import Notification from "../../components/Notification";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
+import { useUpdateUrlWithDataSourceProperties } from "../../../../components/MDSEnabledComponent";
 
-interface CreateSMPolicyProps extends RouteComponentProps {
+interface CreateSMPolicyProps extends RouteComponentProps, DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
   isEdit: boolean;
   notificationService: NotificationService;
   indexService: IndexService;
 }
 
-interface CreateSMPolicyState {
+interface CreateSMPolicyState extends DataSourceMenuProperties {
   policy: SMPolicy;
   policyId: string;
   policySeqNo: number | undefined;
@@ -96,13 +99,14 @@ interface CreateSMPolicyState {
   timezoneError: string;
 }
 
-export default class CreateSnapshotPolicy extends Component<CreateSMPolicyProps, CreateSMPolicyState> {
+export class CreateSnapshotPolicy extends MDSEnabledComponent<CreateSMPolicyProps, CreateSMPolicyState> {
   static contextType = CoreServicesContext;
 
   constructor(props: CreateSMPolicyProps) {
     super(props);
 
     this.state = {
+      ...this.state,
       policy: getDefaultSMPolicy(),
       policyId: "",
       policySeqNo: undefined,
@@ -160,11 +164,67 @@ export default class CreateSnapshotPolicy extends Component<CreateSMPolicyProps,
         BREADCRUMBS.CREATE_SNAPSHOT_POLICY,
       ]);
     }
+    this.updateOptions();
+  }
+
+  async updateOptions() {
     await this.getIndexOptions("");
     await this.getRepos();
     await this.getChannels();
   }
 
+  static getQueryObjectFromState({ dataSourceId, multiDataSourceEnabled }: CreateSMPolicyState) {
+    return {
+      ...(multiDataSourceEnabled ? { dataSourceId } : {}),
+    };
+  }
+
+  resetState() {
+    this.setState({
+      policy: getDefaultSMPolicy(),
+      policyId: "",
+      policySeqNo: undefined,
+      policyPrimaryTerm: undefined,
+
+      isSubmitting: false,
+
+      channels: [],
+      loadingChannels: false,
+
+      indexOptions: DEFAULT_INDEX_OPTIONS,
+      selectedIndexOptions: [],
+
+      repositories: [],
+      selectedRepoValue: "",
+
+      maxAgeNum: 1,
+      maxAgeUnit: "d",
+
+      creationScheduleFrequencyType: "daily",
+      deletionScheduleFrequencyType: "daily",
+
+      deleteConditionEnabled: false,
+      deletionScheduleEnabled: false,
+
+      advancedSettingsOpen: false,
+      showCreateRepoFlyout: false,
+
+      policyIdError: "",
+      repoError: "",
+      minCountError: "",
+      timezoneError: "",
+    });
+  }
+
+  async componentDidUpdate(prevProps: CreateSMPolicyProps, prevState: CreateSMPolicyState) {
+    const prevQuery = CreateSnapshotPolicy.getQueryObjectFromState(prevState);
+    const currQuery = CreateSnapshotPolicy.getQueryObjectFromState(this.state);
+    if (!_.isEqual(prevQuery, currQuery)) {
+      this.resetState();
+      this.updateOptions();
+    }
+  }
+
   getPolicy = async (policyId: string): Promise<void> => {
     try {
       const { snapshotManagementService } = this.props;
@@ -877,3 +937,9 @@ export default class CreateSnapshotPolicy extends Component<CreateSMPolicyProps,
     return _.set(this.state.policy, path, newValue);
   };
 }
+
+export default function (props: Omit<CreateSMPolicyProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  useUpdateUrlWithDataSourceProperties();
+  return <CreateSnapshotPolicy {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/Main/Main.tsx b/public/pages/Main/Main.tsx
index 73498d603..1830d52e9 100644
--- a/public/pages/Main/Main.tsx
+++ b/public/pages/Main/Main.tsx
@@ -177,6 +177,16 @@ const dataSourceEnabledPaths: string[] = [
   ROUTES.CREATE_TRANSFORM,
   ROUTES.EDIT_TRANSFORM,
   ROUTES.TRANSFORM_DETAILS,
+  ROUTES.SNAPSHOT_POLICIES,
+  ROUTES.SNAPSHOT_POLICY_DETAILS,
+  ROUTES.CREATE_SNAPSHOT_POLICY,
+  ROUTES.EDIT_SNAPSHOT_POLICY,
+  ROUTES.SNAPSHOTS,
+  ROUTES.CREATE_SNAPSHOT,
+  ROUTES.EDIT_SNAPSHOT,
+  ROUTES.REPOSITORIES,
+  ROUTES.CREATE_REPOSITORY,
+  ROUTES.EDIT_REPOSITORY,
 ];
 
 export default class Main extends Component<MainProps, MainState> {
@@ -241,6 +251,7 @@ export default class Main extends Component<MainProps, MainState> {
       services.notificationService = new NotificationService(http, this.state.dataSourceId, this.props.multiDataSourceEnabled);
       services.rollupService = new RollupService(http, this.state.dataSourceId, this.props.multiDataSourceEnabled);
       services.transformService = new TransformService(http, this.state.dataSourceId, this.props.multiDataSourceEnabled);
+      services.snapshotManagementService = new SnapshotManagementService(http, this.state.dataSourceId, this.props.multiDataSourceEnabled);
     }
     return services;
   }
@@ -447,6 +458,8 @@ export default class Main extends Component<MainProps, MainState> {
                                 ROUTES.ROLLUP_DETAILS,
                                 ROUTES.TRANSFORM_DETAILS,
                                 ROUTES.EDIT_TRANSFORM,
+                                ROUTES.EDIT_SNAPSHOT_POLICY,
+                                ROUTES.SNAPSHOT_POLICY_DETAILS,
                               ]}
                               render={() => (
                                 <DataSourceViewer
@@ -477,6 +490,16 @@ export default class Main extends Component<MainProps, MainState> {
                                 ROUTES.CHANGE_POLICY,
                                 ROUTES.ROLLUPS,
                                 ROUTES.TRANSFORMS,
+                                ROUTES.SNAPSHOT_POLICIES,
+                                ROUTES.SNAPSHOT_POLICY_DETAILS,
+                                ROUTES.CREATE_SNAPSHOT_POLICY,
+                                ROUTES.EDIT_SNAPSHOT_POLICY,
+                                ROUTES.SNAPSHOTS,
+                                ROUTES.CREATE_SNAPSHOT,
+                                ROUTES.EDIT_SNAPSHOT,
+                                ROUTES.REPOSITORIES,
+                                ROUTES.CREATE_REPOSITORY,
+                                ROUTES.EDIT_REPOSITORY,
                               ]}
                               render={() => (
                                 <DataSourceMenu
diff --git a/public/pages/Repositories/components/CreateRepositoryFlyout/CreateRepositoryFlyout.tsx b/public/pages/Repositories/components/CreateRepositoryFlyout/CreateRepositoryFlyout.tsx
index c3c80787b..1f069efc6 100644
--- a/public/pages/Repositories/components/CreateRepositoryFlyout/CreateRepositoryFlyout.tsx
+++ b/public/pages/Repositories/components/CreateRepositoryFlyout/CreateRepositoryFlyout.tsx
@@ -23,7 +23,7 @@ import {
 } from "@elastic/eui";
 import _ from "lodash";
 import { CreateRepositorySettings } from "../../../../../server/models/interfaces";
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import FlyoutFooter from "../../../VisualCreatePolicy/components/FlyoutFooter";
 import { CoreServicesContext } from "../../../../components/core_services";
 import { SnapshotManagementService } from "../../../../services";
@@ -36,15 +36,17 @@ import {
   S3_REPOSITORY_DOCUMENTATION_URL,
   SNAPSHOT_MANAGEMENT_DOCUMENTATION_URL,
 } from "../../../../utils/constants";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
 
-interface CreateRepositoryProps {
+interface CreateRepositoryProps extends DataSourceMenuProperties {
   service: SnapshotManagementService;
   editRepo: string | null;
   onCloseFlyout: () => void;
   createRepo: (repoName: string, repoType: string, settings: CreateRepositorySettings) => void;
 }
 
-interface CreateRepositoryState {
+interface CreateRepositoryState extends DataSourceMenuProperties {
   repoName: string;
   location: string;
   // repoTypeOptions: EuiComboBoxOptionOption<string>[];
@@ -59,13 +61,14 @@ interface CreateRepositoryState {
   locationError: string;
 }
 
-export default class CreateRepositoryFlyout extends Component<CreateRepositoryProps, CreateRepositoryState> {
+export class CreateRepositoryFlyout extends MDSEnabledComponent<CreateRepositoryProps, CreateRepositoryState> {
   static contextType = CoreServicesContext;
 
   constructor(props: CreateRepositoryProps) {
     super(props);
 
     this.state = {
+      ...this.state,
       repoName: "",
       location: "",
       // repoTypeOptions: [],
@@ -85,6 +88,25 @@ export default class CreateRepositoryFlyout extends Component<CreateRepositoryPr
     }
   }
 
+  async componentDidUpdate(prevProps: CreateRepositoryProps, prevState: CreateRepositoryState) {
+    if (prevState.dataSourceId != this.state.dataSourceId) {
+      const { editRepo } = this.props;
+      this.setState({
+        repoName: "",
+        location: "",
+        selectedRepoTypeOption: REPO_SELECT_OPTIONS[0].value as string,
+        fsSettingsJsonString: JSON.stringify(FS_ADVANCED_SETTINGS, null, 4),
+        customSettingsJsonString: JSON.stringify(CUSTOM_CONFIGURATION, null, 4),
+        repoNameError: "",
+        repoTypeError: "",
+        locationError: "",
+      });
+      if (!!editRepo) {
+        await this.getRepo(editRepo);
+      }
+    }
+  }
+
   getRepo = async (repoName: string) => {
     const { service } = this.props;
     try {
@@ -291,3 +313,8 @@ export default class CreateRepositoryFlyout extends Component<CreateRepositoryPr
     );
   }
 }
+
+export default function (props: Omit<CreateRepositoryProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  return <CreateRepositoryFlyout {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/Repositories/containers/Repositories/Repositories.tsx b/public/pages/Repositories/containers/Repositories/Repositories.tsx
index 12986972c..bebe19e59 100644
--- a/public/pages/Repositories/containers/Repositories/Repositories.tsx
+++ b/public/pages/Repositories/containers/Repositories/Repositories.tsx
@@ -13,7 +13,7 @@ import {
   EuiTextColor,
 } from "@elastic/eui";
 import { getErrorMessage } from "../../../../utils/helpers";
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import { RouteComponentProps } from "react-router-dom";
 import { CatRepository, CreateRepositoryBody, CreateRepositorySettings } from "../../../../../server/models/interfaces";
 import { CoreServicesContext } from "../../../../components/core_services";
@@ -24,12 +24,15 @@ import { FieldValueSelectionFilterConfigType } from "@elastic/eui/src/components
 import { BREADCRUMBS } from "../../../../utils/constants";
 import DeleteModal from "../../components/DeleteModal";
 import { truncateSpan } from "../../../Snapshots/helper";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
+import { useUpdateUrlWithDataSourceProperties } from "../../../../components/MDSEnabledComponent";
 
-interface RepositoriesProps extends RouteComponentProps {
+interface RepositoriesProps extends RouteComponentProps, DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
 }
 
-interface RepositoriesState {
+interface RepositoriesState extends DataSourceMenuProperties {
   repositories: CatRepository[];
   loading: boolean;
   selectedItems: CatRepository[];
@@ -43,7 +46,7 @@ interface RepositoriesState {
   isDeleteModalVisible: boolean;
 }
 
-export default class Repositories extends Component<RepositoriesProps, RepositoriesState> {
+export class Repositories extends MDSEnabledComponent<RepositoriesProps, RepositoriesState> {
   static contextType = CoreServicesContext;
   columns: EuiTableFieldDataColumnType<CatRepository>[];
 
@@ -51,6 +54,7 @@ export default class Repositories extends Component<RepositoriesProps, Repositor
     super(props);
 
     this.state = {
+      ...this.state,
       repositories: [],
       loading: false,
       selectedItems: [],
@@ -93,6 +97,20 @@ export default class Repositories extends Component<RepositoriesProps, Repositor
     await this.getRepos();
   }
 
+  async componentDidUpdate(prevProps: RepositoriesProps, prevState: RepositoriesState) {
+    const prevQuery = Repositories.getQueryObjectFromState(prevState);
+    const currQuery = Repositories.getQueryObjectFromState(this.state);
+    if (!_.isEqual(prevQuery, currQuery)) {
+      await this.getRepos();
+    }
+  }
+
+  static getQueryObjectFromState({ dataSourceId, multiDataSourceEnabled }: RepositoriesState) {
+    return {
+      ...(multiDataSourceEnabled ? { dataSourceId } : {}),
+    };
+  }
+
   getRepos = async () => {
     this.setState({ loading: true });
 
@@ -177,7 +195,6 @@ export default class Repositories extends Component<RepositoriesProps, Repositor
 
   render() {
     const { repositories, loading, selectedItems, showFlyout, isPopoverOpen, editRepo, isDeleteModalVisible } = this.state;
-
     const popoverActionItems = [
       <EuiContextMenuItem
         key="Edit"
@@ -319,3 +336,9 @@ export default class Repositories extends Component<RepositoriesProps, Repositor
       .join(", ");
   };
 }
+
+export default function (props: Omit<RepositoriesProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  useUpdateUrlWithDataSourceProperties();
+  return <Repositories {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/SnapshotPolicies/containers/SnapshotPolicies/SnapshotPolicies.tsx b/public/pages/SnapshotPolicies/containers/SnapshotPolicies/SnapshotPolicies.tsx
index 68ab334c0..789fb89ab 100644
--- a/public/pages/SnapshotPolicies/containers/SnapshotPolicies/SnapshotPolicies.tsx
+++ b/public/pages/SnapshotPolicies/containers/SnapshotPolicies/SnapshotPolicies.tsx
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import _ from "lodash";
 import queryString from "query-string";
 import { RouteComponentProps } from "react-router-dom";
@@ -39,12 +39,15 @@ import DeleteModal from "../../../PolicyDetails/components/DeleteModal";
 import { OnSearchChangeArgs } from "../../../../models/interfaces";
 import { humanCronExpression, parseCronExpression } from "../../../CreateSnapshotPolicy/components/CronSchedule/helper";
 import { truncateSpan } from "../../../Snapshots/helper";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
+import { useUpdateUrlWithDataSourceProperties } from "../../../../components/MDSEnabledComponent";
 
-interface SnapshotPoliciesProps extends RouteComponentProps {
+interface SnapshotPoliciesProps extends RouteComponentProps, DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
 }
 
-interface SnapshotPoliciesState {
+interface SnapshotPoliciesState extends DataSourceMenuProperties {
   policies: SMPolicy[];
   totalPolicies: number;
   loadingPolicies: boolean;
@@ -66,7 +69,7 @@ interface SnapshotPoliciesState {
   isDeleteModalVisible: boolean;
 }
 
-export default class SnapshotPolicies extends Component<SnapshotPoliciesProps, SnapshotPoliciesState> {
+export class SnapshotPolicies extends MDSEnabledComponent<SnapshotPoliciesProps, SnapshotPoliciesState> {
   static contextType = CoreServicesContext;
   columns: EuiTableFieldDataColumnType<SMPolicy>[];
 
@@ -75,6 +78,7 @@ export default class SnapshotPolicies extends Component<SnapshotPoliciesProps, S
 
     const { from, size, sortField, sortOrder } = getSMPoliciesQueryParamsFromURL(this.props.location);
     this.state = {
+      ...this.state,
       policies: [],
       totalPolicies: 0,
       loadingPolicies: false,
@@ -196,8 +200,23 @@ export default class SnapshotPolicies extends Component<SnapshotPoliciesProps, S
     this.setState({ loadingPolicies: false });
   };
 
-  static getQueryObjectFromState({ from, size, sortField, sortOrder, queryString }: SnapshotPoliciesState) {
-    return { from, size, sortField, sortOrder, queryString };
+  static getQueryObjectFromState({
+    from,
+    size,
+    sortField,
+    sortOrder,
+    queryString,
+    multiDataSourceEnabled,
+    dataSourceId,
+  }: SnapshotPoliciesState) {
+    return {
+      from,
+      size,
+      sortField,
+      sortOrder,
+      queryString,
+      ...(multiDataSourceEnabled ? { dataSourceId } : {}),
+    };
   }
 
   onSelectionChange = (selectedItems: SMPolicy[]): void => {
@@ -472,3 +491,9 @@ export default class SnapshotPolicies extends Component<SnapshotPoliciesProps, S
     );
   }
 }
+
+export default function (props: Omit<SnapshotPoliciesProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  useUpdateUrlWithDataSourceProperties();
+  return <SnapshotPolicies {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/SnapshotPolicyDetails/containers/SnapshotPolicyDetails/SnapshotPolicyDetails.tsx b/public/pages/SnapshotPolicyDetails/containers/SnapshotPolicyDetails/SnapshotPolicyDetails.tsx
index 70a75feef..11e7701bb 100644
--- a/public/pages/SnapshotPolicyDetails/containers/SnapshotPolicyDetails/SnapshotPolicyDetails.tsx
+++ b/public/pages/SnapshotPolicyDetails/containers/SnapshotPolicyDetails/SnapshotPolicyDetails.tsx
@@ -5,7 +5,7 @@
 
 import _ from "lodash";
 import { RouteComponentProps } from "react-router-dom";
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import queryString from "query-string";
 import {
   EuiAccordion,
@@ -36,13 +36,16 @@ import InfoModal from "../../components/InfoModal";
 import { getAllowPartial, getIgnoreUnavailabel, getIncludeGlobalState } from "../../../CreateSnapshotPolicy/containers/helper";
 import { truncateSpan } from "../../../Snapshots/helper";
 import { NotificationConfig } from "../../../../../server/models/interfaces";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
+import { useUpdateUrlWithDataSourceProperties } from "../../../../components/MDSEnabledComponent";
 
-interface SnapshotPolicyDetailsProps extends RouteComponentProps {
+interface SnapshotPolicyDetailsProps extends RouteComponentProps, DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
   notificationService: NotificationService;
 }
 
-interface SnapshotPolicyDetailsState {
+interface SnapshotPolicyDetailsState extends DataSourceMenuProperties {
   policyId: string;
   policy: SMPolicy | null;
 
@@ -53,7 +56,7 @@ interface SnapshotPolicyDetailsState {
   channel: NotificationConfig | null;
 }
 
-export default class SnapshotPolicyDetails extends Component<SnapshotPolicyDetailsProps, SnapshotPolicyDetailsState> {
+export class SnapshotPolicyDetails extends MDSEnabledComponent<SnapshotPolicyDetailsProps, SnapshotPolicyDetailsState> {
   static contextType = CoreServicesContext;
   columns: EuiTableFieldDataColumnType<LatestActivities>[];
 
@@ -61,6 +64,7 @@ export default class SnapshotPolicyDetails extends Component<SnapshotPolicyDetai
     super(props);
 
     this.state = {
+      ...this.state,
       policyId: "",
       policy: null,
       metadata: null,
@@ -448,3 +452,9 @@ export default class SnapshotPolicyDetails extends Component<SnapshotPolicyDetai
     );
   }
 }
+
+export default function (props: Omit<SnapshotPolicyDetailsProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  useUpdateUrlWithDataSourceProperties();
+  return <SnapshotPolicyDetails {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/Snapshots/components/CreateSnapshotFlyout/CreateSnapshotFlyout.tsx b/public/pages/Snapshots/components/CreateSnapshotFlyout/CreateSnapshotFlyout.tsx
index 4657ad5f6..7e125a81e 100644
--- a/public/pages/Snapshots/components/CreateSnapshotFlyout/CreateSnapshotFlyout.tsx
+++ b/public/pages/Snapshots/components/CreateSnapshotFlyout/CreateSnapshotFlyout.tsx
@@ -17,7 +17,7 @@ import {
 } from "@elastic/eui";
 import _ from "lodash";
 
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import FlyoutFooter from "../../../VisualCreatePolicy/components/FlyoutFooter";
 import { CoreServicesContext } from "../../../../components/core_services";
 import { IndexService, SnapshotManagementService } from "../../../../services";
@@ -30,15 +30,17 @@ import SnapshotIndicesRepoInput from "../../../CreateSnapshotPolicy/components/S
 import { ChangeEvent } from "react";
 import { getEmptySnapshot } from "./constants";
 import { ERROR_PROMPT } from "../../../CreateSnapshotPolicy/constants";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
 
-interface CreateSnapshotProps {
+interface CreateSnapshotProps extends DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
   indexService: IndexService;
   onCloseFlyout: () => void;
   createSnapshot: (snapshotId: string, repository: string, snapshot: Snapshot) => void;
 }
 
-interface CreateSnapshotState {
+interface CreateSnapshotState extends DataSourceMenuProperties {
   indexOptions: EuiComboBoxOptionOption<IndexItem>[];
   selectedIndexOptions: EuiComboBoxOptionOption<IndexItem>[];
 
@@ -52,12 +54,13 @@ interface CreateSnapshotState {
   snapshotIdError: string;
 }
 
-export default class CreateSnapshotFlyout extends Component<CreateSnapshotProps, CreateSnapshotState> {
+export class CreateSnapshotFlyout extends MDSEnabledComponent<CreateSnapshotProps, CreateSnapshotState> {
   static contextType = CoreServicesContext;
   constructor(props: CreateSnapshotProps) {
     super(props);
 
     this.state = {
+      ...this.state,
       indexOptions: [],
       selectedIndexOptions: [],
       repositories: [],
@@ -74,6 +77,23 @@ export default class CreateSnapshotFlyout extends Component<CreateSnapshotProps,
     await this.getRepos();
   }
 
+  async componentDidUpdate(prevProps: CreateSnapshotProps, prevState: CreateSnapshotState) {
+    if (prevState.dataSourceId != this.state.dataSourceId) {
+      this.setState({
+        indexOptions: [],
+        selectedIndexOptions: [],
+        repositories: [],
+        selectedRepoValue: "",
+        snapshot: getEmptySnapshot(),
+        snapshotId: "",
+        repoError: "",
+        snapshotIdError: "",
+      });
+      await this.getIndexOptions("");
+      await this.getRepos();
+    }
+  }
+
   onClickAction = () => {
     const { createSnapshot } = this.props;
     const { snapshotId, selectedRepoValue, snapshot } = this.state;
@@ -250,3 +270,8 @@ export default class CreateSnapshotFlyout extends Component<CreateSnapshotProps,
     );
   }
 }
+
+export default function (props: Omit<CreateSnapshotProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  return <CreateSnapshotFlyout {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/Snapshots/components/RestoreSnapshotFlyout/RestoreSnapshotFlyout.tsx b/public/pages/Snapshots/components/RestoreSnapshotFlyout/RestoreSnapshotFlyout.tsx
index 32528ad2a..491f5660e 100644
--- a/public/pages/Snapshots/components/RestoreSnapshotFlyout/RestoreSnapshotFlyout.tsx
+++ b/public/pages/Snapshots/components/RestoreSnapshotFlyout/RestoreSnapshotFlyout.tsx
@@ -17,18 +17,18 @@ import {
   EuiAccordion,
   EuiCheckbox,
   EuiCallOut,
-  EuiText
+  EuiText,
 } from "@elastic/eui";
 import _ from "lodash";
-import React, { Component, ChangeEvent } from "react";
+import React, { Component, ChangeEvent, useContext } from "react";
 import FlyoutFooter from "../../../VisualCreatePolicy/components/FlyoutFooter";
 import { CoreServicesContext } from "../../../../components/core_services";
 import { IndexService, SnapshotManagementService } from "../../../../services";
 import { RESTORE_OPTIONS } from "../../../../models/interfaces";
 import { getErrorMessage } from "../../../../utils/helpers";
-import { checkBadJSON, checkBadRegex, checkBadReplacement, checkCustomIgnoreConflict, checkNoSelectedIndices } from "../../helper"
-import { browseIndicesCols } from "../../../../utils/constants"
-import { ERROR_TOAST_TITLE } from "../../constants"
+import { checkBadJSON, checkBadRegex, checkBadReplacement, checkCustomIgnoreConflict, checkNoSelectedIndices } from "../../helper";
+import { browseIndicesCols } from "../../../../utils/constants";
+import { ERROR_TOAST_TITLE } from "../../constants";
 import { IndexItem } from "../../../../../models/interfaces";
 import { CatRepository, GetSnapshot } from "../../../../../server/models/interfaces";
 import CustomLabel from "../../../../components/CustomLabel";
@@ -40,18 +40,20 @@ import RenameInput from "../RenameInput";
 import SnapshotIndicesInput from "../SnapshotIndicesInput";
 import IndexList from "../IndexList";
 import { ERROR_PROMPT } from "../../../CreateSnapshotPolicy/constants";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
 
-interface RestoreSnapshotProps {
+interface RestoreSnapshotProps extends DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
   indexService: IndexService;
   onCloseFlyout: () => void;
-  getRestoreTime: (time: number) => void
+  getRestoreTime: (time: number) => void;
   restoreSnapshot: (snapshotId: string, repository: string, options: object) => void;
   snapshotId: string;
   repository: string;
 }
 
-interface RestoreSnapshotState {
+interface RestoreSnapshotState extends DataSourceMenuProperties {
   indexOptions: EuiComboBoxOptionOption<IndexItem>[];
   selectedIndexOptions: EuiComboBoxOptionOption<IndexItem>[];
   renameIndices: string;
@@ -70,17 +72,18 @@ interface RestoreSnapshotState {
   badPattern: boolean;
   badRename: boolean;
   badJSON: boolean;
-  badIgnore: boolean
+  badIgnore: boolean;
   repoError: string;
   noIndicesSelected: boolean;
   snapshotIdError: string;
 }
 
-export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProps, RestoreSnapshotState> {
+export class RestoreSnapshotFlyout extends MDSEnabledComponent<RestoreSnapshotProps, RestoreSnapshotState> {
   static contextType = CoreServicesContext;
   constructor(props: RestoreSnapshotProps) {
     super(props);
     this.state = {
+      ...this.state,
       indexOptions: [],
       selectedIndexOptions: [],
       renameIndices: "add_prefix",
@@ -110,6 +113,12 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
     await this.getIndexOptions();
   }
 
+  async componentDidUpdate(prevProps: RestoreSnapshotProps, prevState: RestoreSnapshotState) {
+    if (prevState.dataSourceId != this.state.dataSourceId) {
+      await this.getIndexOptions();
+    }
+  }
+
   onClickAction = () => {
     const { restoreSnapshot, snapshotId, repository, onCloseFlyout, getRestoreTime } = this.props;
     const {
@@ -156,14 +165,15 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
     const badIgnore = options.ignore_index_settings ? checkCustomIgnoreConflict(customIndexSettings, ignoreIndexSettings) : false;
     const noIndicesSelected = checkNoSelectedIndices(options.indices, restoreSpecific);
 
-    if ((!badIgnore) && (badPattern || badRename || noIndicesSelected || badJSON)) {
+    if (!badIgnore && (badPattern || badRename || noIndicesSelected || badJSON)) {
       this.context.notifications.toasts.addDanger(null, { title: ERROR_TOAST_TITLE });
     }
 
     if (badIgnore) {
       this.context.notifications.toasts.addDanger(null, {
         title: "Cannot apply and ignore the same index settings",
-        text: "One or more index settings was declared in both Custom index settings and Ignore index settings. Remove the index setting from either fields."
+        text:
+          "One or more index settings was declared in both Custom index settings and Ignore index settings. Remove the index setting from either fields.",
       });
     }
 
@@ -191,11 +201,11 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
     }
 
     if (options.index_settings) {
-      options.index_settings = JSON.parse(options.index_settings)
+      options.index_settings = JSON.parse(options.index_settings);
     }
     getRestoreTime(Date.now());
     restoreSnapshot(snapshotId, repository, options);
-    onCloseFlyout()
+    onCloseFlyout();
   };
 
   onClickIndices = async () => {
@@ -226,11 +236,11 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
 
         this.setState({ snapshot: response.response, indexOptions: [...newOptions] });
       } else {
-        const message = JSON.parse(response.error).error.root_cause[0].reason
+        const message = JSON.parse(response.error).error.root_cause[0].reason;
         const trimmedMessage = message.slice(message.indexOf("]") + 1, message.indexOf(".") + 1);
         this.context.notifications.toasts.addError(response.error, {
           title: `There was a problem getting the snapshot`,
-          toastMessage: `${trimmedMessage} Open browser console & click below for details.`
+          toastMessage: `${trimmedMessage} Open browser console & click below for details.`,
         });
       }
     } catch (err) {
@@ -287,7 +297,11 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
     }
 
     if (e.target.id === restore_all_indices) {
-      this.setState({ restoreSpecific: false, noIndicesSelected: false, snapshot: _.set(this.state.snapshot!, e.target.id, e.target.checked) });
+      this.setState({
+        restoreSpecific: false,
+        noIndicesSelected: false,
+        snapshot: _.set(this.state.snapshot!, e.target.id, e.target.checked),
+      });
       return;
     }
 
@@ -308,8 +322,8 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
       let renameReplacement = this.state.renameReplacement;
 
       if (e.target.id !== "rename_indices") {
-        renamePattern = "(.+)"
-        renameReplacement = "restored_$1"
+        renamePattern = "(.+)";
+        renameReplacement = "restored_$1";
       }
 
       this.setState({
@@ -318,7 +332,7 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
         badPattern: false,
         renameReplacement,
         badRename: false,
-        snapshot: _.set(this.state.snapshot!, e.target.id, e.target.checked)
+        snapshot: _.set(this.state.snapshot!, e.target.id, e.target.checked),
       });
       return;
     }
@@ -340,7 +354,7 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
       badRename,
       badJSON,
       badIgnore,
-      noIndicesSelected
+      noIndicesSelected,
     } = this.state;
 
     const {
@@ -378,7 +392,6 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
             </EuiFlyoutHeader>
 
             <EuiFlyoutBody>
-
               <EuiFlexGroup alignItems="flexStart">
                 <EuiFlexItem>
                   <CustomLabel title="Snapshot name" />
@@ -387,12 +400,17 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
                 </EuiFlexItem>
                 <EuiFlexItem>
                   <CustomLabel title="Status" />
-                  <EuiHealth textSize="m" color={`${status?.toLowerCase()}`} title={`${status} indicator icon`}> {status}</EuiHealth>
+                  <EuiHealth textSize="m" color={`${status?.toLowerCase()}`} title={`${status} indicator icon`}>
+                    {" "}
+                    {status}
+                  </EuiHealth>
                 </EuiFlexItem>
                 <EuiFlexItem>
                   <CustomLabel title="Indices" />
                   <EuiSpacer size="xs" />
-                  <a onClick={this.onClickIndices} style={{ fontSize: "1.1rem" }}>{snapshot?.indices.length}</a>
+                  <a onClick={this.onClickIndices} style={{ fontSize: "1.1rem" }}>
+                    {snapshot?.indices.length}
+                  </a>
                 </EuiFlexItem>
               </EuiFlexGroup>
 
@@ -408,20 +426,18 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
 
               <EuiSpacer size="l" />
 
-              {
-                restoreSpecific && (
-                  <SnapshotIndicesInput
-                    indexOptions={indexOptions}
-                    selectedIndexOptions={selectedIndexOptions}
-                    onIndicesSelectionChange={this.onIndicesSelectionChange}
-                    getIndexOptions={this.getIndexOptions}
-                    onCreateOption={this.onCreateOption}
-                    selectedRepoValue={selectedRepoValue}
-                    showError={noIndicesSelected}
-                    isClearable={true}
-                  />
-                )
-              }
+              {restoreSpecific && (
+                <SnapshotIndicesInput
+                  indexOptions={indexOptions}
+                  selectedIndexOptions={selectedIndexOptions}
+                  onIndicesSelectionChange={this.onIndicesSelectionChange}
+                  getIndexOptions={this.getIndexOptions}
+                  onCreateOption={this.onCreateOption}
+                  selectedRepoValue={selectedRepoValue}
+                  showError={noIndicesSelected}
+                  isClearable={true}
+                />
+              )}
 
               <EuiSpacer size="l" />
 
@@ -436,15 +452,14 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
               />
 
               {renameIndices === add_prefix && <AddPrefixInput getPrefix={this.getPrefix} />}
-              {
-                renameIndices === rename_indices && (
-                  <RenameInput
-                    getRenamePattern={this.getRenamePattern}
-                    getRenameReplacement={this.getRenameReplacement}
-                    showPatternError={badPattern}
-                    showRenameError={badRename} />
-                )
-              }
+              {renameIndices === rename_indices && (
+                <RenameInput
+                  getRenamePattern={this.getRenamePattern}
+                  getRenameReplacement={this.getRenameReplacement}
+                  showPatternError={badPattern}
+                  showRenameError={badRename}
+                />
+              )}
 
               <EuiSpacer size="xxl" />
               <EuiAccordion id="advanced_restore_options" buttonContent="Advanced options">
@@ -470,24 +485,22 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
 
               <EuiSpacer size="l" />
 
-              {snapshot?.failed_shards && <EuiCallOut
-                title="Restoring a partial snapshot"
-                color="warning"
-              >
-                <p>
-                  You are about to restore a partial snapshot. One or more shards may be missing in this
-                  <br />
-                  snapshot. Do you want to continue?
-                </p>
-                <EuiSpacer size="s" />
-                <EuiCheckbox
-                  id={partial}
-                  label={<EuiText size="s">Allow restore partial snapshots</EuiText>}
-                  checked={String(_.get(snapshot, partial, false)) == "true"}
-                  onChange={this.onToggle}
-                />
-              </EuiCallOut>}
-
+              {snapshot?.failed_shards && (
+                <EuiCallOut title="Restoring a partial snapshot" color="warning">
+                  <p>
+                    You are about to restore a partial snapshot. One or more shards may be missing in this
+                    <br />
+                    snapshot. Do you want to continue?
+                  </p>
+                  <EuiSpacer size="s" />
+                  <EuiCheckbox
+                    id={partial}
+                    label={<EuiText size="s">Allow restore partial snapshots</EuiText>}
+                    checked={String(_.get(snapshot, partial, false)) == "true"}
+                    onChange={this.onToggle}
+                  />
+                </EuiCallOut>
+              )}
             </EuiFlyoutBody>
 
             <EuiFlyoutFooter>
@@ -497,7 +510,8 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
                 action=""
                 onClickAction={this.onClickAction}
                 onClickCancel={onCloseFlyout}
-                disabledAction={!!restoreDisabled} />
+                disabledAction={!!restoreDisabled}
+              />
             </EuiFlyoutFooter>
           </>
         )}
@@ -505,3 +519,8 @@ export default class RestoreSnapshotFlyout extends Component<RestoreSnapshotProp
     );
   }
 }
+
+export default function (props: Omit<RestoreSnapshotProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  return <RestoreSnapshotFlyout {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/Snapshots/components/SnapshotFlyout/SnapshotFlyout.tsx b/public/pages/Snapshots/components/SnapshotFlyout/SnapshotFlyout.tsx
index bf3542a36..3141edd0d 100644
--- a/public/pages/Snapshots/components/SnapshotFlyout/SnapshotFlyout.tsx
+++ b/public/pages/Snapshots/components/SnapshotFlyout/SnapshotFlyout.tsx
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import {
   EuiButtonEmpty,
   EuiFlexGrid,
@@ -25,8 +25,10 @@ import * as H from "history";
 import { ROUTES } from "../../../../utils/constants";
 import InfoModal from "../../../SnapshotPolicyDetails/components/InfoModal";
 import { ModalConsumer } from "../../../../components/Modal";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
 
-interface SnapshotFlyoutProps {
+interface SnapshotFlyoutProps extends DataSourceMenuProperties {
   snapshotId: string;
   repository: string;
   snapshotManagementService: SnapshotManagementService;
@@ -34,17 +36,18 @@ interface SnapshotFlyoutProps {
   history: H.History;
 }
 
-interface SnapshotFlyoutState {
+interface SnapshotFlyoutState extends DataSourceMenuProperties {
   snapshot: GetSnapshot | null;
 }
 
-export default class SnapshotFlyout extends Component<SnapshotFlyoutProps, SnapshotFlyoutState> {
+export class SnapshotFlyout extends MDSEnabledComponent<SnapshotFlyoutProps, SnapshotFlyoutState> {
   static contextType = CoreServicesContext;
 
   constructor(props: SnapshotFlyoutProps) {
     super(props);
 
     this.state = {
+      ...this.state,
       snapshot: null,
     };
   }
@@ -54,6 +57,13 @@ export default class SnapshotFlyout extends Component<SnapshotFlyoutProps, Snaps
     await this.getSnapshot(snapshotId, repository);
   }
 
+  async componentDidUpdate(prevProps: SnapshotFlyoutProps, prevState: SnapshotFlyoutState) {
+    if (prevState.dataSourceId != this.state.dataSourceId) {
+      const { snapshotId, repository } = this.props;
+      await this.getSnapshot(snapshotId, repository);
+    }
+  }
+
   getSnapshot = async (snapshotId: string, repository: string) => {
     const { snapshotManagementService } = this.props;
     try {
@@ -153,3 +163,8 @@ export default class SnapshotFlyout extends Component<SnapshotFlyoutProps, Snaps
     );
   }
 }
+
+export default function (props: Omit<SnapshotFlyoutProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  return <SnapshotFlyout {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/pages/Snapshots/containers/Snapshots/Snapshots.tsx b/public/pages/Snapshots/containers/Snapshots/Snapshots.tsx
index 11855028c..ff98d7784 100644
--- a/public/pages/Snapshots/containers/Snapshots/Snapshots.tsx
+++ b/public/pages/Snapshots/containers/Snapshots/Snapshots.tsx
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { Component } from "react";
+import React, { Component, useContext } from "react";
 import _ from "lodash";
 import { RouteComponentProps } from "react-router-dom";
 import {
@@ -36,13 +36,16 @@ import ErrorModal from "../../../Snapshots/components/ErrorModal/ErrorModal";
 import DeleteModal from "../../../Repositories/components/DeleteModal/DeleteModal";
 import { getToasts } from "../../helper";
 import { snapshotStatusRender, truncateSpan } from "../../helper";
+import { DataSourceMenuContext, DataSourceMenuProperties } from "../../../../services/DataSourceMenuContext";
+import MDSEnabledComponent from "../../../../components/MDSEnabledComponent";
+import { useUpdateUrlWithDataSourceProperties } from "../../../../components/MDSEnabledComponent";
 
-interface SnapshotsProps extends RouteComponentProps {
+interface SnapshotsProps extends RouteComponentProps, DataSourceMenuProperties {
   snapshotManagementService: SnapshotManagementService;
   indexService: IndexService;
 }
 
-interface SnapshotsState {
+interface SnapshotsState extends DataSourceMenuProperties {
   snapshots: SnapshotsWithRepoAndPolicy[];
   existingPolicyNames: string[];
   loadingSnapshots: boolean;
@@ -67,7 +70,7 @@ interface SnapshotsState {
   isDeleteModalVisible: boolean;
 }
 
-export default class Snapshots extends Component<SnapshotsProps, SnapshotsState> {
+export class Snapshots extends MDSEnabledComponent<SnapshotsProps, SnapshotsState> {
   static contextType = CoreServicesContext;
   columns: EuiTableFieldDataColumnType<SnapshotsWithRepoAndPolicy>[];
   private tabsRef;
@@ -76,6 +79,7 @@ export default class Snapshots extends Component<SnapshotsProps, SnapshotsState>
     super(props);
 
     this.state = {
+      ...this.state,
       snapshots: [],
       existingPolicyNames: [],
       loadingSnapshots: false,
@@ -161,6 +165,20 @@ export default class Snapshots extends Component<SnapshotsProps, SnapshotsState>
     await this.getSnapshots();
   }
 
+  async componentDidUpdate(prevProps: SnapshotsProps, prevState: SnapshotsState) {
+    const prevQuery = Snapshots.getQueryObjectFromState(prevState);
+    const currQuery = Snapshots.getQueryObjectFromState(this.state);
+    if (!_.isEqual(prevQuery, currQuery)) {
+      await this.getSnapshots();
+    }
+  }
+
+  static getQueryObjectFromState({ dataSourceId, multiDataSourceEnabled }: SnapshotsState) {
+    return {
+      ...(multiDataSourceEnabled ? { dataSourceId } : {}),
+    };
+  }
+
   getSnapshots = async () => {
     this.setState({ loadingSnapshots: true, message: "Loading snapshots..." });
 
@@ -558,3 +576,9 @@ export default class Snapshots extends Component<SnapshotsProps, SnapshotsState>
       .join(", ");
   };
 }
+
+export default function (props: Omit<SnapshotsProps, keyof DataSourceMenuProperties>) {
+  const dataSourceMenuProps = useContext(DataSourceMenuContext);
+  useUpdateUrlWithDataSourceProperties();
+  return <Snapshots {...props} {...dataSourceMenuProps} />;
+}
diff --git a/public/services/SnapshotManagementService.ts b/public/services/SnapshotManagementService.ts
index 74d35b10d..6873fb36c 100644
--- a/public/services/SnapshotManagementService.ts
+++ b/public/services/SnapshotManagementService.ts
@@ -19,59 +19,68 @@ import {
 } from "../../server/models/interfaces";
 import { ServerResponse } from "../../server/models/types";
 import { DocumentSMPolicy, DocumentSMPolicyWithMetadata, SMPolicy, Snapshot } from "../../models/interfaces";
+import { MDSEnabledClientService } from "./MDSEnabledClientService";
 
-export default class SnapshotManagementService {
-  httpClient: HttpSetup;
-
-  constructor(httpClient: HttpSetup) {
-    this.httpClient = httpClient;
-  }
-
+export default class SnapshotManagementService extends MDSEnabledClientService {
   getAllSnapshotsWithPolicy = async (): Promise<ServerResponse<GetSnapshotsResponse>> => {
     let url = `..${NODE_API._SNAPSHOTS}`;
-    const response = (await this.httpClient.get(url)) as ServerResponse<GetSnapshotsResponse>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.get(url, params)) as ServerResponse<GetSnapshotsResponse>;
     return response;
   };
 
   getSnapshot = async (snapshotId: string, repository: string): Promise<ServerResponse<GetSnapshot>> => {
     let url = `..${NODE_API._SNAPSHOTS}/${snapshotId}`;
-    const response = (await this.httpClient.get(url, { query: { repository } })) as ServerResponse<GetSnapshot>;
+    const query = this.patchQueryObjectWithDataSourceId({ repository });
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.get(url, params)) as ServerResponse<GetSnapshot>;
     return response;
   };
 
   deleteSnapshot = async (snapshotId: string, repository: string): Promise<ServerResponse<AcknowledgedResponse>> => {
     let url = `..${NODE_API._SNAPSHOTS}/${snapshotId}`;
-    const response = (await this.httpClient.delete(url, { query: { repository } })) as ServerResponse<AcknowledgedResponse>;
+    const query = this.patchQueryObjectWithDataSourceId({ repository });
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.delete(url, params)) as ServerResponse<AcknowledgedResponse>;
     return response;
   };
 
   createSnapshot = async (snapshotId: string, repository: string, snapshot: Snapshot): Promise<ServerResponse<CreateSnapshotResponse>> => {
     let url = `..${NODE_API._SNAPSHOTS}/${snapshotId}`;
+    const query = this.patchQueryObjectWithDataSourceId({ repository });
+    const params = query ? { query } : {};
     const response = (await this.httpClient.put(url, {
-      query: { repository },
       body: JSON.stringify(snapshot),
+      ...params,
     })) as ServerResponse<CreateSnapshotResponse>;
     return response;
   };
 
   restoreSnapshot = async (snapshotId: string, repository: string, options: object): Promise<ServerResponse<RestoreSnapshotResponse>> => {
     let url = `..${NODE_API._SNAPSHOTS}/${snapshotId}`;
+    const query = this.patchQueryObjectWithDataSourceId({ repository });
+    const params = query ? { query } : {};
     const response = (await this.httpClient.post(url, {
-      query: { repository },
       body: JSON.stringify(options),
+      ...params,
     })) as ServerResponse<RestoreSnapshotResponse>;
     return response;
   };
 
   getIndexRecovery = async (): Promise<ServerResponse<GetIndexRecoveryResponse>> => {
     const url = NODE_API._RECOVERY;
-    const response = (await this.httpClient.get(url)) as ServerResponse<GetIndexRecoveryResponse>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.get(url, params)) as ServerResponse<GetIndexRecoveryResponse>;
     return response;
   };
 
   createPolicy = async (policyId: string, policy: SMPolicy): Promise<ServerResponse<DocumentSMPolicy>> => {
     let url = `..${NODE_API.SMPolicies}/${policyId}`;
-    const response = (await this.httpClient.post(url, { body: JSON.stringify(policy) })) as ServerResponse<DocumentSMPolicy>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.post(url, { body: JSON.stringify(policy), ...params })) as ServerResponse<DocumentSMPolicy>;
     return response;
   };
 
@@ -82,70 +91,92 @@ export default class SnapshotManagementService {
     primaryTerm: number
   ): Promise<ServerResponse<DocumentSMPolicy>> => {
     let url = `..${NODE_API.SMPolicies}/${policyId}`;
+    const query = this.patchQueryObjectWithDataSourceId({ seqNo, primaryTerm });
+    const params = query ? { query } : {};
     const response = (await this.httpClient.put(url, {
-      query: { seqNo, primaryTerm },
       body: JSON.stringify(policy),
+      ...params,
     })) as ServerResponse<DocumentSMPolicy>;
     return response;
   };
 
   getPolicies = async (queryObject: HttpFetchQuery): Promise<ServerResponse<GetSMPoliciesResponse>> => {
     let url = `..${NODE_API.SMPolicies}`;
-    const response = (await this.httpClient.get(url, { query: queryObject })) as ServerResponse<GetSMPoliciesResponse>;
+    const query = this.patchQueryObjectWithDataSourceId(queryObject);
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.get(url, params)) as ServerResponse<GetSMPoliciesResponse>;
     return response;
   };
 
   getPolicy = async (policyId: string): Promise<ServerResponse<DocumentSMPolicyWithMetadata>> => {
     const url = `..${NODE_API.SMPolicies}/${policyId}`;
-    const response = (await this.httpClient.get(url)) as ServerResponse<DocumentSMPolicyWithMetadata>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.get(url, params)) as ServerResponse<DocumentSMPolicyWithMetadata>;
     return response;
   };
 
   deletePolicy = async (policyId: string): Promise<ServerResponse<boolean>> => {
     const url = `..${NODE_API.SMPolicies}/${policyId}`;
-    const response = (await this.httpClient.delete(url)) as ServerResponse<boolean>;
+    const query = this.patchQueryObjectWithDataSourceId();
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.delete(url, params)) as ServerResponse<boolean>;
     return response;
   };
 
   startPolicy = async (policyId: string): Promise<ServerResponse<boolean>> => {
     const url = `..${NODE_API.SMPolicies}/${policyId}/_start`;
-    const response = (await this.httpClient.post(url)) as ServerResponse<boolean>;
+    const query = this.patchQueryObjectWithDataSourceId();
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.post(url, params)) as ServerResponse<boolean>;
     return response;
   };
 
   stopPolicy = async (policyId: string): Promise<ServerResponse<boolean>> => {
     const url = `..${NODE_API.SMPolicies}/${policyId}/_stop`;
-    const response = (await this.httpClient.post(url)) as ServerResponse<boolean>;
+    const query = this.patchQueryObjectWithDataSourceId();
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.post(url, params)) as ServerResponse<boolean>;
     return response;
   };
 
   catRepositories = async (): Promise<ServerResponse<CatRepository[]>> => {
     const url = `..${NODE_API._REPOSITORIES}`;
-    const response = (await this.httpClient.get(url)) as ServerResponse<CatRepository[]>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.get(url, params)) as ServerResponse<CatRepository[]>;
     return response;
   };
 
   catSnapshotIndices = async (indices: string): Promise<ServerResponse<CatIndex[]>> => {
     const url = `..${NODE_API._INDICES}/${indices}`;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
     const response = (await this.httpClient.get(url)) as ServerResponse<CatIndex[]>;
     return response;
   };
 
   getRepository = async (repo: string): Promise<ServerResponse<any>> => {
     const url = `..${NODE_API._REPOSITORIES}/${repo}`;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
     const response = (await this.httpClient.get(url)) as ServerResponse<any>;
     return response;
   };
 
   createRepository = async (repo: string, createRepoBody: CreateRepositoryBody): Promise<ServerResponse<any>> => {
     const url = `..${NODE_API._REPOSITORIES}/${repo}`;
-    const response = (await this.httpClient.put(url, { body: JSON.stringify(createRepoBody) })) as ServerResponse<any>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.put(url, { body: JSON.stringify(createRepoBody), ...params })) as ServerResponse<any>;
     return response;
   };
 
   deleteRepository = async (repo: string): Promise<ServerResponse<any>> => {
     const url = `..${NODE_API._REPOSITORIES}/${repo}`;
-    const response = (await this.httpClient.delete(url)) as ServerResponse<any>;
+    const query = this.patchQueryObjectWithDataSourceId({});
+    const params = query ? { query } : {};
+    const response = (await this.httpClient.delete(url, params)) as ServerResponse<any>;
     return response;
   };
 }
diff --git a/server/plugin.ts b/server/plugin.ts
index 387fada15..af0b69706 100644
--- a/server/plugin.ts
+++ b/server/plugin.ts
@@ -55,7 +55,7 @@ export class IndexPatternManagementPlugin implements Plugin<IndexManagementPlugi
     const rollupService = new RollupService(osDriver, dataSourceEnabled);
     const transformService = new TransformService(osDriver, dataSourceEnabled);
     const notificationService = new NotificationService(osDriver, dataSourceEnabled);
-    const snapshotManagementService = new SnapshotManagementService(osDriver);
+    const snapshotManagementService = new SnapshotManagementService(osDriver, dataSourceEnabled);
     const commonService = new CommonService(osDriver, dataSourceEnabled);
     const aliasService = new AliasServices(osDriver, dataSourceEnabled);
     const services: NodeServices = {
@@ -86,7 +86,7 @@ export class IndexPatternManagementPlugin implements Plugin<IndexManagementPlugi
     rollups(services, router, dataSourceEnabled);
     transforms(services, router, dataSourceEnabled);
     notifications(services, router, dataSourceEnabled);
-    snapshotManagement(services, router);
+    snapshotManagement(services, router, dataSourceEnabled);
     common(services, router, dataSourceEnabled);
     aliases(services, router, dataSourceEnabled);
 
diff --git a/server/routes/snapshotManagement.ts b/server/routes/snapshotManagement.ts
index 5e452eafa..50b9153f9 100644
--- a/server/routes/snapshotManagement.ts
+++ b/server/routes/snapshotManagement.ts
@@ -8,13 +8,17 @@ import { IRouter } from "../../../../src/core/server";
 import { NODE_API } from "../../utils/constants";
 import { schema } from "@osd/config-schema";
 
-export default function (services: NodeServices, router: IRouter) {
+export default function (services: NodeServices, router: IRouter, dataSourceEnabled: boolean) {
   const { snapshotManagementService } = services;
 
   router.get(
     {
       path: NODE_API._SNAPSHOTS,
-      validate: {},
+      validate: {
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
+      },
     },
     snapshotManagementService.getAllSnapshotsWithPolicy
   );
@@ -28,6 +32,7 @@ export default function (services: NodeServices, router: IRouter) {
         }),
         query: schema.object({
           repository: schema.string(),
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
         }),
       },
     },
@@ -43,6 +48,7 @@ export default function (services: NodeServices, router: IRouter) {
         }),
         query: schema.object({
           repository: schema.string(),
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
         }),
       },
     },
@@ -58,6 +64,7 @@ export default function (services: NodeServices, router: IRouter) {
         }),
         query: schema.object({
           repository: schema.string(),
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
         }),
         body: schema.any(),
       },
@@ -74,6 +81,7 @@ export default function (services: NodeServices, router: IRouter) {
         }),
         query: schema.object({
           repository: schema.string(),
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
         }),
         body: schema.any(),
       },
@@ -89,6 +97,9 @@ export default function (services: NodeServices, router: IRouter) {
           id: schema.string(),
         }),
         body: schema.any(),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.createPolicy
@@ -104,6 +115,7 @@ export default function (services: NodeServices, router: IRouter) {
         query: schema.object({
           seqNo: schema.maybe(schema.number()),
           primaryTerm: schema.maybe(schema.number()),
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
         }),
         body: schema.any(),
       },
@@ -122,6 +134,7 @@ export default function (services: NodeServices, router: IRouter) {
           sortField: schema.string(),
           sortOrder: schema.string(),
           queryString: schema.string(),
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
         }),
       },
     },
@@ -135,6 +148,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           id: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.getPolicy
@@ -147,6 +163,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           id: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.deletePolicy
@@ -159,6 +178,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           id: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.startPolicy
@@ -171,6 +193,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           id: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.stopPolicy
@@ -179,7 +204,11 @@ export default function (services: NodeServices, router: IRouter) {
   router.get(
     {
       path: NODE_API._REPOSITORIES,
-      validate: {},
+      validate: {
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
+      },
     },
     snapshotManagementService.catRepositoriesWithSnapshotCount
   );
@@ -191,6 +220,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           indices: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.catSnapshotIndices
@@ -199,7 +231,11 @@ export default function (services: NodeServices, router: IRouter) {
   router.get(
     {
       path: NODE_API._RECOVERY,
-      validate: {},
+      validate: {
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
+      },
     },
     snapshotManagementService.getIndexRecovery
   );
@@ -211,6 +247,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           id: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.deleteRepository
@@ -223,6 +262,9 @@ export default function (services: NodeServices, router: IRouter) {
         params: schema.object({
           id: schema.string(),
         }),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.getRepository
@@ -236,6 +278,9 @@ export default function (services: NodeServices, router: IRouter) {
           id: schema.string(),
         }),
         body: schema.any(),
+        query: schema.object({
+          ...(dataSourceEnabled ? { dataSourceId: schema.string() } : {}),
+        }),
       },
     },
     snapshotManagementService.createRepository
diff --git a/server/services/SnapshotManagementService.ts b/server/services/SnapshotManagementService.ts
index fb236dc8a..b06581fe4 100644
--- a/server/services/SnapshotManagementService.ts
+++ b/server/services/SnapshotManagementService.ts
@@ -28,14 +28,9 @@ import {
   RestoreSnapshotResponse,
 } from "../models/interfaces";
 import { FailedServerResponse, ServerResponse } from "../models/types";
+import { MDSEnabledClientService } from "./MDSEnabledClientService";
 
-export default class SnapshotManagementService {
-  osDriver: ILegacyCustomClusterClient;
-
-  constructor(osDriver: ILegacyCustomClusterClient) {
-    this.osDriver = osDriver;
-  }
-
+export default class SnapshotManagementService extends MDSEnabledClientService {
   getAllSnapshotsWithPolicy = async (
     context: RequestHandlerContext,
     request: OpenSearchDashboardsRequest,
@@ -57,7 +52,7 @@ export default class SnapshotManagementService {
         });
       }
 
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
       let snapshots: CatSnapshotWithRepoAndPolicy[] = [];
       for (let i = 0; i < repositories.length; i++) {
         const res: GetSnapshotResponse = await callWithRequest("snapshot.get", {
@@ -109,12 +104,12 @@ export default class SnapshotManagementService {
       const { repository } = request.query as {
         repository: string;
       };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: GetSnapshotResponse = await callWithRequest("snapshot.get", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: GetSnapshotResponse = (await callWithRequest("snapshot.get", {
         repository: repository,
         snapshot: `${id}`,
         ignore_unavailable: true,
-      });
+      })) as GetSnapshotResponse;
 
       return response.custom({
         statusCode: 200,
@@ -140,11 +135,11 @@ export default class SnapshotManagementService {
       const { repository } = request.query as {
         repository: string;
       };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const resp: AcknowledgedResponse = await callWithRequest("snapshot.delete", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const resp: AcknowledgedResponse = (await callWithRequest("snapshot.delete", {
         repository: repository,
         snapshot: `${id}`,
-      });
+      })) as AcknowledgedResponse;
 
       return response.custom({
         statusCode: 200,
@@ -175,8 +170,8 @@ export default class SnapshotManagementService {
         snapshot: id,
         body: JSON.stringify(request.body),
       };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const resp: CreateSnapshotResponse = await callWithRequest("snapshot.create", params);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const resp: CreateSnapshotResponse = (await callWithRequest("snapshot.create", params)) as CreateSnapshotResponse;
 
       return response.custom({
         statusCode: 200,
@@ -207,8 +202,8 @@ export default class SnapshotManagementService {
         snapshot: id,
         body: JSON.stringify(request.body),
       };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const resp: RestoreSnapshotResponse = await callWithRequest("snapshot.restore", params);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const resp: RestoreSnapshotResponse = (await callWithRequest("snapshot.restore", params)) as RestoreSnapshotResponse;
 
       return response.custom({
         statusCode: 200,
@@ -234,7 +229,7 @@ export default class SnapshotManagementService {
         body: JSON.stringify(request.body),
       };
 
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
       const rawRes = await callWithRequest("ism.createSMPolicy", params);
       const res: DocumentSMPolicy = {
         seqNo: rawRes._seq_no,
@@ -270,7 +265,7 @@ export default class SnapshotManagementService {
         body: JSON.stringify(request.body),
       };
 
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
       const rawRes = await callWithRequest("ism.updateSMPolicy", params);
       const res: DocumentSMPolicy = {
         seqNo: rawRes._seq_no,
@@ -305,7 +300,7 @@ export default class SnapshotManagementService {
         queryString: string;
       };
 
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
       let params = {
         from,
         size,
@@ -351,7 +346,7 @@ export default class SnapshotManagementService {
     try {
       const { id } = request.params as { id: string };
       const params = { id: id };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
       const getResponse = await callWithRequest("ism.getSMPolicy", params);
       const metadata = await callWithRequest("ism.explainSnapshotPolicy", params);
       const documentPolicy = {
@@ -390,8 +385,8 @@ export default class SnapshotManagementService {
     try {
       const { id } = request.params as { id: string };
       const params = { policyId: id };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const deletePolicyResponse: DeletePolicyResponse = await callWithRequest("ism.deleteSMPolicy", params);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const deletePolicyResponse: DeletePolicyResponse = (await callWithRequest("ism.deleteSMPolicy", params)) as DeletePolicyResponse;
       if (deletePolicyResponse.result !== "deleted") {
         return response.custom({
           statusCode: 200,
@@ -421,8 +416,8 @@ export default class SnapshotManagementService {
     try {
       const { id } = request.params as { id: string };
       const params = { id: id };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const resp: AcknowledgedResponse = await callWithRequest("ism.startSnapshotPolicy", params);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const resp: AcknowledgedResponse = (await callWithRequest("ism.startSnapshotPolicy", params)) as AcknowledgedResponse;
       if (resp.acknowledged) {
         return response.custom({
           statusCode: 200,
@@ -447,8 +442,8 @@ export default class SnapshotManagementService {
     try {
       const { id } = request.params as { id: string };
       const params = { id: id };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const resp: AcknowledgedResponse = await callWithRequest("ism.stopSnapshotPolicy", params);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const resp: AcknowledgedResponse = (await callWithRequest("ism.stopSnapshotPolicy", params)) as AcknowledgedResponse;
       if (resp.acknowledged) {
         return response.custom({
           statusCode: 200,
@@ -471,10 +466,10 @@ export default class SnapshotManagementService {
     response: OpenSearchDashboardsResponseFactory
   ): Promise<IOpenSearchDashboardsResponse<ServerResponse<CatRepository[]>>> => {
     try {
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: CatRepository[] = await callWithRequest("cat.repositories", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: CatRepository[] = (await callWithRequest("cat.repositories", {
         format: "json",
-      });
+      })) as CatRepository[];
       return response.custom({
         statusCode: 200,
         body: {
@@ -493,10 +488,10 @@ export default class SnapshotManagementService {
     response: OpenSearchDashboardsResponseFactory
   ): Promise<IOpenSearchDashboardsResponse<ServerResponse<GetIndexRecoveryResponse>>> => {
     try {
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: GetIndexRecoveryResponse = await callWithRequest("indices.recovery", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: GetIndexRecoveryResponse = (await callWithRequest("indices.recovery", {
         format: "json",
-      });
+      })) as GetIndexRecoveryResponse;
       return response.custom({
         statusCode: 200,
         body: {
@@ -515,10 +510,10 @@ export default class SnapshotManagementService {
     response: OpenSearchDashboardsResponseFactory
   ): Promise<IOpenSearchDashboardsResponse<ServerResponse<CatSnapshotIndex[]>>> => {
     try {
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: CatSnapshotIndex[] = await callWithRequest("cat.indices", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: CatSnapshotIndex[] = (await callWithRequest("cat.indices", {
         format: "json",
-      });
+      })) as CatSnapshotIndex[];
 
       return response.custom({
         statusCode: 200,
@@ -538,17 +533,17 @@ export default class SnapshotManagementService {
     response: OpenSearchDashboardsResponseFactory
   ): Promise<IOpenSearchDashboardsResponse<ServerResponse<CatRepository[]>>> => {
     try {
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: CatRepository[] = await callWithRequest("cat.repositories", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: CatRepository[] = (await callWithRequest("cat.repositories", {
         format: "json",
-      });
+      })) as CatRepository[];
 
       for (let i = 0; i < res.length; i++) {
-        const getSnapshotRes: GetSnapshotResponse = await callWithRequest("snapshot.get", {
+        const getSnapshotRes: GetSnapshotResponse = (await callWithRequest("snapshot.get", {
           repository: res[i].id,
           snapshot: "_all",
           ignore_unavailable: true,
-        });
+        })) as GetSnapshotResponse;
         res[i].snapshotCount = getSnapshotRes.snapshots.length;
       }
 
@@ -571,10 +566,10 @@ export default class SnapshotManagementService {
   ): Promise<IOpenSearchDashboardsResponse<ServerResponse<AcknowledgedResponse>>> => {
     try {
       const { id } = request.params as { id: string };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: AcknowledgedResponse = await callWithRequest("snapshot.deleteRepository", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: AcknowledgedResponse = (await callWithRequest("snapshot.deleteRepository", {
         repository: id,
-      });
+      })) as AcknowledgedResponse;
       return response.custom({
         statusCode: 200,
         body: {
@@ -594,10 +589,10 @@ export default class SnapshotManagementService {
   ): Promise<IOpenSearchDashboardsResponse<ServerResponse<GetRepositoryResponse>>> => {
     try {
       const { id } = request.params as { id: string };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: GetRepositoryResponse = await callWithRequest("snapshot.getRepository", {
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: GetRepositoryResponse = (await callWithRequest("snapshot.getRepository", {
         repository: id,
-      });
+      })) as GetRepositoryResponse;
       return response.custom({
         statusCode: 200,
         body: {
@@ -621,8 +616,8 @@ export default class SnapshotManagementService {
         repository: id,
         body: JSON.stringify(request.body),
       };
-      const { callAsCurrentUser: callWithRequest } = this.osDriver.asScoped(request);
-      const res: AcknowledgedResponse = await callWithRequest("snapshot.createRepository", params);
+      const callWithRequest = this.getClientBasedOnDataSource(context, request);
+      const res: AcknowledgedResponse = (await callWithRequest("snapshot.createRepository", params)) as AcknowledgedResponse;
       return response.custom({
         statusCode: 200,
         body: {