Add more builder options for Java Flight servers

cpp/src/arrow/flight/test-integration-client.cc

@@ -102,7 +102,6 @@ int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   std::unique_ptr<arrow::flight::FlightClient> client;
-  std::stringstream uri;
   arrow::flight::Location location;
   ABORT_NOT_OK(arrow::flight::Location::ForGrpcTcp(FLAGS_host, FLAGS_port, &location));
   ABORT_NOT_OK(arrow::flight::FlightClient::Connect(location, &client));

java/flight/pom.xml

@@ -73,6 +73,16 @@
       <groupId>io.netty</groupId>
       <artifactId>netty-buffer</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.netty</groupId>
+      <artifactId>netty-handler</artifactId>
+      <version>${dep.netty.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>io.netty</groupId>
+      <artifactId>netty-transport</artifactId>
+      <version>${dep.netty.version}</version>
+    </dependency>
     <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
@@ -302,4 +312,50 @@
       </plugin>
     </plugins>
   </build>
+  <profiles>
+    <profile>
+      <id>linux-netty-native</id>
+      <activation>
+        <os>
+          <family>linux</family>
+        </os>
+      </activation>
+      <dependencies>
+        <dependency>
+          <groupId>io.netty</groupId>
+          <artifactId>netty-transport-native-unix-common</artifactId>
+          <version>${dep.netty.version}</version>
+          <classifier>${os.detected.name}-${os.detected.arch}</classifier>
+        </dependency>
+        <dependency>
+          <groupId>io.netty</groupId>
+          <artifactId>netty-transport-native-epoll</artifactId>
+          <version>${dep.netty.version}</version>
+          <classifier>${os.detected.name}-${os.detected.arch}</classifier>
+        </dependency>
+      </dependencies>
+    </profile>
+    <profile>
+      <id>mac-netty-native</id>
+      <activation>
+        <os>
+          <family>mac</family>
+        </os>
+      </activation>
+      <dependencies>
+        <dependency>
+          <groupId>io.netty</groupId>
+          <artifactId>netty-transport-native-unix-common</artifactId>
+          <version>${dep.netty.version}</version>
+          <classifier>${os.detected.name}-${os.detected.arch}</classifier>
+        </dependency>
+        <dependency>
+          <groupId>io.netty</groupId>
+          <artifactId>netty-transport-native-kqueue</artifactId>
+          <version>${dep.netty.version}</version>
+          <classifier>${os.detected.name}-${os.detected.arch}</classifier>
+        </dependency>
+      </dependencies>
+    </profile>
+  </profiles>
 </project>

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

@@ -20,11 +20,14 @@ package org.apache.arrow.flight;
 import static io.grpc.stub.ClientCalls.asyncClientStreamingCall;
 import static io.grpc.stub.ClientCalls.asyncServerStreamingCall;
 
+import java.io.InputStream;
 import java.net.URISyntaxException;
 import java.util.Iterator;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
+import javax.net.ssl.SSLException;
+
 import org.apache.arrow.flight.auth.BasicClientAuthHandler;
 import org.apache.arrow.flight.auth.ClientAuthHandler;
 import org.apache.arrow.flight.auth.ClientAuthInterceptor;
@@ -50,13 +53,19 @@ import com.google.common.util.concurrent.SettableFuture;
 import io.grpc.ClientCall;
 import io.grpc.ManagedChannel;
 import io.grpc.MethodDescriptor;
+import io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.NettyChannelBuilder;
 import io.grpc.stub.ClientCallStreamObserver;
 import io.grpc.stub.ClientResponseObserver;
 import io.grpc.stub.StreamObserver;
 
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.ServerChannel;
+import io.netty.handler.ssl.SslContext;
+import io.netty.handler.ssl.SslContextBuilder;
+
 /**
- * Client for flight servers.
+ * Client for Flight services.
  */
 public class FlightClient implements AutoCloseable {
   private static final int PENDING_REQUESTS = 5;
@@ -315,6 +324,13 @@ public class FlightClient implements AutoCloseable {
     allocator.close();
   }
 
+  /**
+   * Create a builder for a Flight client.
+   */
+  public static Builder builder() {
+    return new Builder();
+  }
+
   /**
    * Create a builder for a Flight client.
    * @param allocator The allocator to use for the client.
@@ -329,13 +345,20 @@ public class FlightClient implements AutoCloseable {
    */
   public static final class Builder {
 
-    private final BufferAllocator allocator;
-    private final Location location;
+    private BufferAllocator allocator;
+    private Location location;
     private boolean forceTls = false;
+    private int maxInboundMessageSize = FlightServer.MAX_GRPC_MESSAGE_SIZE;
+    private InputStream trustedCertificates = null;
+    private InputStream clientCertificate = null;
+    private InputStream clientKey = null;
+
+    private Builder() {
+    }
 
     private Builder(BufferAllocator allocator, Location location) {
-      this.allocator = allocator;
-      this.location = location;
+      this.allocator = Preconditions.checkNotNull(allocator);
+      this.location = Preconditions.checkNotNull(location);
     }
 
     /**
@@ -346,6 +369,37 @@ public class FlightClient implements AutoCloseable {
       return this;
     }
 
+    /** Set the maximum inbound message size. */
+    public Builder maxInboundMessageSize(int maxSize) {
+      Preconditions.checkArgument(maxSize > 0);
+      this.maxInboundMessageSize = maxSize;
+      return this;
+    }
+
+    /** Set the trusted TLS certificates. */
+    public Builder trustedCertificates(final InputStream stream) {
+      this.trustedCertificates = Preconditions.checkNotNull(stream);
+      return this;
+    }
+
+    /** Set the trusted TLS certificates. */
+    public Builder clientCertificate(final InputStream clientCertificate, final InputStream clientKey) {
+      Preconditions.checkNotNull(clientKey);
+      this.clientCertificate = Preconditions.checkNotNull(clientCertificate);
+      this.clientKey = Preconditions.checkNotNull(clientKey);
+      return this;
+    }
+
+    public Builder allocator(BufferAllocator allocator) {
+      this.allocator = Preconditions.checkNotNull(allocator);
+      return this;
+    }
+
+    public Builder location(Location location) {
+      this.location = Preconditions.checkNotNull(location);
+      return this;
+    }
+
     /**
      * Create the client from this builder.
      */
@@ -356,7 +410,32 @@ public class FlightClient implements AutoCloseable {
         case LocationSchemes.GRPC:
         case LocationSchemes.GRPC_INSECURE:
         case LocationSchemes.GRPC_TLS: {
-          builder = NettyChannelBuilder.forAddress(location.getUri().getHost(), location.getUri().getPort());
+          builder = NettyChannelBuilder.forAddress(location.toSocketAddress());
+          break;
+        }
+        case LocationSchemes.GRPC_DOMAIN_SOCKET: {
+          // The implementation is platform-specific, so we have to find the classes at runtime
+          builder = NettyChannelBuilder.forAddress(location.toSocketAddress());
+          try {
+            try {
+              // Linux
+              builder.channelType(
+                  (Class<? extends ServerChannel>) Class.forName("io.netty.channel.epoll.EpollDomainSocketChannel"));
+              final EventLoopGroup elg = (EventLoopGroup) Class.forName("io.netty.channel.epoll.EpollEventLoopGroup")
+                  .newInstance();
+              builder.eventLoopGroup(elg);
+            } catch (ClassNotFoundException e) {
+              // BSD
+              builder.channelType(
+                  (Class<? extends ServerChannel>) Class.forName("io.netty.channel.kqueue.KQueueDomainSocketChannel"));
+              final EventLoopGroup elg = (EventLoopGroup) Class.forName("io.netty.channel.kqueue.KQueueEventLoopGroup")
+                  .newInstance();
+              builder.eventLoopGroup(elg);
+            }
+          } catch (ClassNotFoundException | InstantiationException | IllegalAccessException e) {
+            throw new UnsupportedOperationException(
+                "Could not find suitable Netty native transport implementation for domain socket address.");
+          }
           break;
         }
         default:
@@ -365,13 +444,28 @@ public class FlightClient implements AutoCloseable {
 
       if (this.forceTls || LocationSchemes.GRPC_TLS.equals(location.getUri().getScheme())) {
         builder.useTransportSecurity();
+
+        if (this.trustedCertificates != null || this.clientCertificate != null || this.clientKey != null) {
+          final SslContextBuilder sslContextBuilder = GrpcSslContexts.forClient();
+          if (this.trustedCertificates != null) {
+            sslContextBuilder.trustManager(this.trustedCertificates);
+          }
+          if (this.clientCertificate != null && this.clientKey != null) {
+            sslContextBuilder.keyManager(this.clientCertificate, this.clientKey);
+          }
+          try {
+            builder.sslContext(sslContextBuilder.build());
+          } catch (SSLException e) {
+            throw new RuntimeException(e);
+          }
+        }
       } else {
         builder.usePlaintext();
       }
 
       builder
           .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS)
-          .maxInboundMessageSize(FlightServer.MAX_GRPC_MESSAGE_SIZE);
+          .maxInboundMessageSize(maxInboundMessageSize);
       return new FlightClient(allocator, builder.build());
     }
   }

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

@@ -19,22 +19,26 @@ package org.apache.arrow.flight;
 
 import java.io.File;
 import java.io.FileInputStream;
-import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
-import java.net.InetSocketAddress;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.TimeUnit;
 
 import org.apache.arrow.flight.auth.ServerAuthHandler;
 import org.apache.arrow.flight.auth.ServerAuthInterceptor;
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.util.Preconditions;
 
 import io.grpc.Server;
 import io.grpc.ServerInterceptors;
 import io.grpc.netty.NettyServerBuilder;
 
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.ServerChannel;
+
 public class FlightServer implements AutoCloseable {
 
   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(FlightServer.class);
@@ -88,45 +92,80 @@ public class FlightServer implements AutoCloseable {
     }
   }
 
+  /** Create a builder for a Flight server. */
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  /** Create a builder for a Flight server. */
   public static Builder builder(BufferAllocator allocator, Location location, FlightProducer producer) {
     return new Builder(allocator, location, producer);
   }
 
+  /** A builder for Flight servers. */
   public static final class Builder {
-
-    private final BufferAllocator allocator;
-    private final Location location;
-    private final FlightProducer producer;
+    private BufferAllocator allocator;
+    private Location location;
+    private FlightProducer producer;
+    private final Map<String, Object> builderOptions;
     private ServerAuthHandler authHandler = ServerAuthHandler.NO_OP;
-
+    private Executor executor = null;
+    private int maxInboundMessageSize = MAX_GRPC_MESSAGE_SIZE;
     private InputStream certChain;
     private InputStream key;
 
-    Builder(BufferAllocator allocator, Location location, FlightProducer producer) {
-      this.allocator = allocator;
-      this.location = location;
-      this.producer = producer;
+    Builder() {
+      builderOptions = new HashMap<>();
     }
 
+    Builder(BufferAllocator allocator, Location location, FlightProducer producer) {
+      this.allocator = Preconditions.checkNotNull(allocator);
+      this.location = Preconditions.checkNotNull(location);
+      this.producer = Preconditions.checkNotNull(producer);
+      builderOptions = new HashMap<>();
+    }
+
+    /** Create the server for this builder. */
     public FlightServer build() {
       final NettyServerBuilder builder;
       switch (location.getUri().getScheme()) {
         case LocationSchemes.GRPC_DOMAIN_SOCKET: {
-          // TODO: need reflection to check if domain sockets are available
-          throw new UnsupportedOperationException("Domain sockets are not available.");
+          // The implementation is platform-specific, so we have to find the classes at runtime
+          builder = NettyServerBuilder.forAddress(location.toSocketAddress());
+          try {
+            try {
+              // Linux
+              builder.channelType(
+                  (Class<? extends ServerChannel>) Class
+                      .forName("io.netty.channel.epoll.EpollServerDomainSocketChannel"));
+              final EventLoopGroup elg = (EventLoopGroup) Class.forName("io.netty.channel.epoll.EpollEventLoopGroup")
+                  .newInstance();
+              builder.bossEventLoopGroup(elg).workerEventLoopGroup(elg);
+            } catch (ClassNotFoundException e) {
+              // BSD
+              builder.channelType(
+                  (Class<? extends ServerChannel>) Class
+                      .forName("io.netty.channel.kqueue.KQueueServerDomainSocketChannel"));
+              final EventLoopGroup elg = (EventLoopGroup) Class.forName("io.netty.channel.kqueue.KQueueEventLoopGroup")
+                  .newInstance();
+              builder.bossEventLoopGroup(elg).workerEventLoopGroup(elg);
+            }
+          } catch (ClassNotFoundException | InstantiationException | IllegalAccessException e) {
+            throw new UnsupportedOperationException(
+                "Could not find suitable Netty native transport implementation for domain socket address.");
+          }
+          break;
         }
         case LocationSchemes.GRPC:
         case LocationSchemes.GRPC_INSECURE: {
-          builder = NettyServerBuilder
-              .forAddress(new InetSocketAddress(location.getUri().getHost(), location.getUri().getPort()));
+          builder = NettyServerBuilder.forAddress(location.toSocketAddress());
           break;
         }
         case LocationSchemes.GRPC_TLS: {
           if (certChain == null) {
             throw new IllegalArgumentException("Must provide a certificate and key to serve gRPC over TLS");
           }
-          builder = NettyServerBuilder
-              .forAddress(new InetSocketAddress(location.getUri().getHost(), location.getUri().getPort()));
+          builder = NettyServerBuilder.forAddress(location.toSocketAddress());
           break;
         }
         default:
@@ -137,15 +176,33 @@ public class FlightServer implements AutoCloseable {
         builder.useTransportSecurity(certChain, key);
       }
 
-      final Server server = builder
-          .executor(new ForkJoinPool())
-          .maxInboundMessageSize(MAX_GRPC_MESSAGE_SIZE)
+      builder
+          .executor(executor != null ? executor : new ForkJoinPool())
+          .maxInboundMessageSize(maxInboundMessageSize)
           .addService(
               ServerInterceptors.intercept(
                   new FlightBindingService(allocator, producer, authHandler),
-                  new ServerAuthInterceptor(authHandler)))
-          .build();
-      return new FlightServer(server);
+                  new ServerAuthInterceptor(authHandler)));
+
+      // Allow setting some Netty-specific options
+      builderOptions.computeIfPresent("netty.bossEventLoopGroup", (key, elg) -> {
+        builder.bossEventLoopGroup((EventLoopGroup) elg);
+        return null;
+      });
+      builderOptions.computeIfPresent("netty.workerEventLoopGroup", (key, elg) -> {
+        builder.workerEventLoopGroup((EventLoopGroup) elg);
+        return null;
+      });
+
+      return new FlightServer(builder.build());
+    }
+
+    /**
+     * Set the maximum size of a message. Defaults to "unlimited", depending on the underlying transport.
+     */
+    public Builder maxInboundMessageSize(int maxMessageSize) {
+      this.maxInboundMessageSize = maxMessageSize;
+      return this;
     }
 
     /**
@@ -159,15 +216,54 @@ public class FlightServer implements AutoCloseable {
       return this;
     }
 
+    /**
+     * Enable TLS on the server.
+     * @param certChain The certificate chain to use.
+     * @param key The private key to use.
+     */
     public Builder useTls(final InputStream certChain, final InputStream key) {
       this.certChain = certChain;
       this.key = key;
       return this;
     }
 
-    public Builder setAuthHandler(ServerAuthHandler authHandler) {
+    /**
+     * Set the executor used by the server.
+     */
+    public Builder executor(Executor executor) {
+      this.executor = executor;
+      return this;
+    }
+
+    /**
+     * Set the authentication handler.
+     */
+    public Builder authHandler(ServerAuthHandler authHandler) {
       this.authHandler = authHandler;
       return this;
     }
+
+    /**
+     * Provide a transport-specific option. Not guaranteed to have any effect.
+     */
+    public Builder transportHint(final String key, Object option) {
+      builderOptions.put(key, option);
+      return this;
+    }
+
+    public Builder allocator(BufferAllocator allocator) {
+      this.allocator = Preconditions.checkNotNull(allocator);
+      return this;
+    }
+
+    public Builder location(Location location) {
+      this.location = Preconditions.checkNotNull(location);
+      return this;
+    }
+
+    public Builder producer(FlightProducer producer) {
+      this.producer = Preconditions.checkNotNull(producer);
+      return this;
+    }
   }
 }

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

@@ -17,6 +17,9 @@
 
 package org.apache.arrow.flight;
 
+import java.lang.reflect.InvocationTargetException;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
 
@@ -30,21 +33,69 @@ public class Location {
    * Constructs a new instance.
    *
    * @param uri the URI of the Flight service
+   * @throws IllegalArgumentException if the URI scheme is unsupported
    */
   public Location(String uri) throws URISyntaxException {
-    super();
-    this.uri = new URI(uri);
+    this(new URI(uri));
   }
 
+  /**
+   * Construct a new instance from an existing URI.
+   *
+   * @param uri the URI of the Flight service
+   * @throws IllegalArgumentException if the URI scheme is unsupported
+   */
   public Location(URI uri) {
     super();
     this.uri = uri;
+    // Validate the scheme
+    switch (uri.getScheme()) {
+      case LocationSchemes.GRPC:
+      case LocationSchemes.GRPC_DOMAIN_SOCKET:
+      case LocationSchemes.GRPC_INSECURE:
+      case LocationSchemes.GRPC_TLS: {
+        break;
+      }
+      default:
+        throw new IllegalArgumentException("Scheme is not supported: " + this.uri);
+    }
   }
 
   public URI getUri() {
     return uri;
   }
 
+  /**
+   * Helper method to turn this Location into a SocketAddress.
+   *
+   * @return null if could not be converted
+   */
+  SocketAddress toSocketAddress() {
+    switch (uri.getScheme()) {
+      case LocationSchemes.GRPC:
+      case LocationSchemes.GRPC_TLS:
+      case LocationSchemes.GRPC_INSECURE: {
+        return new InetSocketAddress(uri.getHost(), uri.getPort());
+      }
+
+      case LocationSchemes.GRPC_DOMAIN_SOCKET: {
+        try {
+          // This dependency is not available on non-Unix platforms.
+          return (SocketAddress) Class.forName("io.netty.channel.unix.DomainSocketAddress")
+              .getConstructor(String.class)
+              .newInstance(uri.getPath());
+        } catch (InstantiationException | ClassNotFoundException | InvocationTargetException |
+            NoSuchMethodException | IllegalAccessException e) {
+          return null;
+        }
+      }
+
+      default: {
+        return null;
+      }
+    }
+  }
+
   /**
    * Convert this Location into its protocol-level representation.
    */
@@ -69,4 +120,15 @@ public class Location {
       throw new IllegalArgumentException(e);
     }
   }
+
+  /**
+   * Construct a URI for a Flight+gRPC server over a Unix domain socket.
+   */
+  public static Location forGrpcDomainSocket(String path) {
+    try {
+      return new Location(new URI(LocationSchemes.GRPC_DOMAIN_SOCKET, null, path, null));
+    } catch (URISyntaxException e) {
+      throw new IllegalArgumentException(e);
+    }
+  }
 }

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

@@ -20,9 +20,13 @@ package org.apache.arrow.flight;
 /**
  * Constants representing well-known URI schemes for Flight services.
  */
-public class LocationSchemes {
+public final class LocationSchemes {
   public static final String GRPC = "grpc";
   public static final String GRPC_INSECURE = "grpc+tcp";
   public static final String GRPC_DOMAIN_SOCKET = "grpc+unix";
   public static final String GRPC_TLS = "grpc+tls";
+
+  private LocationSchemes() {
+    throw new AssertionError("Do not instantiate this class.");
+  }
 }

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

@@ -62,6 +62,28 @@ public class FlightTestUtil {
     return server;
   }
 
+  static boolean isEpollAvailable() {
+    try {
+      Class<?> epoll = Class.forName("io.netty.channel.epoll.Epoll");
+      return (Boolean) epoll.getMethod("isAvailable").invoke(null);
+    } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
+      return false;
+    }
+  }
+
+  static boolean isKqueueAvailable() {
+    try {
+      Class<?> kqueue = Class.forName("io.netty.channel.kqueue.KQueue");
+      return (Boolean) kqueue.getMethod("isAvailable").invoke(null);
+    } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
+      return false;
+    }
+  }
+
+  static boolean isNativeTransportAvailable() {
+    return isEpollAvailable() || isKqueueAvailable();
+  }
+
   private FlightTestUtil() {
   }
 }

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

