ARROW-6074: [FlightRPC][C++][Python] Flight middleware

cpp/src/arrow/flight/CMakeLists.txt

@@ -83,6 +83,8 @@ set(ARROW_FLIGHT_SRCS
     internal.cc
     protocol_internal.cc
     serialization_internal.cc
+    middleware.cc
+    middleware_internal.cc
     server.cc
     server_auth.cc
     types.cc)

cpp/src/arrow/flight/api.h

@@ -19,6 +19,9 @@
 
 #include "arrow/flight/client.h"
 #include "arrow/flight/client_auth.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/middleware.h"
 #include "arrow/flight/server.h"
 #include "arrow/flight/server_auth.h"
+#include "arrow/flight/server_middleware.h"
 #include "arrow/flight/types.h"

cpp/src/arrow/flight/client.cc

@@ -40,10 +40,14 @@
 #include "arrow/status.h"
 #include "arrow/type.h"
 #include "arrow/util/logging.h"
+#include "arrow/util/stl.h"
 #include "arrow/util/uri.h"
 
 #include "arrow/flight/client_auth.h"
+#include "arrow/flight/client_middleware.h"
 #include "arrow/flight/internal.h"
+#include "arrow/flight/middleware.h"
+#include "arrow/flight/middleware_internal.h"
 #include "arrow/flight/serialization_internal.h"
 #include "arrow/flight/types.h"
 
@@ -86,6 +90,145 @@ struct ClientRpc {
   }
 };
 
+class GrpcAddCallHeaders : public AddCallHeaders {
+ public:
+  explicit GrpcAddCallHeaders(std::multimap<grpc::string, grpc::string>* metadata)
+      : metadata_(metadata) {}
+  ~GrpcAddCallHeaders() = default;
+
+  void AddHeader(const std::string& key, const std::string& value) override {
+    metadata_->insert(std::make_pair(key, value));
+  }
+
+ private:
+  std::multimap<grpc::string, grpc::string>* metadata_;
+};
+
+class GrpcClientInterceptorAdapter : public grpc::experimental::Interceptor {
+ public:
+  GrpcClientInterceptorAdapter(std::vector<std::unique_ptr<ClientMiddleware>> middleware)
+      : middleware_(std::move(middleware)), failed_() {}
+
+  void Intercept(grpc::experimental::InterceptorBatchMethods* methods) {
+    using InterceptionHookPoints = grpc::experimental::InterceptionHookPoints;
+    if (methods->QueryInterceptionHookPoint(
+            InterceptionHookPoints::PRE_SEND_INITIAL_METADATA)) {
+      GrpcAddCallHeaders add_headers(methods->GetSendInitialMetadata());
+      for (const auto& middleware : middleware_) {
+        const Status& status = middleware->SendingHeaders(add_headers);
+        if (!status.ok()) {
+          failed_ = status;
+          break;
+        }
+      }
+    }
+
+    // TODO: also check for trailing metadata if call failed?
+    // (grpc-java and grpc-core differ on this point, it seems)
+
+    // Don't run this hook if an earlier middleware already failed the call
+    if (failed_.ok() && methods->QueryInterceptionHookPoint(
+                            InterceptionHookPoints::POST_RECV_INITIAL_METADATA)) {
+      internal::GrpcCallHeaders headers(methods->GetRecvInitialMetadata());
+      for (const auto& middleware : middleware_) {
+        const Status& status = middleware->ReceivedHeaders(headers);
+        if (!status.ok()) {
+          failed_ = status;
+          break;
+        }
+      }
+    }
+
+    if (methods->QueryInterceptionHookPoint(InterceptionHookPoints::POST_RECV_STATUS)) {
+      DCHECK_NE(nullptr, methods->GetRecvStatus());
+      const Status status =
+          failed_.ok() ? internal::FromGrpcStatus(*methods->GetRecvStatus()) : failed_;
+
+      for (const auto& middleware : middleware_) {
+        const Status& substatus = middleware->CallCompleted(status);
+        if (!substatus.ok()) {
+          // Abandon the rest
+          failed_ = substatus;
+          break;
+        }
+      }
+
+      // Overwrite the gRPC status with the interceptor status
+      if (!failed_.ok()) {
+        *methods->GetRecvStatus() = internal::ToGrpcStatus(failed_);
+      }
+    }
+
+    methods->Proceed();
+  }
+
+ private:
+  std::vector<std::unique_ptr<ClientMiddleware>> middleware_;
+  // Stores an Arrow failure status returned by middleware (if any)
+  Status failed_;
+};
+
+class GrpcClientInterceptorAdapterFactory
+    : public grpc::experimental::ClientInterceptorFactoryInterface {
+ public:
+  GrpcClientInterceptorAdapterFactory(
+      std::vector<std::shared_ptr<ClientMiddlewareFactory>> middleware)
+      : middleware_(middleware) {}
+
+  grpc::experimental::Interceptor* CreateClientInterceptor(
+      grpc::experimental::ClientRpcInfo* info) override {
+    std::vector<std::unique_ptr<ClientMiddleware>> middleware;
+
+    FlightMethod flight_method = FlightMethod::Invalid;
+    util::string_view method(info->method());
+    if (method.ends_with("/Handshake")) {
+      flight_method = FlightMethod::Handshake;
+    } else if (method.ends_with("/ListFlights")) {
+      flight_method = FlightMethod::ListFlights;
+    } else if (method.ends_with("/GetFlightInfo")) {
+      flight_method = FlightMethod::GetFlightInfo;
+    } else if (method.ends_with("/GetSchema")) {
+      flight_method = FlightMethod::GetSchema;
+    } else if (method.ends_with("/DoGet")) {
+      flight_method = FlightMethod::DoGet;
+    } else if (method.ends_with("/DoPut")) {
+      flight_method = FlightMethod::DoPut;
+    } else if (method.ends_with("/DoAction")) {
+      flight_method = FlightMethod::DoAction;
+    } else if (method.ends_with("/ListActions")) {
+      flight_method = FlightMethod::ListActions;
+    }
+
+    const CallInfo flight_info{flight_method};
+    for (auto& factory : middleware_) {
+      std::unique_ptr<ClientMiddleware> instance;
+      const Status& status = factory->StartCall(flight_info, &instance);
+      if (!status.ok()) {
+        // Call was rejected here, but we don't have a formal way to
+        // stop the call with a status.
+        ARROW_LOG(INFO) << "Call cancelled: " << status;
+        info->client_context()->TryCancel();
+        for (auto& existing_instance : middleware) {
+          const Status& sub_status = existing_instance->CallCompleted(status);
+          if (!sub_status.ok()) {
+            // Aborted twice, log and bail
+            ARROW_LOG(INFO) << "Call cancelled while processing cancellation: " << status;
+            break;
+          }
+        }
+        return nullptr;
+      }
+      if (instance) {
+        middleware.push_back(std::move(instance));
+      }
+    }
+    return new GrpcClientInterceptorAdapter(std::move(middleware));
+  }
+
+ private:
+  std::vector<std::shared_ptr<ClientMiddlewareFactory>> middleware_;
+};
+
 class GrpcClientAuthSender : public ClientAuthSender {
  public:
   explicit GrpcClientAuthSender(
@@ -442,8 +585,15 @@ class FlightClient::FlightClientImpl {
       args.SetSslTargetNameOverride(options.override_hostname);
     }
 
+    std::vector<std::unique_ptr<grpc::experimental::ClientInterceptorFactoryInterface>>
+        interceptors;
+    interceptors.emplace_back(
+        new GrpcClientInterceptorAdapterFactory(std::move(options.middleware)));
+
     stub_ = pb::FlightService::NewStub(
-        grpc::CreateCustomChannel(grpc_uri.str(), creds, args));
+        grpc::experimental::CreateCustomChannelWithInterceptors(
+            grpc_uri.str(), creds, args, std::move(interceptors)));
+
     return Status::OK();
   }
 

cpp/src/arrow/flight/client.h

@@ -23,6 +23,7 @@
 #include <chrono>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "arrow/ipc/reader.h"
