Fix duplicate query output column names for dict rows

Author: RamiNoodle733

sqlite_utils/db.py

@@ -82,6 +82,23 @@ def quote_identifier(identifier: str) -> str:
     return '"{}"'.format(identifier.replace('"', '""'))
 
 
+def _row_to_dict(keys: Sequence[str], row: Sequence[Any]) -> Dict[str, Any]:
+    """
+    Convert a row plus column names to a dictionary.
+
+    Duplicate column names are suffixed with ``_2``, ``_3``... so values are
+    preserved instead of overwritten.
+    """
+    counts: Dict[str, int] = {}
+    result: Dict[str, Any] = {}
+    for key, value in zip(keys, row):
+        count = counts.get(key, 0) + 1
+        counts[key] = count
+        final_key = key if count == 1 else "{}_{}".format(key, count)
+        result[final_key] = value
+    return result
+
+
 try:
     import pandas as pd  # type: ignore
 except ImportError:
@@ -548,7 +565,7 @@ def query(
         cursor = self.execute(sql, params or tuple())
         keys = [d[0] for d in cursor.description]
         for row in cursor:
-            yield dict(zip(keys, row))
+            yield _row_to_dict(keys, row)
 
     def execute(
         self, sql: str, parameters: Optional[Union[Sequence, Dict[str, Any]]] = None
@@ -1445,7 +1462,7 @@ def rows_where(
         cursor = self.db.execute(sql, where_args or [])
         columns = [c[0] for c in cursor.description]
         for row in cursor:
-            yield dict(zip(columns, row))
+            yield _row_to_dict(columns, row)
 
     def pks_and_rows_where(
         self,
@@ -2862,7 +2879,7 @@ def search(
         )
         columns = [c[0] for c in cursor.description]
         for row in cursor:
-            yield dict(zip(columns, row))
+            yield _row_to_dict(columns, row)
 
     def value_or_default(self, key: str, value: Any) -> Any:
         return self._defaults[key] if value is DEFAULT else value

tests/test_query.py

@@ -15,3 +15,18 @@ def test_execute_returning_dicts(fresh_db):
     assert fresh_db.execute_returning_dicts("select * from test") == [
         {"id": 1, "bar": 2}
     ]
+
+
+def test_query_duplicate_output_columns_are_suffixed(fresh_db):
+    fresh_db.execute("create table one (id integer, value text)")
+    fresh_db.execute("create table two (id integer, value text)")
+    fresh_db["one"].insert({"id": 1, "value": "left"})
+    fresh_db["two"].insert({"id": 2, "value": "right"})
+
+    rows = list(
+        fresh_db.query(
+            "select one.id, two.id, one.value, two.value from one, two where one.id = 1"
+        )
+    )
+
+    assert rows == [{"id": 1, "id_2": 2, "value": "left", "value_2": "right"}]