def test_do_put_independent_read_write(): """Ensure that separate threads can read/write on a DoPut.""" # ARROW-6063: previously this would cause gRPC to abort when the # writer was closed (due to simultaneous reads), or would hang # forever. data = [pa.array([-10, -5, 0, 5, 10])] table = pa.Table.from_arrays(data, names=['a']) with MetadataFlightServer() as server: client = FlightClient(('localhost', server.port)) writer, metadata_reader = client.do_put( flight.FlightDescriptor.for_path(''), table.schema) count = [0] def _reader_thread(): while metadata_reader.read() is not None: count[0] += 1 thread = threading.Thread(target=_reader_thread) thread.start() batches = table.to_batches(max_chunksize=1) with writer: for idx, batch in enumerate(batches): metadata = struct.pack('<i', idx) writer.write_with_metadata(batch, metadata) # Causes the server to stop writing and end the call writer.done_writing() # Thus reader thread will break out of loop thread.join() # writer.close() won't segfault since reader thread has # stopped assert count[0] == len(batches)
def test_flight_generator_stream(): """Try downloading a flight of RecordBatches in a GeneratorStream.""" data = pa.Table.from_arrays([pa.array(range(0, 10 * 1024))], names=['a']) with EchoStreamFlightServer() as server: client = FlightClient(('localhost', server.port)) writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'), data.schema) writer.write_table(data) writer.close() result = client.do_get(flight.Ticket(b'')).read_all() assert result.equals(data)
def test_flight_do_put_metadata(): """Try a simple do_put call with metadata.""" data = [pa.array([-10, -5, 0, 5, 10])] table = pa.Table.from_arrays(data, names=['a']) with MetadataFlightServer() as server: client = FlightClient(('localhost', server.port)) writer, metadata_reader = client.do_put( flight.FlightDescriptor.for_path(''), table.schema) with writer: for idx, batch in enumerate(table.to_batches(max_chunksize=1)): metadata = struct.pack('<i', idx) writer.write_with_metadata(batch, metadata) buf = metadata_reader.read() assert buf is not None server_idx, = struct.unpack('<i', buf.to_pybytes()) assert idx == server_idx
def test_flight_large_message(): """Try sending/receiving a large message via Flight. See ARROW-4421: by default, gRPC won't allow us to send messages > 4MiB in size. """ data = pa.Table.from_arrays([pa.array(range(0, 10 * 1024 * 1024))], names=['a']) with EchoFlightServer(expected_schema=data.schema) as server: client = FlightClient(('localhost', server.port)) writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'), data.schema) # Write a single giant chunk writer.write_table(data, 10 * 1024 * 1024) writer.close() result = client.do_get(flight.Ticket(b'')).read_all() assert result.equals(data)
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