Merge pull request #59 from ricardomomm/improve-row-data-retrieval

fix issue with duplicate column names

Author: lovelydinosaur

databases/backends/postgres.py

@@ -6,6 +6,7 @@
 from sqlalchemy.dialects.postgresql import pypostgresql
 from sqlalchemy.engine.interfaces import Dialect
 from sqlalchemy.sql import ClauseElement
+from sqlalchemy.sql.schema import Column
 
 from databases.core import DatabaseURL
 from databases.interfaces import ConnectionBackend, DatabaseBackend, TransactionBackend
@@ -73,9 +74,16 @@ def __init__(self, row: tuple, result_columns: tuple, dialect: Dialect) -> None:
             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)
+        }
 
-    def __getitem__(self, key: str) -> typing.Any:
-        idx, datatype = self._column_map[key]
+    def __getitem__(self, key: typing.Any) -> typing.Any:
+        if type(key) is Column:
+            idx, datatype = self._column_map_full[str(key)]
+        else:
+            idx, datatype = self._column_map[key]
         raw = self._row[idx]
         try:
             processor = _result_processors[datatype]

tests/test_databases.py

@@ -178,6 +178,37 @@ async def test_results_support_mapping_interface(database_url):
             assert results_as_dicts[0]["completed"] == True
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_results_support_column_reference(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):
+            now = datetime.datetime.now().replace(microsecond=0)
+            today = datetime.date.today()
+
+            # execute()
+            query = articles.insert()
+            values = {"title": "Hello, world Article", "published": now}
+            await database.execute(query, values)
+
+            query = custom_date.insert()
+            values = {"title": "Hello, world Custom", "published": today}
+            await database.execute(query, values)
+
+            # fetch_all()
+            query = sqlalchemy.select([articles, custom_date])
+            results = await database.fetch_all(query=query)
+            assert len(results) == 1
+            assert results[0][articles.c.title] == "Hello, world Article"
+            assert results[0][articles.c.published] == now
+            assert results[0][custom_date.c.title] == "Hello, world Custom"
+            assert results[0][custom_date.c.published] == today
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_fetch_one_returning_no_results(database_url):