From 5abbcb23cae335341a19fcddf81328d977dfd8b6 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 10 Oct 2018 23:23:56 +0100
Subject: [PATCH] Fixed infinite loop in pq_get_last_result after COPY

There will be an error downstream but we have to get out of this
function first.

Close #781
---
 NEWS                |  3 ++-
 psycopg/pqpath.c    | 11 ++++++++---
 tests/test_async.py |  6 ++++++
 tests/test_green.py |  8 +++++++-
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/NEWS b/NEWS
index e9408c52..ab86b7d5 100644
--- a/NEWS
+++ b/NEWS
@@ -6,7 +6,8 @@ What's new in psycopg 2.7.6
 
 - Close named cursors if exist, even if `~cursor.execute()` wasn't called
   (:ticket:`#746`).
-- Fixed building on modern FreeBSD versions with Python 3.7 (:ticket:`#755`)
+- Fixed building on modern FreeBSD versions with Python 3.7 (:ticket:`#755`).
+- Fixed hang trying to :sql:`COPY` via `~cursor.execute()` (:ticket:`#781`).
 
 
 What's new in psycopg 2.7.5
diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c
index 204a6b00..b490d4f7 100644
--- a/psycopg/pqpath.c
+++ b/psycopg/pqpath.c
@@ -1106,12 +1106,13 @@ pq_send_query(connectionObject *conn, const char *query)
  * The function will block only if a command is active and the
  * necessary response data has not yet been read by PQconsumeInput.
  *
- * The result should be disposed using PQclear()
+ * The result should be disposed of using PQclear()
  */
 PGresult *
 pq_get_last_result(connectionObject *conn)
 {
     PGresult *result = NULL, *res;
+    ExecStatusType status;
 
     /* Read until PQgetResult gives a NULL */
     while (NULL != (res = PQgetResult(conn->pgconn))) {
@@ -1124,11 +1125,15 @@ pq_get_last_result(connectionObject *conn)
             PQclear(result);
         }
         result = res;
+        status = PQresultStatus(result);
+        Dprintf("pq_get_last_result: got result %s", PQresStatus(status));
 
-        /* After entering copy both mode, libpq will make a phony
+        /* After entering copy mode, libpq will make a phony
          * PGresult for us every time we query for it, so we need to
          * break out of this endless loop. */
-        if (PQresultStatus(result) == PGRES_COPY_BOTH) {
+        if (status == PGRES_COPY_BOTH
+                || status == PGRES_COPY_OUT
+                || status == PGRES_COPY_IN) {
             break;
         }
     }
diff --git a/tests/test_async.py b/tests/test_async.py
index 4eb5e6a2..3043678f 100755
--- a/tests/test_async.py
+++ b/tests/test_async.py
@@ -450,6 +450,12 @@ class AsyncTests(ConnectingTestCase):
         else:
             self.fail("no exception raised")
 
+    @skip_before_postgres(8, 2)
+    def test_copy_no_hang(self):
+        cur = self.conn.cursor()
+        cur.execute("copy (select 1) to stdout")
+        self.assertRaises(psycopg2.ProgrammingError, self.wait, self.conn)
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)
diff --git a/tests/test_green.py b/tests/test_green.py
index 8c1c20ce..57dd8053 100755
--- a/tests/test_green.py
+++ b/tests/test_green.py
@@ -27,7 +27,7 @@ import psycopg2
 import psycopg2.extensions
 import psycopg2.extras
 
-from testutils import ConnectingTestCase, slow
+from testutils import ConnectingTestCase, skip_before_postgres, slow
 
 
 class ConnectionStub(object):
@@ -111,6 +111,12 @@ class GreenTestCase(ConnectingTestCase):
         curs.execute("select 1")
         self.assertEqual(curs.fetchone()[0], 1)
 
+    @skip_before_postgres(8, 2)
+    def test_copy_no_hang(self):
+        cur = self.conn.cursor()
+        self.assertRaises(psycopg2.ProgrammingError,
+            cur.execute, "copy (select 1) to stdout")
+
 
 class CallbackErrorTestCase(ConnectingTestCase):
     def setUp(self):