Preserve duplicate column names in query results, closes #624

docs/cli.rst

@@ -45,6 +45,8 @@ The default format returned for queries is JSON:
     [{"id": 1, "age": 4, "name": "Cleo"},
      {"id": 2, "age": 2, "name": "Pancakes"}]
 
+If the query returns more than one column with the same name, later occurrences are renamed with a numeric suffix - ``select 1 as id, 2 as id`` returns ``[{"id": 1, "id_2": 2}]``. This only applies to JSON output: :ref:`CSV and TSV <cli_query_csv>` and :ref:`table <cli_query_table>` output keep the duplicate column headers unchanged.
+
 .. _cli_query_nl:
 
 Newline-delimited JSON

docs/python-api.rst

@@ -233,6 +233,17 @@ The SQL query is executed as soon as ``db.query()`` is called. The resulting row
 
 ``db.query()`` can only be used with SQL that returns rows. Passing a statement that returns no rows - an ``INSERT`` or ``UPDATE`` without a ``RETURNING`` clause, for example - will raise a ``ValueError``. The rejected statement is rolled back, so it has no effect on the database. Use :ref:`db.execute() <python_api_execute>` for those statements instead.
 
+If a query returns more than one column with the same name - a join between two tables that share column names, for example - later occurrences are renamed with a numeric suffix, so every value is included in the dictionary:
+
+.. code-block:: python
+
+    row = next(db.query("select 1 as id, 2 as id, 3 as id"))
+    print(row)
+    # Outputs:
+    # {'id': 1, 'id_2': 2, 'id_3': 3}
+
+A suffix that would collide with another column in the query is skipped - ``select 1 as id, 2 as id, 3 as id_2`` returns ``{'id': 1, 'id_3': 2, 'id_2': 3}``. The same renaming is applied by ``table.rows_where()`` and ``table.search()``.
+
 .. _python_api_execute:
 
 db.execute(sql, params)

sqlite_utils/cli.py

@@ -34,6 +34,7 @@
     OperationalError,
     _compile_code,
     chunks,
+    dedupe_keys,
     file_progress,
     find_spatialite,
     flatten as _flatten,
@@ -3493,6 +3494,10 @@ def __init__(self, exception, path):
 
 
 def output_rows(iterator, headers, nl, arrays, json_cols):
+    # Duplicate column names would collide as dictionary keys, so rename
+    # later occurrences id, id -> id, id_2 - CSV and table output keep
+    # the original duplicate headers since they never build dictionaries
+    headers = dedupe_keys(headers)
     # We have to iterate two-at-a-time so we can know if we
     # should output a trailing comma or if we have reached
     # the last row.

sqlite_utils/db.py

