Skip to content

Commit

Permalink
Read metadata before returning
Browse files Browse the repository at this point in the history
The ambition is to make sure to have received an ACK from the server
before continuing to send a DONE message.

Co-authored-by: Max Kießling <max.kiessling@neotechnology.com>
  • Loading branch information
Mats-SX and DarthMax committed Oct 18, 2024
1 parent ec86a75 commit a53d10b
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/neo4j_arrow/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,13 +291,14 @@ def _write_batches(
upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(desc).encode("utf-8"))
n_rows, n_bytes = 0, 0
try:
writer, _ = client.do_put(upload_descriptor, schema, self.call_opts)
writer, metadata_reader = client.do_put(upload_descriptor, schema, self.call_opts)
with writer:
for batch in batches:
mapped_batch = fn(batch)
self._write_batch_with_retries(mapped_batch, writer)
n_rows += batch.num_rows
n_bytes += batch.get_total_buffer_size()
metadata_reader.read()
except Exception as e:
raise error.interpret(e)
return n_rows, n_bytes
Expand Down

0 comments on commit a53d10b

Please sign in to comment.