Merge pull request #74 from encode/fetch_val

Add `fetch_val()` interface

Author: tomchristie

README.md

@@ -97,6 +97,10 @@ rows = await database.fetch_all(query=query)
 query = notes.select()
 row = await database.fetch_one(query=query)
 
+# Fetch single value, defaults to `column=0`.
+query = notes.select()
+value = await database.fetch_val(query=query)
+
 # Fetch multiple rows without loading them all into memory at once
 query = notes.select()
 async for row in database.iterate(query=query):

databases/backends/postgres.py

@@ -7,6 +7,7 @@
 from sqlalchemy.engine.interfaces import Dialect
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.sql.schema import Column
+from sqlalchemy.types import TypeEngine
 
 from databases.core import DatabaseURL
 from databases.interfaces import ConnectionBackend, DatabaseBackend, TransactionBackend
@@ -72,20 +73,27 @@ def __init__(
         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)
-        }
-        self._column_map_full = {
-            str(column[0]): (idx, datatype)
-            for idx, (_, _, column, datatype) in enumerate(self._result_columns)
-        }
+        self._column_map = (
+            {}
+        )  # type: typing.Mapping[str, typing.Tuple[int, TypeEngine]]
+        self._column_map_int = (
+            {}
+        )  # type: typing.Mapping[int, typing.Tuple[int, TypeEngine]]
+        self._column_map_full = (
+            {}
+        )  # type: typing.Mapping[str, typing.Tuple[int, TypeEngine]]
+        for idx, (column_name, _, column, datatype) in enumerate(self._result_columns):
+            self._column_map[column_name] = (idx, datatype)
+            self._column_map_int[idx] = (idx, datatype)
+            self._column_map_full[str(column[0])] = (idx, datatype)
 
     def __getitem__(self, key: typing.Any) -> typing.Any:
         if len(self._column_map) == 0:  # raw query
             return self._row[tuple(self._row.keys()).index(key)]
         elif type(key) is Column:
             idx, datatype = self._column_map_full[str(key)]
+        elif type(key) is int:
+            idx, datatype = self._column_map_int[key]
         else:
             idx, datatype = self._column_map[key]
         raw = self._row[idx]

databases/core.py

@@ -6,7 +6,6 @@
 from urllib.parse import SplitResult, parse_qsl, urlsplit
 
 from sqlalchemy import text
-from sqlalchemy.engine import RowProxy
 from sqlalchemy.sql import ClauseElement
 
 from databases.importer import import_from_string
@@ -91,16 +90,25 @@ async def __aexit__(
 
     async def fetch_all(
         self, query: typing.Union[ClauseElement, str], values: dict = None
-    ) -> typing.List[RowProxy]:
+    ) -> typing.List[typing.Mapping]:
         async with self.connection() as connection:
             return await connection.fetch_all(query, values)
 
     async def fetch_one(
         self, query: typing.Union[ClauseElement, str], values: dict = None
-    ) -> RowProxy:
+    ) -> typing.Optional[typing.Mapping]:
         async with self.connection() as connection:
             return await connection.fetch_one(query, values)
 
+    async def fetch_val(
+        self,
+        query: typing.Union[ClauseElement, str],
+        values: dict = None,
+        column: typing.Any = 0,
+    ) -> typing.Any:
+        async with self.connection() as connection:
+            return await connection.fetch_val(query, values, column=column)
+
     async def execute(
         self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.Any:
@@ -115,7 +123,7 @@ async def execute_many(
 
     async def iterate(
         self, query: typing.Union[ClauseElement, str], values: dict = None
-    ) -> typing.AsyncGenerator[RowProxy, None]:
+    ) -> typing.AsyncGenerator[typing.Mapping, None]:
         async with self.connection() as connection:
             async for record in connection.iterate(query, values):
                 yield record
@@ -167,14 +175,23 @@ async def __aexit__(
 
     async def fetch_all(
         self, query: typing.Union[ClauseElement, str], values: dict = None
-    ) -> typing.Any:
+    ) -> typing.List[typing.Mapping]:
         return await self._connection.fetch_all(self._build_query(query, values))
 
     async def fetch_one(
         self, query: typing.Union[ClauseElement, str], values: dict = None
-    ) -> typing.Any:
+    ) -> typing.Optional[typing.Mapping]:
         return await self._connection.fetch_one(self._build_query(query, values))
 
+    async def fetch_val(
+        self,
+        query: typing.Union[ClauseElement, str],
+        values: dict = None,
+        column: typing.Any = 0,
+    ) -> typing.Any:
+        row = await self._connection.fetch_one(self._build_query(query, values))
+        return None if row is None else row[column]
+
     async def execute(
         self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.Any:

tests/test_databases.py

@@ -137,6 +137,11 @@ async def test_queries(database_url):
             assert result["text"] == "example1"
             assert result["completed"] == True
 
+            # fetch_val()
+            query = sqlalchemy.sql.select([notes.c.text])
+            result = await database.fetch_val(query=query)
+            assert result == "example1"
+
             # iterate()
             query = notes.select()
             iterate_results = []