From: Andrew Kane <andrew@ankane.org>
Date: Mon, 4 Mar 2024 22:34:56 +0000 (-0800)
Subject: Improve performance of copy
X-Git-Tag: 3.1.19~3^2~1
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=963c5f8a1b169eccee05fbf0b9738308df86706e;p=thirdparty%2Fpsycopg.git

Improve performance of copy
---

diff --git a/psycopg/psycopg/generators.py b/psycopg/psycopg/generators.py
index 4f2ec878b..aa01f1d62 100644
--- a/psycopg/psycopg/generators.py
+++ b/psycopg/psycopg/generators.py
@@ -293,6 +293,16 @@ def copy_to(pgconn: PGconn, buffer: Buffer) -> PQGen[None]:
     while pgconn.put_copy_data(buffer) == 0:
         yield WAIT_W
 
+    # Repeat until it the message is flushed to the server
+    while True:
+        while True:
+            ready = yield WAIT_W
+            if ready:
+                break
+        f = pgconn.flush()
+        if f == 0:
+            break
+
 
 def copy_end(pgconn: PGconn, error: Optional[bytes]) -> PQGen[PGresult]:
     # Retry enqueuing end copy message until successful