# Copyright (C) 2020 The Psycopg Team
from typing import Any, Dict, Optional, Sequence, Tuple, Type, Union
-from typing import cast
-from psycopg.pq.abc import PGconn, PGresult
-from psycopg.pq._enums import DiagnosticField
+from .pq.abc import PGconn, PGresult
+from .pq._enums import DiagnosticField
+from ._compat import TypeGuard
ErrorInfo = Union[None, PGresult, Dict[int, Optional[bytes]]]
__module__ = "psycopg"
sqlstate: Optional[str] = None
- pgconn: Optional[PGconn] = None
def __init__(
self,
super().__init__(*args)
self._info = info
self._encoding = encoding
- self.pgconn = pgconn
+ self._pgconn = pgconn
# Handle sqlstate codes for which we don't have a class.
if not self.sqlstate and info:
self.sqlstate = self.diag.sqlstate
+ @property
+ def pgconn(self) -> Optional[PGconn]:
+ return self._pgconn if self._pgconn else None
+
+ @property
+ def pgresult(self) -> Optional[PGresult]:
+ return self._info if _is_pgresult(self._info) else None
+
@property
def diag(self) -> "Diagnostic":
"""
def __reduce__(self) -> Union[str, Tuple[Any, ...]]:
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
+ res[2]["_info"] = _info_to_dict(self._info)
+ res[2]["_pgconn"] = None
return res
- @classmethod
- def _info_to_dict(cls, info: ErrorInfo) -> ErrorInfo:
- """
- Convert a PGresult to a dictionary to make the info picklable.
- """
- # PGresult is a protocol, can't use isinstance
- if hasattr(info, "error_field"):
- info = cast(PGresult, info)
- return {v: info.error_field(v) for v in DiagnosticField}
- else:
- return info
-
class InterfaceError(Error):
"""
def __reduce__(self) -> Union[str, Tuple[Any, ...]]:
res = super().__reduce__()
if isinstance(res, tuple) and len(res) >= 3:
- res[2]["_info"] = Error._info_to_dict(self._info)
+ res[2]["_info"] = _info_to_dict(self._info)
return res
+def _info_to_dict(info: ErrorInfo) -> ErrorInfo:
+ """
+ Convert a PGresult to a dictionary to make the info picklable.
+ """
+ # PGresult is a protocol, can't use isinstance
+ if _is_pgresult(info):
+ return {v: info.error_field(v) for v in DiagnosticField}
+ else:
+ return info
+
+
def lookup(sqlstate: str) -> Type[Error]:
"""Lookup an error code or `constant name`__ and return its exception class.
)
+def _is_pgresult(info: ErrorInfo) -> TypeGuard[PGresult]:
+ """Return True if an ErrorInfo is a PGresult instance."""
+ # PGresult is a protocol, can't use isinstance
+ return hasattr(info, "error_field")
+
+
def _class_for_state(sqlstate: str) -> Type[Error]:
try:
return lookup(sqlstate)
psycopg.connect("dbname=nosuchdb")
exc = pickle.loads(pickle.dumps(excinfo.value))
- assert not exc.pgconn
+ assert exc.pgconn is None
+
+
+def test_pgresult(conn):
+ with pytest.raises(e.DatabaseError) as excinfo:
+ conn.execute("select 1 from wat")
+
+ exc = excinfo.value
+ assert exc.pgresult
+ assert exc.pgresult.error_field(pq.DiagnosticField.SQLSTATE) == b"42P01"
+
+
+def test_pgresult_pickle(conn):
+ with pytest.raises(e.DatabaseError) as excinfo:
+ conn.execute("select 1 from wat")
+
+ exc = pickle.loads(pickle.dumps(excinfo.value))
+ assert exc.pgresult is None
+ assert exc.diag.sqlstate == "42P01"