From 08cb85c9365106cf6e9c9b3c8163dbee486812bd Mon Sep 17 00:00:00 2001
From: Ashish Singh <ssashish@amazon.com>
Date: Tue, 19 Mar 2024 13:20:55 +0530
Subject: [PATCH] Use cluster default remote store path type as fallback during
 snapshot restore

Signed-off-by: Ashish Singh <ssashish@amazon.com>
---
 .../metadata/MetadataCreateIndexService.java  | 22 ++++++++++++++-----
 .../opensearch/snapshots/RestoreService.java  |  2 +-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java b/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java
index f117d1a4a11a2..f6a14d8ec9d63 100644
--- a/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java
+++ b/server/src/main/java/org/opensearch/cluster/metadata/MetadataCreateIndexService.java
@@ -553,11 +553,7 @@ IndexMetadata buildAndValidateTemporaryIndexMetadata(
         tmpImdBuilder.setRoutingNumShards(routingNumShards);
         tmpImdBuilder.settings(indexSettings);
         tmpImdBuilder.system(isSystem);
-
-        if (remoteStorePathResolver != null) {
-            String pathType = remoteStorePathResolver.resolveType().toString();
-            tmpImdBuilder.putCustom(IndexMetadata.REMOTE_STORE_CUSTOM_KEY, Map.of(RemoteStorePathType.NAME, pathType));
-        }
+        addRemoteCustomData(tmpImdBuilder);
 
         // Set up everything, now locally create the index to see that things are ok, and apply
         IndexMetadata tempMetadata = tmpImdBuilder.build();
@@ -566,6 +562,22 @@ IndexMetadata buildAndValidateTemporaryIndexMetadata(
         return tempMetadata;
     }
 
+    public void addRemoteCustomData(IndexMetadata.Builder tmpImdBuilder) {
+        if (remoteStorePathResolver != null) {
+            // It is possible that remote custom data exists already. In such cases, we need to only update the path type
+            // in the remote store custom data map.
+            Map<String, String> existingRemoteCustomData = tmpImdBuilder.removeCustom(IndexMetadata.REMOTE_STORE_CUSTOM_KEY);
+            Map<String, String> remoteCustomData = existingRemoteCustomData == null
+                ? new HashMap<>()
+                : new HashMap<>(existingRemoteCustomData);
+            // Determine the path type for use using the remoteStorePathResolver.
+            String newPathType = remoteStorePathResolver.resolveType().toString();
+            String oldPathType = remoteCustomData.put(RemoteStorePathType.NAME, newPathType);
+            logger.trace(() -> new ParameterizedMessage("Added new path type {}, replaced old path type {}", newPathType, oldPathType));
+            tmpImdBuilder.putCustom(IndexMetadata.REMOTE_STORE_CUSTOM_KEY, remoteCustomData);
+        }
+    }
+
     private ClusterState applyCreateIndexRequestWithV1Templates(
         final ClusterState currentState,
         final CreateIndexClusterStateUpdateRequest request,
diff --git a/server/src/main/java/org/opensearch/snapshots/RestoreService.java b/server/src/main/java/org/opensearch/snapshots/RestoreService.java
index bf2c7fc74be92..e5ac604e0a5e3 100644
--- a/server/src/main/java/org/opensearch/snapshots/RestoreService.java
+++ b/server/src/main/java/org/opensearch/snapshots/RestoreService.java
@@ -216,7 +216,6 @@ public RestoreService(
 
         // Task is onboarded for throttling, it will get retried from associated TransportClusterManagerNodeAction.
         restoreSnapshotTaskKey = clusterService.registerClusterManagerTask(ClusterManagerTaskKeys.RESTORE_SNAPSHOT_KEY, true);
-
     }
 
     /**
@@ -452,6 +451,7 @@ public ClusterState execute(ClusterState currentState) {
                                             .put(snapshotIndexMetadata.getSettings())
                                             .put(IndexMetadata.SETTING_INDEX_UUID, UUIDs.randomBase64UUID())
                                     );
+                                    createIndexService.addRemoteCustomData(indexMdBuilder);
                                     shardLimitValidator.validateShardLimit(
                                         renamedIndexName,
                                         snapshotIndexMetadata.getSettings(),