fix: fix infinite loop and OOM in bad executemany

The bad condition is only reached using COPY into executemany in
pipeline mode and with prepared statements disabled. It should probably
never happen outside the unit test.

diff --git a/psycopg/psycopg/generators.py b/psycopg/psycopg/generators.py
index 81b52a2afc5d4dafe2618133df15146e3d522748..4f2ec878bb9cc70d06c392f884cba07962b9669e 100644 (file)
--- a/psycopg/psycopg/generators.py
+++ b/psycopg/psycopg/generators.py
@@ -205,11 +205,22 @@ def _pipeline_communicate(
                         break
                     results.append(res)
                     res = []
-                elif r.status == PIPELINE_SYNC:
-                    assert not res
-                    results.append([r])
                 else:
-                    res.append(r)
+                    status = r.status
+                    if status == PIPELINE_SYNC:
+                        assert not res
+                        results.append([r])
+                    elif status == COPY_IN or status == COPY_OUT or status == COPY_BOTH:
+                        # This shouldn't happen, but insisting hard enough, it will.
+                        # For instance, in test_executemany_badquery(), with the COPY
+                        # statement and the AsyncClientCursor, which disables
+                        # prepared statements).
+                        # Bail out from the resulting infinite loop.
+                        raise e.NotSupportedError(
+                            "COPY cannot be used in pipeline mode"
+                        )
+                    else:
+                        res.append(r)
 
         if ready & READY_W:
             pgconn.flush()

diff --git a/psycopg_c/psycopg_c/_psycopg/generators.pyx b/psycopg_c/psycopg_c/_psycopg/generators.pyx
index fbd901385992b3c83e7d4f22d82cb187fa61db0e..a51fce5e28d5f2117e1a161fa6c3c2a8663ede9c 100644 (file)
--- a/psycopg_c/psycopg_c/_psycopg/generators.pyx
+++ b/psycopg_c/psycopg_c/_psycopg/generators.pyx
@@ -241,6 +241,19 @@ def pipeline_communicate(
                     if status == libpq.PGRES_PIPELINE_SYNC:
                         results.append([r])
                         break
+                    elif (
+                        status == libpq.PGRES_COPY_IN
+                        or status == libpq.PGRES_COPY_OUT
+                        or status == libpq.PGRES_COPY_BOTH
+                    ):
+                        # This shouldn't happen, but insisting hard enough, it will.
+                        # For instance, in test_executemany_badquery(), with the COPY
+                        # statement and the AsyncClientCursor, which disables
+                        # prepared statements).
+                        # Bail out from the resulting infinite loop.
+                        raise e.NotSupportedError(
+                            "COPY cannot be used in pipeline mode"
+                        )
                     else:
                         res.append(r)
 

diff --git a/tests/test_client_cursor.py b/tests/test_client_cursor.py
index 5ac793b643e16db29791a06dba90081b8d63af5d..18941a54efa0cd1544c0f07620e423a3d609c18c 100644 (file)
--- a/tests/test_client_cursor.py
+++ b/tests/test_client_cursor.py
@@ -380,10 +380,7 @@ def test_executemany_rowcount_no_hit(conn, execmany):
     "query",
     [
         "insert into nosuchtable values (%s, %s)",
-        # This fails, but only because we try to copy in pipeline mode,
-        # crashing the connection. Which would be even fine, but with
-        # the async cursor it's worse... See test_client_cursor_async.py.
-        # "copy (select %s, %s) to stdout",
+        "copy (select %s, %s) to stdout",
         "wat (%s, %s)",
     ],
 )

diff --git a/tests/test_client_cursor_async.py b/tests/test_client_cursor_async.py
index ec6504d07906881ede7a36639ee9a0fea478047c..67b314acb3a78ccee7a56058505a082941ecf9f4 100644 (file)
--- a/tests/test_client_cursor_async.py
+++ b/tests/test_client_cursor_async.py
@@ -372,11 +372,7 @@ async def test_executemany_rowcount_no_hit(aconn, execmany):
     "query",
     [
         "insert into nosuchtable values (%s, %s)",
-        # This fails because we end up trying to copy in pipeline mode.
-        # However, sometimes (and pretty regularly if we enable pgconn.trace())
-        # something goes in a loop and only terminates by OOM. Strace shows
-        # an allocation loop. I think it's in the libpq.
-        # "copy (select %s, %s) to stdout",
+        "copy (select %s, %s) to stdout",
         "wat (%s, %s)",
     ],
 )