diff --git a/databases/backends/common/records.py b/databases/backends/common/records.py index 5f6ba263..65032fc8 100644 --- a/databases/backends/common/records.py +++ b/databases/backends/common/records.py @@ -39,9 +39,6 @@ def __init__( @property def _mapping(self) -> typing.Mapping: - if hasattr(self._row, "_asdict"): - return self._row._asdict() - return self._row def keys(self) -> typing.KeysView: @@ -52,10 +49,7 @@ def values(self) -> typing.ValuesView: def __getitem__(self, key: typing.Any) -> typing.Any: if len(self._column_map) == 0: - try: - return self._row[key] - except TypeError: - return self._mapping[key] + return self._row[key] elif isinstance(key, Column): idx, datatype = self._column_map_full[str(key)] elif isinstance(key, int): diff --git a/databases/backends/dialects/psycopg.py b/databases/backends/dialects/psycopg.py index 7740d976..07bd1880 100644 --- a/databases/backends/dialects/psycopg.py +++ b/databases/backends/dialects/psycopg.py @@ -7,9 +7,10 @@ import typing +from sqlalchemy import types, util from sqlalchemy.dialects.postgresql.base import PGDialect, PGExecutionContext from sqlalchemy.engine import processors -from sqlalchemy.types import Numeric +from sqlalchemy.types import Float, Numeric class PGExecutionContext_psycopg(PGExecutionContext): @@ -32,10 +33,13 @@ def result_processor( class PGDialect_psycopg(PGDialect): - colspecs = { - **PGDialect.colspecs, - Numeric: PGNumeric, - } + colspecs = util.update_copy( + PGDialect.colspecs, + { + types.Numeric: PGNumeric, + types.Float: Float, + }, + ) execution_ctx_cls = PGExecutionContext_psycopg diff --git a/databases/backends/psycopg.py b/databases/backends/psycopg.py index da0a6718..eb3cec16 100644 --- a/databases/backends/psycopg.py +++ b/databases/backends/psycopg.py @@ -1,6 +1,9 @@ import typing +import orjson import psycopg +import psycopg.adapt +import psycopg.types import psycopg_pool from psycopg.rows import namedtuple_row from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg @@ -18,6 +21,16 @@ ) +class JsonLoader(psycopg.adapt.Loader): + def load(self, data): + return orjson.loads(data) + + +class JsonDumper(psycopg.adapt.Dumper): + def dump(self, data): + return orjson.dumps(data) + + class PsycopgBackend(DatabaseBackend): _database_url: DatabaseURL _options: typing.Dict[str, typing.Any] @@ -73,8 +86,13 @@ async def acquire(self) -> None: raise RuntimeError("PsycopgBackend is not running") # TODO: Add configurable timeouts - self._connection = await self._database._pool.getconn() - await self._connection.set_autocommit(True) + connection = await self._database._pool.getconn() + connection.adapters.register_loader("json", JsonLoader) + connection.adapters.register_loader("jsonb", JsonLoader) + connection.adapters.register_dumper(dict, JsonDumper) + connection.adapters.register_dumper(list, JsonDumper) + await connection.set_autocommit(True) + self._connection = connection async def release(self) -> None: if self._connection is None: