A server for testing DoExchange.
| 455 | |
| 456 | |
| 457 | class ExchangeFlightServer(FlightServerBase): |
| 458 | """A server for testing DoExchange.""" |
| 459 | |
| 460 | def __init__(self, options=None, **kwargs): |
| 461 | super().__init__(**kwargs) |
| 462 | self.options = options |
| 463 | |
| 464 | def do_exchange(self, context, descriptor, reader, writer): |
| 465 | assert reader.stats.num_messages == 0 |
| 466 | if descriptor.descriptor_type != flight.DescriptorType.CMD: |
| 467 | raise pa.ArrowInvalid("Must provide a command descriptor") |
| 468 | elif descriptor.command == b"echo": |
| 469 | return self.exchange_echo(context, reader, writer) |
| 470 | elif descriptor.command == b"get": |
| 471 | return self.exchange_do_get(context, reader, writer) |
| 472 | elif descriptor.command == b"put": |
| 473 | return self.exchange_do_put(context, reader, writer) |
| 474 | elif descriptor.command == b"transform": |
| 475 | return self.exchange_transform(context, reader, writer) |
| 476 | else: |
| 477 | raise pa.ArrowInvalid( |
| 478 | f"Unknown command: {descriptor.command}") |
| 479 | |
| 480 | def exchange_do_get(self, context, reader, writer): |
| 481 | """Emulate DoGet with DoExchange.""" |
| 482 | data = pa.Table.from_arrays([ |
| 483 | pa.array(range(0, 10 * 1024)) |
| 484 | ], names=["a"]) |
| 485 | writer.begin(data.schema) |
| 486 | writer.write_table(data) |
| 487 | |
| 488 | def exchange_do_put(self, context, reader, writer): |
| 489 | """Emulate DoPut with DoExchange.""" |
| 490 | num_batches = 0 |
| 491 | for chunk in reader: |
| 492 | if not chunk.data: |
| 493 | raise pa.ArrowInvalid("All chunks must have data.") |
| 494 | assert reader.stats.num_messages != 0 |
| 495 | num_batches += 1 |
| 496 | assert reader.stats.num_record_batches == num_batches |
| 497 | writer.write_metadata(str(num_batches).encode("utf-8")) |
| 498 | |
| 499 | def exchange_echo(self, context, reader, writer): |
| 500 | """Run a simple echo server.""" |
| 501 | assert reader.stats.num_messages == 0 |
| 502 | started = False |
| 503 | for chunk in reader: |
| 504 | if not started and chunk.data: |
| 505 | writer.begin(chunk.data.schema, options=self.options) |
| 506 | started = True |
| 507 | if chunk.app_metadata and chunk.data: |
| 508 | writer.write_with_metadata(chunk.data, chunk.app_metadata) |
| 509 | elif chunk.app_metadata: |
| 510 | writer.write_metadata(chunk.app_metadata) |
| 511 | elif chunk.data: |
| 512 | assert reader.stats.num_messages != 0 |
| 513 | writer.write_batch(chunk.data) |
| 514 | else: |
no outgoing calls