Make Python Flight bindings more complete

cpp/src/arrow/flight/types.cc

@@ -112,6 +112,10 @@ std::string Location::scheme() const {
   return scheme;
 }
 
+bool Location::Equals(const Location& other) const {
+  return ToString() == other.ToString();
+}
+
 SimpleFlightListing::SimpleFlightListing(const std::vector<FlightInfo>& flights)
     : position_(0), flights_(flights) {}
 

cpp/src/arrow/flight/types.h

@@ -157,8 +157,10 @@ struct Location {
   /// \brief Get the scheme of this URI.
   std::string scheme() const;
 
+  bool Equals(const Location& other) const;
+
   friend bool operator==(const Location& left, const Location& right) {
-    return left.ToString() == right.ToString();
+    return left.Equals(right);
   }
   friend bool operator!=(const Location& left, const Location& right) {
     return !(left == right);

cpp/src/arrow/python/flight.cc

@@ -217,7 +217,7 @@ Status PyGeneratorFlightDataStream::Next(FlightPayload* payload) {
 Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,
                         const arrow::flight::FlightDescriptor& descriptor,
                         const std::vector<arrow::flight::FlightEndpoint>& endpoints,
-                        uint64_t total_records, uint64_t total_bytes,
+                        int64_t total_records, int64_t total_bytes,
                         std::unique_ptr<arrow::flight::FlightInfo>* out) {
   arrow::flight::FlightInfo::Data flight_data;
   RETURN_NOT_OK(arrow::flight::internal::SchemaToString(*schema, &flight_data.schema));

cpp/src/arrow/python/flight.h

@@ -197,7 +197,7 @@ ARROW_PYTHON_EXPORT
 Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,
                         const arrow::flight::FlightDescriptor& descriptor,
                         const std::vector<arrow::flight::FlightEndpoint>& endpoints,
-                        uint64_t total_records, uint64_t total_bytes,
+                        int64_t total_records, int64_t total_bytes,
                         std::unique_ptr<arrow::flight::FlightInfo>* out);
 
 }  // namespace flight

python/pyarrow/_flight.pyx

@@ -184,6 +184,10 @@ cdef class Ticket:
     def __init__(self, ticket):
         self.ticket.ticket = tobytes(ticket)
 
+    @property
+    def ticket(self):
+        return self.ticket.ticket
+
     def __repr__(self):
         return '<Ticket {}>'.format(self.ticket.ticket)
 
@@ -199,6 +203,18 @@ cdef class Location:
     def __repr__(self):
         return '<Location {}>'.format(self.location.ToString())
 
+    @property
+    def uri(self):
+        return self.location.ToString()
+
+    def equals(self, Location other):
+        return self == other
+
+    def __eq__(self, other):
+        if not isinstance(other, Location):
+            return NotImplemented
+        return self.location.Equals((<Location> other).location)
+
     @staticmethod
     def for_grpc_tcp(host, port):
         """Create a Location for a TCP-based gRPC service."""
@@ -265,8 +281,11 @@ cdef class FlightEndpoint:
             self.endpoint.ticket.ticket = tobytes(ticket)
 
         for location in locations:
-            c_location = CLocation()
-            check_status(CLocation.Parse(tobytes(location), &c_location))
+            if isinstance(location, Location):
+                c_location = (<Location> location).location
+            else:
+                c_location = CLocation()
+                check_status(CLocation.Parse(tobytes(location), &c_location))
             self.endpoint.locations.push_back(c_location)
 
     @property

python/pyarrow/includes/libarrow_flight.pxd

@@ -81,6 +81,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
     cdef cppclass CLocation" arrow::flight::Location":
         CLocation()
         c_string ToString()
+        c_bool Equals(const CLocation& other)
 
         @staticmethod
         CStatus Parse(c_string& uri_string, CLocation* location)
@@ -97,8 +98,8 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
 
     cdef cppclass CFlightInfo" arrow::flight::FlightInfo":
         CFlightInfo(CFlightInfo info)
-        uint64_t total_records()
-        uint64_t total_bytes()
+        int64_t total_records()
+        int64_t total_bytes()
         CStatus GetSchema(CDictionaryMemo* memo, shared_ptr[CSchema]* out)
         CFlightDescriptor& descriptor()
         const vector[CFlightEndpoint]& endpoints()
@@ -274,8 +275,8 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
         shared_ptr[CSchema] schema,
         CFlightDescriptor& descriptor,
         vector[CFlightEndpoint] endpoints,
-        uint64_t total_records,
-        uint64_t total_bytes,
+        int64_t total_records,
+        int64_t total_bytes,
         unique_ptr[CFlightInfo]* out)
 
 cdef extern from "<utility>" namespace "std":

