Move getRequestOptions to ClientUtil from DriverRemoteConnection

Bit-Quill · Mar 9, 2024 · 3532289 · 3532289
1 parent 9f5d222
commit 3532289
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 31 deletions.
diff --git a/gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/Client.java b/gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/Client.java
@@ -22,11 +22,12 @@
 import org.apache.tinkerpop.gremlin.driver.exception.ConnectionException;
 import org.apache.tinkerpop.gremlin.driver.exception.NoHostAvailableException;
 import org.apache.tinkerpop.gremlin.driver.message.RequestMessage;
-import org.apache.tinkerpop.gremlin.driver.remote.DriverRemoteConnection;
 import org.apache.tinkerpop.gremlin.process.traversal.Bytecode;
 import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
 import org.apache.tinkerpop.gremlin.process.traversal.TraversalSource;
 import org.apache.tinkerpop.gremlin.process.traversal.Traverser;
+import org.apache.tinkerpop.gremlin.process.traversal.strategy.decoration.OptionsStrategy;
+import org.apache.tinkerpop.gremlin.process.traversal.util.BytecodeHelper;
 import org.apache.tinkerpop.gremlin.structure.Graph;
 import org.apache.tinkerpop.gremlin.util.iterator.IteratorUtils;
 import org.slf4j.Logger;
@@ -53,6 +54,12 @@
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
+import static org.apache.tinkerpop.gremlin.driver.Client.ClientUtil.getRequestOptions;
+import static org.apache.tinkerpop.gremlin.driver.Tokens.ARGS_BATCH_SIZE;
+import static org.apache.tinkerpop.gremlin.driver.Tokens.ARGS_EVAL_TIMEOUT;
+import static org.apache.tinkerpop.gremlin.driver.Tokens.ARGS_USER_AGENT;
+import static org.apache.tinkerpop.gremlin.driver.Tokens.REQUEST_ID;
+
 /**
  * A {@code Client} is constructed from a {@link Cluster} and represents a way to send messages to Gremlin Server.
  * This class itself is a base class as there are different implementations that provide differing kinds of
@@ -642,7 +649,7 @@ public static class AliasClusteredClient extends Client {
 
         @Override
         public CompletableFuture<ResultSet> submitAsync(final Bytecode bytecode) {
-            return submitAsync(bytecode, DriverRemoteConnection.getRequestOptions(bytecode));
+            return submitAsync(bytecode, getRequestOptions(bytecode));
         }
 
         @Override
@@ -914,6 +921,26 @@ public Settings create() {
         }
     }
 
+    public static class ClientUtil {
+        public static RequestOptions getRequestOptions(final Bytecode bytecode) {
+            final Iterator<OptionsStrategy> itty = BytecodeHelper.findStrategies(bytecode, OptionsStrategy.class);
+            final RequestOptions.Builder builder = RequestOptions.build();
+            while (itty.hasNext()) {
+                final OptionsStrategy optionsStrategy = itty.next();
+                final Map<String,Object> options = optionsStrategy.getOptions();
+                if (options.containsKey(ARGS_EVAL_TIMEOUT))
+                    builder.timeout(((Number) options.get(ARGS_EVAL_TIMEOUT)).longValue());
+                if (options.containsKey(REQUEST_ID))
+                    builder.overrideRequestId((UUID) options.get(REQUEST_ID));
+                if (options.containsKey(ARGS_BATCH_SIZE))
+                    builder.batchSize(((Number) options.get(ARGS_BATCH_SIZE)).intValue());
+                if (options.containsKey(ARGS_USER_AGENT))
+                    builder.userAgent((String) options.get(ARGS_USER_AGENT));
+            }
+            return builder.create();
+        }
+    }
+
     /**
      * Settings for a {@link Client} that involve a session.
      */

diff --git a/...iver/src/main/java/org/apache/tinkerpop/gremlin/driver/remote/DriverRemoteConnection.java b/...iver/src/main/java/org/apache/tinkerpop/gremlin/driver/remote/DriverRemoteConnection.java
@@ -21,27 +21,19 @@
 import org.apache.commons.configuration2.Configuration;
 import org.apache.tinkerpop.gremlin.driver.Client;
 import org.apache.tinkerpop.gremlin.driver.Cluster;
