S01E10

Author: ansipunk

databases/backends/aiopg.py

@@ -5,13 +5,14 @@
 import uuid
 
 import aiopg
-from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
 from sqlalchemy.engine.cursor import CursorResultMetaData
 from sqlalchemy.engine.interfaces import Dialect, ExecutionContext
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.sql.ddl import DDLElement
 
 from databases.backends.common.records import Record, Row, create_column_maps
+from databases.backends.compilers.psycopg import PGCompiler_psycopg
+from databases.backends.dialects.psycopg import PGDialect_psycopg
 from databases.core import LOG_EXTRA, DatabaseURL
 from databases.interfaces import (
     ConnectionBackend,
@@ -36,10 +37,12 @@ def _get_dialect(self) -> Dialect:
         dialect = PGDialect_psycopg(
             json_serializer=json.dumps, json_deserializer=lambda x: x
         )
+        dialect.statement_compiler = PGCompiler_psycopg
         dialect.implicit_returning = True
         dialect.supports_native_enum = True
         dialect.supports_smallserial = True  # 9.2+
         dialect._backslash_escapes = False
+        dialect.supports_sane_multi_rowcount = True  # psycopg 2.0.9+
         dialect._has_native_hstore = True
         dialect.supports_native_decimal = True
 

databases/backends/asyncpg.py

@@ -2,12 +2,12 @@
 import typing
 
 import asyncpg
-from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
 from sqlalchemy.engine.interfaces import Dialect
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.sql.ddl import DDLElement
 
 from databases.backends.common.records import Record, create_column_maps
+from databases.backends.dialects.psycopg import dialect as psycopg_dialect
 from databases.core import LOG_EXTRA, DatabaseURL
 from databases.interfaces import (
     ConnectionBackend,
@@ -29,14 +29,16 @@ def __init__(
         self._pool = None
 
     def _get_dialect(self) -> Dialect:
-        dialect = PGDialect_psycopg(paramstyle="pyformat")
+        dialect = psycopg_dialect(paramstyle="pyformat")
+
         dialect.implicit_returning = True
         dialect.supports_native_enum = True
         dialect.supports_smallserial = True  # 9.2+
         dialect._backslash_escapes = False
         dialect.supports_sane_multi_rowcount = True  # psycopg 2.0.9+
         dialect._has_native_hstore = True
         dialect.supports_native_decimal = True
+
         return dialect
 
     def _get_connection_kwargs(self) -> dict:
@@ -189,10 +191,10 @@ def _compile(self, query: ClauseElement) -> typing.Tuple[str, list, tuple]:
         )
         return compiled_query, args, result_map
 
-        @property
-        def raw_connection(self) -> asyncpg.connection.Connection:
-            assert self._connection is not None, "Connection is not acquired"
-            return self._connection
+    @property
+    def raw_connection(self) -> asyncpg.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
 
 
 class AsyncpgTransaction(TransactionBackend):

databases/backends/compilers/__init__.py

@@ -0,0 +1,17 @@
+from sqlalchemy.dialects.postgresql.psycopg import PGCompiler_psycopg
+
+
+class APGCompiler_psycopg2(PGCompiler_psycopg):
+    def construct_params(self, *args, **kwargs):
+        pd = super().construct_params(*args, **kwargs)
+
+        for column in self.prefetch:
+            pd[column.key] = self._exec_default(column.default)
+
+        return pd
+
+    def _exec_default(self, default):
+        if default.is_callable:
+            return default.arg(self.dialect)
+        else:
+            return default.arg

databases/backends/compilers/psycopg.py

@@ -0,0 +1,42 @@
+"""
+All the unique changes for the databases package
+with the custom Numeric as the deprecated pypostgresql
+for backwards compatibility and to make sure the
+package can go to SQLAlchemy 2.0+.
+"""
+
+import typing
+
+from sqlalchemy.dialects.postgresql.base import PGDialect, PGExecutionContext
+from sqlalchemy.engine import processors
+from sqlalchemy.types import Numeric
+
+
+class PGExecutionContext_psycopg(PGExecutionContext):
+    ...
+
+
+class PGNumeric(Numeric):
+    def bind_processor(
+        self, dialect: typing.Any
+    ) -> typing.Union[str, None]:  # pragma: no cover
+        return processors.to_str
+
+    def result_processor(
+        self, dialect: typing.Any, coltype: typing.Any
+    ) -> typing.Union[float, None]:  # pragma: no cover
+        if self.asdecimal:
+            return None
+        else:
+            return processors.to_float
+
+
+class PGDialect_psycopg(PGDialect):
+    colspecs = {
+        **PGDialect.colspecs,
+        Numeric: PGNumeric,
+    }
+    execution_ctx_cls = PGExecutionContext_psycopg
+
+
+dialect = PGDialect_psycopg