ARROW-10467: [FlightRPC][Java] Add the ability to pass arbitrary client headers.

java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightConstants.java

@@ -23,5 +23,6 @@
 public interface FlightConstants {
 
   String SERVICE = "arrow.flight.protocol.FlightService";
+  String PROPERTY_HEADER = "Arrow-Properties";
 
 }

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

@@ -168,6 +168,7 @@ public static final class Builder {
     private final Map<String, Object> builderOptions;
     private ServerAuthHandler authHandler = ServerAuthHandler.NO_OP;
     private CallHeaderAuthenticator headerAuthenticator = CallHeaderAuthenticator.NO_OP;
+    private ServerHeaderHandler headerHandler = ServerHeaderHandler.NO_OP;
     private ExecutorService executor = null;
     private int maxInboundMessageSize = MAX_GRPC_MESSAGE_SIZE;
     private InputStream certChain;
@@ -196,6 +197,13 @@ public FlightServer build() {
         this.middleware(FlightServerMiddleware.Key.of(Auth2Constants.AUTHORIZATION_HEADER),
             new ServerCallHeaderAuthMiddleware.Factory(headerAuthenticator));
       }
+
+      // Add the property handler middleware if applicable.
+      if (headerHandler != ServerHeaderHandler.NO_OP) {
+        this.middleware(FlightServerMiddleware.Key.of(FlightConstants.PROPERTY_HEADER),
+            new ServerHeaderMiddleware.Factory(headerHandler));
+      }
+
       final NettyServerBuilder builder;
       switch (location.getUri().getScheme()) {
         case LocationSchemes.GRPC_DOMAIN_SOCKET: {
@@ -351,6 +359,14 @@ public Builder headerAuthenticator(CallHeaderAuthenticator headerAuthenticator)
       return this;
     }
 
+    /**
+     * Set the header handler.
+     */
+    public Builder headerHandler(ServerHeaderHandler headerHandler) {
+      this.headerHandler = headerHandler;
+      return this;
+    }
+
     /**
      * Provide a transport-specific option. Not guaranteed to have any effect.
      */

java/flight/flight-core/src/main/java/org/apache/arrow/flight/HeaderCallOption.java

@@ -0,0 +1,54 @@
+/*
+ * 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 io.grpc.Metadata;
+import io.grpc.stub.AbstractStub;
+import io.grpc.stub.MetadataUtils;
+
+/**
+ * Method option for supplying headers to method calls.
+ */
+public class HeaderCallOption implements CallOptions.GrpcCallOption {
+  private final Metadata propertiesMetadata = new Metadata();
+
+  /**
+   * Header property constructor.
+   *
+   * @param headers the headers that should be sent across. If a header is a string, it should only be valid ASCII
+   *                characters. Binary headers should end in "-bin".
+   */
+  public HeaderCallOption(CallHeaders headers) {
+    for (String key : headers.keys()) {
+      if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+        propertiesMetadata.put(
+            Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER),
+            headers.getByte(key));
+      } else {
+        propertiesMetadata.put(
+            Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER),
+            headers.get(key));
+      }
+    }
+  }
+
+  @Override
+  public <T extends AbstractStub<T>> T wrapStub(T stub) {
+    return MetadataUtils.attachHeaders(stub, propertiesMetadata);
+  }
+}

java/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerHeaderHandler.java

@@ -0,0 +1,30 @@
+/*
+ * 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.function.Consumer;
+
+/**
+ * Interface for server side property handling.
+ */
+public interface ServerHeaderHandler extends Consumer<CallHeaders> {
+  /**
+   * A property handler that does nothing.
+   */
+  ServerHeaderHandler NO_OP = (s) -> { };
+}

java/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerHeaderMiddleware.java

@@ -0,0 +1,60 @@
+/*
+ * 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;
+
+/**
+ * Middleware that's used to extract and pass properties to the server during requests.
+ */
+public class ServerHeaderMiddleware implements FlightServerMiddleware {
+  /**
+   * Factory for accessing ServerCallPropertyMiddleware.
+   */
+  public static class Factory implements FlightServerMiddleware.Factory<ServerHeaderMiddleware> {
+    private final ServerHeaderHandler propertyHandler;
+
+    /**
+     * Construct a factory with the given header handler.
+     * @param headerHandler The header handler that will be used to pass properties.
+     */
+    public Factory(ServerHeaderHandler headerHandler) {
+      this.propertyHandler = headerHandler;
+    }
+
+    @Override
+    public ServerHeaderMiddleware onCallStarted(CallInfo callInfo, CallHeaders incomingHeaders,
+                                                RequestContext context) {
+      propertyHandler.accept(incomingHeaders);
+      return new ServerHeaderMiddleware();
+    }
+  }
+
+  private ServerHeaderMiddleware() {
+  }
+
+  @Override
+  public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
+  }
+
+  @Override
+  public void onCallCompleted(CallStatus status) {
+  }
+
+  @Override
+  public void onCallErrored(Throwable err) {
+  }
+}

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

@@ -24,12 +24,15 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
+import org.apache.arrow.flight.grpc.MetadataAdapter;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.junit.Assert;
 import org.junit.Ignore;
 import org.junit.Test;
 
+import io.grpc.Metadata;
+
 public class TestCallOptions {
 
   @Test
@@ -64,6 +67,61 @@ public void underTimeout() {
     });
   }
 
+  @Test
+  public void singleProperty() {
+    final MetadataAdapter headers = new MetadataAdapter(new Metadata());
+    headers.insert("key", "value");
+    testHeaders(headers);
+  }
+
+  @Test
+  public void multipleProperties() {
+    final MetadataAdapter headers = new MetadataAdapter(new Metadata());
+    headers.insert("key", "value");
+    headers.insert("key2", "value2");
+    testHeaders(headers);
+  }
+
+  @Test
+  public void binaryProperties() {
+    final MetadataAdapter headers = new MetadataAdapter(new Metadata());
+    headers.insert("key-bin", "value".getBytes());
+    headers.insert("key3-bin", "ëfßæ".getBytes());
+    testHeaders(headers);
+  }
+
+  @Test
+  public void mixedProperties() {
+    final MetadataAdapter headers = new MetadataAdapter(new Metadata());
+    headers.insert("key", "value");
+    headers.insert("key3-bin", "ëfßæ".getBytes());
+    testHeaders(headers);
+  }
+
+  private void testHeaders(MetadataAdapter headers) {
+    final ServerHeaderHandler propertyHandler = (incomingHeaders) -> {
+      for (String key : headers.keys()) {
+        if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+          Assert.assertArrayEquals(headers.getByte(key), incomingHeaders.getByte(key));
+        } else {
+          Assert.assertEquals(headers.get(key), incomingHeaders.get(key));
+        }
+      }
+    };
+
+    try (
+        BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
+        Producer producer = new Producer(a);
+        FlightServer s =
+            FlightTestUtil.getStartedServer((location) ->
+                FlightServer.builder(a, location, producer).headerHandler(propertyHandler).build());
+        FlightClient client = FlightClient.builder(a, s.getLocation()).build()) {
+      client.doAction(new Action("fast"), new HeaderCallOption(headers)).hasNext();
+    } catch (InterruptedException | IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   void test(Consumer<FlightClient> testFn) {
     try (
         BufferAllocator a = new RootAllocator(Long.MAX_VALUE);