@@ -1,5 +1,6 @@
 from .utils import (
     chunks,
+    dedupe_keys,
     hash_record,
     sqlite3,
     OperationalError,
@@ -786,7 +787,7 @@ def query(
                     cursor = self.conn.execute(sql, *args)
             if cursor.description is None:
                 raise ValueError(message)
-            keys = [d[0] for d in cursor.description]
+            keys = dedupe_keys(d[0] for d in cursor.description)
             return (dict(zip(keys, row)) for row in cursor)
         # Execute inside a savepoint, so a statement that turns out not to
         # return rows can be rolled back before the ValueError is raised
@@ -796,7 +797,7 @@ def query(
             cursor = self.conn.execute(sql, *args)
             if cursor.description is None:
                 raise ValueError(message)
-            keys = [d[0] for d in cursor.description]
+            keys = dedupe_keys(d[0] for d in cursor.description)
             try:
                 self.conn.execute('RELEASE "sqlite_utils_query"')
                 released = True
@@ -1865,7 +1866,7 @@ def rows_where(
         if offset is not None:
             sql += " offset {}".format(offset)
         cursor = self.db.execute(sql, where_args or [])
-        columns = [c[0] for c in cursor.description]
+        columns = dedupe_keys(c[0] for c in cursor.description)
         for row in cursor:
             yield dict(zip(columns, row))
 
@@ -3398,7 +3399,7 @@ def search(
             ),
             args,
         )
-        columns = [c[0] for c in cursor.description]
+        columns = dedupe_keys(c[0] for c in cursor.description)
         for row in cursor:
             yield dict(zip(columns, row))
 

sqlite_utils/utils.py

@@ -613,6 +613,37 @@ def hash_record(record: Dict[str, Any], keys: Optional[Iterable[str]] = None) ->
     ).hexdigest()
 
 
+def dedupe_keys(keys: Iterable[str]) -> List[str]:
+    """
+    Rename duplicates in a list of column names so every name is unique,
+    by appending ``_2``, ``_3``... to later occurrences - skipping any
+    suffix that would collide with another column in the list.
+
+    Used when converting SQL query rows to dictionaries, where duplicate
+    column names would otherwise silently overwrite each other.
+
+    :param keys: List of column names, possibly containing duplicates
+    """
+    keys = list(keys)
+    taken = set(keys)
+    if len(taken) == len(keys):
+        # No duplicates - the common case
+        return keys
+    seen: set = set()
+    result = []
+    for key in keys:
+        if key in seen:
+            new_key = key
+            suffix = 2
+            while new_key in seen or new_key in taken:
+                new_key = "{}_{}".format(key, suffix)
+                suffix += 1
+            key = new_key
+        seen.add(key)
+        result.append(key)
+    return result
+
+
 def _flatten(d: Dict[str, Any]) -> Generator[Tuple[str, Any], None, None]:
     for key, value in d.items():
         if isinstance(value, dict):

tests/test_cli.py

@@ -746,6 +746,26 @@ def test_query_json_empty(db_path):
     assert result.output.strip() == "[]"
 
 
+def test_query_json_duplicate_columns_are_deduped(db_path):
+    # https://github.com/simonw/sqlite-utils/issues/624
+    result = CliRunner().invoke(
+        cli.cli,
+        [db_path, "select 1 as id, 2 as id, 'x' as value, 'y' as value"],
+    )
+    assert result.output.strip() == (
+        '[{"id": 1, "id_2": 2, "value": "x", "value_2": "y"}]'
+    )
+
+
+def test_query_csv_duplicate_columns_are_preserved(db_path):
+    # CSV output should keep the duplicate headers, not rename them
+    result = CliRunner().invoke(
+        cli.cli,
+        [db_path, "select 1 as id, 2 as id", "--csv"],
+    )
+    assert result.output.replace("\r", "").strip() == "id,id\n1,2"
+
+
 def test_query_invalid_function(db_path):
     result = CliRunner().invoke(
         cli.cli, [db_path, "select bad()", "--functions", "def invalid_python"]

tests/test_fts.py

@@ -83,6 +83,20 @@ def test_enable_fts_escape_table_names(fresh_db):
     assert [] == list(table.search("bar"))
 
 
+def test_search_duplicate_columns_are_deduped(fresh_db):
+    # https://github.com/simonw/sqlite-utils/issues/624
+    table = fresh_db["t"]
+    table.insert_all(search_records)
+    table.enable_fts(["text", "country"], fts_version="FTS4")
+    rows = list(table.search("tanuki", columns=["text", "text"]))
+    assert rows == [
+        {
+            "text": "tanuki are running tricksters",
+            "text_2": "tanuki are running tricksters",
+        }
+    ]
+
+
 def test_search_limit_offset(fresh_db):
     table = fresh_db["t"]
     table.insert_all(search_records)

tests/test_query.py

@@ -218,6 +218,22 @@ def test_query_insert_returning_respects_explicit_transaction(fresh_db):
     assert [row["name"] for row in fresh_db["dogs"].rows] == ["Cleo"]
 
 
+def test_query_duplicate_column_names_are_deduped(fresh_db):
+    # https://github.com/simonw/sqlite-utils/issues/624
+    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")
+    )
+    assert rows == [{"id": 1, "id_2": 2, "value": "left", "value_2": "right"}]
+
+
+def test_query_deduped_column_avoids_existing_names(fresh_db):
+    # The renamed duplicate must not overwrite a real column called id_2
+    rows = list(fresh_db.query("select 1 as id, 2 as id, 3 as id_2"))
+    assert rows == [{"id": 1, "id_3": 2, "id_2": 3}]
+
+
 def test_execute_returning_dicts(fresh_db):
     # Like db.query() but returns a list, included for backwards compatibility
     # see https://github.com/simonw/sqlite-utils/issues/290

tests/test_rows.py

@@ -104,3 +104,10 @@ def test_pks_and_rows_where_compound_pk(fresh_db):
         (("number", 1), {"type": "number", "number": 1, "plusone": 2}),
         (("number", 2), {"type": "number", "number": 2, "plusone": 3}),
     ]
+
+
+def test_rows_where_duplicate_select_columns_are_deduped(fresh_db):
+    # https://github.com/simonw/sqlite-utils/issues/624
+    fresh_db["t"].insert({"id": 1, "name": "Cleo"})
+    rows = list(fresh_db["t"].rows_where(select="id, id, name"))
+    assert rows == [{"id": 1, "id_2": 1, "name": "Cleo"}]

tests/test_utils.py

@@ -83,3 +83,20 @@ def test_maximize_csv_field_size_limit():
 )
 def test_flatten(input, expected):
     assert utils.flatten(input) == expected
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    (
+        ([], []),
+        (["id", "name"], ["id", "name"]),
+        (["id", "id"], ["id", "id_2"]),
+        (["id", "id", "id"], ["id", "id_2", "id_3"]),
+        # A renamed duplicate must not clobber a real column called id_2
+        (["id", "id", "id_2"], ["id", "id_3", "id_2"]),
+        (["id_2", "id", "id"], ["id_2", "id", "id_3"]),
+        (["id", "id", "id_2", "id_2"], ["id", "id_3", "id_2", "id_2_2"]),
+    ),
+)
+def test_dedupe_keys(input, expected):
+    assert utils.dedupe_keys(input) == expected