@@ -41,6 +42,8 @@ class Schema;
 namespace flight {
 
 class ClientAuthHandler;
+class ClientMiddleware;
+class ClientMiddlewareFactory;
 
 /// \brief A duration type for Flight call timeouts.
 typedef std::chrono::duration<double, std::chrono::seconds::period> TimeoutDuration;
@@ -64,6 +67,8 @@ class ARROW_FLIGHT_EXPORT FlightClientOptions {
   std::string tls_root_certs;
   /// \brief Override the hostname checked by TLS. Use with caution.
   std::string override_hostname;
+  /// \brief A list of client middleware to apply.
+  std::vector<std::shared_ptr<ClientMiddlewareFactory>> middleware;
 };
 
 /// \brief A RecordBatchReader exposing Flight metadata and cancel

cpp/src/arrow/flight/client_middleware.h

@@ -0,0 +1,61 @@
+// 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.
+
+// Interfaces for defining middleware for Flight clients. Currently
+// experimental.
+
+#pragma once
+
+#include "arrow/flight/middleware.h"
+#include "arrow/flight/visibility.h"  // IWYU pragma: keep
+#include "arrow/status.h"
+
+namespace arrow {
+
+namespace flight {
+
+class ARROW_FLIGHT_EXPORT ClientMiddleware {
+ public:
+  virtual ~ClientMiddleware() = default;
+
+  /// While a non-OK status will do its best to stop the call, this
+  /// isn't guaranteed.
+  virtual Status SendingHeaders(AddCallHeaders& outgoing_headers) = 0;
+  virtual Status ReceivedHeaders(const CallHeaders& incoming_headers) = 0;
+  virtual Status CallCompleted(const Status& status) = 0;
+};
+
+class ARROW_FLIGHT_EXPORT ClientMiddlewareFactory {
+ public:
+  virtual ~ClientMiddlewareFactory() = default;
+
+  /// \brief A callback for the start of a new call.
+  ///
+  /// \param info Information about the call.
+  /// \param[out] middleware The middleware instance for this call. If
+  ///     unset, will not add middleware to this call instance from
+  ///     this factory.
+  /// \return Status a status. A non-OK status will try to stop the
+  /// call, but the underlying implementation may not be able to
+  /// support this.
+  virtual Status StartCall(const CallInfo& info,
+                           std::unique_ptr<ClientMiddleware>* middleware) = 0;
+};
+
+}  // namespace flight
+
+}  // namespace arrow

cpp/src/arrow/flight/flight_test.cc

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <atomic>
 #include <chrono>
 #include <cstdint>
 #include <cstdio>
@@ -41,6 +42,7 @@
 #endif
 
 #include "arrow/flight/internal.h"
+#include "arrow/flight/middleware_internal.h"
 #include "arrow/flight/test_util.h"
 
 namespace pb = arrow::flight::protocol;
@@ -581,6 +583,279 @@ class TestTls : public ::testing::Test {
   std::unique_ptr<InProcessTestServer> server_;
 };
 
+class RejectServerMiddlewareFactory : public ServerMiddlewareFactory {
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    return MakeFlightError(FlightStatusCode::Unauthenticated, "All calls are rejected");
+  }
+};
+
+class CountingServerMiddleware : public ServerMiddleware {
+ public:
+  CountingServerMiddleware(std::atomic<int>* successful, std::atomic<int>* failed)
+      : successful_(successful), failed_(failed) {}
+  Status SendingHeaders(AddCallHeaders& outgoing_headers) { return Status::OK(); }
+  Status CallCompleted(const Status& status) {
+    if (status.ok()) {
+      ARROW_IGNORE_EXPR((*successful_)++);
+    } else {
+      ARROW_IGNORE_EXPR((*failed_)++);
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::atomic<int>* successful_;
+  std::atomic<int>* failed_;
+};
+
+class CountingServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  CountingServerMiddlewareFactory() : successful_(0), failed_(0) {}
+
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    *middleware = std::make_shared<CountingServerMiddleware>(&successful_, &failed_);
+    return Status::OK();
+  }
+
+  std::atomic<int> successful_;
+  std::atomic<int> failed_;
+};
+
+class TracingServerMiddleware : public ServerMiddleware {
+ public:
+  TracingServerMiddleware(const std::string& span_id) : span_id_(span_id) {}
+  Status SendingHeaders(AddCallHeaders& outgoing_headers) { return Status::OK(); }
+  Status CallCompleted(const Status& status) { return Status::OK(); }
+
+  std::string span_id_;
+};
+
+static thread_local std::string current_span_id = "";
+
+class TracingServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  TracingServerMiddlewareFactory() {}
+
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const auto& iter_pair = incoming_headers.GetHeaders("x-tracing-span-id");
+    if (iter_pair.first != iter_pair.second) {
+      const auto& value = (*iter_pair.first).second;
+      const std::string copy(value);
+      current_span_id = copy;
+      *middleware = std::make_shared<TracingServerMiddleware>(copy);
+    }
+    return Status::OK();
+  }
+};
+
+class PropagatingClientMiddleware : public ClientMiddleware {
+ public:
+  explicit PropagatingClientMiddleware(std::atomic<int>* received_headers,
+                                       std::vector<Status>* recorded_status)
+      : received_headers_(received_headers), recorded_status_(recorded_status) {}
+
+  Status SendingHeaders(AddCallHeaders& outgoing_headers) {
+    outgoing_headers.AddHeader("x-tracing-span-id", current_span_id);
+    return Status::OK();
+  }
+
+  Status ReceivedHeaders(const CallHeaders& incoming_headers) {
+    (*received_headers_)++;
+    return Status::OK();
+  }
+
+  Status CallCompleted(const Status& status) {
+    recorded_status_->push_back(status);
+    return Status::OK();
+  }
+
+ private:
+  std::atomic<int>* received_headers_;
+  std::vector<Status>* recorded_status_;
+};
+
+class PropagatingClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    recorded_calls_.push_back(info.method);
+    *middleware = arrow::internal::make_unique<PropagatingClientMiddleware>(
+        &received_headers_, &recorded_status_);
+    return Status::OK();
+  }
+
+  void Reset() {
+    recorded_calls_.clear();
+    recorded_status_.clear();
+    received_headers_.fetch_and(0);
+  }
+
+  std::vector<FlightMethod> recorded_calls_;
+  std::vector<Status> recorded_status_;
+  std::atomic<int> received_headers_;
+};
+
+class ReportContextTestServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    std::shared_ptr<TracingServerMiddleware> middleware =
+        std::static_pointer_cast<TracingServerMiddleware>(
+            context.GetMiddleware("tracing"));
+    std::shared_ptr<Buffer> buf;
+    if (middleware) {
+      RETURN_NOT_OK(Buffer::FromString(middleware->span_id_, &buf));
+    } else {
+      RETURN_NOT_OK(Buffer::FromString("", &buf));
+    }
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+class PropagatingTestServer : public FlightServerBase {
+ public:
+  explicit PropagatingTestServer(std::unique_ptr<FlightClient> client)
+      : client_(std::move(client)) {}
+
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    return client_->DoAction(action, result);
+  }
+
+ private:
+  std::unique_ptr<FlightClient> client_;
+};
+
+class TestServerMiddleware : public ::testing::Test {
+ public:
+  void SetUp() {
+    Location location;
+    std::unique_ptr<FlightServerBase> server(new MetadataTestServer);
+
+    ASSERT_OK(Location::ForGrpcTcp("localhost", 0, &location));
+    FlightServerOptions options(location);
+    options.middleware.push_back(
+        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
+            "test", std::make_shared<RejectServerMiddlewareFactory>()));
+    ASSERT_OK(server->Init(options));
+
+    Location real_location;
+    ASSERT_OK(Location::ForGrpcTcp("localhost", server->port(), &real_location));
+    server_.reset(new InProcessTestServer(std::move(server), real_location));
+    ASSERT_OK(server_->Start());
+    ASSERT_OK(ConnectClient());
+  }
+
+  void TearDown() { server_->Stop(); }
+
+  Status ConnectClient() { return FlightClient::Connect(server_->location(), &client_); }
+
+ protected:
+  std::unique_ptr<FlightClient> client_;
+  std::unique_ptr<InProcessTestServer> server_;
+};
+
+class TestCountingServerMiddleware : public ::testing::Test {
+ public:
+  void SetUp() {
+    Location location;
+    std::unique_ptr<FlightServerBase> server(new MetadataTestServer);
+
+    request_counter_ = std::make_shared<CountingServerMiddlewareFactory>();
+
+    ASSERT_OK(Location::ForGrpcTcp("localhost", 0, &location));
+    FlightServerOptions options(location);
+    options.middleware.push_back(
+        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
+            "counting", request_counter_));
+    ASSERT_OK(server->Init(options));
+
+    Location real_location;
+    ASSERT_OK(Location::ForGrpcTcp("localhost", server->port(), &real_location));
+    server_.reset(new InProcessTestServer(std::move(server), real_location));
+    ASSERT_OK(server_->Start());
+    ASSERT_OK(ConnectClient());
+  }
+
+  void TearDown() { server_->Stop(); }
+
+  Status ConnectClient() { return FlightClient::Connect(server_->location(), &client_); }
+
+ protected:
+  std::unique_ptr<FlightClient> client_;
+  std::unique_ptr<InProcessTestServer> server_;
+  std::shared_ptr<CountingServerMiddlewareFactory> request_counter_;
+};
+
+// Setup for this test is 2 servers
+// 1. Client makes request to server A with a request ID set
+// 2. server A extracts the request ID and makes a request to server B
+//    with the same request ID set
+// 3. server B extracts the request ID and sends it back
+// 4. server A returns the response of server B
+// 5. Client validates the response
+class TestPropagatingMiddleware : public ::testing::Test {
+ public:
+  void SetUp() {
+    Location location;
+    ASSERT_OK(Location::ForGrpcTcp("localhost", 0, &location));
+
+    server_middleware_ = std::make_shared<TracingServerMiddlewareFactory>();
+    second_client_middleware_ = std::make_shared<PropagatingClientMiddlewareFactory>();
+    client_middleware_ = std::make_shared<PropagatingClientMiddlewareFactory>();
+
+    second_server_ = arrow::internal::make_unique<ReportContextTestServer>();
+    FlightServerOptions second_server_options(location);
+    second_server_options.middleware.push_back(
+        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
+            "tracing", server_middleware_));
+    ASSERT_OK(second_server_->Init(second_server_options));
+
+    Location real_location;
+    ASSERT_OK(Location::ForGrpcTcp("localhost", second_server_->port(), &real_location));
+    std::unique_ptr<FlightClient> server_client;
+    FlightClientOptions second_client_options{};
+    second_client_options.middleware.push_back(second_client_middleware_);
+    ASSERT_OK(
+        FlightClient::Connect(real_location, second_client_options, &server_client));
+
+    first_server_ =
+        arrow::internal::make_unique<PropagatingTestServer>(std::move(server_client));
+
+    FlightServerOptions first_server_options(location);
+    first_server_options.middleware.push_back(
+        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
+            "tracing", server_middleware_));
+    ASSERT_OK(first_server_->Init(first_server_options));
+
+    ASSERT_OK(Location::ForGrpcTcp("localhost", first_server_->port(), &real_location));
+    FlightClientOptions first_client_options{};
+    first_client_options.middleware.push_back(client_middleware_);
+    ASSERT_OK(FlightClient::Connect(real_location, first_client_options, &client_));
+  }
+
+  void ValidateStatus(const Status& status, const FlightMethod& method) {
+    ASSERT_EQ(1, client_middleware_->received_headers_);
+    ASSERT_EQ(method, client_middleware_->recorded_calls_.at(0));
+    ASSERT_EQ(status.code(), client_middleware_->recorded_status_.at(0).code());
+  }
+
+  void TearDown() {
+    ASSERT_OK(first_server_->Shutdown());
+    ASSERT_OK(second_server_->Shutdown());
+  }
+
+ protected:
+  std::unique_ptr<FlightClient> client_;
+  std::unique_ptr<FlightServerBase> first_server_;
+  std::unique_ptr<FlightServerBase> second_server_;
+  std::shared_ptr<TracingServerMiddlewareFactory> server_middleware_;
+  std::shared_ptr<PropagatingClientMiddlewareFactory> second_client_middleware_;
+  std::shared_ptr<PropagatingClientMiddlewareFactory> client_middleware_;
+};
+
 TEST_F(TestFlightClient, ListFlights) {
   std::unique_ptr<FlightListing> listing;
   ASSERT_OK(client_->ListFlights(&listing));
@@ -1077,5 +1352,172 @@ TEST_F(TestMetadata, DoPutReadMetadata) {
   ASSERT_OK(writer->Close());
 }
 
+TEST_F(TestServerMiddleware, Rejected) {
+  std::unique_ptr<FlightInfo> info;
+  const auto& status = client_->GetFlightInfo(FlightDescriptor{}, &info);
+  ASSERT_RAISES(IOError, status);
+  ASSERT_THAT(status.message(), ::testing::HasSubstr("All calls are rejected"));
+}
+
+TEST_F(TestCountingServerMiddleware, Count) {
+  std::unique_ptr<FlightInfo> info;
+  const auto& status = client_->GetFlightInfo(FlightDescriptor{}, &info);
+  ASSERT_RAISES(NotImplemented, status);
+
+  Ticket ticket{""};
+  std::unique_ptr<FlightStreamReader> stream;
+  ASSERT_OK(client_->DoGet(ticket, &stream));
+
+  // The DoGet hasn't finished yet here.
+  ASSERT_EQ(0, request_counter_->successful_);
+  ASSERT_EQ(1, request_counter_->failed_);
+
+  while (true) {
+    FlightStreamChunk chunk;
+    ASSERT_OK(stream->Next(&chunk));
+    if (chunk.data == nullptr) {
+      break;
+    }
+  }
+
+  ASSERT_EQ(1, request_counter_->successful_);
+  ASSERT_EQ(1, request_counter_->failed_);
+}
+
+TEST_F(TestPropagatingMiddleware, Propagate) {
+  Action action;
+  std::unique_ptr<ResultStream> stream;
+  std::unique_ptr<Result> result;
+
+  current_span_id = "trace-id";
+  client_middleware_->Reset();
+
+  action.type = "action1";
+  const std::string action1_value = "action1-content";
+  ASSERT_OK(Buffer::FromString(action1_value, &action.body));
+  ASSERT_OK(client_->DoAction(action, &stream));
+
+  ASSERT_OK(stream->Next(&result));
+  ASSERT_EQ("trace-id", result->body->ToString());
+  ValidateStatus(Status::OK(), FlightMethod::DoAction);
+}
+
+// For each method, make sure that the client middleware received
+// headers from the server and that the proper method enum value was
+// passed to the interceptor
+TEST_F(TestPropagatingMiddleware, ListFlights) {
+  client_middleware_->Reset();
+  std::unique_ptr<FlightListing> listing;
+  const Status status = client_->ListFlights(&listing);
+  ASSERT_RAISES(NotImplemented, status);
+  ValidateStatus(status, FlightMethod::ListFlights);
+}
+
+TEST_F(TestPropagatingMiddleware, GetFlightInfo) {
+  client_middleware_->Reset();
+  auto descr = FlightDescriptor::Path({"examples", "ints"});
+  std::unique_ptr<FlightInfo> info;
+  const Status status = client_->GetFlightInfo(descr, &info);
+  ASSERT_RAISES(NotImplemented, status);
+  ValidateStatus(status, FlightMethod::GetFlightInfo);
+}
+
+TEST_F(TestPropagatingMiddleware, GetSchema) {
+  client_middleware_->Reset();
+  auto descr = FlightDescriptor::Path({"examples", "ints"});
+  std::unique_ptr<SchemaResult> result;
+  const Status status = client_->GetSchema(descr, &result);
+  ASSERT_RAISES(NotImplemented, status);
+  ValidateStatus(status, FlightMethod::GetSchema);
+}
+
+TEST_F(TestPropagatingMiddleware, ListActions) {
+  client_middleware_->Reset();
+  std::vector<ActionType> actions;
+  const Status status = client_->ListActions(&actions);
+  ASSERT_RAISES(NotImplemented, status);
+  ValidateStatus(status, FlightMethod::ListActions);
+}
+
+TEST_F(TestPropagatingMiddleware, DoGet) {
+  client_middleware_->Reset();
+  Ticket ticket1{"ARROW-5095-fail"};
+  std::unique_ptr<FlightStreamReader> stream;
+  Status status = client_->DoGet(ticket1, &stream);
+  ASSERT_RAISES(NotImplemented, status);
+  ValidateStatus(status, FlightMethod::DoGet);
+}
+
+TEST_F(TestPropagatingMiddleware, DoPut) {
+  client_middleware_->Reset();
+  auto descr = FlightDescriptor::Path({"ints"});
+  auto a1 = ArrayFromJSON(int32(), "[4, 5, 6, null]");
+  auto schema = arrow::schema({field("f1", a1->type())});
+
+  std::unique_ptr<FlightStreamWriter> stream;
+  std::unique_ptr<FlightMetadataReader> reader;
+  ASSERT_OK(client_->DoPut(descr, schema, &stream, &reader));
+  const Status status = stream->Close();
+  ASSERT_RAISES(NotImplemented, status);
+  ValidateStatus(status, FlightMethod::DoPut);
+}
+
+// Validate the implementation of the call headers wrapper
+TEST_F(TestPropagatingMiddleware, CallHeaders) {
+  std::string header1 = "x-header";
+  std::string header1_value1 = "value";
+  std::string header1_value2 = "value";
+  std::string header2 = "x-header-two";
+  std::string header2_value1 = "value-two";
+
+  internal::GrpcMetadataMap headers;
+  headers.insert({header1, header1_value1});
+  headers.insert({header2, header2_value1});
+  headers.insert({header1, header1_value2});
+
+  internal::GrpcCallHeaders incoming_headers(&headers);
+
+  ASSERT_EQ(incoming_headers.cbegin(), incoming_headers.cbegin());
+  ASSERT_EQ(incoming_headers.cend(), incoming_headers.cend());
+  ASSERT_EQ(2, incoming_headers.Count(header1));
+  ASSERT_EQ(1, incoming_headers.Count(header2));
+  {
+    const auto& header_pair = incoming_headers.GetHeaders("x-invalid");
+    ASSERT_EQ(header_pair.first, header_pair.second);
+  }
+  {
+    std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator> header_pair =
+        incoming_headers.GetHeaders(header1);
+    auto& current = header_pair.first;
+    ASSERT_NE(current, header_pair.second);
+    ASSERT_EQ(header1, (*current).first);
+    ASSERT_EQ(header1_value1, (*current).second);
+    ++current;
+    ASSERT_NE(current, header_pair.second);
+    ASSERT_EQ(header1, (*current).first);
+    ASSERT_EQ(header1_value2, (*current).second);
+    ++current;
+    ASSERT_EQ(current, header_pair.second);
+  }
+  {
+    std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator> header_pair =
+        incoming_headers.GetHeaders(header2);
+    auto& current = header_pair.first;
+    ASSERT_NE(current, header_pair.second);
+    ASSERT_EQ(header2, (*current).first);
+    ASSERT_EQ(header2_value1, (*current).second);
+    ++current;
+    ASSERT_EQ(current, header_pair.second);
+  }
+  CallHeaders::const_iterator it = incoming_headers.cbegin();
+  ASSERT_NE(it, incoming_headers.cend());
+  it++;
+  ASSERT_NE(it, incoming_headers.cend());
+  it++;
+  ASSERT_NE(it, incoming_headers.cend());
+  it++;
+  ASSERT_EQ(it, incoming_headers.cend());
+}
+
 }  // namespace flight
 }  // namespace arrow

cpp/src/arrow/flight/middleware.cc

@@ -0,0 +1,77 @@
+// 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.
+
+// Platform-specific defines
+#include "arrow/flight/platform.h"
+
+#include "arrow/flight/middleware.h"
+
+#include <map>
+
+#ifdef GRPCPP_PP_INCLUDE
+#include <grpcpp/grpcpp.h>
+#else
+#include <grpc++/grpc++.h>
+#endif
+
+#include "arrow/util/logging.h"
+#include "arrow/util/stl.h"
+
+namespace arrow {
+namespace flight {
+
+HeaderIterator::HeaderIterator() : impl_(nullptr) {}
+
+HeaderIterator::HeaderIterator(std::unique_ptr<Impl> impl) : impl_(impl.release()) {}
+
+HeaderIterator::HeaderIterator(const HeaderIterator& copy)
+    : HeaderIterator(copy.impl_->Clone()) {}
+
+const HeaderIterator::value_type& HeaderIterator::operator*() const {
+  ARROW_CHECK_NE(impl_, nullptr);
+  return impl_->Dereference();
+}
+
+HeaderIterator& HeaderIterator::operator=(const HeaderIterator& other) {
+  impl_ = other.impl_->Clone();
+  return *this;
+}
+
+HeaderIterator& HeaderIterator::operator++() {
+  ARROW_CHECK_NE(impl_, nullptr);
+  impl_->Next();
+  return *this;
+}
+
+HeaderIterator HeaderIterator::operator++(int) {
+  HeaderIterator ret = *this;
+  this->operator++();
+  return ret;
+}
+
+bool HeaderIterator::operator==(const HeaderIterator& r) const {
+  ARROW_CHECK_NE(impl_, nullptr);
+  return impl_->Equals(r.impl_.get());
+}
+
+bool HeaderIterator::operator!=(const HeaderIterator& r) const {
+  ARROW_CHECK_NE(impl_, nullptr);
+  return !(*this == r);
+}
+
+}  // namespace flight
+}  // namespace arrow

cpp/src/arrow/flight/middleware.h

@@ -0,0 +1,111 @@
+// 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.
+
+// Interfaces for defining middleware for Flight clients and
+// servers. Currently experimental.
+
+#pragma once
+
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "arrow/flight/visibility.h"  // IWYU pragma: keep
+#include "arrow/status.h"
+#include "arrow/util/string_view.h"
+
+namespace arrow {
+
+namespace flight {
+
+struct ARROW_FLIGHT_EXPORT HeaderIterator {
+ public:
+  typedef std::pair<util::string_view, util::string_view> value_type;
+
+  class Impl {
+   public:
+    virtual void Next() = 0;
+    virtual const value_type& Dereference() const = 0;
+    virtual bool Equals(const void* other) const = 0;
+    virtual std::unique_ptr<Impl> Clone() const = 0;
+    virtual ~Impl() = default;
+  };
+
+  /// \brief Nullary constructor to allow usage in Cython.
+  /// Do NOT use this constructor yourself.
+  HeaderIterator();
+  explicit HeaderIterator(std::unique_ptr<Impl> impl);
+  HeaderIterator(const HeaderIterator& copy);
+
+  ~HeaderIterator() = default;
+
+  const value_type& operator*() const;
+  HeaderIterator& operator++();
+  HeaderIterator operator++(int);
+  bool operator==(const HeaderIterator& r) const;
+  bool operator!=(const HeaderIterator& r) const;
+  HeaderIterator& operator=(const HeaderIterator& other);
+
+ private:
+  std::unique_ptr<Impl> impl_;
+};
+
+/// \brief A read-only wrapper around headers for an RPC call.
+class ARROW_FLIGHT_EXPORT CallHeaders {
+ public:
+  typedef HeaderIterator const_iterator;
+
+  virtual ~CallHeaders() = default;
+
+  virtual std::pair<const_iterator, const_iterator> GetHeaders(
+      const std::string& key) const = 0;
+  virtual std::size_t Count(const std::string& key) const = 0;
+  virtual const_iterator cbegin() const noexcept = 0;
+  virtual const_iterator cend() const noexcept = 0;
+};
+
+/// \brief A write-only wrapper around headers for an RPC call.
+class ARROW_FLIGHT_EXPORT AddCallHeaders {
+ public:
+  virtual ~AddCallHeaders() = default;
+
+  virtual void AddHeader(const std::string& key, const std::string& value) = 0;
+};
+
+/// \brief An enumeration of the RPC methods Flight implements.
+enum class FlightMethod : char {
+  Invalid = 0,
+  Handshake = 1,
+  ListFlights = 2,
+  GetFlightInfo = 3,
+  GetSchema = 4,
+  DoGet = 5,
+  DoPut = 6,
+  DoAction = 7,
+  ListActions = 8,
+};
+
+/// \brief Information about an instance of a Flight RPC.
+class ARROW_FLIGHT_EXPORT CallInfo {
+ public:
+  FlightMethod method;
+};
+
+}  // namespace flight
+
+}  // namespace arrow

