Decode the "query" column on Python side

We no longer convert the "query" column on Postgres side. Instead, we
let the database adapter load TEXT values as bytes, and finally decode
those when loading each row by using the database encoding.

Adapter configuration is done through the new 'text_as_bytes' flag,
passed to fetch*() functions, triggering the registration of a bytes
loader. Row decoding is done in BaseProcess.from_bytes() class method.

Decoding is made "robust" (not crashing) by simply passing
errors="replace" to bytes.decode().

Test test_data.py::test_encoding() remains valid for regression. On the
hand, test_InvalidTextRepresentation() is no longer needed as we don't
use convert_from() anymore.

Author: dlax

CHANGELOG.md

@@ -14,6 +14,9 @@
 
 ### Fixed
 
+* Rework decoding of the `query` column to (hopefully) make it more robust
+  (see #149 for the original report, #302 for a new problem raised while fixing
+  the previous one and #332 for the latest update).
 * Fix a few typos in the man page.
 
 ### Misc.

pgactivity/data.py

@@ -412,7 +412,8 @@ def pg_get_activities(self, duration_mode: int = 1) -> List[RunningProcess]:
                 "min_duration": self.min_duration,
                 "dbname_filter": self.filters.dbname,
             },
-            mkrow=RunningProcess,
+            mkrow=RunningProcess.from_bytes,
+            text_as_bytes=True,
         )
 
     def pg_get_waiting(self, duration_mode: int = 1) -> List[WaitingProcess]:
@@ -438,7 +439,8 @@ def pg_get_waiting(self, duration_mode: int = 1) -> List[WaitingProcess]:
                 "min_duration": self.min_duration,
                 "dbname_filter": self.filters.dbname,
             },
-            mkrow=WaitingProcess,
+            mkrow=WaitingProcess.from_bytes,
+            text_as_bytes=True,
         )
 
     def pg_get_blocking(self, duration_mode: int = 1) -> List[BlockingProcess]:
@@ -466,7 +468,8 @@ def pg_get_blocking(self, duration_mode: int = 1) -> List[BlockingProcess]:
                 "min_duration": self.min_duration,
                 "dbname_filter": self.filters.dbname,
             },
-            mkrow=BlockingProcess,
+            mkrow=BlockingProcess.from_bytes,
+            text_as_bytes=True,
         )
 
     def pg_is_local(self) -> bool:

pgactivity/pg.py

@@ -19,6 +19,7 @@
 
     import psycopg
     from psycopg import sql as sql
