Refactoring toards using SQLAlchemy's standard RowProxy for results

Author: lovelydinosaur

databases/backends/mysql.py

@@ -5,7 +5,8 @@
 
 import aiomysql
 from sqlalchemy.dialects.mysql import pymysql
-from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.engine.interfaces import Dialect, ExecutionContext
+from sqlalchemy.engine.result import ResultMetaData, RowProxy
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.types import TypeEngine
 
@@ -45,28 +46,9 @@ def connection(self) -> "MySQLConnection":
         return MySQLConnection(self._pool, self._dialect)
 
 
-class Record:
-    def __init__(self, row: tuple, result_columns: tuple, dialect: Dialect) -> None:
-        self._row = row
-        self._result_columns = result_columns
-        self._dialect = dialect
-        self._column_map = {
-            column_name: (idx, datatype)
-            for idx, (column_name, _, _, datatype) in enumerate(self._result_columns)
-        }
-
-    def __getitem__(self, key: str) -> typing.Any:
-        idx, datatype = self._column_map[key]
-        raw = self._row[idx]
-        try:
-            processor = _result_processors[datatype]
-        except KeyError:
-            processor = datatype.result_processor(self._dialect, None)
-            _result_processors[datatype] = processor
-
-        if processor is not None:
-            return processor(raw)
-        return raw
+class CompilationContext:
+    def __init__(self, context: ExecutionContext):
+        self.context = context
 
 
 class MySQLConnection(ConnectionBackend):
@@ -84,33 +66,38 @@ async def release(self) -> None:
         await self._pool.release(self._connection)
         self._connection = None
 
-    async def fetch_all(self, query: ClauseElement) -> typing.Any:
+    async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:
         assert self._connection is not None, "Connection is not acquired"
-        query, args, result_columns = self._compile(query)
+        query, args, context = self._compile(query)
         cursor = await self._connection.cursor()
         try:
             await cursor.execute(query, args)
             rows = await cursor.fetchall()
-            return [Record(row, result_columns, self._dialect) for row in rows]
+            metadata = ResultMetaData(context, cursor.description)
+            return [
+                RowProxy(metadata, row, metadata._processors, metadata._keymap)
+                for row in rows
+            ]
         finally:
             await cursor.close()
 
-    async def fetch_one(self, query: ClauseElement) -> typing.Any:
+    async def fetch_one(self, query: ClauseElement) -> RowProxy:
         assert self._connection is not None, "Connection is not acquired"
-        query, args, result_columns = self._compile(query)
+        query, args, context = self._compile(query)
         cursor = await self._connection.cursor()
         try:
             await cursor.execute(query, args)
             row = await cursor.fetchone()
-            return Record(row, result_columns, self._dialect)
+            metadata = ResultMetaData(context, cursor.description)
+            return RowProxy(metadata, row, metadata._processors, metadata._keymap)
         finally:
             await cursor.close()
 
     async def execute(self, query: ClauseElement, values: dict = None) -> None:
         assert self._connection is not None, "Connection is not acquired"
         if values is not None:
             query = query.values(values)
-        query, args, result_columns = self._compile(query)
+        query, args, context = self._compile(query)
         cursor = await self._connection.cursor()
         try:
             await cursor.execute(query, args)
@@ -123,7 +110,7 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         try:
             for item in values:
                 single_query = query.values(item)
-                single_query, args, result_columns = self._compile(single_query)
+                single_query, args, context = self._compile(single_query)
                 await cursor.execute(single_query, args)
         finally:
             await cursor.close()
@@ -132,26 +119,38 @@ async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Any, None]:
         assert self._connection is not None, "Connection is not acquired"
-        query, args, result_columns = self._compile(query)
+        query, args, context = self._compile(query)
         cursor = await self._connection.cursor()
         try:
             await cursor.execute(query, args)
+            metadata = ResultMetaData(context, cursor.description)
             async for row in cursor:
-                yield Record(row, result_columns, self._dialect)
+                yield RowProxy(metadata, row, metadata._processors, metadata._keymap)
         finally:
             await cursor.close()
 
     def transaction(self) -> TransactionBackend:
         return MySQLTransaction(self)
 
