S01E11

Author: ansipunk

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):

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
 
 

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: