[Python] Implement auth API for Flight

cpp/src/arrow/python/flight.cc

@@ -26,6 +26,49 @@ namespace arrow {
 namespace py {
 namespace flight {
 
+PyServerAuthHandler::PyServerAuthHandler(PyObject* handler,
+                                         PyServerAuthHandlerVtable vtable)
+    : vtable_(vtable) {
+  Py_INCREF(handler);
+  handler_.reset(handler);
+}
+
+Status PyServerAuthHandler::Authenticate(
+    const arrow::flight::ServerAuthSender& outgoing,
+    const arrow::flight::ServerAuthReader& incoming) {
+  PyAcquireGIL lock;
+  vtable_.authenticate(handler_.obj(), outgoing, incoming);
+  return CheckPyError();
+}
+
+Status PyServerAuthHandler::IsValid(const std::string& token,
+                                    std::string* peer_identity) {
+  PyAcquireGIL lock;
+  vtable_.is_valid(handler_.obj(), token, peer_identity);
+  return CheckPyError();
+}
+
+PyClientAuthHandler::PyClientAuthHandler(PyObject* handler,
+                                         PyClientAuthHandlerVtable vtable)
+    : vtable_(vtable) {
+  Py_INCREF(handler);
+  handler_.reset(handler);
+}
+
+Status PyClientAuthHandler::Authenticate(
+    const arrow::flight::ClientAuthSender& outgoing,
+    const arrow::flight::ClientAuthReader& incoming) {
+  PyAcquireGIL lock;
+  vtable_.authenticate(handler_.obj(), outgoing, incoming);
+  return CheckPyError();
+}
+
+Status PyClientAuthHandler::GetToken(std::string* token) {
+  PyAcquireGIL lock;
+  vtable_.get_token(handler_.obj(), token);
+  return CheckPyError();
+}
+
 PyFlightServer::PyFlightServer(PyObject* server, PyFlightServerVtable vtable)
     : vtable_(vtable) {
   Py_INCREF(server);
@@ -37,40 +80,45 @@ Status PyFlightServer::ListFlights(
     const arrow::flight::Criteria* criteria,
     std::unique_ptr<arrow::flight::FlightListing>* listings) {
   PyAcquireGIL lock;
-  vtable_.list_flights(server_.obj(), criteria, listings);
+  vtable_.list_flights(server_.obj(), context, criteria, listings);
   return CheckPyError();
 }
 
-Status PyFlightServer::GetFlightInfo(const arrow::flight::FlightDescriptor& request,
+Status PyFlightServer::GetFlightInfo(const arrow::flight::ServerCallContext& context,
+                                     const arrow::flight::FlightDescriptor& request,
                                      std::unique_ptr<arrow::flight::FlightInfo>* info) {
   PyAcquireGIL lock;
-  vtable_.get_flight_info(server_.obj(), request, info);
+  vtable_.get_flight_info(server_.obj(), context, request, info);
   return CheckPyError();
 }
 
-Status PyFlightServer::DoGet(const arrow::flight::Ticket& request,
+Status PyFlightServer::DoGet(const arrow::flight::ServerCallContext& context,
+                             const arrow::flight::Ticket& request,
                              std::unique_ptr<arrow::flight::FlightDataStream>* stream) {
   PyAcquireGIL lock;
-  vtable_.do_get(server_.obj(), request, stream);
+  vtable_.do_get(server_.obj(), context, request, stream);
   return CheckPyError();
 }
 
-Status PyFlightServer::DoPut(std::unique_ptr<arrow::flight::FlightMessageReader> reader) {
+Status PyFlightServer::DoPut(const arrow::flight::ServerCallContext& context,
+                             std::unique_ptr<arrow::flight::FlightMessageReader> reader) {
   PyAcquireGIL lock;
-  vtable_.do_put(server_.obj(), std::move(reader));
+  vtable_.do_put(server_.obj(), context, std::move(reader));
   return CheckPyError();
 }
 
-Status PyFlightServer::DoAction(const arrow::flight::Action& action,
+Status PyFlightServer::DoAction(const arrow::flight::ServerCallContext& context,
+                                const arrow::flight::Action& action,
                                 std::unique_ptr<arrow::flight::ResultStream>* result) {
   PyAcquireGIL lock;
-  vtable_.do_action(server_.obj(), action, result);
+  vtable_.do_action(server_.obj(), context, action, result);
   return CheckPyError();
 }
 
-Status PyFlightServer::ListActions(std::vector<arrow::flight::ActionType>* actions) {
+Status PyFlightServer::ListActions(const arrow::flight::ServerCallContext& context,
+                                   std::vector<arrow::flight::ActionType>* actions) {
   PyAcquireGIL lock;
-  vtable_.list_actions(server_.obj(), actions);
+  vtable_.list_actions(server_.obj(), context, actions);
   return CheckPyError();
 }
 

cpp/src/arrow/python/flight.h

@@ -19,6 +19,7 @@
 #define PYARROW_FLIGHT_H
 
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "arrow/flight/api.h"
@@ -35,21 +36,70 @@ namespace flight {
 /// Python.
 class ARROW_PYTHON_EXPORT PyFlightServerVtable {
  public:
-  std::function<void(PyObject*, const arrow::flight::Criteria*,
+  std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
+                     const arrow::flight::Criteria*,
                      std::unique_ptr<arrow::flight::FlightListing>*)>
       list_flights;
-  std::function<void(PyObject*, const arrow::flight::FlightDescriptor&,
+  std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
+                     const arrow::flight::FlightDescriptor&,
                      std::unique_ptr<arrow::flight::FlightInfo>*)>
       get_flight_info;
-  std::function<void(PyObject*, const arrow::flight::Ticket&,
+  std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
+                     const arrow::flight::Ticket&,
                      std::unique_ptr<arrow::flight::FlightDataStream>*)>
       do_get;
-  std::function<void(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
+  std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
+                     std::unique_ptr<arrow::flight::FlightMessageReader>)>
       do_put;
-  std::function<void(PyObject*, const arrow::flight::Action&,
+  std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
+                     const arrow::flight::Action&,
                      std::unique_ptr<arrow::flight::ResultStream>*)>
       do_action;
-  std::function<void(PyObject*, std::vector<arrow::flight::ActionType>*)> list_actions;
+  std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
+                     std::vector<arrow::flight::ActionType>*)>
+      list_actions;
+};
+
+class ARROW_PYTHON_EXPORT PyServerAuthHandlerVtable {
+ public:
+  std::function<void(PyObject*, const arrow::flight::ServerAuthSender&,
+                     const arrow::flight::ServerAuthReader&)>
+      authenticate;
+  std::function<void(PyObject*, const std::string&, std::string*)> is_valid;
+};
+
+class ARROW_PYTHON_EXPORT PyClientAuthHandlerVtable {
+ public:
+  std::function<void(PyObject*, const arrow::flight::ClientAuthSender&,
+                     const arrow::flight::ClientAuthReader&)>
+      authenticate;
+  std::function<void(PyObject*, std::string*)> get_token;
+};
+
+/// \brief A helper to implement an auth mechanism in Python.
+class ARROW_PYTHON_EXPORT PyServerAuthHandler : public arrow::flight::ServerAuthHandler {
+ public:
+  explicit PyServerAuthHandler(PyObject* handler, PyServerAuthHandlerVtable vtable);
+  Status Authenticate(const arrow::flight::ServerAuthSender& outgoing,
+                      const arrow::flight::ServerAuthReader& incoming) override;
+  Status IsValid(const std::string& token, std::string* peer_identity) override;
+
+ private:
+  OwnedRefNoGIL handler_;
+  PyServerAuthHandlerVtable vtable_;
+};
+
+/// \brief A helper to implement an auth mechanism in Python.
+class ARROW_PYTHON_EXPORT PyClientAuthHandler : public arrow::flight::ClientAuthHandler {
+ public:
+  explicit PyClientAuthHandler(PyObject* handler, PyClientAuthHandlerVtable vtable);
+  Status Authenticate(const arrow::flight::ClientAuthSender& outgoing,
+                      const arrow::flight::ClientAuthReader& incoming) override;
+  Status GetToken(std::string* token) override;
+
+ private:
+  OwnedRefNoGIL handler_;
+  PyClientAuthHandlerVtable vtable_;
 };
 
 class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBase {
@@ -63,14 +113,19 @@ class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBas
   Status ListFlights(const arrow::flight::ServerCallContext& context,
                      const arrow::flight::Criteria* criteria,
                      std::unique_ptr<arrow::flight::FlightListing>* listings) override;
-  Status GetFlightInfo(const arrow::flight::FlightDescriptor& request,
+  Status GetFlightInfo(const arrow::flight::ServerCallContext& context,
+                       const arrow::flight::FlightDescriptor& request,
                        std::unique_ptr<arrow::flight::FlightInfo>* info) override;
-  Status DoGet(const arrow::flight::Ticket& request,
+  Status DoGet(const arrow::flight::ServerCallContext& context,
+               const arrow::flight::Ticket& request,
                std::unique_ptr<arrow::flight::FlightDataStream>* stream) override;
-  Status DoPut(std::unique_ptr<arrow::flight::FlightMessageReader> reader) override;
-  Status DoAction(const arrow::flight::Action& action,
+  Status DoPut(const arrow::flight::ServerCallContext& context,
+               std::unique_ptr<arrow::flight::FlightMessageReader> reader) override;
+  Status DoAction(const arrow::flight::ServerCallContext& context,
+                  const arrow::flight::Action& action,
                   std::unique_ptr<arrow::flight::ResultStream>* result) override;
-  Status ListActions(std::vector<arrow::flight::ActionType>* actions) override;
+  Status ListActions(const arrow::flight::ServerCallContext& context,
+                     std::vector<arrow::flight::ActionType>* actions) override;
 
  private:
   OwnedRefNoGIL server_;

python/examples/flight/server.py

@@ -35,7 +35,7 @@ class FlightServer(pyarrow.flight.FlightServerBase):
         return (descriptor.descriptor_type.value, descriptor.command,
                 tuple(descriptor.path or tuple()))
 
-    def list_flights(self, criteria):
+    def list_flights(self, context, criteria):
         for key, table in self.flights.items():
             if key[1] is not None:
                 descriptor = \
@@ -51,7 +51,7 @@ class FlightServer(pyarrow.flight.FlightServerBase):
                                             descriptor, endpoints,
                                             table.num_rows, 0)
 
-    def get_flight_info(self, descriptor):
+    def get_flight_info(self, context, descriptor):
         key = FlightServer.descriptor_to_key(descriptor)
         if key in self.flights:
             table = self.flights[key]
@@ -64,25 +64,25 @@ class FlightServer(pyarrow.flight.FlightServerBase):
                                              table.num_rows, 0)
         raise KeyError('Flight not found.')
 
-    def do_put(self, descriptor, reader):
+    def do_put(self, context, descriptor, reader):
         key = FlightServer.descriptor_to_key(descriptor)
         print(key)
         self.flights[key] = reader.read_all()
         print(self.flights[key])
 
-    def do_get(self, ticket):
+    def do_get(self, context, ticket):
         key = ast.literal_eval(ticket.ticket.decode())
         if key not in self.flights:
             return None
         return pyarrow.flight.RecordBatchStream(self.flights[key])
 
-    def list_actions(self):
+    def list_actions(self, context):
         return [
             ("clear", "Clear the stored flights."),
             ("shutdown", "Shut down this server."),
         ]
 
-    def do_action(self, action):
+    def do_action(self, context, action):
         if action.type == "clear":
             raise NotImplementedError(
                 "{} is not implemented.".format(action.type))

python/pyarrow/_flight.pyx

@@ -312,6 +312,17 @@ cdef class FlightClient:
 
         return result
 
+    def authenticate(self, auth_handler):
+        """Authenticate to the server."""
+        cdef unique_ptr[CClientAuthHandler] handler
+        if not isinstance(auth_handler, ClientAuthHandler):
+            raise TypeError(
+                "FlightClient.authenticate takes a ClientAuthHandler, "
+                "not '{}'".format(type(auth_handler)))
+        handler.reset((<ClientAuthHandler> auth_handler).to_handler())
+        with nogil:
+            check_status(self.client.get().Authenticate(move(handler)))
+
     def list_actions(self):
         """List the actions available on a service."""
         cdef:
@@ -474,6 +485,148 @@ cdef class GeneratorStream(FlightDataStream):
         return new CPyGeneratorFlightDataStream(self, self.schema, callback)
 
 
+cdef class ServerCallContext:
+    """Per-call state/context."""
+    cdef:
+        const CServerCallContext* context
+
+    def peer_identity(self):
+        """Get the identity of the authenticated peer.
+
+        May be the empty string.
+        """
+        return tobytes(self.context.peer_identity())
+
+    @staticmethod
+    cdef ServerCallContext wrap(const CServerCallContext& context):
+        cdef ServerCallContext result = \
+            ServerCallContext.__new__(ServerCallContext)
+        result.context = &context
+        return result
+
+
+cdef class ServerAuthReader:
+    """A reader for messages from the client during an auth handshake."""
+    cdef:
+        const CServerAuthReader* reader
+
+    def read(self):
+        cdef c_string token
+        if not self.reader:
+            raise ValueError("Cannot use ServerAuthReader outside "
+                             "ServerAuthHandler.authenticate")
+        with nogil:
+            check_status(self.reader.Read(&token))
+        return token
+
+    cdef void poison(self):
+        """Prevent further usage of this object.
+
+        This object is constructed by taking a pointer to a reference,
+        so we want to make sure Python users do not access this after
+        the reference goes away.
+        """
+        self.reader = NULL
+
+    @staticmethod
+    cdef ServerAuthReader wrap(const CServerAuthReader& reader):
+        cdef ServerAuthReader result = \
+            ServerAuthReader.__new__(ServerAuthReader)
+        result.reader = &reader
+        return result
+
+
+cdef class ServerAuthSender:
+    """A writer for messages to the client during an auth handshake."""
+    cdef:
+        const CServerAuthSender* sender
+
+    def write(self, message):
+        cdef c_string c_message = tobytes(message)
+        if not self.sender:
+            raise ValueError("Cannot use ServerAuthSender outside "
+                             "ServerAuthHandler.authenticate")
+        with nogil:
+            check_status(self.sender.Write(c_message))
+
+    cdef void poison(self):
+        """Prevent further usage of this object.
+
+        This object is constructed by taking a pointer to a reference,
+        so we want to make sure Python users do not access this after
+        the reference goes away.
+        """
+        self.sender = NULL
+
+    @staticmethod
+    cdef ServerAuthSender wrap(const CServerAuthSender& sender):
+        cdef ServerAuthSender result = \
+            ServerAuthSender.__new__(ServerAuthSender)
+        result.sender = &sender
+        return result
+
+
+cdef class ClientAuthReader:
+    """A reader for messages from the server during an auth handshake."""
+    cdef:
+        const CClientAuthReader* reader
+
+    def read(self):
+        cdef c_string token
+        if not self.reader:
+            raise ValueError("Cannot use ClientAuthReader outside "
+                             "ClientAuthHandler.authenticate")
+        with nogil:
+            check_status(self.reader.Read(&token))
+        return token
+
+    cdef void poison(self):
+        """Prevent further usage of this object.
+
+        This object is constructed by taking a pointer to a reference,
+        so we want to make sure Python users do not access this after
+        the reference goes away.
+        """
+        self.reader = NULL
+
+    @staticmethod
+    cdef ClientAuthReader wrap(const CClientAuthReader& reader):
+        cdef ClientAuthReader result = \
+            ClientAuthReader.__new__(ClientAuthReader)
+        result.reader = &reader
+        return result
+
+
+cdef class ClientAuthSender:
+    """A writer for messages to the server during an auth handshake."""
+    cdef:
+        const CClientAuthSender* sender
+
+    def write(self, message):
+        cdef c_string c_message = tobytes(message)
+        if not self.sender:
+            raise ValueError("Cannot use ClientAuthSender outside "
+                             "ClientAuthHandler.authenticate")
+        with nogil:
+            check_status(self.sender.Write(c_message))
+
+    cdef void poison(self):
+        """Prevent further usage of this object.
+
+        This object is constructed by taking a pointer to a reference,
+        so we want to make sure Python users do not access this after
+        the reference goes away.
+        """
+        self.sender = NULL
+
+    @staticmethod
+    cdef ClientAuthSender wrap(const CClientAuthSender& sender):
+        cdef ClientAuthSender result = \
+            ClientAuthSender.__new__(ClientAuthSender)
+        result.sender = &sender
+        return result
+
+
 cdef void _data_stream_next(void* self, CFlightPayload* payload) except *:
     """Callback for implementing FlightDataStream in Python."""
     cdef:
@@ -534,12 +687,14 @@ cdef void _data_stream_next(void* self, CFlightPayload* payload) except *:
                         "not {}.".format(type(result)))
 
 
-cdef void _list_flights(void* self, const CCriteria* c_criteria,
+cdef void _list_flights(void* self, const CServerCallContext& context,
+                        const CCriteria* c_criteria,
                         unique_ptr[CFlightListing]* listing) except *:
     """Callback for implementing ListFlights in Python."""
     cdef:
         vector[CFlightInfo] flights
-    result = (<object> self).list_flights(c_criteria.expression)
+    result = (<object> self).list_flights(ServerCallContext.wrap(context),
+                                          c_criteria.expression)
     for info in result:
         if not isinstance(info, FlightInfo):
             raise TypeError("FlightServerBase.list_flights must return "
@@ -549,14 +704,16 @@ cdef void _list_flights(void* self, const CCriteria* c_criteria,
     listing.reset(new CSimpleFlightListing(flights))
 
 
-cdef void _get_flight_info(void* self, CFlightDescriptor c_descriptor,
+cdef void _get_flight_info(void* self, const CServerCallContext& context,
+                           CFlightDescriptor c_descriptor,
                            unique_ptr[CFlightInfo]* info) except *:
     """Callback for implementing Flight servers in Python."""
     cdef:
         FlightDescriptor py_descriptor = \
             FlightDescriptor.__new__(FlightDescriptor)
     py_descriptor.descriptor = c_descriptor
-    result = (<object> self).get_flight_info(py_descriptor)
+    result = (<object> self).get_flight_info(ServerCallContext.wrap(context),
+                                             py_descriptor)
     if not isinstance(result, FlightInfo):
         raise TypeError("FlightServerBase.get_flight_info must return "
                         "a FlightInfo instance, but got {}".format(
@@ -564,7 +721,7 @@ cdef void _get_flight_info(void* self, CFlightDescriptor c_descriptor,
     info.reset(new CFlightInfo(deref((<FlightInfo> result).info.get())))
 
 
-cdef void _do_put(void* self,
+cdef void _do_put(void* self, const CServerCallContext& context,
                   unique_ptr[CFlightMessageReader] reader) except *:
     """Callback for implementing Flight servers in Python."""
     cdef:
@@ -574,17 +731,20 @@ cdef void _do_put(void* self,
 
     descriptor.descriptor = reader.get().descriptor()
     py_reader.reader.reset(reader.release())
-    (<object> self).do_put(descriptor, py_reader)
+    (<object> self).do_put(ServerCallContext.wrap(context), descriptor,
+                           py_reader)
 
 
-cdef void _do_get(void* self, CTicket ticket,
+cdef void _do_get(void* self, const CServerCallContext& context,
+                  CTicket ticket,
                   unique_ptr[CFlightDataStream]* stream) except *:
     """Callback for implementing Flight servers in Python."""
     cdef:
         unique_ptr[CFlightDataStream] data_stream
 
     py_ticket = Ticket(ticket.ticket)
-    result = (<object> self).do_get(py_ticket)
+    result = (<object> self).do_get(ServerCallContext.wrap(context),
+                                    py_ticket)
     if not isinstance(result, FlightDataStream):
         raise TypeError("FlightServerBase.do_get must return "
                         "a FlightDataStream")
@@ -607,17 +767,20 @@ cdef void _do_action_result_next(void* self,
         result.reset(nullptr)
 
 
-cdef void _do_action(void* self, const CAction& action,
+cdef void _do_action(void* self, const CServerCallContext& context,
+                     const CAction& action,
                      unique_ptr[CResultStream]* result) except *:
     """Callback for implementing Flight servers in Python."""
     cdef:
         function[cb_result_next] ptr = &_do_action_result_next
     py_action = Action(action.type, pyarrow_wrap_buffer(action.body))
-    responses = (<object> self).do_action(py_action)
+    responses = (<object> self).do_action(ServerCallContext.wrap(context),
+                                          py_action)
     result.reset(new CPyFlightResultStream(responses, ptr))
 
 
-cdef void _list_actions(void* self, vector[CActionType]* actions) except *:
+cdef void _list_actions(void* self, const CServerCallContext& context,
+                        vector[CActionType]* actions) except *:
     """Callback for implementing Flight servers in Python."""
     cdef:
         CActionType action_type
@@ -629,17 +792,106 @@ cdef void _list_actions(void* self, vector[CActionType]* actions) except *:
         actions.push_back(action_type)
 
 
+cdef void _server_authenticate(void* self, const CServerAuthSender& outgoing,
+                               const CServerAuthReader& incoming) except *:
+    """Callback for implementing authentication in Python."""
+    sender = ServerAuthSender.wrap(outgoing)
+    reader = ServerAuthReader.wrap(incoming)
+    try:
+        (<object> self).authenticate(sender, reader)
+    finally:
+        sender.poison()
+        reader.poison()
+
+
+cdef void _is_valid(void* self, const c_string& token,
+                    c_string* peer_identity) except *:
+    """Callback for implementing authentication in Python."""
+    cdef c_string c_result
+    c_result = tobytes((<object> self).is_valid(token))
+    peer_identity[0] = c_result
+
+
+cdef void _client_authenticate(void* self, const CClientAuthSender& outgoing,
+                               const CClientAuthReader& incoming) except *:
+    """Callback for implementing authentication in Python."""
+    sender = ClientAuthSender.wrap(outgoing)
+    reader = ClientAuthReader.wrap(incoming)
+    try:
+        (<object> self).authenticate(sender, reader)
+    finally:
+        sender.poison()
+        reader.poison()
+
+
+cdef void _get_token(void* self, c_string* token) except *:
+    """Callback for implementing authentication in Python."""
+    cdef c_string c_result
+    c_result = tobytes((<object> self).get_token())
+    token[0] = c_result
+
+
+cdef class ServerAuthHandler:
+    """Authentication middleware for a server."""
+
+    def authenticate(self, outgoing, incoming):
+        """Conduct the handshake with the client."""
+        raise NotImplementedError
+
+    def is_valid(self, token):
+        """Validate a client token, returning their identity.
+
+        May return an empty string (if the auth mechanism does not
+        name the peer) or raise an exception (if the token is
+        invalid).
+
+        """
+        raise NotImplementedError
+
+    cdef PyServerAuthHandler* to_handler(self):
+        cdef PyServerAuthHandlerVtable vtable
+        vtable.authenticate = _server_authenticate
+        vtable.is_valid = _is_valid
+        return new PyServerAuthHandler(self, vtable)
+
+
+cdef class ClientAuthHandler:
+    """Authentication plugin for a client."""
+
+    def authenticate(self, outgoing, incoming):
+        """Conduct the handshake with the server."""
+        raise NotImplementedError
+
+    def get_token(self):
+        """Get the auth token for a call."""
+        raise NotImplementedError
+
+    cdef PyClientAuthHandler* to_handler(self):
+        cdef PyClientAuthHandlerVtable vtable
+        vtable.authenticate = _client_authenticate
+        vtable.get_token = _get_token
+        return new PyClientAuthHandler(self, vtable)
+
+
 cdef class FlightServerBase:
     """A Flight service definition."""
 
     cdef:
         unique_ptr[PyFlightServer] server
 
-    def run(self, port):
+    def run(self, port, auth_handler=None):
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             int c_port = port
             PyFlightServer* c_server
+            unique_ptr[CServerAuthHandler] c_auth_handler
+
+        if auth_handler:
+            if not isinstance(auth_handler, ServerAuthHandler):
+                raise TypeError("auth_handler must be a ServerAuthHandler, "
+                                "not a '{}'".format(type(auth_handler)))
+            c_auth_handler.reset(
+                (<ServerAuthHandler> auth_handler).to_handler())
 
         vtable.list_flights = &_list_flights
         vtable.get_flight_info = &_get_flight_info
@@ -651,25 +903,25 @@ cdef class FlightServerBase:
         c_server = new PyFlightServer(self, vtable)
         self.server.reset(c_server)
         with nogil:
-            check_status(c_server.Init(c_port))
+            check_status(c_server.Init(move(c_auth_handler), c_port))
             check_status(c_server.ServeWithSignals())
 
-    def list_flights(self, criteria):
+    def list_flights(self, context, criteria):
         raise NotImplementedError
 
-    def get_flight_info(self, descriptor):
+    def get_flight_info(self, context, descriptor):
         raise NotImplementedError
 
-    def do_put(self, descriptor, reader):
+    def do_put(self, context, descriptor, reader):
         raise NotImplementedError
 
-    def do_get(self, ticket):
+    def do_get(self, context, ticket):
         raise NotImplementedError
 
-    def list_actions(self):
+    def list_actions(self, context):
         raise NotImplementedError
 
-    def do_action(self, action):
+    def do_action(self, context, action):
         raise NotImplementedError
 
     def shutdown(self):

python/pyarrow/flight.py

@@ -15,16 +15,20 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from pyarrow._flight import (Action,  # noqa
-                             ActionType,
-                             DescriptorType,
-                             FlightClient,
-                             FlightDescriptor,
-                             FlightEndpoint,
-                             FlightInfo,
-                             FlightServerBase,
-                             GeneratorStream,
-                             Location,
-                             Ticket,
-                             RecordBatchStream,
-                             Result)
+from pyarrow._flight import (  # noqa
+    Action,
+    ActionType,
+    DescriptorType,
+    FlightClient,
+    FlightDescriptor,
+    FlightEndpoint,
+    FlightInfo,
+    FlightServerBase,
+    GeneratorStream,
+    Location,
+    Ticket,
+    RecordBatchStream,
+    Result,
+    ClientAuthHandler,
+    ServerAuthHandler,
+)

python/pyarrow/includes/libarrow_flight.pxd

@@ -120,11 +120,34 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
             " arrow::flight::RecordBatchStream"(CFlightDataStream):
         CRecordBatchStream(shared_ptr[CRecordBatchReader]& reader)
 
+    cdef cppclass CServerAuthReader" arrow::flight::ServerAuthReader":
+        CStatus Read(c_string* token)
+
+    cdef cppclass CServerAuthSender" arrow::flight::ServerAuthSender":
+        CStatus Write(c_string& token)
+
+    cdef cppclass CClientAuthReader" arrow::flight::ClientAuthReader":
+        CStatus Read(c_string* token)
+
+    cdef cppclass CClientAuthSender" arrow::flight::ClientAuthSender":
+        CStatus Write(c_string& token)
+
+    cdef cppclass CServerAuthHandler" arrow::flight::ServerAuthHandler":
+        pass
+
+    cdef cppclass CClientAuthHandler" arrow::flight::ClientAuthHandler":
+        pass
+
+    cdef cppclass CServerCallContext" arrow::flight::ServerCallContext":
+        c_string& peer_identity()
+
     cdef cppclass CFlightClient" arrow::flight::FlightClient":
         @staticmethod
         CStatus Connect(const c_string& host, int port,
                         unique_ptr[CFlightClient]* client)
 
+        CStatus Authenticate(unique_ptr[CClientAuthHandler] auth_handler)
+
         CStatus DoAction(CAction& action, unique_ptr[CResultStream]* results)
         CStatus ListActions(vector[CActionType]* actions)
 
@@ -141,18 +164,29 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
 
 # Callbacks for implementing Flight servers
 # Use typedef to emulate syntax for std::function<void(...)>
-ctypedef void cb_list_flights(object, const CCriteria*,
+ctypedef void cb_list_flights(object, const CServerCallContext&,
+                              const CCriteria*,
                               unique_ptr[CFlightListing]*)
-ctypedef void cb_get_flight_info(object, const CFlightDescriptor&,
+ctypedef void cb_get_flight_info(object, const CServerCallContext&,
+                                 const CFlightDescriptor&,
                                  unique_ptr[CFlightInfo]*)
-ctypedef void cb_do_put(object, unique_ptr[CFlightMessageReader])
-ctypedef void cb_do_get(object, const CTicket&,
+ctypedef void cb_do_put(object, const CServerCallContext&,
+                        unique_ptr[CFlightMessageReader])
+ctypedef void cb_do_get(object, const CServerCallContext&,
+                        const CTicket&,
                         unique_ptr[CFlightDataStream]*)
-ctypedef void cb_do_action(object, const CAction&,
+ctypedef void cb_do_action(object, const CServerCallContext&, const CAction&,
                            unique_ptr[CResultStream]*)
-ctypedef void cb_list_actions(object, vector[CActionType]*)
+ctypedef void cb_list_actions(object, const CServerCallContext&,
+                              vector[CActionType]*)
 ctypedef void cb_result_next(object, unique_ptr[CResult]*)
 ctypedef void cb_data_stream_next(object, CFlightPayload*)
+ctypedef void cb_server_authenticate(object, const CServerAuthSender&,
+                                     const CServerAuthReader&)
+ctypedef void cb_is_valid(object, const c_string&, c_string*)
+ctypedef void cb_client_authenticate(object, const CClientAuthSender&,
+                                     const CClientAuthReader&)
+ctypedef void cb_get_token(object, c_string*)
 
 cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
     cdef cppclass PyFlightServerVtable:
@@ -164,13 +198,31 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
         function[cb_do_action] do_action
         function[cb_list_actions] list_actions
 
+    cdef cppclass PyServerAuthHandlerVtable:
+        PyServerAuthHandlerVtable()
+        function[cb_server_authenticate] authenticate
+        function[cb_is_valid] is_valid
+
+    cdef cppclass PyClientAuthHandlerVtable:
+        PyClientAuthHandlerVtable()
+        function[cb_client_authenticate] authenticate
+        function[cb_get_token] get_token
+
     cdef cppclass PyFlightServer:
         PyFlightServer(object server, PyFlightServerVtable vtable)
 
-        CStatus Init(int port)
+        CStatus Init(unique_ptr[CServerAuthHandler] auth_handler, int port)
         CStatus ServeWithSignals() except *
         void Shutdown()
 
+    cdef cppclass PyServerAuthHandler\
+            " arrow::py::flight::PyServerAuthHandler"(CServerAuthHandler):
+        PyServerAuthHandler(object handler, PyServerAuthHandlerVtable vtable)
+
+    cdef cppclass PyClientAuthHandler\
+            " arrow::py::flight::PyClientAuthHandler"(CClientAuthHandler):
+        PyClientAuthHandler(object handler, PyClientAuthHandlerVtable vtable)
+
     cdef cppclass CPyFlightResultStream\
             " arrow::py::flight::PyFlightResultStream"(CResultStream):
         CPyFlightResultStream(object generator,
@@ -197,4 +249,6 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
         unique_ptr[CFlightInfo]* out)
 
 cdef extern from "<utility>" namespace "std":
-    unique_ptr[CFlightDataStream] move(unique_ptr[CFlightDataStream])
+    unique_ptr[CFlightDataStream] move(unique_ptr[CFlightDataStream]) nogil
+    unique_ptr[CServerAuthHandler] move(unique_ptr[CServerAuthHandler]) nogil
+    unique_ptr[CClientAuthHandler] move(unique_ptr[CClientAuthHandler]) nogil

python/pyarrow/tests/test_flight.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import base64
 import contextlib
 import socket
 import threading
@@ -24,6 +25,8 @@ import pytest
 
 import pyarrow as pa
 
+from pyarrow.compat import tobytes
+
 
 flight = pytest.importorskip("pyarrow.flight")
 
@@ -35,7 +38,7 @@ class ConstantFlightServer(flight.FlightServerBase):
     does not properly hold a reference to the Table object.
     """
 
-    def do_get(self, ticket):
+    def do_get(self, context, ticket):
         data = [
             pa.array([-10, -5, 0, 5, 10])
         ]
@@ -50,28 +53,36 @@ class EchoFlightServer(flight.FlightServerBase):
         super(EchoFlightServer, self).__init__()
         self.last_message = None
 
-    def do_get(self, ticket):
+    def do_get(self, context, ticket):
         return flight.RecordBatchStream(self.last_message)
 
-    def do_put(self, descriptor, reader):
+    def do_put(self, context, descriptor, reader):
         self.last_message = reader.read_all()
 
 
 class EchoStreamFlightServer(EchoFlightServer):
     """An echo server that streams individual record batches."""
 
-    def do_get(self, ticket):
+    def do_get(self, context, ticket):
         return flight.GeneratorStream(
             self.last_message.schema,
             self.last_message.to_batches(chunksize=1024))
 
+    def list_actions(self, context):
+        return []
+
+    def do_action(self, context, action):
+        if action.type == "who-am-i":
+            return iter([flight.Result(context.peer_identity())])
+        raise NotImplementedError
+
 
 class InvalidStreamFlightServer(flight.FlightServerBase):
     """A Flight server that tries to return messages with differing schemas."""
 
     schema = pa.schema([('a', pa.int32())])
 
-    def do_get(self, ticket):
+    def do_get(self, context, ticket):
         data1 = [pa.array([-10, -5, 0, 5, 10], type=pa.int32())]
         data2 = [pa.array([-10.0, -5.0, 0.0, 5.0, 10.0], type=pa.float64())]
         assert data1.type != data2.type
@@ -82,6 +93,83 @@ class InvalidStreamFlightServer(flight.FlightServerBase):
         return flight.GeneratorStream(self.schema, [table1, table2])
 
 
+class HttpBasicServerAuthHandler(flight.ServerAuthHandler):
+    """An example implementation of HTTP basic authentication."""
+
+    def __init__(self, creds):
+        super().__init__()
+        self.creds = creds
+
+    def authenticate(self, outgoing, incoming):
+        pass
+
+    def is_valid(self, token):
+        if not token:
+            raise ValueError("unauthenticated: token not provided")
+        token = base64.b64decode(token)
+        username, password = token.split(b':')
+        if username not in self.creds:
+            raise ValueError("unknown user")
+        if self.creds[username] != password:
+            raise ValueError("wrong password")
+        return username
+
+
+class HttpBasicClientAuthHandler(flight.ClientAuthHandler):
+    """An example implementation of HTTP basic authentication."""
+
+    def __init__(self, username, password):
+        super().__init__()
+        self.username = tobytes(username)
+        self.password = tobytes(password)
+
+    def authenticate(self, outgoing, incoming):
+        pass
+
+    def get_token(self):
+        return base64.b64encode(self.username + b':' + self.password)
+
+
+class TokenServerAuthHandler(flight.ServerAuthHandler):
+    """An example implementation of authentication via handshake."""
+
+    def __init__(self, creds):
+        super().__init__()
+        self.creds = creds
+
+    def authenticate(self, outgoing, incoming):
+        username = incoming.read()
+        password = incoming.read()
+        if username in self.creds and self.creds[username] == password:
+            outgoing.write(base64.b64encode(b'secret:' + username))
+        else:
+            raise ValueError("unauthenticated: invalid username/password")
+
+    def is_valid(self, token):
+        token = base64.b64decode(token)
+        if not token.startswith(b'secret:'):
+            raise ValueError("unauthenticated: invalid token")
+        return token[7:]
+
+
+class TokenClientAuthHandler(flight.ClientAuthHandler):
+    """An example implementation of authentication via handshake."""
+
+    def __init__(self, username, password):
+        super().__init__()
+        self.username = username
+        self.password = password
+        self.token = b''
+
+    def authenticate(self, outgoing, incoming):
+        outgoing.write(self.username)
+        outgoing.write(self.password)
+        self.token = incoming.read()
+
+    def get_token(self):
+        return self.token
+
+
 @contextlib.contextmanager
 def flight_server(server_base, *args, **kwargs):
     """Spawn a Flight server on a free port, shutting it down when done."""
@@ -92,10 +180,14 @@ def flight_server(server_base, *args, **kwargs):
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         port = sock.getsockname()[1]
 
-    server_instance = server_base(*args, **kwargs)
+    auth_handler = kwargs.get('auth_handler')
+    ctor_kwargs = kwargs
+    if auth_handler:
+        del ctor_kwargs['auth_handler']
+    server_instance = server_base(*args, **ctor_kwargs)
 
     def _server_thread():
-        server_instance.run(port)
+        server_instance.run(port, auth_handler=auth_handler)
 
     thread = threading.Thread(target=_server_thread, daemon=True)
     thread.start()
@@ -163,3 +255,64 @@ def test_flight_invalid_generator_stream():
         client = flight.FlightClient.connect('localhost', server_port)
         with pytest.raises(pa.ArrowException):
             client.do_get(flight.Ticket(b'')).read_all()
+
+
+basic_auth_handler = HttpBasicServerAuthHandler(creds={
+    b"test": b"p4ssw0rd",
+})
+
+token_auth_handler = TokenServerAuthHandler(creds={
+    b"test": b"p4ssw0rd",
+})
+
+
+def test_http_basic_unauth():
+    """Test that auth fails when not authenticated."""
+    with flight_server(EchoStreamFlightServer,
+                       auth_handler=basic_auth_handler) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        action = flight.Action("who-am-i", b"")
+        with pytest.raises(pa.ArrowException, match=".*unauthenticated.*"):
+            list(client.do_action(action))
+
+
+def test_http_basic_auth():
+    """Test a Python implementation of HTTP basic authentication."""
+    with flight_server(EchoStreamFlightServer,
+                       auth_handler=basic_auth_handler) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        action = flight.Action("who-am-i", b"")
+        client.authenticate(HttpBasicClientAuthHandler('test', 'p4ssw0rd'))
+        identity = next(client.do_action(action))
+        assert identity.body.to_pybytes() == b'test'
+
+
+def test_http_basic_auth_invalid_password():
+    """Test that auth fails with the wrong password."""
+    with flight_server(EchoStreamFlightServer,
+                       auth_handler=basic_auth_handler) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        action = flight.Action("who-am-i", b"")
+        client.authenticate(HttpBasicClientAuthHandler('test', 'wrong'))
+        with pytest.raises(pa.ArrowException, match=".*wrong password.*"):
+            next(client.do_action(action))
+
+
+def test_token_auth():
+    """Test an auth mechanism that uses a handshake."""
+    with flight_server(EchoStreamFlightServer,
+                       auth_handler=token_auth_handler) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        action = flight.Action("who-am-i", b"")
+        client.authenticate(TokenClientAuthHandler('test', 'p4ssw0rd'))
+        identity = next(client.do_action(action))
+        assert identity.body.to_pybytes() == b'test'
+
+
+def test_token_auth_invalid():
+    """Test an auth mechanism that uses a handshake."""
+    with flight_server(EchoStreamFlightServer,
+                       auth_handler=token_auth_handler) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        with pytest.raises(pa.ArrowException, match=".*unauthenticated.*"):
+            client.authenticate(TokenClientAuthHandler('test', 'wrong'))