Use varargs for Java Flight CallOptions

java/flight/src/main/java/org/apache/arrow/flight/CallOption.java

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight;
+
+/**
+ * Per-call RPC options. These are hints to the underlying RPC layer and may not be respected.
+ */
+public interface CallOption {
+}

java/flight/src/main/java/org/apache/arrow/flight/CallOptions.java

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight;
+
+import java.util.concurrent.TimeUnit;
+
+import io.grpc.stub.AbstractStub;
+
+/**
+ * Common call options.
+ */
+public class CallOptions {
+  public static CallOption timeout(long duration, TimeUnit unit) {
+    return new Timeout(duration, unit);
+  }
+
+  static <T extends AbstractStub<T>> T wrapStub(T stub, CallOption[] options) {
+    for (CallOption option : options) {
+      if (option instanceof GrpcCallOption) {
+        stub = ((GrpcCallOption) option).wrapStub(stub);
+      }
+    }
+    return stub;
+  }
+
+  private static class Timeout implements GrpcCallOption {
+    long timeout;
+    TimeUnit timeoutUnit;
+
+    Timeout(long timeout, TimeUnit timeoutUnit) {
+      this.timeout = timeout;
+      this.timeoutUnit = timeoutUnit;
+    }
+
+    @Override
+    public <T extends AbstractStub<T>> T wrapStub(T stub) {
+      return stub.withDeadlineAfter(timeout, timeoutUnit);
+    }
+  }
+
+  interface GrpcCallOption extends CallOption {
+    <T extends AbstractStub<T>> T wrapStub(T stub);
+  }
+}

java/flight/src/main/java/org/apache/arrow/flight/FlightClient.java

