From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 5 Aug 2023 21:15:05 +0000 (+0100)
Subject: fix: fix infinite loop and OOM in bad executemany
X-Git-Tag: pool-3.1.8~23
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ce4bab206d0c176c78c6dc95560e39c14cf436ad;p=thirdparty%2Fpsycopg.git

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 81b52a2af..4f2ec878b 100644
--- 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 fbd901385..a51fce5e2 100644
--- 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 5ac793b64..18941a54e 100644
--- 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 ec6504d07..67b314acb 100644
--- 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)",
     ],
 )