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

diff --git a/docs/news.rst b/docs/news.rst
index a72290adc9fd407ab7ca14edb75f192ff85f1ffa..f6c92f49cdc7d4df8048a805fc83976c68435598 100644 (file)
--- 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 f0140e4a8c07a7ef80a08633895b28c98079835c..902695b4bfd55595f79de207ccf2559bb86ee630 100644 (file)
--- 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]: