Conforms to changes in core

Signed-off-by: Darshit Chanpura <[email protected]>
DarshitChanpura · Dec 31, 2024 · 413bb0b · 413bb0b
1 parent d8969e5
commit 413bb0b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 16 deletions.
diff --git a/...rc/main/java/org/opensearch/sample/actions/access/revoke/RevokeResourceAccessRequest.java b/...rc/main/java/org/opensearch/sample/actions/access/revoke/RevokeResourceAccessRequest.java
@@ -12,7 +12,7 @@
 import java.util.Map;
 import java.util.Set;
 
-import org.opensearch.accesscontrol.resources.EntityType;
+import org.opensearch.accesscontrol.resources.RecipientType;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -22,18 +22,18 @@
 public class RevokeResourceAccessRequest extends ActionRequest {
 
     private final String resourceId;
-    private final Map<EntityType, Set<String>> revokeAccess;
+    private final Map<RecipientType, Set<String>> revokeAccess;
     private final Set<String> scopes;
 
-    public RevokeResourceAccessRequest(String resourceId, Map<EntityType, Set<String>> revokeAccess, Set<String> scopes) {
+    public RevokeResourceAccessRequest(String resourceId, Map<RecipientType, Set<String>> revokeAccess, Set<String> scopes) {
         this.resourceId = resourceId;
         this.revokeAccess = revokeAccess;
         this.scopes = scopes;
     }
 
     public RevokeResourceAccessRequest(StreamInput in) throws IOException {
         this.resourceId = in.readString();
-        this.revokeAccess = in.readMap(input -> EntityType.valueOf(input.readString()), input -> input.readSet(StreamInput::readString));
+        this.revokeAccess = in.readMap(input -> new RecipientType(input.readString()), input -> input.readSet(StreamInput::readString));
         this.scopes = in.readSet(StreamInput::readString);
     }
 
@@ -42,7 +42,7 @@ public void writeTo(final StreamOutput out) throws IOException {
         out.writeString(resourceId);
         out.writeMap(
             revokeAccess,
-            (streamOutput, entityType) -> streamOutput.writeString(entityType.name()),
+            (streamOutput, recipientType) -> streamOutput.writeString(recipientType.getType()),
             StreamOutput::writeStringCollection
         );
         out.writeStringCollection(scopes);
@@ -62,7 +62,7 @@ public String getResourceId() {
         return resourceId;
     }
 
-    public Map<EntityType, Set<String>> getRevokeAccess() {
+    public Map<RecipientType, Set<String>> getRevokeAccess() {
         return revokeAccess;
     }
 

diff --git a/...main/java/org/opensearch/sample/actions/access/revoke/RevokeResourceAccessRestAction.java b/...main/java/org/opensearch/sample/actions/access/revoke/RevokeResourceAccessRestAction.java
@@ -12,7 +12,8 @@
 import java.util.*;
 import java.util.stream.Collectors;
 
-import org.opensearch.accesscontrol.resources.EntityType;
+import org.opensearch.accesscontrol.resources.RecipientType;
+import org.opensearch.accesscontrol.resources.RecipientTypeRegistry;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.rest.BaseRestHandler;
@@ -46,15 +47,9 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         String resourceId = (String) source.get("resource_id");
         @SuppressWarnings("unchecked")
         Map<String, Set<String>> revokeSource = (Map<String, Set<String>>) source.get("entities");
-        Map<EntityType, Set<String>> revoke = revokeSource.entrySet().stream().collect(Collectors.toMap(entry -> {
-            try {
-                return EntityType.fromValue(entry.getKey());
-            } catch (IllegalArgumentException e) {
-                throw new IllegalArgumentException(
-                    "Invalid entity type: " + entry.getKey() + ". Valid values are: " + Arrays.toString(EntityType.values())
-                );
-            }
-        }, Map.Entry::getValue));
+        Map<RecipientType, Set<String>> revoke = revokeSource.entrySet()
+            .stream()
+            .collect(Collectors.toMap(entry -> RecipientTypeRegistry.fromValue(entry.getKey()), Map.Entry::getValue));
         @SuppressWarnings("unchecked")
         Set<String> scopes = new HashSet<>(source.containsKey("scopes") ? (List<String>) source.get("scopes") : List.of());
         final RevokeResourceAccessRequest revokeResourceAccessRequest = new RevokeResourceAccessRequest(resourceId, revoke, scopes);