Skip to content

Commit

Permalink
S01E11
Browse files Browse the repository at this point in the history
  • Loading branch information
ansipunk committed Mar 3, 2024
1 parent 6910288 commit 09c1b2b
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 14 deletions.
8 changes: 1 addition & 7 deletions databases/backends/common/records.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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):
Expand Down
14 changes: 9 additions & 5 deletions databases/backends/dialects/psycopg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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


Expand Down
22 changes: 20 additions & 2 deletions databases/backends/psycopg.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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]
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 09c1b2b

Please sign in to comment.