@@ -88,79 +88,53 @@ public class FlightClient implements AutoCloseable {
    * Get a list of available flights.
    *
    * @param criteria Criteria for selecting flights
+   * @param options RPC-layer hints for the call.
    * @return FlightInfo Iterable
    */
-  public Iterable<FlightInfo> listFlights(Criteria criteria) {
-    return listFlights(null, criteria);
-  }
-
-  /**
-   * Get a list of available flights.
-   *
-   * @param options RPC-layer hints for the call. May be null for default options.
-   * @param criteria Criteria for selecting flights
-   * @return FlightInfo Iterable
-   */
-  public Iterable<FlightInfo> listFlights(CallOptions options, Criteria criteria) {
-    return ImmutableList.copyOf(CallOptions.wrapBlockingStub(blockingStub, options).listFlights(criteria.asCriteria()))
+  public Iterable<FlightInfo> listFlights(Criteria criteria, CallOption... options) {
+    return ImmutableList.copyOf(CallOptions.wrapStub(blockingStub, options).listFlights(criteria.asCriteria()))
         .stream()
         .map(FlightInfo::new)
         .collect(Collectors.toList());
   }
 
-  /**
-   * List actions available on the Flight service.
-   */
-  public Iterable<ActionType> listActions() {
-    return listActions(null);
-  }
-
   /**
    * List actions available on the Flight service.
    *
-   * @param options RPC-layer hints for the call. May be null for default options.
+   * @param options RPC-layer hints for the call.
    */
-  public Iterable<ActionType> listActions(CallOptions options) {
-    return ImmutableList.copyOf(CallOptions.wrapBlockingStub(blockingStub, options)
+  public Iterable<ActionType> listActions(CallOption... options) {
+    return ImmutableList.copyOf(CallOptions.wrapStub(blockingStub, options)
         .listActions(Empty.getDefaultInstance()))
         .stream()
         .map(ActionType::new)
         .collect(Collectors.toList());
   }
 
-  /**
-   * Perform an action on the Flight service.
-   * @param action The action to perform.
-   * @return An iterator of results.
-   */
-  public Iterator<Result> doAction(Action action) {
-    return doAction(null, action);
-  }
-
   /**
    * Perform an action on the Flight service.
    *
-   * @param options RPC-layer hints for this call. May be null for default options.
    * @param action The action to perform.
+   * @param options RPC-layer hints for this call.
    * @return An iterator of results.
    */
-  public Iterator<Result> doAction(CallOptions options, Action action) {
+  public Iterator<Result> doAction(Action action, CallOption... options) {
     return Iterators
-        .transform(CallOptions.wrapBlockingStub(blockingStub, options).doAction(action.toProtocol()), Result::new);
+        .transform(CallOptions.wrapStub(blockingStub, options).doAction(action.toProtocol()), Result::new);
   }
 
   public void authenticateBasic(String username, String password) {
     BasicClientAuthHandler basicClient = new BasicClientAuthHandler(username, password);
-    authenticate(null, basicClient);
+    authenticate(basicClient);
   }
 
   /**
    * Authenticate against the Flight service.
    *
-   * @param options RPC-layer hints for this call. May be null for default options.
+   * @param options RPC-layer hints for this call.
    * @param handler The auth mechanism to use.
    */
-  public void authenticate(CallOptions options, ClientAuthHandler handler) {
+  public void authenticate(ClientAuthHandler handler, CallOption... options) {
     Preconditions.checkArgument(!authInterceptor.hasAuthHandler(), "Auth already completed.");
     ClientAuthWrapper.doClientAuth(handler, CallOptions.wrapStub(asyncStub, options));
     authInterceptor.setAuthHandler(handler);
@@ -170,21 +144,11 @@ public class FlightClient implements AutoCloseable {
    * Create or append a descriptor with another stream.
    * @param descriptor FlightDescriptor
    * @param root VectorSchemaRoot
+   * @param options RPC-layer hints for this call.
    * @return ClientStreamListener
    */
-  public ClientStreamListener startPut(FlightDescriptor descriptor, VectorSchemaRoot root) {
-    return startPut(null, descriptor, root);
-  }
-
-  /**
-   * Create or append a descriptor with another stream.
-   * @param options RPC-layer hints for this call. May be null for default options.
-   * @param descriptor FlightDescriptor
-   * @param root VectorSchemaRoot
-   * @return ClientStreamListener
-   */
-  public ClientStreamListener startPut(CallOptions options,
-      FlightDescriptor descriptor, VectorSchemaRoot root) {
+  public ClientStreamListener startPut(
+      FlightDescriptor descriptor, VectorSchemaRoot root, CallOption... options) {
     Preconditions.checkNotNull(descriptor);
     Preconditions.checkNotNull(root);
 
@@ -204,32 +168,18 @@ public class FlightClient implements AutoCloseable {
   /**
    * Get info on a stream.
    * @param descriptor The descriptor for the stream.
+   * @param options RPC-layer hints for this call.
    */
-  public FlightInfo getInfo(FlightDescriptor descriptor) {
-    return getInfo(null, descriptor);
-  }
-
-  /**
-   * Get info on a stream.
-   * @param options RPC-layer hints for this call. May be null for default options.
-   * @param descriptor The descriptor for the stream.
-   */
-  public FlightInfo getInfo(CallOptions options, FlightDescriptor descriptor) {
-    return new FlightInfo(CallOptions.wrapBlockingStub(blockingStub, options).getFlightInfo(descriptor.toProtocol()));
+  public FlightInfo getInfo(FlightDescriptor descriptor, CallOption... options) {
+    return new FlightInfo(CallOptions.wrapStub(blockingStub, options).getFlightInfo(descriptor.toProtocol()));
   }
 
   /**
    * Retrieve a stream from the server.
+   * @param ticket The ticket granting access to the data stream.
+   * @param options RPC-layer hints for this call.
    */
-  public FlightStream getStream(Ticket ticket) {
-    return getStream(null, ticket);
-  }
-
-  /**
-   * Retrieve a stream from the server.
-   * @param options RPC-layer hints for this call. May be null for default options.
-   */
-  public FlightStream getStream(CallOptions options, Ticket ticket) {
+  public FlightStream getStream(Ticket ticket, CallOption... options) {
     final io.grpc.CallOptions callOptions = CallOptions.wrapStub(asyncStub, options).getCallOptions();
     ClientCall<Flight.Ticket, ArrowMessage> call =
             authInterceptor.interceptCall(doGetDescriptor, callOptions, channel);
@@ -353,44 +303,4 @@ public class FlightClient implements AutoCloseable {
     allocator.close();
   }
 
-  /**
-   * Per-call RPC options. These are hints to the underlying RPC layer and may not be respected.
-   */
-  public static class CallOptions {
-
-    long timeout = 0;
-    TimeUnit timeoutUnit = null;
-
-    public CallOptions() {
-    }
-
-    /**
-     * Set a timeout on a call.
-     */
-    public CallOptions setTimeout(long duration, TimeUnit unit) {
-      timeout = duration;
-      timeoutUnit = unit;
-      return this;
-    }
-
-    static FlightServiceBlockingStub wrapBlockingStub(FlightServiceBlockingStub stub, CallOptions options) {
-      if (options == null) {
-        return stub;
-      } else if (options.timeoutUnit != null) {
-        return stub.withDeadlineAfter(options.timeout, options.timeoutUnit);
-      } else {
-        return stub;
-      }
-    }
-
-    static FlightServiceStub wrapStub(FlightServiceStub stub, CallOptions options) {
-      if (options == null) {
-        return stub;
-      } else if (options.timeoutUnit != null) {
-        return stub.withDeadlineAfter(options.timeout, options.timeoutUnit);
-      } else {
-        return stub;
-      }
-    }
-  }
 }

java/flight/src/test/java/org/apache/arrow/flight/TestCallOptions.java

@@ -24,7 +24,6 @@ import java.util.Iterator;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
-import org.apache.arrow.flight.FlightClient.CallOptions;
 import org.apache.arrow.flight.auth.ServerAuthHandler;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
@@ -37,7 +36,7 @@ public class TestCallOptions {
   public void timeoutFires() {
     test((client) -> {
       Instant start = Instant.now();
-      Iterator<Result> results = client.doAction(new CallOptions().setTimeout(1, TimeUnit.SECONDS), new Action("hang"));
+      Iterator<Result> results = client.doAction(new Action("hang"), CallOptions.timeout(1, TimeUnit.SECONDS));
       try {
         results.next();
         Assert.fail("Call should have failed");
@@ -54,7 +53,7 @@ public class TestCallOptions {
     test((client) -> {
       Instant start = Instant.now();
       // This shouldn't fail and it should complete within the timeout
-      Iterator<Result> results = client.doAction(new CallOptions().setTimeout(2, TimeUnit.SECONDS), new Action("fast"));
+      Iterator<Result> results = client.doAction(new Action("fast"), CallOptions.timeout(2, TimeUnit.SECONDS));
       Assert.assertArrayEquals(new byte[]{42, 42}, results.next().getBody());
       Instant end = Instant.now();
       Assert.assertTrue("Call took over 2500 ms despite timeout", Duration.between(start, end).toMillis() < 2500);