cpp/src/arrow/flight/middleware_internal.cc

@@ -0,0 +1,82 @@
+// 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.
+
+// Platform-specific defines
+#include "arrow/flight/platform.h"
+
+#include "arrow/flight/middleware_internal.h"
+
+#include <map>
+
+#include "arrow/util/stl.h"
+
+namespace arrow {
+namespace flight {
+
+namespace internal {
+
+class MultimapIterator : public HeaderIterator::Impl {
+ public:
+  explicit MultimapIterator(GrpcMetadataMap::const_iterator it) : it_(it) {}
+
+  void Next() override { it_++; }
+  virtual const HeaderIterator::value_type& Dereference() const override {
+    current_header_ = util::string_view(it_->first.data(), it_->first.length());
+    current_value_ = util::string_view(it_->second.data(), it_->second.length());
+    current_ = std::make_pair(current_header_, current_value_);
+    return current_;
+  }
+  virtual bool Equals(const void* other) const override {
+    return it_ == static_cast<const MultimapIterator*>(other)->it_;
+  }
+  virtual std::unique_ptr<HeaderIterator::Impl> Clone() const override {
+    return arrow::internal::make_unique<MultimapIterator>(it_);
+  }
+  ~MultimapIterator() = default;
+
+ private:
+  GrpcMetadataMap::const_iterator it_;
+  mutable util::string_view current_header_;
+  mutable util::string_view current_value_;
+  mutable std::pair<util::string_view, util::string_view> current_;
+};
+
+std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>
+GrpcCallHeaders::GetHeaders(const std::string& key) const {
+  const auto& result = metadata_->equal_range(key);
+  return std::make_pair<const_iterator, const_iterator>(
+      HeaderIterator(arrow::internal::make_unique<MultimapIterator>(result.first)),
+      HeaderIterator(arrow::internal::make_unique<MultimapIterator>(result.second)));
+}
+
+std::size_t GrpcCallHeaders::Count(const std::string& key) const {
+  return metadata_->count(key);
+}
+
+CallHeaders::const_iterator GrpcCallHeaders::cbegin() const noexcept {
+  return HeaderIterator(
+      arrow::internal::make_unique<MultimapIterator>(metadata_->cbegin()));
+}
+
+CallHeaders::const_iterator GrpcCallHeaders::cend() const noexcept {
+  return HeaderIterator(
+      arrow::internal::make_unique<MultimapIterator>(metadata_->cend()));
+}
+
+}  // namespace internal
+}  // namespace flight
+}  // namespace arrow

cpp/src/arrow/flight/middleware_internal.h

@@ -0,0 +1,62 @@
+// 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.
+
+// Interfaces for defining middleware for Flight clients and
+// servers. Currently experimental.
+
+#pragma once
+
+#include "arrow/flight/platform.h"
+
+#include <map>
+
+#ifdef GRPCPP_PP_INCLUDE
+#include <grpcpp/grpcpp.h>
+#else
+#include <grpc++/grpc++.h>
+#endif
+
+#include "arrow/flight/middleware.h"
+
+namespace arrow {
+
+namespace flight {
+
+namespace internal {
+
+typedef std::multimap<grpc::string_ref, grpc::string_ref> GrpcMetadataMap;
+
+class GrpcCallHeaders : public CallHeaders {
+ public:
+  explicit GrpcCallHeaders(const GrpcMetadataMap* metadata) : metadata_(metadata) {}
+  ~GrpcCallHeaders() = default;
+
+  std::pair<const_iterator, const_iterator> GetHeaders(
+      const std::string& key) const override;
+  std::size_t Count(const std::string& key) const override;
+  const_iterator cbegin() const noexcept override;
+  const_iterator cend() const noexcept override;
+
+ private:
+  const GrpcMetadataMap* metadata_;
+};
+
+}  // namespace internal
+
+}  // namespace flight
+
+}  // namespace arrow

cpp/src/arrow/flight/server.cc

@@ -26,6 +26,7 @@
 #include <memory>
 #include <sstream>
 #include <string>
+#include <unordered_map>
 #include <utility>
 
 #ifdef GRPCPP_PP_INCLUDE
@@ -48,8 +49,11 @@
 #include "arrow/util/uri.h"
 
 #include "arrow/flight/internal.h"
+#include "arrow/flight/middleware.h"
+#include "arrow/flight/middleware_internal.h"
 #include "arrow/flight/serialization_internal.h"
 #include "arrow/flight/server_auth.h"
+#include "arrow/flight/server_middleware.h"
 #include "arrow/flight/types.h"
 
 using FlightService = arrow::flight::protocol::FlightService;
