def __init__(self, **kwargs):
super().__init__(**kwargs)
if self.dbapi:
m = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", self.dbapi.__version__)
if m:
self.psycopg_version = tuple(
int(x) for x in m.group(1, 2, 3) if x is not None)
if self.psycopg_version < (3, 0, 2):
raise ImportError(
"psycopg version 3.0.2 or higher is required.")
from psycopg.adapt import AdaptersMap
self._psycopg_adapters_map = adapters_map = AdaptersMap(
self.dbapi.adapters)
if self._json_deserializer:
from psycopg.types.json import set_json_loads
set_json_loads(self._json_deserializer, adapters_map)
if self._json_serializer:
from psycopg.types.json import set_json_dumps
set_json_dumps(self._json_serializer, adapters_map)
def test_json_load_customise(conn, binary, pgtype):
cur = conn.cursor(binary=binary)
set_json_loads(my_loads)
try:
cur.execute(f"""select '{{"foo": "bar"}}'::{pgtype}""")
obj = cur.fetchone()[0]
assert obj["foo"] == "bar"
assert obj["answer"] == 42
finally:
set_json_loads(json.loads)
def test_json_load_customise_context(conn, binary, pgtype):
cur1 = conn.cursor(binary=binary)
cur2 = conn.cursor(binary=binary)
set_json_loads(my_loads, cur2)
cur1.execute(f"""select '{{"foo": "bar"}}'::{pgtype}""")
got = cur1.fetchone()[0]
assert got["foo"] == "bar"
assert "answer" not in got
cur2.execute(f"""select '{{"foo": "bar"}}'::{pgtype}""")
got = cur2.fetchone()[0]
assert got["foo"] == "bar"
assert got["answer"] == 42
from psycopg_pool import ConnectionPool
import atexit
import logging
from pydantic import BaseSettings
from tenacity import retry, retry_if_exception_type, stop_after_attempt
logger = logging.getLogger(__name__)
def dumps(data: dict) -> str:
"""Dump dictionary as string."""
return orjson.dumps(data).decode()
set_json_dumps(dumps)
set_json_loads(orjson.loads)
def pg_notice_handler(notice: psycopg.errors.Diagnostic) -> None:
"""Add PG messages to logging."""
msg = f"{notice.severity} - {notice.message_primary}"
logger.info(msg)
class Settings(BaseSettings):
"""Base Settings for Database Connection."""
db_min_conn_size: int = 0
db_max_conn_size: int = 1
db_max_queries: int = 5
db_max_idle: int = 5
