Bind Flight DoGet server call to Python

cpp/src/arrow/python/flight.cc

@@ -33,7 +33,9 @@ PyFlightServer::PyFlightServer(PyObject* server, PyFlightServerVtable vtable)
 Status PyFlightServer::ListFlights(
     const arrow::flight::Criteria* criteria,
     std::unique_ptr<arrow::flight::FlightListing>* listings) {
-  return Status::NotImplemented("NYI");
+  PyAcquireGIL lock;
+  vtable_.list_flights(server_.obj(), criteria, listings);
+  return CheckPyError();
 }
 
 Status PyFlightServer::GetFlightInfo(const arrow::flight::FlightDescriptor& request,
@@ -58,11 +60,15 @@ Status PyFlightServer::DoPut(std::unique_ptr<arrow::flight::FlightMessageReader>
 
 Status PyFlightServer::DoAction(const arrow::flight::Action& action,
                                 std::unique_ptr<arrow::flight::ResultStream>* result) {
-  return Status::NotImplemented("NYI");
+  PyAcquireGIL lock;
+  vtable_.do_action(server_.obj(), action, result);
+  return CheckPyError();
 }
 
 Status PyFlightServer::ListActions(std::vector<arrow::flight::ActionType>* actions) {
-  return Status::NotImplemented("NYI");
+  PyAcquireGIL lock;
+  vtable_.list_actions(server_.obj(), actions);
+  return CheckPyError();
 }
 
 Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,

cpp/src/arrow/python/flight.h

@@ -35,14 +35,21 @@ namespace flight {
 /// Python.
 class ARROW_PYTHON_EXPORT PyFlightServerVtable {
  public:
+  std::function<void(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>*)>
       get_flight_info;
-  std::function<void(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
-      do_put;
   std::function<void(PyObject*, const arrow::flight::Ticket&,
                      std::unique_ptr<arrow::flight::FlightDataStream>*)>
       do_get;
+  std::function<void(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
+      do_put;
+  std::function<void(PyObject*, const arrow::flight::Action&,
+                     std::unique_ptr<arrow::flight::ResultStream>*)>
+      do_action;
+  std::function<void(PyObject*, std::vector<arrow::flight::ActionType>*)> list_actions;
 };
 
 class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBase {

python/examples/flight/server.py

@@ -17,6 +17,8 @@
 
 """An example Flight Python server."""
 
+import ast
+
 import pyarrow
 import pyarrow.flight
 
@@ -28,15 +30,18 @@ class FlightServer(pyarrow.flight.FlightServerBase):
 
     @classmethod
     def descriptor_to_key(self, descriptor):
-        return (descriptor.descriptor_type, descriptor.command,
+        return (descriptor.descriptor_type.value, descriptor.command,
                 tuple(descriptor.path or tuple()))
 
     def get_flight_info(self, descriptor):
         key = FlightServer.descriptor_to_key(descriptor)
         if key in self.flights:
             table = self.flights[key]
+            endpoints = [
+                pyarrow.flight.FlightEndpoint(repr(key), [('localhost', 5005)]),
+            ]
             return pyarrow.flight.FlightInfo(table.schema,
-                                             descriptor, [],
+                                             descriptor, endpoints,
                                              table.num_rows, 0)
         raise KeyError('Flight not found.')
 
@@ -46,6 +51,13 @@ class FlightServer(pyarrow.flight.FlightServerBase):
         self.flights[key] = reader.read_all()
         print(self.flights[key])
 
+    def do_get(self, ticket):
+        print(ticket.ticket)
+        key = ast.literal_eval(ticket.ticket.decode())
+        if key not in self.flights:
+            return None
+        return pyarrow.flight.RecordBatchStream(self.flights[key])
+
 
 def main():
     server = FlightServer()

python/pyarrow/_flight.pyx

@@ -140,7 +140,7 @@ cdef class FlightDescriptor:
         return self.descriptor.path
 
     def __repr__(self):
-        return "<FlightDescriptor type: {!r}>".format(self.descriptor_type())
+        return "<FlightDescriptor type: {!r}>".format(self.descriptor_type)
 
 
 class Ticket:
@@ -175,9 +175,9 @@ cdef class FlightEndpoint:
             CLocation c_location = CLocation()
 
         if isinstance(ticket, Ticket):
-            self.endpoint.ticket.ticket = ticket.ticket
+            self.endpoint.ticket.ticket = tobytes(ticket.ticket)
         else:
-            self.endpoint.ticket.ticket = ticket
+            self.endpoint.ticket.ticket = tobytes(ticket)
 
         for location in locations:
             # Accepts Location namedtuple or tuple
@@ -270,7 +270,7 @@ cdef class FlightInfo:
 
         result = []
         for endpoint in endpoints:
-            py_endpoint = FlightEndpoint.__new__()
+            py_endpoint = FlightEndpoint.__new__(FlightEndpoint)
             py_endpoint.endpoint = endpoint
             result.append(py_endpoint)
         return result
@@ -413,15 +413,38 @@ cdef class FlightDataStream:
 
 
 cdef class RecordBatchStream(FlightDataStream):
-    def __init__(self, reader):
-        # TODO: we don't really expose the readers in Python.
-        pass
+    """A Flight data stream backed by RecordBatches."""
+    def __init__(self, data_source):
+        """Create a RecordBatchStream from a data source.
+
+        Parameters
+        ----------
+        data_source : RecordBatchReader or Table
+        """
+        cdef shared_ptr[CRecordBatchReader] reader
+        if isinstance(data_source, _CRecordBatchReader):
+            stream = new CRecordBatchStream((<_CRecordBatchReader> data_source).reader)
+            self.stream.reset(stream)
+        elif isinstance(data_source, lib.Table):
+            reader.reset(new TableBatchReader(deref((<Table> data_source).table)))
+            stream = new CRecordBatchStream(reader)
+            self.stream.reset(stream)
+        else:
+            raise TypeError("Expected RecordBatchReader or Table, but got: {}".format(type(data_source)))
 
 
 cdef void _get_flight_info(void* self, CFlightDescriptor c_descriptor,
                            unique_ptr[CFlightInfo]* info):
     """Callback for implementing Flight servers in Python."""
-    raise NotImplementedError("GetFlightInfo is not implemented")
+    cdef:
+        FlightDescriptor py_descriptor = \
+            FlightDescriptor.__new__(FlightDescriptor)
+    py_descriptor.descriptor = c_descriptor
+    result = (<object> self).get_flight_info(py_descriptor)
+    if not isinstance(result, FlightInfo):
+        raise TypeError("FlightServerBase.get_flight_info must return "
+                        "a FlightInfo instance")
+    info[0] = move((<FlightInfo> result).info)
 
 
 cdef void _do_put(void* self, unique_ptr[CFlightMessageReader] reader):
@@ -439,8 +462,7 @@ cdef void _do_put(void* self, unique_ptr[CFlightMessageReader] reader):
 cdef void _do_get(void* self, CTicket ticket,
                   unique_ptr[CFlightDataStream]* stream):
     """Callback for implementing Flight servers in Python."""
-    py_ticket = Ticket()
-    py_ticket.ticket = ticket.ticket
+    py_ticket = Ticket(ticket.ticket)
     result = (<object> self).do_get(py_ticket)
     if not isinstance(result, FlightDataStream):
         raise TypeError("FlightServerBase.do_get must return "
@@ -471,6 +493,9 @@ cdef class FlightServerBase:
     def do_put(self, descriptor, reader):
         raise NotImplementedError
 
+    def do_get(self, ticket):
+        raise NotImplementedError
+
     def shutdown(self):
         if self.server.get() != NULL:
             self.server.get().Shutdown()

python/pyarrow/flight.py

@@ -16,5 +16,7 @@
 # under the License.
 
 from pyarrow._flight import (FlightClient, Action, ActionType,  # noqa
-                             FlightDescriptor, FlightInfo, Ticket, Location,
-                             FlightServerBase, DescriptorType)
+                             FlightDescriptor, FlightEndpoint,
+                             FlightInfo, Ticket, Location,
+                             FlightServerBase, DescriptorType,
+                             RecordBatchStream)

python/pyarrow/includes/libarrow_flight.pxd

@@ -139,4 +139,7 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
         unique_ptr[CFlightInfo]* out)
 
 cdef extern from "<utility>" namespace "std":
+    # Declare explicit overloads of std::move for use in server
+    # bindings
     unique_ptr[CFlightDataStream] move(unique_ptr[CFlightDataStream])
+    unique_ptr[CFlightInfo] move(unique_ptr[CFlightInfo])