@@ -63,11 +67,29 @@ namespace pb = arrow::flight::protocol;
 namespace arrow {
 namespace flight {
 
-#define CHECK_ARG_NOT_NULL(VAL, MESSAGE)                              \
-  if (VAL == nullptr) {                                               \
-    return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, MESSAGE); \
+// Macro that runs interceptors before returning the given status
+#define RETURN_WITH_MIDDLEWARE(CONTEXT, STATUS) \
+  do {                                          \
+    const auto& __s = (STATUS);                 \
+    return CONTEXT.finish_request(__s);         \
+  } while (false)
+
+#define CHECK_ARG_NOT_NULL(CONTEXT, VAL, MESSAGE)                                      \
+  if (VAL == nullptr) {                                                                \
+    RETURN_WITH_MIDDLEWARE(CONTEXT,                                                    \
+                           grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, MESSAGE)); \
   }
 
+// Same as RETURN_NOT_OK, but accepts either Arrow or gRPC status, and
+// will run interceptors
+#define SERVICE_RETURN_NOT_OK(CONTEXT, expr) \
+  do {                                       \
+    const auto& _s = (expr);                 \
+    if (ARROW_PREDICT_FALSE(!_s.ok())) {     \
+      return CONTEXT.finish_request(_s);     \
+    }                                        \
+  } while (false)
+
 namespace {
 
 // A MessageReader implementation that reads from a gRPC ServerReader
@@ -215,19 +237,64 @@ class FlightServiceImpl;
 class GrpcServerCallContext : public ServerCallContext {
   const std::string& peer_identity() const override { return peer_identity_; }
 
+  std::shared_ptr<ServerMiddleware> GetMiddleware(const std::string& key) const override {
+    auto result = middleware_map_.find(key);
+    if (result == middleware_map_.end()) {
+      return nullptr;
+    }
+    return result->second;
+  }
+
+  // Helper method that runs interceptors given the result of an RPC,
+  // then returns the final gRPC status to send to the client
+  grpc::Status finish_request(const grpc::Status& status) {
+    // Don't double-convert status - return the original one here
+    finish_request(internal::FromGrpcStatus(status));
+    return status;
+  }
+
+  grpc::Status finish_request(const arrow::Status& status) {
+    for (const auto& instance : middleware_) {
+      // If an error occurs here, remainder of middleware are aborted
+      GRPC_RETURN_NOT_OK(instance->CallCompleted(status));
+    }
+    return internal::ToGrpcStatus(status);
+  }
+
  private:
   friend class FlightServiceImpl;
   ServerContext* context_;
   std::string peer_identity_;
+  std::vector<std::shared_ptr<ServerMiddleware>> middleware_;
+  mutable std::unordered_map<std::string, std::shared_ptr<ServerMiddleware>>
+      middleware_map_;
+};
+
+class GrpcAddCallHeaders : public AddCallHeaders {
+ public:
+  explicit GrpcAddCallHeaders(grpc::ServerContext* context) : context_(context) {}
+  ~GrpcAddCallHeaders() = default;
+
+  void AddHeader(const std::string& key, const std::string& value) override {
+    context_->AddInitialMetadata(key, value);
+  }
+
+ private:
+  grpc::ServerContext* context_;
 };
 
 // This class glues an implementation of FlightServerBase together with the
 // gRPC service definition, so the latter is not exposed in the public API
 class FlightServiceImpl : public FlightService::Service {
  public:
-  explicit FlightServiceImpl(std::shared_ptr<ServerAuthHandler> auth_handler,
-                             FlightServerBase* server)
-      : auth_handler_(auth_handler), server_(server) {}
+  explicit FlightServiceImpl(
+      std::shared_ptr<ServerAuthHandler> auth_handler,
+      std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
+          middleware,
+      FlightServerBase* server)
+      : auth_handler_(auth_handler),
+        middleware_(std::move(middleware)),
+        server_(server) {}
 
   template <typename UserType, typename Iterator, typename ProtoType>
   grpc::Status WriteStream(Iterator* iterator, ServerWriter<ProtoType>* writer) {
@@ -270,196 +337,251 @@ class FlightServiceImpl : public FlightService::Service {
   }
 
   // Authenticate the client (if applicable) and construct the call context
-  grpc::Status CheckAuth(ServerContext* context, GrpcServerCallContext& flight_context) {
+  grpc::Status CheckAuth(const FlightMethod& method, ServerContext* context,
+                         GrpcServerCallContext& flight_context) {
     flight_context.context_ = context;
     if (!auth_handler_) {
       flight_context.peer_identity_ = "";
-      return grpc::Status::OK;
+    } else {
+      const auto client_metadata = context->client_metadata();
+      const auto auth_header = client_metadata.find(internal::kGrpcAuthHeader);
+      std::string token;
+      if (auth_header == client_metadata.end()) {
+        token = "";
+      } else {
+        token = std::string(auth_header->second.data(), auth_header->second.length());
+      }
+      GRPC_RETURN_NOT_OK(auth_handler_->IsValid(token, &flight_context.peer_identity_));
     }
 
-    const auto client_metadata = context->client_metadata();
-    const auto auth_header = client_metadata.find(internal::kGrpcAuthHeader);
-    std::string token;
-    if (auth_header == client_metadata.end()) {
-      token = "";
-    } else {
-      token = std::string(auth_header->second.data(), auth_header->second.length());
+    return MakeCallContext(method, context, flight_context);
+  }
+
+  // Authenticate the client (if applicable) and construct the call context
+  grpc::Status MakeCallContext(const FlightMethod& method, ServerContext* context,
+                               GrpcServerCallContext& flight_context) {
+    // Run server middleware
+    const CallInfo info{method};
+    const internal::GrpcCallHeaders incoming_headers(&context->client_metadata());
+    GrpcAddCallHeaders outgoing_headers(context);
+    for (const auto& factory : middleware_) {
+      std::shared_ptr<ServerMiddleware> instance;
+      Status result = factory.second->StartCall(info, incoming_headers, &instance);
+      if (!result.ok()) {
+        // Interceptor rejected call, end the request on all existing
+        // interceptors
+        return flight_context.finish_request(result);
+      }
+      if (instance) {
+        flight_context.middleware_.push_back(instance);
+        flight_context.middleware_map_[factory.first] = instance;
+        result = instance->SendingHeaders(outgoing_headers);
+        if (!result.ok()) {
+          // Interceptor rejected call, end the request on all existing
+          // interceptors
+          return flight_context.finish_request(result);
+        }
+      }
     }
-    GRPC_RETURN_NOT_OK(auth_handler_->IsValid(token, &flight_context.peer_identity_));
+
     return grpc::Status::OK;
   }
 
   grpc::Status Handshake(
       ServerContext* context,
       grpc::ServerReaderWriter<pb::HandshakeResponse, pb::HandshakeRequest>* stream) {
+    GrpcServerCallContext flight_context;
+    GRPC_RETURN_NOT_GRPC_OK(
+        MakeCallContext(FlightMethod::Handshake, context, flight_context));
+
     if (!auth_handler_) {
-      return grpc::Status(
-          grpc::StatusCode::UNIMPLEMENTED,
-          "This service does not have an authentication mechanism enabled.");
+      RETURN_WITH_MIDDLEWARE(
+          flight_context,
+          grpc::Status(
+              grpc::StatusCode::UNIMPLEMENTED,
+              "This service does not have an authentication mechanism enabled."));
     }
     GrpcServerAuthSender outgoing{stream};
     GrpcServerAuthReader incoming{stream};
-    GRPC_RETURN_NOT_OK(auth_handler_->Authenticate(&outgoing, &incoming));
-    return grpc::Status::OK;
+    RETURN_WITH_MIDDLEWARE(flight_context,
+                           auth_handler_->Authenticate(&outgoing, &incoming));
   }
 
   grpc::Status ListFlights(ServerContext* context, const pb::Criteria* request,
                            ServerWriter<pb::FlightInfo>* writer) {
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
+    GRPC_RETURN_NOT_GRPC_OK(
+        CheckAuth(FlightMethod::ListFlights, context, flight_context));
 
     // Retrieve the listing from the implementation
     std::unique_ptr<FlightListing> listing;
 
     Criteria criteria;
     if (request) {
-      GRPC_RETURN_NOT_OK(internal::FromProto(*request, &criteria));
+      SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &criteria));
     }
-    GRPC_RETURN_NOT_OK(server_->ListFlights(flight_context, &criteria, &listing));
+    SERVICE_RETURN_NOT_OK(flight_context,
+                          server_->ListFlights(flight_context, &criteria, &listing));
     if (!listing) {
       // Treat null listing as no flights available
-      return grpc::Status::OK;
+      RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::OK);
     }
-    return WriteStream<FlightInfo>(listing.get(), writer);
+    RETURN_WITH_MIDDLEWARE(flight_context,
+                           WriteStream<FlightInfo>(listing.get(), writer));
   }
 
   grpc::Status GetFlightInfo(ServerContext* context, const pb::FlightDescriptor* request,
                              pb::FlightInfo* response) {
-    CHECK_ARG_NOT_NULL(request, "FlightDescriptor cannot be null");
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
+    GRPC_RETURN_NOT_GRPC_OK(
+        CheckAuth(FlightMethod::GetFlightInfo, context, flight_context));
+
+    CHECK_ARG_NOT_NULL(flight_context, request, "FlightDescriptor cannot be null");
 
     FlightDescriptor descr;
-    GRPC_RETURN_NOT_OK(internal::FromProto(*request, &descr));
+    SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &descr));
 
     std::unique_ptr<FlightInfo> info;
-    GRPC_RETURN_NOT_OK(server_->GetFlightInfo(flight_context, descr, &info));
+    SERVICE_RETURN_NOT_OK(flight_context,
+                          server_->GetFlightInfo(flight_context, descr, &info));
 
     if (!info) {
       // Treat null listing as no flights available
-      return grpc::Status(grpc::StatusCode::NOT_FOUND, "Flight not found");
+      RETURN_WITH_MIDDLEWARE(
+          flight_context, grpc::Status(grpc::StatusCode::NOT_FOUND, "Flight not found"));
     }
 
-    GRPC_RETURN_NOT_OK(internal::ToProto(*info, response));
-    return grpc::Status::OK;
+    SERVICE_RETURN_NOT_OK(flight_context, internal::ToProto(*info, response));
+    RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::OK);
   }
 
   grpc::Status GetSchema(ServerContext* context, const pb::FlightDescriptor* request,
                          pb::SchemaResult* response) {
-    CHECK_ARG_NOT_NULL(request, "FlightDescriptor cannot be null");
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
+    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::GetSchema, context, flight_context));
+
+    CHECK_ARG_NOT_NULL(flight_context, request, "FlightDescriptor cannot be null");
 
     FlightDescriptor descr;
-    GRPC_RETURN_NOT_OK(internal::FromProto(*request, &descr));
+    SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &descr));
 
     std::unique_ptr<SchemaResult> result;
-    GRPC_RETURN_NOT_OK(server_->GetSchema(flight_context, descr, &result));
+    SERVICE_RETURN_NOT_OK(flight_context,
+                          server_->GetSchema(flight_context, descr, &result));
 
     if (!result) {
       // Treat null listing as no flights available
-      return grpc::Status(grpc::StatusCode::NOT_FOUND, "Flight not found");
+      RETURN_WITH_MIDDLEWARE(
+          flight_context, grpc::Status(grpc::StatusCode::NOT_FOUND, "Flight not found"));
     }
 
-    GRPC_RETURN_NOT_OK(internal::ToProto(*result, response));
-    return grpc::Status::OK;
+    SERVICE_RETURN_NOT_OK(flight_context, internal::ToProto(*result, response));
+    RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::OK);
   }
 
   grpc::Status DoGet(ServerContext* context, const pb::Ticket* request,
                      ServerWriter<pb::FlightData>* writer) {
-    CHECK_ARG_NOT_NULL(request, "ticket cannot be null");
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
+    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoGet, context, flight_context));
+
+    CHECK_ARG_NOT_NULL(flight_context, request, "ticket cannot be null");
 
     Ticket ticket;
-    GRPC_RETURN_NOT_OK(internal::FromProto(*request, &ticket));
+    SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &ticket));
 
     std::unique_ptr<FlightDataStream> data_stream;
-    GRPC_RETURN_NOT_OK(server_->DoGet(flight_context, ticket, &data_stream));
+    SERVICE_RETURN_NOT_OK(flight_context,
+                          server_->DoGet(flight_context, ticket, &data_stream));
 
     if (!data_stream) {
-      return grpc::Status(grpc::StatusCode::NOT_FOUND, "No data in this flight");
+      RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status(grpc::StatusCode::NOT_FOUND,
+                                                          "No data in this flight"));
     }
 
     // Write the schema as the first message in the stream
     FlightPayload schema_payload;
-    GRPC_RETURN_NOT_OK(data_stream->GetSchemaPayload(&schema_payload));
+    SERVICE_RETURN_NOT_OK(flight_context, data_stream->GetSchemaPayload(&schema_payload));
     if (!internal::WritePayload(schema_payload, writer)) {
       // Connection terminated?  XXX return error code?
-      return grpc::Status::OK;
+      RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::OK);
     }
 
     // Consume data stream and write out payloads
     while (true) {
       FlightPayload payload;
-      GRPC_RETURN_NOT_OK(data_stream->Next(&payload));
+      SERVICE_RETURN_NOT_OK(flight_context, data_stream->Next(&payload));
       if (payload.ipc_message.metadata == nullptr ||
           !internal::WritePayload(payload, writer))
         // No more messages to write, or connection terminated for some other
         // reason
         break;
     }
-    return grpc::Status::OK;
+    RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::OK);
   }
 
   grpc::Status DoPut(ServerContext* context,
                      grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader) {
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
+    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoPut, context, flight_context));
 
     auto message_reader =
         std::unique_ptr<FlightMessageReaderImpl>(new FlightMessageReaderImpl(reader));
-    GRPC_RETURN_NOT_OK(message_reader->Init());
+    SERVICE_RETURN_NOT_OK(flight_context, message_reader->Init());
     auto metadata_writer =
         std::unique_ptr<FlightMetadataWriter>(new GrpcMetadataWriter(reader));
-    return internal::ToGrpcStatus(server_->DoPut(
-        flight_context, std::move(message_reader), std::move(metadata_writer)));
+    RETURN_WITH_MIDDLEWARE(flight_context,
+                           server_->DoPut(flight_context, std::move(message_reader),
+                                          std::move(metadata_writer)));
   }
 
   grpc::Status ListActions(ServerContext* context, const pb::Empty* request,
                            ServerWriter<pb::ActionType>* writer) {
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
+    GRPC_RETURN_NOT_GRPC_OK(
+        CheckAuth(FlightMethod::ListActions, context, flight_context));
     // Retrieve the listing from the implementation
     std::vector<ActionType> types;
-    GRPC_RETURN_NOT_OK(server_->ListActions(flight_context, &types));
-    return WriteStream<ActionType>(types, writer);
+    SERVICE_RETURN_NOT_OK(flight_context, server_->ListActions(flight_context, &types));
+    RETURN_WITH_MIDDLEWARE(flight_context, WriteStream<ActionType>(types, writer));
   }
 
   grpc::Status DoAction(ServerContext* context, const pb::Action* request,
                         ServerWriter<pb::Result>* writer) {
     GrpcServerCallContext flight_context;
-    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
-    CHECK_ARG_NOT_NULL(request, "Action cannot be null");
+    GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoAction, context, flight_context));
+    CHECK_ARG_NOT_NULL(flight_context, request, "Action cannot be null");
     Action action;
