Skip to content

Commit

Permalink
Use server_encoding as a fall back to database encoding
Browse files Browse the repository at this point in the history
(We need an assertion to pass static type checker as parameter_status()
would return None for unknown parameters, but we're querying a known one
here.)
  • Loading branch information
dlax committed Feb 27, 2023
1 parent b0c3b34 commit fe64960
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 9 deletions.
11 changes: 8 additions & 3 deletions pgactivity/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import re
from argparse import Namespace
from functools import partial
from typing import Dict, List, Optional

import attr
Expand Down Expand Up @@ -68,6 +69,7 @@ class Data:
pg_conn: Connection
pg_version: str
pg_num_version: int
server_encoding: str
min_duration: float
filters: Filters
dsn_parameters: Dict[str, str]
Expand Down Expand Up @@ -103,10 +105,13 @@ def pg_connect(
if pg.server_version(pg_conn) >= 130000:
pg.execute(pg_conn, queries.get("disable_log_min_duration_sample"))
pg_version = pg_get_short_version(pg_get_version(pg_conn))
server_encoding = pg_conn.info.parameter_status("server_encoding")
assert server_encoding is not None
return cls(
pg_conn,
pg_version,
pg.server_version(pg_conn),
server_encoding,
min_duration=min_duration,
failed_queries=FailedQueriesInfo(),
filters=filters,
Expand Down Expand Up @@ -412,7 +417,7 @@ def pg_get_activities(self, duration_mode: int = 1) -> List[RunningProcess]:
"min_duration": self.min_duration,
"dbname_filter": self.filters.dbname,
},
mkrow=RunningProcess.from_bytes,
mkrow=partial(RunningProcess.from_bytes, self.server_encoding),
text_as_bytes=True,
)

Expand All @@ -439,7 +444,7 @@ def pg_get_waiting(self, duration_mode: int = 1) -> List[WaitingProcess]:
"min_duration": self.min_duration,
"dbname_filter": self.filters.dbname,
},
mkrow=WaitingProcess.from_bytes,
mkrow=partial(WaitingProcess.from_bytes, self.server_encoding),
text_as_bytes=True,
)

Expand Down Expand Up @@ -468,7 +473,7 @@ def pg_get_blocking(self, duration_mode: int = 1) -> List[BlockingProcess]:
"min_duration": self.min_duration,
"dbname_filter": self.filters.dbname,
},
mkrow=BlockingProcess.from_bytes,
mkrow=partial(BlockingProcess.from_bytes, self.server_encoding),
text_as_bytes=True,
)

Expand Down
18 changes: 12 additions & 6 deletions pgactivity/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

import attr
import psutil
from attr import validators, converters
from attr import validators

from . import compat, colors, utils

Expand Down Expand Up @@ -908,23 +908,29 @@ class BaseProcess:
duration: Optional[float]
state: str
query: Optional[str]
encoding: str = attr.ib(converter=converters.default_if_none("utf-8")) # type: ignore[misc]
encoding: Optional[str]
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
cls: Type[_P],
server_encoding: str,
*,
encoding: Optional[Union[str, bytes]],
**kwargs: Any,
) -> _P:
if encoding is None:
encoding = "utf-8"
enc = server_encoding
elif isinstance(encoding, bytes): # psycopg2
encoding = encoding.decode()
enc = encoding = encoding.decode()
else:
enc = encoding
for name, value in kwargs.items():
if isinstance(value, bytes):
kwargs[name] = value.decode(encoding, errors="replace")
kwargs[name] = value.decode(enc, errors="replace")
return cls(encoding=encoding, **kwargs)


Expand Down

0 comments on commit fe64960

Please sign in to comment.