def test_doexchange_echo(): """Try a DoExchange echo server.""" data = pa.Table.from_arrays([pa.array(range(0, 10 * 1024))], names=["a"]) batches = data.to_batches(max_chunksize=512) with ExchangeFlightServer() as server: client = FlightClient(("localhost", server.port)) descriptor = flight.FlightDescriptor.for_command(b"echo") writer, reader = client.do_exchange(descriptor) with writer: # Read/write metadata before starting data. for i in range(10): buf = str(i).encode("utf-8") writer.write_metadata(buf) chunk = reader.read_chunk() assert chunk.data is None assert chunk.app_metadata == buf # Now write data without metadata. writer.begin(data.schema) for batch in batches: writer.write_batch(batch) assert reader.schema == data.schema chunk = reader.read_chunk() assert chunk.data == batch assert chunk.app_metadata is None # And write data with metadata. for i, batch in enumerate(batches): buf = str(i).encode("utf-8") writer.write_with_metadata(batch, buf) chunk = reader.read_chunk() assert chunk.data == batch assert chunk.app_metadata == buf
def test_doexchange_get(): """Emulate DoGet with DoExchange.""" expected = pa.Table.from_arrays([pa.array(range(0, 10 * 1024))], names=["a"]) with ExchangeFlightServer() as server: client = FlightClient(("localhost", server.port)) descriptor = flight.FlightDescriptor.for_command(b"get") writer, reader = client.do_exchange(descriptor) with writer: table = reader.read_all() assert expected == table
def test_doexchange_put(): """Emulate DoPut with DoExchange.""" data = pa.Table.from_arrays([pa.array(range(0, 10 * 1024))], names=["a"]) batches = data.to_batches(max_chunksize=512) with ExchangeFlightServer() as server: client = FlightClient(("localhost", server.port)) descriptor = flight.FlightDescriptor.for_command(b"put") writer, reader = client.do_exchange(descriptor) with writer: writer.begin(data.schema) for batch in batches: writer.write_batch(batch) writer.done_writing() chunk = reader.read_chunk() assert chunk.data is None expected_buf = str(len(batches)).encode("utf-8") assert chunk.app_metadata == expected_buf
def test_middleware_mapping(): """Test that middleware records methods correctly.""" server_middleware = RecordingServerMiddlewareFactory() client_middleware = RecordingClientMiddlewareFactory() with FlightServerBase(middleware={"test": server_middleware}) as server: client = FlightClient( ('localhost', server.port), middleware=[client_middleware] ) descriptor = flight.FlightDescriptor.for_command(b"") with pytest.raises(NotImplementedError): list(client.list_flights()) with pytest.raises(NotImplementedError): client.get_flight_info(descriptor) with pytest.raises(NotImplementedError): client.get_schema(descriptor) with pytest.raises(NotImplementedError): client.do_get(flight.Ticket(b"")) with pytest.raises(NotImplementedError): writer, _ = client.do_put(descriptor, pa.schema([])) writer.close() with pytest.raises(NotImplementedError): list(client.do_action(flight.Action(b"", b""))) with pytest.raises(NotImplementedError): list(client.list_actions()) with pytest.raises(NotImplementedError): writer, _ = client.do_exchange(descriptor) writer.close() expected = [ flight.FlightMethod.LIST_FLIGHTS, flight.FlightMethod.GET_FLIGHT_INFO, flight.FlightMethod.GET_SCHEMA, flight.FlightMethod.DO_GET, flight.FlightMethod.DO_PUT, flight.FlightMethod.DO_ACTION, flight.FlightMethod.LIST_ACTIONS, flight.FlightMethod.DO_EXCHANGE, ] assert server_middleware.methods == expected assert client_middleware.methods == expected
def test_doexchange_transform(): """Transform a table with a service.""" data = pa.Table.from_arrays([ pa.array(range(0, 1024)), pa.array(range(1, 1025)), pa.array(range(2, 1026)), ], names=["a", "b", "c"]) expected = pa.Table.from_arrays([ pa.array(range(3, 1024 * 3 + 3, 3)), ], names=["sum"]) with ExchangeFlightServer() as server: client = FlightClient(("localhost", server.port)) descriptor = flight.FlightDescriptor.for_command(b"transform") writer, reader = client.do_exchange(descriptor) with writer: writer.begin(data.schema) writer.write_table(data) writer.done_writing() table = reader.read_all() assert expected == table