-    GRPC_RETURN_NOT_OK(internal::FromProto(*request, &action));
+    SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &action));
 
     std::unique_ptr<ResultStream> results;
-    GRPC_RETURN_NOT_OK(server_->DoAction(flight_context, action, &results));
+    SERVICE_RETURN_NOT_OK(flight_context,
+                          server_->DoAction(flight_context, action, &results));
 
     if (!results) {
-      return grpc::Status::CANCELLED;
+      RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::CANCELLED);
     }
 
     while (true) {
       std::unique_ptr<Result> result;
-      GRPC_RETURN_NOT_OK(results->Next(&result));
+      SERVICE_RETURN_NOT_OK(flight_context, results->Next(&result));
       if (!result) {
         // No more results
         break;
       }
       pb::Result pb_result;
-      GRPC_RETURN_NOT_OK(internal::ToProto(*result, &pb_result));
+      SERVICE_RETURN_NOT_OK(flight_context, internal::ToProto(*result, &pb_result));
       if (!writer->Write(pb_result)) {
         // Stream may be closed
         break;
       }
     }
-    return grpc::Status::OK;
+    RETURN_WITH_MIDDLEWARE(flight_context, grpc::Status::OK);
   }
 
  private:
   std::shared_ptr<ServerAuthHandler> auth_handler_;
+  std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
+      middleware_;
   FlightServerBase* server_;
 };
 
@@ -518,7 +640,13 @@ thread_local std::atomic<FlightServerBase::Impl*>
 #endif
 
 FlightServerOptions::FlightServerOptions(const Location& location_)
-    : location(location_), auth_handler(nullptr), tls_certificates(), builder_hook() {}
+    : location(location_),
+      auth_handler(nullptr),
+      tls_certificates(),
+      middleware(),
+      builder_hook() {}
+
+FlightServerOptions::~FlightServerOptions() = default;
 
 FlightServerBase::FlightServerBase() { impl_.reset(new Impl); }
 
