From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 28 Oct 2020 19:26:01 +0000 (+0100)
Subject: Errors can be pickled
X-Git-Tag: 3.0.dev0~421
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=726972a5471a1b68a4ff3bac708c7860ab931e1a;p=thirdparty%2Fpsycopg.git

Errors can be pickled
---

diff --git a/psycopg3/psycopg3/errors.py b/psycopg3/psycopg3/errors.py
index f69109a0b..96279fc17 100644
--- a/psycopg3/psycopg3/errors.py
+++ b/psycopg3/psycopg3/errors.py
@@ -18,7 +18,7 @@ DBAPI-defined Exceptions are defined in the following hierarchy::
 
 # Copyright (C) 2020 The Psycopg Team
 
-from typing import Any, Callable, Dict, Optional, Sequence, Type
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Type, Union
 from psycopg3.pq.proto import PGresult
 from psycopg3.pq.enums import DiagnosticField
 
@@ -50,6 +50,13 @@ class Error(Exception):
     def diag(self) -> "Diagnostic":
         return Diagnostic(self.pgresult, encoding=self._encoding)
 
+    def __reduce__(self) -> Union[str, Tuple[Any, ...]]:
+        res = super().__reduce__()
+        if isinstance(res, tuple) and len(res) >= 3:
+            res[2]['pgresult'] = None
+
+        return res
+
 
 class InterfaceError(Error):
     """
diff --git a/tests/test_connection.py b/tests/test_connection.py
index eb00c6b3d..2d741a020 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -31,10 +31,16 @@ def test_close(conn):
     conn.close()
     assert conn.closed
     assert conn.status == conn.ConnStatus.BAD
+
+    cur = conn.cursor()
+
     conn.close()
     assert conn.closed
     assert conn.status == conn.ConnStatus.BAD
 
+    with pytest.raises(psycopg3.InterfaceError):
+        cur.execute("select 1")
+
 
 def test_weakref(dsn):
     conn = psycopg3.connect(dsn)
diff --git a/tests/test_connection_async.py b/tests/test_connection_async.py
index 0e23dc751..daa8aad03 100644
--- a/tests/test_connection_async.py
+++ b/tests/test_connection_async.py
@@ -34,10 +34,16 @@ async def test_close(aconn):
     await aconn.close()
     assert aconn.closed
     assert aconn.status == aconn.ConnStatus.BAD
+
+    cur = aconn.cursor()
+
     await aconn.close()
     assert aconn.closed
     assert aconn.status == aconn.ConnStatus.BAD
 
+    with pytest.raises(psycopg3.InterfaceError):
+        await cur.execute("select 1")
+
 
 async def test_weakref(dsn):
     conn = await psycopg3.AsyncConnection.connect(dsn)
diff --git a/tests/test_errors.py b/tests/test_errors.py
index d7f58cf69..b4b34750c 100644
--- a/tests/test_errors.py
+++ b/tests/test_errors.py
@@ -1,4 +1,7 @@
+import pickle
+
 import pytest
+
 from psycopg3 import pq
 from psycopg3 import errors as e
 
@@ -103,3 +106,13 @@ def test_lookup():
 
     with pytest.raises(KeyError):
         e.lookup("XXXXX")
+
+
+def test_error_pickle(conn):
+    cur = conn.cursor()
+    with pytest.raises(e.DatabaseError) as excinfo:
+        cur.execute("select 1 from wat")
+
+    exc = pickle.loads(pickle.dumps(excinfo.value))
+    assert isinstance(exc, e.UndefinedTable)
+    assert exc.pgresult is None