-    def _compile(self, query: ClauseElement) -> typing.Tuple[str, list, tuple]:
+    def _compile(
+        self, query: ClauseElement
+    ) -> typing.Tuple[str, list, CompilationContext]:
         compiled = query.compile(dialect=self._dialect)
         args = compiled.construct_params()
         logger.debug(compiled.string, args)
         for key, val in args.items():
             if key in compiled._bind_processors:
                 args[key] = compiled._bind_processors[key](val)
-        return compiled.string, args, compiled._result_columns
+
+        execution_context = self._dialect.execution_ctx_cls()
+        execution_context.dialect = self._dialect
+        execution_context.result_column_struct = (
+            compiled._result_columns,
+            compiled._ordered_columns,
+            compiled._textual_ordered_columns,
+        )
+
+        return compiled.string, args, CompilationContext(execution_context)
 
 
 class MySQLTransaction(TransactionBackend):

databases/core.py

@@ -5,6 +5,7 @@
 from types import TracebackType
 from urllib.parse import SplitResult, urlsplit
 
+from sqlalchemy.engine import RowProxy
 from sqlalchemy.sql import ClauseElement
 
 from databases.importer import import_from_string
@@ -27,7 +28,6 @@ def __init__(
     ):
         self._url = DatabaseURL(url)
         self._force_rollback = force_rollback
-
         self.is_connected = False
 
         backend_str = self.SUPPORTED_BACKENDS[self._url.dialect]
@@ -44,6 +44,9 @@ def __init__(
         self._global_transaction = None  # type: typing.Optional[Transaction]
 
     async def connect(self) -> None:
+        """
+        Establish the connection pool.
+        """
         assert not self.is_connected, "Already connected."
 
         await self._backend.connect()
@@ -57,6 +60,9 @@ async def connect(self) -> None:
             await self._global_transaction.__aenter__()
 
     async def disconnect(self) -> None:
+        """
+        Close all connections in the connection pool.
+        """
         assert self.is_connected, "Already disconnected."
 
         if self._force_rollback:
@@ -80,11 +86,11 @@ async def __aexit__(
     ) -> None:
         await self.disconnect()
 
-    async def fetch_all(self, query: ClauseElement) -> typing.Any:
+    async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:
         async with self.connection() as connection:
             return await connection.fetch_all(query=query)
 
-    async def fetch_one(self, query: ClauseElement) -> typing.Any:
+    async def fetch_one(self, query: ClauseElement) -> RowProxy:
         async with self.connection() as connection:
             return await connection.fetch_one(query=query)
 
@@ -98,7 +104,7 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
 
     async def iterate(
         self, query: ClauseElement
-    ) -> typing.AsyncGenerator[typing.Any, None]:
+    ) -> typing.AsyncGenerator[RowProxy, None]:
         async with self.connection() as connection:
             async for record in connection.iterate(query):
                 yield record

databases/interfaces.py

@@ -1,5 +1,6 @@
 import typing
 
+from sqlalchemy.engine import RowProxy
 from sqlalchemy.sql import ClauseElement
 
 
@@ -21,10 +22,10 @@ async def acquire(self) -> None:
     async def release(self) -> None:
         raise NotImplementedError()  # pragma: no cover
 
-    async def fetch_all(self, query: ClauseElement) -> typing.Any:
+    async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:
         raise NotImplementedError()  # pragma: no cover
 
-    async def fetch_one(self, query: ClauseElement) -> typing.Any:
+    async def fetch_one(self, query: ClauseElement) -> RowProxy:
         raise NotImplementedError()  # pragma: no cover
 
     async def execute(self, query: ClauseElement, values: dict = None) -> None:
@@ -35,7 +36,7 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
 
     async def iterate(
         self, query: ClauseElement
-    ) -> typing.AsyncGenerator[typing.Any, None]:
+    ) -> typing.AsyncGenerator[RowProxy, None]:
         raise NotImplementedError()  # pragma: no cover
         # mypy needs async iterators to contain a `yield`
         # https://github.com/python/mypy/issues/5385#issuecomment-407281656