From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 11 Mar 2022 19:48:08 +0000 (+0000)
Subject: feat: add `Error.pgconn` attribute
X-Git-Tag: 3.1~168^2~3
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=9220293dc023b757f2a57702c14b1460ff8f25b0;p=thirdparty%2Fpsycopg.git

feat: add `Error.pgconn` attribute

Might be useful to investigate errors such as bad/missing password. See #242
for details.
---

diff --git a/psycopg/psycopg/errors.py b/psycopg/psycopg/errors.py
index d23ce7506..e5a94913a 100644
--- a/psycopg/psycopg/errors.py
+++ b/psycopg/psycopg/errors.py
@@ -21,7 +21,7 @@ DBAPI-defined Exceptions are defined in the following hierarchy::
 from typing import Any, Dict, Optional, Sequence, Tuple, Type, Union
 from typing import cast
 
-from psycopg.pq.abc import PGresult
+from psycopg.pq.abc import PGconn, PGresult
 from psycopg.pq._enums import DiagnosticField
 
 ErrorInfo = Union[None, PGresult, Dict[int, Optional[bytes]]]
@@ -52,13 +52,19 @@ class Error(Exception):
     __module__ = "psycopg"
 
     sqlstate: Optional[str] = None
+    pgconn: Optional[PGconn] = None
 
     def __init__(
-        self, *args: Sequence[Any], info: ErrorInfo = None, encoding: str = "utf-8"
+        self,
+        *args: Sequence[Any],
+        info: ErrorInfo = None,
+        encoding: str = "utf-8",
+        pgconn: Optional[PGconn] = None
     ):
         super().__init__(*args)
         self._info = info
         self._encoding = encoding
+        self.pgconn = pgconn
 
         # Handle sqlstate codes for which we don't have a class.
         if not self.sqlstate and info:
@@ -75,6 +81,8 @@ class Error(Exception):
         res = super().__reduce__()
         if isinstance(res, tuple) and len(res) >= 3:
             res[2]["_info"] = self._info_to_dict(self._info)
+            # To make the exception picklable
+            res[2]["pgconn"] = None
 
         return res
 
diff --git a/psycopg/psycopg/generators.py b/psycopg/psycopg/generators.py
index 679763107..f0140e4a8 100644
--- a/psycopg/psycopg/generators.py
+++ b/psycopg/psycopg/generators.py
@@ -39,7 +39,8 @@ def connect(conninfo: str) -> PQGenConn[PGconn]:
         if conn.status == ConnStatus.BAD:
             encoding = conninfo_encoding(conninfo)
             raise e.OperationalError(
-                f"connection is bad: {pq.error_message(conn, encoding=encoding)}"
+                f"connection is bad: {pq.error_message(conn, encoding=encoding)}",
+                pgconn=conn,
             )
 
         status = conn.connect_poll()
@@ -52,10 +53,11 @@ def connect(conninfo: str) -> PQGenConn[PGconn]:
         elif status == PollingStatus.FAILED:
             encoding = conninfo_encoding(conninfo)
             raise e.OperationalError(
-                f"connection failed: {pq.error_message(conn, encoding=encoding)}"
+                f"connection failed: {pq.error_message(conn, encoding=encoding)}",
+                pgconn=conn,
             )
         else:
-            raise e.InternalError(f"unexpected poll status: {status}")
+            raise e.InternalError(f"unexpected poll status: {status}", pgconn=conn)
 
     conn.nonblocking = 1
     return conn
diff --git a/psycopg_c/psycopg_c/_psycopg/generators.pyx b/psycopg_c/psycopg_c/_psycopg/generators.pyx
index 4c8405743..951785e74 100644
--- a/psycopg_c/psycopg_c/_psycopg/generators.pyx
+++ b/psycopg_c/psycopg_c/_psycopg/generators.pyx
@@ -35,7 +35,8 @@ def connect(conninfo: str) -> PQGenConn[abc.PGconn]:
         if conn_status == libpq.CONNECTION_BAD:
             encoding = conninfo_encoding(conninfo)
             raise e.OperationalError(
-                f"connection is bad: {error_message(conn, encoding=encoding)}"
+                f"connection is bad: {error_message(conn, encoding=encoding)}",
+                pgconn=conn
             )
 
         poll_status = libpq.PQconnectPoll(pgconn_ptr)
@@ -49,10 +50,13 @@ def connect(conninfo: str) -> PQGenConn[abc.PGconn]:
         elif poll_status == libpq.PGRES_POLLING_FAILED:
             encoding = conninfo_encoding(conninfo)
             raise e.OperationalError(
-                f"connection failed: {error_message(conn, encoding=encoding)}"
+                f"connection failed: {error_message(conn, encoding=encoding)}",
+                pgconn=conn
             )
         else:
-            raise e.InternalError(f"unexpected poll status: {poll_status}")
+            raise e.InternalError(
+                f"unexpected poll status: {poll_status}", pgconn=conn
+            )
 
     conn.nonblocking = 1
     return conn
diff --git a/tests/test_errors.py b/tests/test_errors.py
index 5ae3b51b9..c54a0e43f 100644
--- a/tests/test_errors.py
+++ b/tests/test_errors.py
@@ -4,6 +4,7 @@ from weakref import ref
 
 import pytest
 
+import psycopg
 from psycopg import pq
 from psycopg import errors as e
 
@@ -260,3 +261,20 @@ def test_unknown_sqlstate(conn):
     # Survives pickling too
     pexc = pickle.loads(pickle.dumps(exc))
     assert pexc.sqlstate == code
+
+
+def test_pgconn_error():
+    with pytest.raises(psycopg.OperationalError) as excinfo:
+        psycopg.connect("dbname=nosuchdb")
+
+    exc = excinfo.value
+    assert exc.pgconn
+    assert exc.pgconn.db == b"nosuchdb"
+
+
+def test_pgconn_error_pickle():
+    with pytest.raises(psycopg.OperationalError) as excinfo:
+        psycopg.connect("dbname=nosuchdb")
+
+    exc = pickle.loads(pickle.dumps(excinfo.value))
+    assert not exc.pgconn