@@ -157,7 +157,7 @@ public class TestBasicOperation {
   /**
    * An example FlightProducer for test purposes.
    */
-  public class Producer implements FlightProducer, AutoCloseable {
+  public static class Producer implements FlightProducer, AutoCloseable {
 
     private final BufferAllocator allocator;
 

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

@@ -0,0 +1,63 @@
+/*
+ * 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.io.File;
+
+import org.apache.arrow.flight.TestBasicOperation.Producer;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class TestServerOptions {
+
+  @Test
+  public void domainSocket() throws Exception {
+    Assume.assumeTrue("We have a native transport available", FlightTestUtil.isNativeTransportAvailable());
+    final File domainSocket = File.createTempFile("flight-unit-test-", ".sock");
+    Assert.assertTrue(domainSocket.delete());
+    final Location location = Location.forGrpcDomainSocket(domainSocket.getAbsolutePath());
+    try (
+        BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
+        Producer producer = new Producer(a);
+        FlightServer s =
+            FlightTestUtil.getStartedServer(
+                (port) -> FlightServer.builder(a, location, producer).build()
+            )) {
+      try (FlightClient c = FlightClient.builder(a, location).build()) {
+        FlightStream stream = c.getStream(new Ticket(new byte[0]));
+        VectorSchemaRoot root = stream.getRoot();
+        IntVector iv = (IntVector) root.getVector("c1");
+        int value = 0;
+        while (stream.next()) {
+          for (int i = 0; i < root.getRowCount(); i++) {
+            Assert.assertEquals(value, iv.get(i));
+            value++;
+          }
+        }
+      }
+    }
+  }
+}

java/flight/src/test/java/org/apache/arrow/flight/auth/TestAuth.java

@@ -149,7 +149,7 @@ public class TestAuth {
             root.clear();
             listener.completed();
           }
-        }).setAuthHandler(new BasicServerAuthHandler(validator)).build());
+        }).authHandler(new BasicServerAuthHandler(validator)).build());
     client = FlightClient.builder(allocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
         .build();
   }

python/pyarrow/_flight.pyx

@@ -400,16 +400,26 @@ cdef class FlightClient:
                         .format(self.__class__.__name__))
 
     @staticmethod
-    def connect(location, **kwargs):
-        """Connect to a Flight service on the given host and port."""
+    def connect(location, tls_root_certs=None):
+        """
+        Connect to a Flight service on the given host and port.
+
+        Parameters
+        ----------
+        location : Location
+            location to connect to
+
+        tls_root_certs : bytes
+            PEM-encoded
+        """
         cdef:
             FlightClient result = FlightClient.__new__(FlightClient)
             int c_port = 0
             CLocation c_location = Location.unwrap(location)
             CFlightClientOptions c_options
 
-        if "tls_root_certs" in kwargs:
-            c_options.tls_root_certs = tobytes(kwargs["tls_root_certs"])
+        if tls_root_certs:
+            c_options.tls_root_certs = tobytes(tls_root_certs)
 
         with nogil:
             check_status(CFlightClient.Connect(c_location, c_options,
@@ -1004,7 +1014,8 @@ cdef class FlightServerBase:
     cdef:
         unique_ptr[PyFlightServer] server
 
-    def run(self, location, auth_handler=None, **kwargs):
+    def run(self, location, auth_handler=None,
+            tls_cert_chain=None, tls_private_key=None):
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             PyFlightServer* c_server
@@ -1019,14 +1030,12 @@ cdef class FlightServerBase:
             c_options.get().auth_handler.reset(
                 (<ServerAuthHandler> auth_handler).to_handler())
 
-        if "tls_cert_chain" in kwargs:
-            if "tls_private_key" not in kwargs:
+        if tls_cert_chain:
+            if not tls_private_key:
                 raise ValueError(
                     "Must provide both cert chain and private key")
-            c_options.get().tls_cert_chain = tobytes(
-                kwargs["tls_cert_chain"])
-            c_options.get().tls_private_key = tobytes(
-                kwargs["tls_private_key"])
+            c_options.get().tls_cert_chain = tobytes(tls_cert_chain)
+            c_options.get().tls_private_key = tobytes(tls_private_key)
 
         vtable.list_flights = &_list_flights
         vtable.get_flight_info = &_get_flight_info

python/pyarrow/tests/test_flight.py

@@ -465,9 +465,9 @@ def test_token_auth_invalid():
 
 def test_location_invalid():
     """Test constructing invalid URIs."""
-    with pytest.raises(pa.ArrowException, match=".*Cannot parse URI:.*"):
+    with pytest.raises(pa.ArrowInvalid, match=".*Cannot parse URI:.*"):
         flight.FlightClient.connect("%")
 
     server = ConstantFlightServer()
-    with pytest.raises(pa.ArrowException, match=".*Cannot parse URI:.*"):
+    with pytest.raises(pa.ArrowInvalid, match=".*Cannot parse URI:.*"):
         server.run("%")