Respond to code review feedback

Signed-off-by: Craig Perkins <[email protected]>
cwperks · Jan 13, 2025 · c8fab69 · c8fab69
1 parent 9439a0e
commit c8fab69
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 15 deletions.
diff --git a/plugins/identity-shiro/src/main/java/org/opensearch/identity/shiro/ShiroIdentityPlugin.java b/plugins/identity-shiro/src/main/java/org/opensearch/identity/shiro/ShiroIdentityPlugin.java
@@ -23,8 +23,8 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.env.Environment;
 import org.opensearch.env.NodeEnvironment;
+import org.opensearch.identity.RunAsSystemClient;
 import org.opensearch.identity.Subject;
-import org.opensearch.identity.noop.RunAsSystemClient;
 import org.opensearch.identity.tokens.AuthToken;
 import org.opensearch.identity.tokens.TokenManager;
 import org.opensearch.plugins.ActionPlugin;

diff --git a/...arch/identity/noop/RunAsSystemClient.java → ...pensearch/identity/RunAsSystemClient.java b/...arch/identity/noop/RunAsSystemClient.java → ...pensearch/identity/RunAsSystemClient.java
@@ -6,7 +6,7 @@
  * compatible open source license.
  */
 
-package org.opensearch.identity.noop;
+package org.opensearch.identity;
 
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionType;
@@ -36,21 +36,11 @@ public RunAsSystemClient(Client delegate) {
     protected <Request extends ActionRequest, Response extends ActionResponse> void doExecute(
         ActionType<Response> action,
         Request request,
-        ActionListener<Response> actionListener
+        ActionListener<Response> listener
     ) {
         ThreadContext threadContext = threadPool().getThreadContext();
-
         try (ThreadContext.StoredContext ctx = threadContext.stashContext()) {
-
-            ActionListener<Response> wrappedListener = ActionListener.wrap(r -> {
-                ctx.restore();
-                actionListener.onResponse(r);
-            }, e -> {
-                ctx.restore();
-                actionListener.onFailure(e);
-            });
-
-            super.doExecute(action, request, wrappedListener);
+            super.doExecute(action, request, ActionListener.runBefore(listener, ctx::restore));
         }
     }
 }
diff --git a/server/src/main/java/org/opensearch/identity/noop/NoopIdentityPlugin.java b/server/src/main/java/org/opensearch/identity/noop/NoopIdentityPlugin.java
@@ -9,6 +9,7 @@
 package org.opensearch.identity.noop;
 
 import org.opensearch.client.Client;
+import org.opensearch.identity.RunAsSystemClient;
 import org.opensearch.identity.Subject;
 import org.opensearch.identity.tokens.TokenManager;
 import org.opensearch.plugins.IdentityPlugin;

diff --git a/server/src/test/java/org/opensearch/identity/RunAsSystemClientTests.java b/server/src/test/java/org/opensearch/identity/RunAsSystemClientTests.java
@@ -12,7 +12,6 @@
 import org.opensearch.action.admin.cluster.health.ClusterHealthResponse;
 import org.opensearch.client.Client;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.identity.noop.RunAsSystemClient;
 import org.opensearch.test.OpenSearchSingleNodeTestCase;
 import org.junit.Before;