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]]]
__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:
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
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()
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
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)
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
import pytest
+import psycopg
from psycopg import pq
from psycopg import errors as e
# 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