python/pyarrow/tests/test_flight.py

@@ -103,6 +103,42 @@ class EchoStreamFlightServer(EchoFlightServer):
         raise NotImplementedError
 
 
+class GetInfoFlightServer(flight.FlightServerBase):
+    """A Flight server that tests GetFlightInfo."""
+
+    def get_flight_info(self, context, descriptor):
+        return flight.FlightInfo(
+            pa.schema([('a', pa.int32())]),
+            descriptor,
+            [
+                flight.FlightEndpoint(b'', ['grpc://test']),
+                flight.FlightEndpoint(
+                    b'',
+                    [flight.Location.for_grpc_tcp('localhost', 5005)],
+                ),
+            ],
+            -1,
+            -1,
+        )
+
+
+class CheckTicketFlightServer(flight.FlightServerBase):
+    """A Flight server that compares the given ticket to an expected value."""
+
+    def __init__(self, expected_ticket):
+        super(CheckTicketFlightServer, self).__init__()
+        self.expected_ticket = expected_ticket
+
+    def do_get(self, context, ticket):
+        assert self.expected_ticket == ticket.ticket
+        data1 = [pa.array([-10, -5, 0, 5, 10], type=pa.int32())]
+        table = pa.Table.from_arrays(data1, names=['a'])
+        return flight.RecordBatchStream(table)
+
+    def do_put(self, context, descriptor, reader):
+        self.last_message = reader.read_all()
+
+
 class InvalidStreamFlightServer(flight.FlightServerBase):
     """A Flight server that tries to return messages with differing schemas."""
 
@@ -255,6 +291,34 @@ def test_flight_do_get_dicts():
         assert data.equals(table)
 
 
+def test_flight_do_get_ticket():
+    """Make sure Tickets get passed to the server."""
+    data1 = [pa.array([-10, -5, 0, 5, 10], type=pa.int32())]
+    table = pa.Table.from_arrays(data1, names=['a'])
+    with flight_server(
+            CheckTicketFlightServer,
+            expected_ticket=b'the-ticket',
+    ) as server_location:
+        client = flight.FlightClient.connect(server_location)
+        data = client.do_get(flight.Ticket(b'the-ticket')).read_all()
+        assert data.equals(table)
+
+
+def test_flight_get_info():
+    """Make sure FlightEndpoint accepts string and object URIs."""
+    with flight_server(GetInfoFlightServer) as server_location:
+        client = flight.FlightClient.connect(server_location)
+        info = client.get_flight_info(flight.FlightDescriptor.for_command(b''))
+        assert info.total_records == -1
+        assert info.total_bytes == -1
+        assert info.schema == pa.schema([('a', pa.int32())])
+        assert len(info.endpoints) == 2
+        assert len(info.endpoints[0].locations) == 1
+        assert info.endpoints[0].locations[0] == flight.Location('grpc://test')
+        assert info.endpoints[1].locations[0] == \
+            flight.Location.for_grpc_tcp('localhost', 5005)
+
+
 def test_flight_domain_socket():
     """Try a simple do_get call over a domain socket."""
     data = [