-import org.apache.tinkerpop.gremlin.driver.RequestOptions;
 import org.apache.tinkerpop.gremlin.process.remote.RemoteConnection;
 import org.apache.tinkerpop.gremlin.process.remote.RemoteConnectionException;
 import org.apache.tinkerpop.gremlin.process.remote.traversal.RemoteTraversal;
 import org.apache.tinkerpop.gremlin.process.traversal.Bytecode;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
-import org.apache.tinkerpop.gremlin.process.traversal.strategy.decoration.OptionsStrategy;
-import org.apache.tinkerpop.gremlin.process.traversal.util.BytecodeHelper;
 import org.apache.tinkerpop.gremlin.structure.Transaction;
 import org.apache.tinkerpop.gremlin.util.iterator.IteratorUtils;
 
-import java.util.Iterator;
-import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 
-import static org.apache.tinkerpop.gremlin.driver.Tokens.ARGS_BATCH_SIZE;
-import static org.apache.tinkerpop.gremlin.driver.Tokens.ARGS_EVAL_TIMEOUT;
-import static org.apache.tinkerpop.gremlin.driver.Tokens.ARGS_USER_AGENT;
-import static org.apache.tinkerpop.gremlin.driver.Tokens.REQUEST_ID;
+import static org.apache.tinkerpop.gremlin.driver.Client.ClientUtil.getRequestOptions;
 
 /**
  * A {@link RemoteConnection} implementation for Gremlin Server. Each {@code DriverServerConnection} is bound to one
@@ -244,24 +236,6 @@ Optional<String> getSessionId() {
         return Optional.empty();
     }
 
-    public static RequestOptions getRequestOptions(final Bytecode bytecode) {
-        final Iterator<OptionsStrategy> itty = BytecodeHelper.findStrategies(bytecode, OptionsStrategy.class);
-        final RequestOptions.Builder builder = RequestOptions.build();
-        while (itty.hasNext()) {
-            final OptionsStrategy optionsStrategy = itty.next();
-            final Map<String,Object> options = optionsStrategy.getOptions();
-            if (options.containsKey(ARGS_EVAL_TIMEOUT))
-                builder.timeout(((Number) options.get(ARGS_EVAL_TIMEOUT)).longValue());
-            if (options.containsKey(REQUEST_ID))
-                builder.overrideRequestId((UUID) options.get(REQUEST_ID));
-            if (options.containsKey(ARGS_BATCH_SIZE))
-                builder.batchSize(((Number) options.get(ARGS_BATCH_SIZE)).intValue());
-            if (options.containsKey(ARGS_USER_AGENT))
-                builder.userAgent((String) options.get(ARGS_USER_AGENT));
-        }
-        return builder.create();
-    }
-
     @Override
     public void close() throws Exception {
         try {

diff --git a/.../src/test/java/org/apache/tinkerpop/gremlin/driver/remote/DriverRemoteConnectionTest.java b/.../src/test/java/org/apache/tinkerpop/gremlin/driver/remote/DriverRemoteConnectionTest.java
@@ -26,6 +26,7 @@
 
 import java.util.UUID;
 
+import static org.apache.tinkerpop.gremlin.driver.Client.ClientUtil.getRequestOptions;
 import static org.junit.Assert.assertEquals;
 
 /**
@@ -37,7 +38,7 @@ public class DriverRemoteConnectionTest {
     @Test
     public void shouldBuildRequestOptions() {
         final UUID requestId = UUID.fromString("34a9f45f-8854-4d33-8b40-92a8171ee495");
-        final RequestOptions options = DriverRemoteConnection.getRequestOptions(
+        final RequestOptions options = getRequestOptions(
                 g.with("x").
                         with("y", 100).
                         with(Tokens.ARGS_BATCH_SIZE, 1000).
@@ -53,7 +54,7 @@ public void shouldBuildRequestOptions() {
 
     @Test
     public void shouldBuildRequestOptionsWithNumerics() {
-        final RequestOptions options = DriverRemoteConnection.getRequestOptions(
+        final RequestOptions options = getRequestOptions(
                 g.with(Tokens.ARGS_BATCH_SIZE, 100).
                   with(Tokens.ARGS_EVAL_TIMEOUT, 1000).
                   V().asAdmin().getBytecode());