@@ -526,7 +654,8 @@ FlightServerBase::~FlightServerBase() {}
 
 Status FlightServerBase::Init(FlightServerOptions& options) {
   std::shared_ptr<ServerAuthHandler> handler = std::move(options.auth_handler);
-  impl_->service_.reset(new FlightServiceImpl(handler, this));
+  impl_->service_.reset(
+      new FlightServiceImpl(handler, std::move(options.middleware), this));
 
   grpc::ServerBuilder builder;
   // Allow uploading messages of any length

cpp/src/arrow/flight/server.h

@@ -23,6 +23,7 @@
 #include <functional>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "arrow/flight/server_auth.h"
@@ -40,6 +41,9 @@ class Status;
 
 namespace flight {
 
+class ServerMiddleware;
+class ServerMiddlewareFactory;
+
 /// \brief Interface that produces a sequence of IPC payloads to be sent in
 /// FlightData protobuf messages
 class ARROW_FLIGHT_EXPORT FlightDataStream {
@@ -98,12 +102,18 @@ class ARROW_FLIGHT_EXPORT ServerCallContext {
   virtual ~ServerCallContext() = default;
   /// \brief The name of the authenticated peer (may be the empty string)
   virtual const std::string& peer_identity() const = 0;
+  /// \brief Get a middleware instance for this call.
+  /// \return nullptr if not found
+  virtual std::shared_ptr<ServerMiddleware> GetMiddleware(
+      const std::string& key) const = 0;
 };
 
 class ARROW_FLIGHT_EXPORT FlightServerOptions {
  public:
   explicit FlightServerOptions(const Location& location_);
 
+  ~FlightServerOptions();
+
   /// \brief The host & port (or domain socket path) to listen on.
   /// Use port 0 to bind to an available port.
   Location location;
@@ -111,8 +121,16 @@ class ARROW_FLIGHT_EXPORT FlightServerOptions {
   std::unique_ptr<ServerAuthHandler> auth_handler;
   /// \brief A list of TLS certificate+key pairs to use.
   std::vector<CertKeyPair> tls_certificates;
+
+  /// \brief A list of server middleware to apply.
+  ///
+  /// It is a vector rather than a map since the order of middleware matters.
+  std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
+      middleware;
+
   /// \brief A Flight implementation-specific callback to customize
   /// transport-specific options.
+  ///
   /// Not guaranteed to be called. The type of the parameter is
   /// specific to the Flight implementation. Users should take care to
   /// link to the same transport implementation as Flight to avoid

cpp/src/arrow/flight/server_middleware.h

@@ -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.
+
+// Interfaces for defining middleware for Flight servers. Currently
+// experimental.
+
+#pragma once
+
+#include "arrow/flight/middleware.h"
+#include "arrow/flight/visibility.h"  // IWYU pragma: keep
+#include "arrow/status.h"
+
+namespace arrow {
+
+namespace flight {
+
+class ARROW_FLIGHT_EXPORT ServerMiddleware {
+ public:
+  virtual ~ServerMiddleware() = default;
+
+  virtual Status SendingHeaders(AddCallHeaders& outgoing_headers) = 0;
+  virtual Status CallCompleted(const Status& status) = 0;
+};
+
+class ARROW_FLIGHT_EXPORT ServerMiddlewareFactory {
+ public:
+  virtual ~ServerMiddlewareFactory() = default;
+
+  /// \brief A callback for the start of a new call.
+  ///
+  /// Return a non-OK status to reject the call with the given status.
+  ///
+  /// \param info Information about the call.
+  /// \param incoming_headers Headers sent by the client for this call.
+  ///     Do not retain a reference to this object.
+  /// \param[out] middleware The middleware instance for this call. If
+  ///     unset, will not add middleware to this call instance from
+  ///     this factory.
+  virtual Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                           std::shared_ptr<ServerMiddleware>* middleware) = 0;
+};
+
+}  // namespace flight
+
+}  // namespace arrow

cpp/src/arrow/python/flight.cc

@@ -239,6 +239,100 @@ Status PyGeneratorFlightDataStream::Next(FlightPayload* payload) {
   });
 }
 
+// Flight Server Middleware
+
+PyServerMiddlewareFactory::PyServerMiddlewareFactory(PyObject* factory,
+                                                     StartCallCallback start_call)
+    : start_call_(start_call) {
+  Py_INCREF(factory);
+  factory_.reset(factory);
+}
+
+Status PyServerMiddlewareFactory::StartCall(
+    const arrow::flight::CallInfo& info,
+    const arrow::flight::CallHeaders& incoming_headers,
+    std::shared_ptr<arrow::flight::ServerMiddleware>* middleware) {
+  return SafeCallIntoPython([&] {
+    const Status status = start_call_(factory_.obj(), info, incoming_headers, middleware);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
+PyServerMiddleware::PyServerMiddleware(PyObject* middleware, Vtable vtable)
+    : vtable_(vtable) {
+  Py_INCREF(middleware);
+  middleware_.reset(middleware);
+}
+
+Status PyServerMiddleware::SendingHeaders(
+    arrow::flight::AddCallHeaders& outgoing_headers) {
+  return SafeCallIntoPython([&] {
+    const Status status = vtable_.sending_headers(middleware_.obj(), outgoing_headers);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
+Status PyServerMiddleware::CallCompleted(const Status& call_status) {
+  return SafeCallIntoPython([&] {
+    const Status status = vtable_.call_completed(middleware_.obj(), call_status);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
+// Flight Client Middleware
+
+PyClientMiddlewareFactory::PyClientMiddlewareFactory(PyObject* factory,
+                                                     StartCallCallback start_call)
+    : start_call_(start_call) {
+  Py_INCREF(factory);
+  factory_.reset(factory);
+}
+
+Status PyClientMiddlewareFactory::StartCall(
+    const arrow::flight::CallInfo& info,
+    std::unique_ptr<arrow::flight::ClientMiddleware>* middleware) {
+  return SafeCallIntoPython([&] {
+    const Status status = start_call_(factory_.obj(), info, middleware);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
+PyClientMiddleware::PyClientMiddleware(PyObject* middleware, Vtable vtable)
+    : vtable_(vtable) {
+  Py_INCREF(middleware);
+  middleware_.reset(middleware);
+}
+
+Status PyClientMiddleware::SendingHeaders(
+    arrow::flight::AddCallHeaders& outgoing_headers) {
+  return SafeCallIntoPython([&] {
+    const Status status = vtable_.sending_headers(middleware_.obj(), outgoing_headers);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
+Status PyClientMiddleware::ReceivedHeaders(
+    const arrow::flight::CallHeaders& incoming_headers) {
+  return SafeCallIntoPython([&] {
+    const Status status = vtable_.received_headers(middleware_.obj(), incoming_headers);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
+Status PyClientMiddleware::CallCompleted(const Status& call_status) {
+  return SafeCallIntoPython([&] {
+    const Status status = vtable_.call_completed(middleware_.obj(), call_status);
+    RETURN_NOT_OK(CheckPyError());
+    return status;
+  });
+}
+
 Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,
                         const arrow::flight::FlightDescriptor& descriptor,
                         const std::vector<arrow::flight::FlightEndpoint>& endpoints,

cpp/src/arrow/python/flight.h

@@ -180,6 +180,99 @@ class ARROW_PYTHON_EXPORT PyFlightDataStream : public arrow::flight::FlightDataS
   std::unique_ptr<arrow::flight::FlightDataStream> stream_;
 };
 
+class ARROW_PYTHON_EXPORT PyServerMiddlewareFactory
+    : public arrow::flight::ServerMiddlewareFactory {
+ public:
+  /// \brief A callback to create the middleware instance in Python
+  typedef std::function<Status(
+      PyObject*, const arrow::flight::CallInfo& info,
+      const arrow::flight::CallHeaders& incoming_headers,
+      std::shared_ptr<arrow::flight::ServerMiddleware>* middleware)>
+      StartCallCallback;
+
+  /// \brief Must only be called while holding the GIL.
+  explicit PyServerMiddlewareFactory(PyObject* factory, StartCallCallback start_call);
+
+  Status StartCall(const arrow::flight::CallInfo& info,
+                   const arrow::flight::CallHeaders& incoming_headers,
+                   std::shared_ptr<arrow::flight::ServerMiddleware>* middleware) override;
+
+ private:
+  OwnedRefNoGIL factory_;
+  StartCallCallback start_call_;
+};
+
+class ARROW_PYTHON_EXPORT PyServerMiddleware : public arrow::flight::ServerMiddleware {
+ public:
+  typedef std::function<Status(PyObject*,
+                               arrow::flight::AddCallHeaders& outgoing_headers)>
+      SendingHeadersCallback;
+  typedef std::function<Status(PyObject*, const Status& status)> CallCompletedCallback;
+
+  struct Vtable {
+    SendingHeadersCallback sending_headers;
+    CallCompletedCallback call_completed;
+  };
+
+  /// \brief Must only be called while holding the GIL.
+  explicit PyServerMiddleware(PyObject* middleware, Vtable vtable);
+
+  Status SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) override;
+  Status CallCompleted(const Status& status) override;
+
+ private:
+  OwnedRefNoGIL middleware_;
+  Vtable vtable_;
+};
+
+class ARROW_PYTHON_EXPORT PyClientMiddlewareFactory
+    : public arrow::flight::ClientMiddlewareFactory {
+ public:
+  /// \brief A callback to create the middleware instance in Python
+  typedef std::function<Status(
+      PyObject*, const arrow::flight::CallInfo& info,
+      std::unique_ptr<arrow::flight::ClientMiddleware>* middleware)>
+      StartCallCallback;
+
+  /// \brief Must only be called while holding the GIL.
+  explicit PyClientMiddlewareFactory(PyObject* factory, StartCallCallback start_call);
+
+  Status StartCall(const arrow::flight::CallInfo& info,
+                   std::unique_ptr<arrow::flight::ClientMiddleware>* middleware) override;
+
+ private:
+  OwnedRefNoGIL factory_;
+  StartCallCallback start_call_;
+};
+
+class ARROW_PYTHON_EXPORT PyClientMiddleware : public arrow::flight::ClientMiddleware {
+ public:
+  typedef std::function<Status(PyObject*,
+                               arrow::flight::AddCallHeaders& outgoing_headers)>
+      SendingHeadersCallback;
+  typedef std::function<Status(PyObject*,
+                               const arrow::flight::CallHeaders& incoming_headers)>
+      ReceivedHeadersCallback;
+  typedef std::function<Status(PyObject*, const Status& status)> CallCompletedCallback;
+
+  struct Vtable {
+    SendingHeadersCallback sending_headers;
+    ReceivedHeadersCallback received_headers;
+    CallCompletedCallback call_completed;
+  };
+
+  /// \brief Must only be called while holding the GIL.
+  explicit PyClientMiddleware(PyObject* factory, Vtable vtable);
+
+  Status SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) override;
+  Status ReceivedHeaders(const arrow::flight::CallHeaders& incoming_headers) override;
+  Status CallCompleted(const Status& status) override;
+
+ private:
+  OwnedRefNoGIL middleware_;
+  Vtable vtable_;
+};
+
 /// \brief A callback that obtains the next payload from a Flight result stream.
 typedef std::function<Status(PyObject*, arrow::flight::FlightPayload*)>
     PyGeneratorFlightDataStreamCallback;

python/pyarrow/_flight.pyx

@@ -31,6 +31,7 @@ import warnings
 import six
 
 from cython.operator cimport dereference as deref
+from cython.operator cimport postincrement
 
 from pyarrow.compat import frombytes, tobytes
 from pyarrow.lib cimport *
@@ -282,6 +283,40 @@ class DescriptorType(enum.Enum):
     CMD = 2
 
 
+class FlightMethod(enum.Enum):
+    """The implemented methods in Flight."""
+
+    INVALID = 0
+    HANDSHAKE = 1
+    LIST_FLIGHTS = 2
+    GET_FLIGHT_INFO = 3
+    GET_SCHEMA = 4
+    DO_GET = 5
+    DO_PUT = 6
+    DO_ACTION = 7
+    LIST_ACTIONS = 8
+
+
+cdef wrap_flight_method(CFlightMethod method):
+    if method == CFlightMethodHandshake:
+        return FlightMethod.HANDSHAKE
+    elif method == CFlightMethodListFlights:
+        return FlightMethod.LIST_FLIGHTS
+    elif method == CFlightMethodGetFlightInfo:
+        return FlightMethod.GET_FLIGHT_INFO
+    elif method == CFlightMethodGetSchema:
+        return FlightMethod.GET_SCHEMA
+    elif method == CFlightMethodDoGet:
+        return FlightMethod.DO_GET
+    elif method == CFlightMethodDoPut:
+        return FlightMethod.DO_PUT
+    elif method == CFlightMethodDoAction:
+        return FlightMethod.DO_ACTION
+    elif method == CFlightMethodListActions:
+        return FlightMethod.LIST_ACTIONS
+    return FlightMethod.INVALID
+
+
 cdef class FlightDescriptor:
     """A description of a data stream available from a Flight service."""
     cdef:
@@ -855,11 +890,14 @@ cdef class FlightClient:
         PEM-encoded
     override_hostname : str or None
         Override the hostname checked by TLS. Insecure, use with caution.
+    middleware : list optional, default None
+        A list of ClientMiddlewareFactory instances.
     """
     cdef:
         unique_ptr[CFlightClient] client
 
-    def __init__(self, location, tls_root_certs=None, override_hostname=None):
+    def __init__(self, location, tls_root_certs=None, override_hostname=None,
+                 middleware=None):
         if isinstance(location, six.string_types):
             location = Location(location)
         elif isinstance(location, tuple):
@@ -871,18 +909,27 @@ cdef class FlightClient:
         elif not isinstance(location, Location):
             raise TypeError('`location` argument must be a string, tuple or a '
                             'Location instance')
-        self.init(location, tls_root_certs, override_hostname)
+        self.init(location, tls_root_certs, override_hostname, middleware)
 
-    cdef init(self, Location location, tls_root_certs, override_hostname):
+    cdef init(self, Location location, tls_root_certs, override_hostname,
+              middleware):
         cdef:
             int c_port = 0
             CLocation c_location = Location.unwrap(location)
             CFlightClientOptions c_options
+            function[cb_client_middleware_start_call] start_call = \
+                &_client_middleware_start_call
 
         if tls_root_certs:
             c_options.tls_root_certs = tobytes(tls_root_certs)
         if override_hostname:
             c_options.override_hostname = tobytes(override_hostname)
+        if middleware:
+            for factory in middleware:
+                c_options.middleware.push_back(
+                    <shared_ptr[CClientMiddlewareFactory]>
+                    make_shared[CPyClientMiddlewareFactory](
+                        <PyObject*> factory, start_call))
 
         with nogil:
             check_flight_status(CFlightClient.Connect(c_location, c_options,
@@ -1597,6 +1644,104 @@ cdef CStatus _get_token(void* self, c_string* token) except *:
     return CStatus_OK()
 
 
+cdef CStatus _middleware_sending_headers(
+        void* self, CAddCallHeaders& add_headers) except *:
+    """Callback for implementing middleware."""
+    try:
+        headers = (<object> self).sending_headers()
+    except FlightError as flight_error:
+        return (<FlightError> flight_error).to_status()
+
+    if headers:
+        for header, values in headers.items():
+            if isinstance(values, (six.text_type, six.binary_type)):
+                values = (values,)
+            for value in values:
+                add_headers.AddHeader(tobytes(header), tobytes(value))
+
+    return CStatus_OK()
+
+
+cdef CStatus _middleware_call_completed(
+        void* self,
+        const CStatus& call_status) except *:
+    """Callback for implementing middleware."""
+    try:
+        try:
+            check_flight_status(call_status)
+        except Exception as e:
+            (<object> self).call_completed(e)
+        else:
+            (<object> self).call_completed(None)
+    except FlightError as flight_error:
+        return (<FlightError> flight_error).to_status()
+    return CStatus_OK()
+
+
+cdef CStatus _middleware_received_headers(
+        void* self,
+        const CCallHeaders& c_headers) except *:
+    """Callback for implementing middleware."""
+    try:
+        headers = convert_headers(c_headers)
+        (<object> self).received_headers(headers)
+    except FlightError as flight_error:
+        return (<FlightError> flight_error).to_status()
+    return CStatus_OK()
+
+
+cdef dict convert_headers(const CCallHeaders& c_headers):
+    cdef:
+        CHeaderIterator header_iter = c_headers.cbegin()
+    headers = {}
+    while header_iter != c_headers.cend():
+        header = frombytes(c_string(deref(header_iter).first))
+        value = frombytes(c_string(deref(header_iter).second))
+        if header not in headers:
+            headers[header] = []
+        headers[header].append(value)
+        postincrement(header_iter)
+    return headers
+
+
+cdef CStatus _server_middleware_start_call(
+        void* self,
+        const CCallInfo& c_info,
+        const CCallHeaders& c_headers,
+        shared_ptr[CServerMiddleware]* c_instance) except *:
+    """Callback for implementing server middleware."""
+    instance = None
+    try:
+        call_info = wrap_call_info(c_info)
+        headers = convert_headers(c_headers)
+        instance = (<object> self).start_call(call_info, headers)
+    except FlightError as flight_error:
+        return (<FlightError> flight_error).to_status()
+
+    if instance:
+        ServerMiddleware.wrap(instance, c_instance)
+
+    return CStatus_OK()
+
+
+cdef CStatus _client_middleware_start_call(
+        void* self,
+        const CCallInfo& c_info,
+        unique_ptr[CClientMiddleware]* c_instance) except *:
+    """Callback for implementing client middleware."""
+    instance = None
+    try:
+        call_info = wrap_call_info(c_info)
+        instance = (<object> self).start_call(call_info)
+    except FlightError as flight_error:
+        return (<FlightError> flight_error).to_status()
+
+    if instance:
+        ClientMiddleware.wrap(instance, c_instance)
+
+    return CStatus_OK()
+
+
 cdef class ServerAuthHandler:
     """Authentication middleware for a server.
 
@@ -1667,6 +1812,64 @@ cdef class ClientAuthHandler:
         return new PyClientAuthHandler(self, vtable)
 
 
+_CallInfo = collections.namedtuple("_CallInfo", ["method"])
+
+
+class CallInfo(_CallInfo):
+    """Information about a particular RPC for Flight middleware."""
+
+
+cdef wrap_call_info(const CCallInfo& c_info):
+    method = wrap_flight_method(c_info.method)
+    return CallInfo(method=method)
+
+
+cdef class ClientMiddlewareFactory:
+    def start_call(self, info):
+        pass
+
+
+cdef class ClientMiddleware:
+    def sending_headers(self):
+        pass
+
+    def received_headers(self, headers):
+        pass
+
+    def call_completed(self, exception):
+        pass
+
+    @staticmethod
+    cdef void wrap(object py_middleware,
+                   unique_ptr[CClientMiddleware]* c_instance):
+        cdef PyClientMiddlewareVtable vtable
+        vtable.sending_headers = _middleware_sending_headers
+        vtable.received_headers = _middleware_received_headers
+        vtable.call_completed = _middleware_call_completed
+        c_instance[0].reset(new CPyClientMiddleware(py_middleware, vtable))
+
+
+cdef class ServerMiddlewareFactory:
+    def start_call(self, info, headers):
+        pass
+
+
+cdef class ServerMiddleware:
+    def sending_headers(self):
+        pass
+
+    def call_completed(self, exception):
+        pass
+
+    @staticmethod
+    cdef void wrap(object py_middleware,
+                   shared_ptr[CServerMiddleware]* c_instance):
+        cdef PyServerMiddlewareVtable vtable
+        vtable.sending_headers = _middleware_sending_headers
+        vtable.call_completed = _middleware_call_completed
+        c_instance[0].reset(new CPyServerMiddleware(py_middleware, vtable))
+
+
 cdef class FlightServerBase:
     """A Flight service definition.
 
@@ -1683,13 +1886,15 @@ cdef class FlightServerBase:
         An authentication mechanism to use. May be None.
     tls_certificates : list optional, default None
         A list of (certificate, key) pairs.
+    middleware : dict optional, default None
+        A dictionary of (middleware key, middleware factory) items.
     """
 
     cdef:
         unique_ptr[PyFlightServer] server
 
     def __init__(self, location=None, auth_handler=None,
-                 tls_certificates=None):
+                 tls_certificates=None, middleware=None):
         if isinstance(location, six.string_types):
             location = Location(location)
         elif isinstance(location, (tuple, type(None))):
@@ -1703,15 +1908,19 @@ cdef class FlightServerBase:
         elif not isinstance(location, Location):
             raise TypeError('`location` argument must be a string, tuple or a '
                             'Location instance')
-        self.init(location, auth_handler, tls_certificates)
+        self.init(location, auth_handler, tls_certificates, middleware)
 
     cdef init(self, Location location, ServerAuthHandler auth_handler,
-              list tls_certificates):
+              list tls_certificates, dict middleware):
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             PyFlightServer* c_server
             unique_ptr[CFlightServerOptions] c_options
             CCertKeyPair c_cert
+            pair[c_string, shared_ptr[CServerMiddlewareFactory]] \
+                middleware_pair
+            function[cb_server_middleware_start_call] start_call = \
+                &_server_middleware_start_call
 
         c_options.reset(new CFlightServerOptions(Location.unwrap(location)))
 
@@ -1728,6 +1937,14 @@ cdef class FlightServerBase:
                 c_cert.pem_key = tobytes(key)
                 c_options.get().tls_certificates.push_back(c_cert)
 
+        if middleware:
+            # Takes advantage of dict ordering in Python>=3.6
+            for key, factory in middleware.items():
+                middleware_pair.first = tobytes(key)
+                middleware_pair.second.reset(
+                    new CPyServerMiddlewareFactory(factory, start_call))
+                c_options.get().middleware.push_back(middleware_pair)
+
         vtable.list_flights = &_list_flights
         vtable.get_flight_info = &_get_flight_info
         vtable.get_schema = &_get_schema
@@ -1820,7 +2037,8 @@ cdef class FlightServerBase:
         self.wait()
 
 
-def connect(location, tls_root_certs=None, override_hostname=None):
+def connect(location, tls_root_certs=None, override_hostname=None,
+            middleware=None):
     """
     Connect to the Flight server
 
@@ -1833,10 +2051,13 @@ def connect(location, tls_root_certs=None, override_hostname=None):
         PEM-encoded
     override_hostname : str or None
         Override the hostname checked by TLS. Insecure, use with caution.
+    middleware : list or None
+        A list of ClientMiddlewareFactory instances to apply.
 
     Returns
     -------
     client : FlightClient
     """
     return FlightClient(location, tls_root_certs=tls_root_certs,
-                        override_hostname=override_hostname)
+                        override_hostname=override_hostname,
+                        middleware=middleware)

python/pyarrow/flight.py

@@ -49,5 +49,9 @@ from pyarrow._flight import (  # noqa
     Result,
     ClientAuthHandler,
     ServerAuthHandler,
+    ClientMiddleware,
+    ClientMiddlewareFactory,
+    ServerMiddleware,
+    ServerMiddlewareFactory,
     BasicAuth
 )

python/pyarrow/includes/libarrow_flight.pxd

@@ -201,16 +201,71 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
         c_string pem_cert
         c_string pem_key
 
+    cdef cppclass CFlightMethod" arrow::flight::FlightMethod":
+        bint operator==(CFlightMethod)
+
+    CFlightMethod CFlightMethodInvalid\
+        " arrow::flight::FlightMethod::Invalid"
+    CFlightMethod CFlightMethodHandshake\
+        " arrow::flight::FlightMethod::Handshake"
+    CFlightMethod CFlightMethodListFlights\
+        " arrow::flight::FlightMethod::ListFlights"
+    CFlightMethod CFlightMethodGetFlightInfo\
+        " arrow::flight::FlightMethod::GetFlightInfo"
+    CFlightMethod CFlightMethodGetSchema\
+        " arrow::flight::FlightMethod::GetSchema"
+    CFlightMethod CFlightMethodDoGet\
+        " arrow::flight::FlightMethod::DoGet"
+    CFlightMethod CFlightMethodDoPut\
+        " arrow::flight::FlightMethod::DoPut"
+    CFlightMethod CFlightMethodDoAction\
+        " arrow::flight::FlightMethod::DoAction"
+    CFlightMethod CFlightMethodListActions\
+        " arrow::flight::FlightMethod::ListActions"
+
+    cdef cppclass CCallInfo" arrow::flight::CallInfo":
+        CFlightMethod method
+
+    cdef cppclass CHeaderIterator" arrow::flight::HeaderIterator":
+        CHeaderIterator& operator++()
+        const pair[c_string, c_string] operator*()
+        bint operator==(const CHeaderIterator& r)
+        bint operator!=(const CHeaderIterator& r)
+
+    cdef cppclass CCallHeaders" arrow::flight::CallHeaders":
+        size_t Count(const c_string& key)
+        CHeaderIterator cbegin()
+        CHeaderIterator cend()
+
+    cdef cppclass CAddCallHeaders" arrow::flight::AddCallHeaders":
+        void AddHeader(const c_string& key, const c_string& value)
+
+    cdef cppclass CServerMiddleware" arrow::flight::ServerMiddleware":
+        pass
+
+    cdef cppclass CServerMiddlewareFactory\
+            " arrow::flight::ServerMiddlewareFactory":
+        pass
+
+    cdef cppclass CClientMiddleware" arrow::flight::ClientMiddleware":
+        pass
+
+    cdef cppclass CClientMiddlewareFactory\
+            " arrow::flight::ClientMiddlewareFactory":
+        pass
+
     cdef cppclass CFlightServerOptions" arrow::flight::FlightServerOptions":
         CFlightServerOptions(const CLocation& location)
         CLocation location
         unique_ptr[CServerAuthHandler] auth_handler
         vector[CCertKeyPair] tls_certificates
+        vector[pair[c_string, shared_ptr[CServerMiddlewareFactory]]] middleware
 
     cdef cppclass CFlightClientOptions" arrow::flight::FlightClientOptions":
         CFlightClientOptions()
         c_string tls_root_certs
         c_string override_hostname
+        vector[shared_ptr[CClientMiddlewareFactory]] middleware
 
     cdef cppclass CFlightClient" arrow::flight::FlightClient":
         @staticmethod
@@ -298,6 +353,20 @@ ctypedef CStatus cb_client_authenticate(object, CClientAuthSender*,
                                         CClientAuthReader*)
 ctypedef CStatus cb_get_token(object, c_string*)
 
+ctypedef CStatus cb_middleware_sending_headers(object, CAddCallHeaders&)
+ctypedef CStatus cb_middleware_call_completed(object, const CStatus&)
+ctypedef CStatus cb_client_middleware_received_headers(
+    object, const CCallHeaders&)
+ctypedef CStatus cb_server_middleware_start_call(
+    object,
+    const CCallInfo&,
+    const CCallHeaders&,
+    shared_ptr[CServerMiddleware]*)
+ctypedef CStatus cb_client_middleware_start_call(
+    object,
+    const CCallInfo&,
+    unique_ptr[CClientMiddleware]*)
+
 cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
     cdef cppclass PyFlightServerVtable:
         PyFlightServerVtable()
@@ -353,6 +422,41 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
                                      shared_ptr[CSchema] schema,
                                      function[cb_data_stream_next] callback)
 
+    cdef cppclass PyServerMiddlewareVtable\
+            " arrow::py::flight::PyServerMiddleware::Vtable":
+        PyServerMiddlewareVtable()
+        function[cb_middleware_sending_headers] sending_headers
+        function[cb_middleware_call_completed] call_completed
+
+    cdef cppclass PyClientMiddlewareVtable\
+            " arrow::py::flight::PyClientMiddleware::Vtable":
+        PyClientMiddlewareVtable()
+        function[cb_middleware_sending_headers] sending_headers
+        function[cb_client_middleware_received_headers] received_headers
+        function[cb_middleware_call_completed] call_completed
+
+    cdef cppclass CPyServerMiddleware\
+            " arrow::py::flight::PyServerMiddleware"(CServerMiddleware):
+        CPyServerMiddleware(object middleware, PyServerMiddlewareVtable vtable)
+
+    cdef cppclass CPyServerMiddlewareFactory\
+            " arrow::py::flight::PyServerMiddlewareFactory"\
+            (CServerMiddlewareFactory):
+        CPyServerMiddlewareFactory(
+            object factory,
+            function[cb_server_middleware_start_call] start_call)
+
+    cdef cppclass CPyClientMiddleware\
+            " arrow::py::flight::PyClientMiddleware"(CClientMiddleware):
+        CPyClientMiddleware(object middleware, PyClientMiddlewareVtable vtable)
+
+    cdef cppclass CPyClientMiddlewareFactory\
+            " arrow::py::flight::PyClientMiddlewareFactory"\
+            (CClientMiddlewareFactory):
+        CPyClientMiddlewareFactory(
+            object factory,
+            function[cb_client_middleware_start_call] start_call)
+
     cdef CStatus CreateFlightInfo" arrow::py::flight::CreateFlightInfo"(
         shared_ptr[CSchema] schema,
         CFlightDescriptor& descriptor,