Stricter interface definitions of fetch_all, fetch_one

Author: lovelydinosaur

databases/backends/mysql.py

@@ -98,7 +98,7 @@ async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:
         finally:
             await cursor.close()
 
-    async def fetch_one(self, query: ClauseElement) -> RowProxy:
+    async def fetch_one(self, query: ClauseElement) -> typing.Optional[RowProxy]:
         assert self._connection is not None, "Connection is not acquired"
         query, args, context = self._compile(query)
         cursor = await self._connection.cursor()

databases/backends/postgres.py

@@ -1,5 +1,6 @@
 import logging
 import typing
+from collections.abc import Mapping
 
 import asyncpg
 from sqlalchemy.dialects.postgresql import pypostgresql
@@ -63,7 +64,7 @@ def connection(self) -> "PostgresConnection":
         return PostgresConnection(self, self._dialect)
 
 
-class Record:
+class Record(Mapping):
     def __init__(self, row: tuple, result_columns: tuple, dialect: Dialect) -> None:
         self._row = row
         self._result_columns = result_columns
@@ -86,6 +87,12 @@ def __getitem__(self, key: str) -> typing.Any:
             return processor(raw)
         return raw
 
+    def __iter__(self) -> typing.Iterator:
+        return iter(self._column_map)
+
+    def __len__(self) -> int:
+        return len(self._column_map)
+
 
 class PostgresConnection(ConnectionBackend):
     def __init__(self, database: PostgresBackend, dialect: Dialect):
@@ -104,13 +111,13 @@ async def release(self) -> None:
         self._connection = await self._database._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[Record]:
         assert self._connection is not None, "Connection is not acquired"
         query, args, result_columns = self._compile(query)
         rows = await self._connection.fetch(query, *args)
         return [Record(row, result_columns, self._dialect) for row in rows]
 
-    async def fetch_one(self, query: ClauseElement) -> typing.Any:
+    async def fetch_one(self, query: ClauseElement) -> typing.Optional[Record]:
         assert self._connection is not None, "Connection is not acquired"
         query, args, result_columns = self._compile(query)
         row = await self._connection.fetchrow(query, *args)

databases/backends/sqlite.py

@@ -91,7 +91,7 @@ async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:
                 for row in rows
             ]
 
-    async def fetch_one(self, query: ClauseElement) -> RowProxy:
+    async def fetch_one(self, query: ClauseElement) -> typing.Optional[RowProxy]:
         assert self._connection is not None, "Connection is not acquired"
         query, args, context = self._compile(query)
 

databases/interfaces.py

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

scripts/lint

@@ -10,3 +10,4 @@ set -x
 ${PREFIX}autoflake --in-place --recursive databases tests
 ${PREFIX}black databases tests
 ${PREFIX}isort --multi-line=3 --trailing-comma --force-grid-wrap=0 --combine-as --line-width 88 --recursive --apply databases tests
+${PREFIX}mypy databases --ignore-missing-imports --disallow-untyped-defs

tests/test_databases.py

@@ -151,6 +151,46 @@ async def test_queries(database_url):
             assert iterate_results[2]["completed"] == True
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_results_support_mapping_interface(database_url):
+    """
+    Casting results to a dict should work, since the interface defines them
+    as supporting the mapping interface.
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            # execute()
+            query = notes.insert()
+            values = {"text": "example1", "completed": True}
+            await database.execute(query, values)
+
+            # fetch_all()
+            query = notes.select()
+            results = await database.fetch_all(query=query)
+            results_dict = dict(results[0])
+
+            assert len(results_dict) == 3
+            assert isinstance(results_dict["id"], int)
+            assert results_dict["text"] == "example1"
+            assert results_dict["completed"] == True
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_fetch_one_returning_no_results(database_url):
+    """
+    fetch_one should return `None` when no results match.
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            # fetch_all()
+            query = notes.select()
+            result = await database.fetch_one(query=query)
+            assert result is None
+
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_execute_return_val(database_url):