From be2e9ae6f75aa7dbdbd3c6370d382f5c316a8042 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 19 Mar 2020 20:55:10 +1300
Subject: [PATCH] Added error_message function instead of message_str method

Keeping the pq interface simple and low level and moving the policy a
bit aside.
---
 psycopg3/pq/__init__.py  |  3 +++
 psycopg3/pq/misc.py      | 42 ++++++++++++++++++++++++++++++++++++++++
 psycopg3/pq/pq_ctypes.py | 17 +++++-----------
 tests/pq/test_misc.py    | 15 ++++++++++++++
 4 files changed, 65 insertions(+), 12 deletions(-)
 create mode 100644 psycopg3/pq/misc.py
 create mode 100644 tests/pq/test_misc.py

diff --git a/psycopg3/pq/__init__.py b/psycopg3/pq/__init__.py
index 3bbe1587e..05c7a9225 100644
--- a/psycopg3/pq/__init__.py
+++ b/psycopg3/pq/__init__.py
@@ -18,10 +18,12 @@ from .enums import (
     DiagnosticField,
 )
 from .encodings import py_codecs
+from .misc import error_message
 
 from . import pq_ctypes as pq_module
 
 PGconn = pq_module.PGconn
+PGresult = pq_module.PGresult
 PQerror = pq_module.PQerror
 Conninfo = pq_module.Conninfo
 
@@ -35,5 +37,6 @@ __all__ = (
     "PGconn",
     "Conninfo",
     "PQerror",
+    "error_message",
     "py_codecs",
 )
diff --git a/psycopg3/pq/misc.py b/psycopg3/pq/misc.py
new file mode 100644
index 000000000..9385cc7ac
--- /dev/null
+++ b/psycopg3/pq/misc.py
@@ -0,0 +1,42 @@
+"""
+Various functionalities to make easier to work with the libpq.
+"""
+
+# Copyright (C) 2020 The Psycopg Team
+
+
+def error_message(obj):
+    """
+    Return an error message from a PGconn or PGresult.
+
+    The return value is a str (unlike pq data which is usually bytes).
+    """
+    from psycopg3 import pq
+
+    if isinstance(obj, pq.PGconn):
+        msg = obj.error_message
+
+        # strip severity and whitespaces
+        if msg:
+            msg = msg.splitlines()[0].split(b":", 1)[-1].strip()
+
+    elif isinstance(obj, pq.PGresult):
+        msg = obj.error_field(pq.DiagnosticField.PG_DIAG_MESSAGE_PRIMARY)
+        if not msg:
+            msg = obj.error_message
+
+            # strip severity and whitespaces
+            if msg:
+                msg = msg.splitlines()[0].split(b":", 1)[-1].strip()
+
+    else:
+        raise TypeError(
+            f"PGconn or PGresult expected, got {type(obj).__name__}"
+        )
+
+    if msg:
+        msg = msg.decode("utf8", "replace")  # TODO: or in connection encoding?
+    else:
+        msg = "no details available"
+
+    return msg
diff --git a/psycopg3/pq/pq_ctypes.py b/psycopg3/pq/pq_ctypes.py
index 31f2b9189..7e4549abd 100644
--- a/psycopg3/pq/pq_ctypes.py
+++ b/psycopg3/pq/pq_ctypes.py
@@ -19,6 +19,7 @@ from .enums import (
     TransactionStatus,
     Ping,
 )
+from .misc import error_message
 from . import _pq_ctypes as impl
 from ..exceptions import OperationalError
 
@@ -158,14 +159,6 @@ class PGconn:
     def error_message(self):
         return impl.PQerrorMessage(self.pgconn_ptr)
 
-    @property
-    def error_str(self):
-        rv = self.error_message
-        if rv:
-            return rv.encode('utf8', 'replace').rstrip()
-        else:
-            return "no details available"
-
     @property
     def socket(self):
         return impl.PQsocket(self.pgconn_ptr)
@@ -351,7 +344,7 @@ class PGconn:
 
     def consume_input(self):
         if 1 != impl.PQconsumeInput(self.pgconn_ptr):
-            raise PQerror(f"consuming input failed: {self.error_str}")
+            raise PQerror(f"consuming input failed: {error_message(self)}")
 
     def is_busy(self):
         return impl.PQisBusy(self.pgconn_ptr)
@@ -363,12 +356,12 @@ class PGconn:
     @nonblocking.setter
     def nonblocking(self, arg):
         if 0 > impl.PQsetnonblocking(self.pgconn_ptr, arg):
-            raise PQerror(f"setting nonblocking failed: {self.error_str}")
+            raise PQerror(f"setting nonblocking failed: {error_message(self)}")
 
     def flush(self):
         rv = impl.PQflush(self.pgconn_ptr)
         if rv < 0:
-            raise PQerror(f"flushing failed: {self.error_str}")
+            raise PQerror(f"flushing failed: {error_message(self)}")
         return rv
 
 
@@ -479,7 +472,7 @@ class Conninfo:
             if not errmsg:
                 raise MemoryError("couldn't allocate on conninfo parse")
             else:
-                exc = PQerror(errmsg.value)
+                exc = PQerror(errmsg.value.decode("utf8", "replace"))
                 impl.PQfreemem(errmsg)
                 raise exc
 
diff --git a/tests/pq/test_misc.py b/tests/pq/test_misc.py
new file mode 100644
index 000000000..c193c87f0
--- /dev/null
+++ b/tests/pq/test_misc.py
@@ -0,0 +1,15 @@
+import pytest
+
+
+def test_error_message(pq, pgconn):
+    res = pgconn.exec_(b"wat")
+    assert res.status == pq.ExecStatus.PGRES_FATAL_ERROR
+    msg = pq.error_message(pgconn)
+    assert msg == 'syntax error at or near "wat"'
+    assert msg == pq.error_message(res)
+    assert msg == res.error_field(
+        pq.DiagnosticField.PG_DIAG_MESSAGE_PRIMARY
+    ).decode("ascii")
+
+    with pytest.raises(TypeError):
+        pq.error_message(None)
-- 
2.47.3