+    from psycopg.adapt import Buffer, Loader
     from psycopg.rows import dict_row
     from psycopg.errors import (
         FeatureNotSupported as FeatureNotSupported,
@@ -34,6 +35,12 @@
 
     Connection = psycopg.Connection[Dict[str, Any]]
 
+    class BytesLoader(Loader):
+        def load(self, data: Buffer) -> bytes:
+            if isinstance(data, memoryview):
+                return bytes(data)
+            return data
+
     def connect(*args: Any, **kwargs: Any) -> Connection:
         return psycopg.connect(*args, autocommit=True, row_factory=dict_row, **kwargs)
 
@@ -51,19 +58,27 @@ def execute(
         conn.execute(query, args, prepare=True)
 
     @overload
-    def cursor(conn: Connection, mkrow: Callable[..., Row]) -> psycopg.Cursor[Row]:
+    def cursor(
+        conn: Connection, mkrow: Callable[..., Row], text_as_bytes: bool
+    ) -> psycopg.Cursor[Row]:
         ...
 
     @overload
-    def cursor(conn: Connection, mkrow: None) -> psycopg.Cursor[psycopg.rows.DictRow]:
+    def cursor(
+        conn: Connection, mkrow: None, text_as_bytes: bool
+    ) -> psycopg.Cursor[psycopg.rows.DictRow]:
         ...
 
     def cursor(
-        conn: Connection, mkrow: Optional[Callable[..., Row]]
+        conn: Connection, mkrow: Optional[Callable[..., Row]], text_as_bytes: bool
     ) -> Union[psycopg.Cursor[psycopg.rows.DictRow], psycopg.Cursor[Row]]:
         if mkrow is not None:
-            return conn.cursor(row_factory=psycopg.rows.kwargs_row(mkrow))
-        return conn.cursor()
+            cur = conn.cursor(row_factory=psycopg.rows.kwargs_row(mkrow))
+        else:
+            cur = conn.cursor()  # type: ignore[assignment]
+        if text_as_bytes:
+            cur.adapters.register_loader("text", BytesLoader)
+        return cur
 
     @overload
     def fetchone(
@@ -72,6 +87,7 @@ def fetchone(
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
         *,
         mkrow: Callable[..., Row],
+        text_as_bytes: bool = False,
     ) -> Row:
         ...
 
@@ -80,6 +96,8 @@ def fetchone(
         conn: Connection,
         query: Union[str, sql.Composed],
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
+        *,
+        text_as_bytes: bool = False,
     ) -> Dict[str, Any]:
         ...
 
@@ -89,8 +107,9 @@ def fetchone(
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
         *,
         mkrow: Optional[Callable[..., Row]] = None,
+        text_as_bytes: bool = False,
     ) -> Union[Dict[str, Any], Row]:
-        with cursor(conn, mkrow) as cur:
+        with cursor(conn, mkrow, text_as_bytes) as cur:
             row = cur.execute(query, args, prepare=True).fetchone()
         assert row is not None
         return row
@@ -102,6 +121,7 @@ def fetchall(
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
         *,
         mkrow: Callable[..., Row],
+        text_as_bytes: bool = False,
     ) -> List[Row]:
         ...
 
@@ -110,6 +130,8 @@ def fetchall(
         conn: Connection,
         query: Union[str, sql.Composed],
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
+        *,
+        text_as_bytes: bool = False,
     ) -> List[Dict[str, Any]]:
         ...
 
@@ -118,13 +140,15 @@ def fetchall(
         query: Union[str, sql.Composed],
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
         *,
+        text_as_bytes: bool = False,
         mkrow: Optional[Callable[..., Row]] = None,
     ) -> Union[List[Dict[str, Any]], List[Row]]:
-        with cursor(conn, mkrow) as cur:
+        with cursor(conn, mkrow, text_as_bytes) as cur:
             return cur.execute(query, args, prepare=True).fetchall()
 
 except ImportError:
     import psycopg2
+    import psycopg2.extensions
     from psycopg2.extras import DictCursor
     from psycopg2 import sql as sql  # type: ignore[no-redef]
     from psycopg2.errors import (  # type: ignore[no-redef]
@@ -170,8 +194,11 @@ def fetchone(  # type: ignore[no-redef]
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
         *,
         mkrow: Optional[Callable[..., Row]] = None,
+        text_as_bytes: bool = False,
     ) -> Union[Dict[str, Any], Row]:
         with conn.cursor() as cur:
+            if text_as_bytes:
+                psycopg2.extensions.register_type(psycopg2.extensions.BYTES, cur)
             cur.execute(query, args)
             row = cur.fetchone()
         assert row is not None
@@ -185,8 +212,11 @@ def fetchall(  # type: ignore[no-redef]
         args: Union[None, Sequence[Any], Dict[str, Any]] = None,
         *,
         mkrow: Optional[Callable[..., Row]] = None,
+        text_as_bytes: bool = False,
     ) -> Union[List[Dict[str, Any]], List[Row]]:
         with conn.cursor() as cur:
+            if text_as_bytes:
+                psycopg2.extensions.register_type(psycopg2.extensions.BYTES, cur)
             cur.execute(query, args)
             rows = cur.fetchall()
         if mkrow is not None:

pgactivity/queries/get_blocking_oldest.sql

@@ -17,7 +17,7 @@ SELECT
       END AS state,
       CASE WHEN sq.query LIKE '<IDLE>%%'
           THEN NULL
-          ELSE convert_from(replace(sq.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8'))
+          ELSE sq.query
       END AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       waiting AS wait

pgactivity/queries/get_blocking_post_090200.sql

@@ -13,7 +13,7 @@ SELECT
       locktype AS type,
       duration,
       state,
-      convert_from(replace(sq.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      sq.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       waiting as wait
   FROM

pgactivity/queries/get_blocking_post_090600.sql

@@ -11,7 +11,7 @@ SELECT
       locktype AS type,
       duration,
       state,
-      convert_from(replace(sq.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      sq.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       wait_event as wait
   FROM

pgactivity/queries/get_pg_activity_oldest.sql

@@ -18,7 +18,7 @@ SELECT
       END AS state,
       CASE
           WHEN a.current_query LIKE '<IDLE>%%' THEN NULL
-          ELSE convert_from(replace(a.current_query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8'))
+          ELSE a.current_query
       END AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       NULL AS query_leader_pid,

pgactivity/queries/get_pg_activity_post_090200.sql

@@ -13,7 +13,7 @@ SELECT
       a.waiting AS wait,
       a.usename AS user,
       a.state AS state,
-      convert_from(replace(a.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      a.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       NULL AS query_leader_pid,
       false AS is_parallel_worker

pgactivity/queries/get_pg_activity_post_090600.sql

@@ -13,7 +13,7 @@ SELECT
       a.wait_event AS wait,
       a.usename AS user,
       a.state AS state,
-      convert_from(replace(a.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      a.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       NULL AS query_leader_pid,
       false AS is_parallel_worker

pgactivity/queries/get_pg_activity_post_100000.sql

@@ -13,7 +13,7 @@ SELECT
       a.wait_event as wait,
       a.usename AS user,
       a.state AS state,
-      convert_from(replace(a.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      a.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       NULL AS query_leader_pid,
       (   a.backend_type = 'background worker'

pgactivity/queries/get_pg_activity_post_110000.sql

@@ -12,7 +12,7 @@ SELECT
       a.wait_event AS wait,
       a.usename AS user,
       a.state AS state,
-      convert_from(replace(a.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      a.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       NULL AS query_leader_pid,
       a.backend_type = 'parallel worker' AS is_parallel_worker

pgactivity/queries/get_pg_activity_post_130000.sql

@@ -12,7 +12,7 @@ SELECT
       a.wait_event AS wait,
       a.usename AS user,
       a.state AS state,
-      convert_from(replace(a.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      a.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding,
       coalesce(a.leader_pid, a.pid) AS query_leader_pid,
       a.backend_type = 'parallel worker' AS is_parallel_worker

pgactivity/queries/get_waiting_oldest.sql

@@ -20,7 +20,7 @@ SELECT
       END AS state,
       CASE WHEN a.current_query LIKE '<IDLE>%%'
           THEN NULL
-          ELSE convert_from(replace(a.current_query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8'))
+          ELSE a.current_query
       END AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding
   FROM

pgactivity/queries/get_waiting_post_090200.sql

@@ -16,7 +16,7 @@ SELECT
       pg_locks.relation::regclass AS relation,
       EXTRACT(epoch FROM (NOW() - a.{duration_column})) AS duration,
       a.state as state,
-      convert_from(replace(a.query, '\', '\\')::bytea, coalesce(pg_catalog.pg_encoding_to_char(b.encoding), 'UTF8')) AS query,
+      a.query AS query,
       pg_catalog.pg_encoding_to_char(b.encoding) AS encoding
   FROM
       pg_catalog.pg_locks

pgactivity/types.py

@@ -14,6 +14,7 @@
     overload,
     Sequence,
     Tuple,
+    Type,
     TypeVar,
     Union,
 )
@@ -911,6 +912,21 @@ class BaseProcess:
     query_leader_pid: Optional[int]
     is_parallel_worker: bool
 
+    _P = TypeVar("_P", bound="BaseProcess")
+
+    @classmethod
+    def from_bytes(
+        cls: Type[_P], *, encoding: Optional[Union[str, bytes]], **kwargs: Any
+    ) -> _P:
+        if encoding is None:
+            encoding = "utf-8"
+        elif isinstance(encoding, bytes):  # psycopg2
+            encoding = encoding.decode()
+        for name, value in kwargs.items():
+            if isinstance(value, bytes):
+                kwargs[name] = value.decode(encoding, errors="replace")
+        return cls(encoding=encoding, **kwargs)
+
 
 @attr.s(auto_attribs=True, frozen=True, slots=True)
 class RunningProcess(BaseProcess):

tests/test_data.py

@@ -107,7 +107,7 @@ def test_terminate_backend(postgresql, data):
 
 
 def test_encoding(postgresql, data, execute):
-    """Test for issue #149"""
+    """Test for issue #149, #332."""
     conninfo = postgresql.info.dsn
     conn = psycopg.connect(conninfo)
     conn.autocommit = True
@@ -137,13 +137,6 @@ def test_encoding(postgresql, data, execute):
     assert "blocking éléphant" in blocking.query
 
 
-def test_InvalidTextRepresentation(postgresql, data, execute):
-    """Test for issue #275"""
-    postgresql.execute("select '123' ~ '\\d+', 'Hello world!\n', pg_sleep(3)")
-    running = data.pg_get_activities()
-    assert "123" in running[0].query
-
-
 def test_filters_dbname(data, execute):
     data_filtered = attr.evolve(data, filters=types.Filters(dbname="temp"))
     execute("SELECT pg_sleep(2)", dbname="template1", autocommit=True)