| 27 | |
| 28 | |
| 29 | class FlightServer(pyarrow.flight.FlightServerBase): |
| 30 | def __init__(self, host="localhost", location=None, |
| 31 | tls_certificates=None, verify_client=False, |
| 32 | root_certificates=None, auth_handler=None): |
| 33 | super(FlightServer, self).__init__( |
| 34 | location, auth_handler, tls_certificates, verify_client, |
| 35 | root_certificates) |
| 36 | self.flights = {} |
| 37 | self.host = host |
| 38 | self.tls_certificates = tls_certificates |
| 39 | |
| 40 | @classmethod |
| 41 | def descriptor_to_key(self, descriptor): |
| 42 | return (descriptor.descriptor_type.value, descriptor.command, |
| 43 | tuple(descriptor.path or tuple())) |
| 44 | |
| 45 | def _make_flight_info(self, key, descriptor, table): |
| 46 | if self.tls_certificates: |
| 47 | location = pyarrow.flight.Location.for_grpc_tls( |
| 48 | self.host, self.port) |
| 49 | else: |
| 50 | location = pyarrow.flight.Location.for_grpc_tcp( |
| 51 | self.host, self.port) |
| 52 | endpoints = [pyarrow.flight.FlightEndpoint(repr(key), [location]), ] |
| 53 | |
| 54 | mock_sink = pyarrow.MockOutputStream() |
| 55 | stream_writer = pyarrow.RecordBatchStreamWriter( |
| 56 | mock_sink, table.schema) |
| 57 | stream_writer.write_table(table) |
| 58 | stream_writer.close() |
| 59 | data_size = mock_sink.size() |
| 60 | |
| 61 | return pyarrow.flight.FlightInfo(table.schema, |
| 62 | descriptor, endpoints, |
| 63 | table.num_rows, data_size) |
| 64 | |
| 65 | def list_flights(self, context, criteria): |
| 66 | for key, table in self.flights.items(): |
| 67 | if key[1] is not None: |
| 68 | descriptor = \ |
| 69 | pyarrow.flight.FlightDescriptor.for_command(key[1]) |
| 70 | else: |
| 71 | descriptor = pyarrow.flight.FlightDescriptor.for_path(*key[2]) |
| 72 | |
| 73 | yield self._make_flight_info(key, descriptor, table) |
| 74 | |
| 75 | def get_flight_info(self, context, descriptor): |
| 76 | key = FlightServer.descriptor_to_key(descriptor) |
| 77 | if key in self.flights: |
| 78 | table = self.flights[key] |
| 79 | return self._make_flight_info(key, descriptor, table) |
| 80 | raise KeyError('Flight not found.') |
| 81 | |
| 82 | def do_put(self, context, descriptor, reader, writer): |
| 83 | key = FlightServer.descriptor_to_key(descriptor) |
| 84 | print(key) |
| 85 | self.flights[key] = reader.read_all() |
| 86 | print(self.flights[key]) |