From 5ccdf4d01cdfb6c8b62a1850b220d48a080f165f Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Thu, 24 Mar 2022 14:01:58 +0100 Subject: [PATCH] fix(copy): split large buffers before sending them to PQputCopyData Sending excessively large data doesn't allow the libpq to flush its content to the server, ending up in an infinite loop. Close #255 --- docs/news.rst | 1 + psycopg/psycopg/generators.py | 13 ++++++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/docs/news.rst b/docs/news.rst index a72290adc..f6c92f49c 100644 --- a/docs/news.rst +++ b/docs/news.rst @@ -26,6 +26,7 @@ Psycopg 3.0.11 (unreleased) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ - Fix `DataError` loading arrays with dimensions information (:ticket:`#253`). +- Fix hanging during COPY in case of memory error (:ticket:`#255`). - Fix error propagation from COPY worker thread (mentioned in :ticket:`#255`). diff --git a/psycopg/psycopg/generators.py b/psycopg/psycopg/generators.py index f0140e4a8..902695b4b 100644 --- a/psycopg/psycopg/generators.py +++ b/psycopg/psycopg/generators.py @@ -203,9 +203,16 @@ def copy_from(pgconn: PGconn) -> PQGen[Union[memoryview, PGresult]]: def copy_to(pgconn: PGconn, buffer: bytes) -> PQGen[None]: - # Retry enqueuing data until successful - while pgconn.put_copy_data(buffer) == 0: - yield Wait.W + # Split the data to send in chunks not larger than 1Mb. + # + # The libpq docs say to retry on 0. What they don't say is that they return + # 0 pretty much only on palloc error, not on socket EWOULDBLOCK. Passing a + # block too big will always fail and create an infinite loop (See #255). + COPY_BUFFER_SIZE = 2**20 + for i in range(0, len(buffer), COPY_BUFFER_SIZE): + # Retry enqueuing data until successful + while pgconn.put_copy_data(buffer[i : i + COPY_BUFFER_SIZE]) == 0: + yield Wait.W def copy_end(pgconn: PGconn, error: Optional[bytes]) -> PQGen[PGresult]: -- 2.47.2