From e30fc3981b313f932a35cfd06dff700cefc75488 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 1 Jul 2020 05:58:08 +1200
Subject: [PATCH] Handle a bad command passed to copy

---
 psycopg3/cursor.py       | 12 +++++++++---
 tests/test_copy.py       | 15 +++++++++++++++
 tests/test_copy_async.py | 15 +++++++++++++++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/psycopg3/cursor.py b/psycopg3/cursor.py
index 8ee944658..fd4e60252 100644
--- a/psycopg3/cursor.py
+++ b/psycopg3/cursor.py
@@ -261,10 +261,16 @@ class BaseCursor:
 
         result = results[0]
         status = result.status
-        if status not in (pq.ExecStatus.COPY_IN, pq.ExecStatus.COPY_OUT):
+        if status in (pq.ExecStatus.COPY_IN, pq.ExecStatus.COPY_OUT):
+            return
+        elif status == pq.ExecStatus.FATAL_ERROR:
+            raise e.error_from_result(
+                result, encoding=self.connection.codec.name
+            )
+        else:
             raise e.ProgrammingError(
-                "copy() should be used only with COPY ... TO STDOUT"
-                " or COPY ... FROM STDIN statements"
+                "copy() should be used only with COPY ... TO STDOUT or COPY ..."
+                f" FROM STDIN statements, got {pq.ExecStatus(status).name}"
             )
 
 
diff --git a/tests/test_copy.py b/tests/test_copy.py
index 3e256df6a..e739f1985 100644
--- a/tests/test_copy.py
+++ b/tests/test_copy.py
@@ -69,6 +69,21 @@ def test_copy_in_buffers_pg_error(conn):
     assert conn.pgconn.transaction_status == conn.TransactionStatus.INERROR
 
 
+def test_copy_bad_result(conn):
+    conn.autocommit = True
+
+    cur = conn.cursor()
+
+    with pytest.raises(e.SyntaxError):
+        cur.copy("wat")
+
+    with pytest.raises(e.ProgrammingError):
+        cur.copy("select 1")
+
+    with pytest.raises(e.ProgrammingError):
+        cur.copy("reset timezone")
+
+
 @pytest.mark.parametrize(
     "format, buffer",
     [(Format.TEXT, "sample_text"), (Format.BINARY, "sample_binary")],
diff --git a/tests/test_copy_async.py b/tests/test_copy_async.py
index cce63aab5..289a6c428 100644
--- a/tests/test_copy_async.py
+++ b/tests/test_copy_async.py
@@ -72,6 +72,21 @@ async def test_copy_in_buffers_pg_error(aconn):
     assert aconn.pgconn.transaction_status == aconn.TransactionStatus.INERROR
 
 
+async def test_copy_bad_result(conn):
+    conn.autocommit = True
+
+    cur = conn.cursor()
+
+    with pytest.raises(e.SyntaxError):
+        await cur.copy("wat")
+
+    with pytest.raises(e.ProgrammingError):
+        await cur.copy("select 1")
+
+    with pytest.raises(e.ProgrammingError):
+        await cur.copy("reset timezone")
+
+
 @pytest.mark.parametrize(
     "format, buffer",
     [(Format.TEXT, "sample_text"), (Format.BINARY, "sample_binary")],
-- 
2.47.2