Bind Flight action endpoints to Python

cpp/src/arrow/flight/server.cc

@@ -264,6 +264,10 @@ class FlightServiceImpl : public FlightService::Service {
     std::unique_ptr<ResultStream> results;
     GRPC_RETURN_NOT_OK(server_->DoAction(action, &results));
 
+    if (!results) {
+      return grpc::Status::CANCELLED;
+    }
+
     std::unique_ptr<Result> result;
     pb::Result pb_result;
     while (true) {

cpp/src/arrow/python/flight.cc

@@ -71,6 +71,19 @@ Status PyFlightServer::ListActions(std::vector<arrow::flight::ActionType>* actio
   return CheckPyError();
 }
 
+PyFlightResultStream::PyFlightResultStream(PyObject* generator,
+                                           PyFlightResultStreamCallback callback)
+    : callback_(callback) {
+  Py_INCREF(generator);
+  generator_.reset(generator);
+}
+
+Status PyFlightResultStream::Next(std::unique_ptr<arrow::flight::Result>* result) {
+  PyAcquireGIL lock;
+  callback_(generator_.obj(), result);
+  return CheckPyError();
+}
+
 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

@@ -34,22 +34,24 @@ namespace flight {
 /// \brief A table of function pointers for calling from C++ into
 /// Python.
 class ARROW_PYTHON_EXPORT PyFlightServerVtable {
+
  public:
-  std::function<void(PyObject*, const arrow::flight::Criteria*,
-                     std::unique_ptr<arrow::flight::FlightListing>*)>
+  // Must return int so Cython will propagate exceptions
+  std::function<int(PyObject*, const arrow::flight::Criteria*,
+                    std::unique_ptr<arrow::flight::FlightListing>*)>
       list_flights;
-  std::function<void(PyObject*, const arrow::flight::FlightDescriptor&,
-                     std::unique_ptr<arrow::flight::FlightInfo>*)>
+  std::function<int(PyObject*, const arrow::flight::FlightDescriptor&,
+                    std::unique_ptr<arrow::flight::FlightInfo>*)>
       get_flight_info;
-  std::function<void(PyObject*, const arrow::flight::Ticket&,
-                     std::unique_ptr<arrow::flight::FlightDataStream>*)>
+  std::function<int(PyObject*, const arrow::flight::Ticket&,
+                    std::unique_ptr<arrow::flight::FlightDataStream>*)>
       do_get;
-  std::function<void(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
+  std::function<int(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
       do_put;
-  std::function<void(PyObject*, const arrow::flight::Action&,
-                     std::unique_ptr<arrow::flight::ResultStream>*)>
+  std::function<int(PyObject*, 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<int(PyObject*, std::vector<arrow::flight::ActionType>*)> list_actions;
 };
 
 class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBase {
@@ -72,6 +74,20 @@ class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBas
   PyFlightServerVtable vtable_;
 };
 
+/// \brief A callback that obtains the next result from a Flight action.
+typedef std::function<int(PyObject*, std::unique_ptr<arrow::flight::Result>*)> PyFlightResultStreamCallback;
+
+/// \brief A ResultStream built around a Python callback.
+class ARROW_PYTHON_EXPORT PyFlightResultStream : public arrow::flight::ResultStream {
+ public:
+  explicit PyFlightResultStream(PyObject* generator, PyFlightResultStreamCallback callback);
+  Status Next(std::unique_ptr<arrow::flight::Result>* result) override;
+
+ private:
+  OwnedRefNoGIL generator_;
+  PyFlightResultStreamCallback callback_;
+};
+
 ARROW_PYTHON_EXPORT
 Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,
                         const arrow::flight::FlightDescriptor& descriptor,

python/examples/flight/server.py

@@ -72,6 +72,18 @@ class FlightServer(pyarrow.flight.FlightServerBase):
             return None
         return pyarrow.flight.RecordBatchStream(self.flights[key])
 
+    def list_actions(self):
+        return [
+            ("clear", "Clear the stored flights."),
+            ("shutdown", "Shut down this server."),
+        ]
+
+    def do_action(self, action):
+        if action.type == "clear":
+            raise NotImplementedError("{} is not implemented.".format(action.type))
+        else:
+            yield pyarrow.flight.Result(pyarrow.py_buffer(b'Shutdown!'))
+
 
 def main():
     server = FlightServer()

python/pyarrow/_flight.pyx

@@ -47,33 +47,23 @@ cdef class Action:
         return pyarrow_wrap_buffer(self.action.body)
 
 
-cdef class ActionType:
+class ActionType(collections.namedtuple('ActionType', ['type', 'description'])):
     """A type of action executable on a Flight service."""
-    cdef:
-        CActionType action_type
-
-    @property
-    def type(self):
-        return frombytes(self.action_type.type)
-
-    @property
-    def description(self):
-        return frombytes(self.action_type.description)
 
     def make_action(self, buf):
         """Create an Action with this type."""
         return Action(self.type, buf)
 
-    def __repr__(self):
-        return '<ActionType type={} description={}>'.format(
-            self.type, self.description)
-
 
 cdef class Result:
     """A result from executing an Action."""
     cdef:
         unique_ptr[CResult] result
 
+    def __init__(self, buf):
+        self.result.reset(new CResult())
+        self.result.get().body = pyarrow_unwrap_buffer(as_buffer(buf))
+
     @property
     def body(self):
         """Get the Buffer containing the result."""
@@ -329,8 +319,8 @@ cdef class FlightClient:
 
         result = []
         for action_type in results:
-            py_action = ActionType()
-            py_action.action_type = action_type
+            py_action = ActionType(frombytes(action_type.type),
+                                   frombytes(action_type.description))
             result.append(py_action)
 
         return result
@@ -339,11 +329,12 @@ cdef class FlightClient:
         """Execute an action on a service."""
         cdef:
             unique_ptr[CResultStream] results
+            Result result
         with nogil:
             check_status(self.client.get().DoAction(action.action, &results))
 
         while True:
-            result = Result()
+            result = Result.__new__(Result)
             with nogil:
                 check_status(results.get().Next(&result.result))
                 if result.result == NULL:
@@ -433,8 +424,8 @@ cdef class RecordBatchStream(FlightDataStream):
             raise TypeError("Expected RecordBatchReader or Table, but got: {}".format(type(data_source)))
 
 
-cdef void _list_flights(void* self, const CCriteria* c_criteria,
-                        unique_ptr[CFlightListing]* listing):
+cdef int _list_flights(void* self, const CCriteria* c_criteria,
+                        unique_ptr[CFlightListing]* listing) except -1:
     """Callback for implementing ListFlights in Python."""
     cdef:
         vector[CFlightInfo] flights
@@ -445,10 +436,11 @@ cdef void _list_flights(void* self, const CCriteria* c_criteria,
                             "FlightInfo instances, but got {}".format(type(info)))
         flights.push_back(CFlightInfo(deref((<FlightInfo> info).info.get())))
     listing.reset(new CSimpleFlightListing(flights))
+    return 0
 
 
-cdef void _get_flight_info(void* self, CFlightDescriptor c_descriptor,
-                           unique_ptr[CFlightInfo]* info):
+cdef int _get_flight_info(void* self, CFlightDescriptor c_descriptor,
+                           unique_ptr[CFlightInfo]* info) except -1:
     """Callback for implementing Flight servers in Python."""
     cdef:
         FlightDescriptor py_descriptor = \
@@ -459,9 +451,10 @@ cdef void _get_flight_info(void* self, CFlightDescriptor c_descriptor,
         raise TypeError("FlightServerBase.get_flight_info must return "
                         "a FlightInfo instance")
     info[0] = move((<FlightInfo> result).info)
+    return 0
 
 
-cdef void _do_put(void* self, unique_ptr[CFlightMessageReader] reader):
+cdef int _do_put(void* self, unique_ptr[CFlightMessageReader] reader) except -1:
     """Callback for implementing Flight servers in Python."""
     cdef:
         FlightRecordBatchReader py_reader = FlightRecordBatchReader()
@@ -471,10 +464,11 @@ cdef void _do_put(void* self, unique_ptr[CFlightMessageReader] reader):
     descriptor.descriptor = reader.get().descriptor()
     py_reader.reader.reset(reader.release())
     (<object> self).do_put(descriptor, py_reader)
+    return 0
 
 
-cdef void _do_get(void* self, CTicket ticket,
-                  unique_ptr[CFlightDataStream]* stream):
+cdef int _do_get(void* self, CTicket ticket,
+                  unique_ptr[CFlightDataStream]* stream) except -1:
     """Callback for implementing Flight servers in Python."""
     py_ticket = Ticket(ticket.ticket)
     result = (<object> self).do_get(py_ticket)
@@ -482,8 +476,46 @@ cdef void _do_get(void* self, CTicket ticket,
         raise TypeError("FlightServerBase.do_get must return "
                         "a FlightDataStream")
     stream[0] = move((<FlightDataStream> result).stream)
+    return 0
 
 
+cdef int _do_action_result_next(void* self,
+                                 unique_ptr[CResult]* result) except -1:
+    """Callback for implementing Flight servers in Python."""
+    try:
+        action_result = next(<object> self)
+        if not isinstance(action_result, Result):
+            raise TypeError("Result of FlightServerBase.do_action must "
+                            "return a Result during iteration")
+        result[0] = move((<Result> action_result).result)
+    except StopIteration:
+        result.reset(nullptr)
+    return 0
+
+
+cdef int _do_action(void* self, const CAction& action,
+                    unique_ptr[CResultStream]* result) except -1:
+    """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)
+    result[0].reset(new CPyFlightResultStream(responses, ptr))
+    return 0
+
+
+cdef int _list_actions(void* self, vector[CActionType]* actions) except -1:
+    """Callback for implementing Flight servers in Python."""
+    cdef:
+        CActionType action_type
+    # Method should return a list of ActionTypes or similar tuple
+    result = (<object> self).list_actions()
+    for action in result:
+        action_type.type = tobytes(action[0])
+        action_type.description = tobytes(action[1])
+        actions.push_back(action_type)
+    return 0
+
 cdef class FlightServerBase:
     """A Flight service definition."""
 
@@ -498,6 +530,8 @@ cdef class FlightServerBase:
         vtable.get_flight_info = &_get_flight_info
         vtable.do_put = &_do_put
         vtable.do_get = &_do_get
+        vtable.list_actions = &_list_actions
+        vtable.do_action = &_do_action
         self.server.reset(new PyFlightServer(self, vtable))
         with nogil:
             self.server.get().Run(c_port)
@@ -514,6 +548,12 @@ cdef class FlightServerBase:
     def do_get(self, ticket):
         raise NotImplementedError
 
+    def list_actions(self):
+        raise NotImplementedError
+
+    def do_action(self, action):
+        raise NotImplementedError
+
     def shutdown(self):
         if self.server.get() != NULL:
             self.server.get().Shutdown()

python/pyarrow/flight.py

@@ -19,4 +19,4 @@ from pyarrow._flight import (FlightClient, Action, ActionType,  # noqa
                              FlightDescriptor, FlightEndpoint,
                              FlightInfo, Ticket, Location,
                              FlightServerBase, DescriptorType,
-                             RecordBatchStream)
+                             RecordBatchStream, Result)

python/pyarrow/includes/libarrow_flight.pxd

@@ -120,13 +120,17 @@ 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*,
-                              unique_ptr[CFlightListing]*)
-ctypedef void cb_get_flight_info(object, const CFlightDescriptor&,
-                                 unique_ptr[CFlightInfo]*)
-ctypedef void cb_do_put(object, unique_ptr[CFlightMessageReader])
-ctypedef void cb_do_get(object, const CTicket&,
-                        unique_ptr[CFlightDataStream]*)
+ctypedef int cb_list_flights(object, const CCriteria*,
+                             unique_ptr[CFlightListing]*)
+ctypedef int cb_get_flight_info(object, const CFlightDescriptor&,
+                                unique_ptr[CFlightInfo]*)
+ctypedef int cb_do_put(object, unique_ptr[CFlightMessageReader])
+ctypedef int cb_do_get(object, const CTicket&,
+                       unique_ptr[CFlightDataStream]*)
+ctypedef int cb_do_action(object, const CAction&,
+                          unique_ptr[CResultStream]*)
+ctypedef int cb_list_actions(object, vector[CActionType]*)
+ctypedef int cb_result_next(object, unique_ptr[CResult]*)
 
 cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
     cdef cppclass PyFlightServerVtable:
@@ -135,12 +139,17 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
         function[cb_get_flight_info] get_flight_info
         function[cb_do_put] do_put
         function[cb_do_get] do_get
+        function[cb_do_action] do_action
+        function[cb_list_actions] list_actions
 
     cdef cppclass PyFlightServer:
         PyFlightServer(object server, PyFlightServerVtable vtable)
         void Run(int port)
         void Shutdown()
 
+    cdef cppclass CPyFlightResultStream" arrow::py::flight::PyFlightResultStream"(CResultStream):
+        CPyFlightResultStream(object generator, function[cb_result_next] callback)
+
     cdef CStatus CreateFlightInfo" arrow::py::flight::CreateFlightInfo"(
         shared_ptr[CSchema] schema,
         CFlightDescriptor& descriptor,
@@ -154,3 +163,4 @@ cdef extern from "<utility>" namespace "std":
     # bindings
     unique_ptr[CFlightDataStream] move(unique_ptr[CFlightDataStream])
     unique_ptr[CFlightInfo] move(unique_ptr[CFlightInfo])
+    unique_ptr[CResult] move(unique_ptr[CResult])