diff --git a/docs/cli.rst b/docs/cli.rst index c9389d88c..84a65d950 100644 --- a/docs/cli.rst +++ b/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 ` and :ref:`table ` output keep the duplicate column headers unchanged. + .. _cli_query_nl: Newline-delimited JSON diff --git a/docs/python-api.rst b/docs/python-api.rst index 6723efb2e..7ac39513d 100644 --- a/docs/python-api.rst +++ b/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() ` 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) diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index a30f79acb..d5a4b1e75 100644 --- a/sqlite_utils/cli.py +++ b/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. diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 3f23c7d65..d009f3531 100644 --- a/sqlite_utils/db.py +++ b/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)) diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 865ee796c..b39b11764 100644 --- a/sqlite_utils/utils.py +++ b/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): diff --git a/tests/test_cli.py b/tests/test_cli.py index d26e4dd5d..f19a00cec 100644 --- a/tests/test_cli.py +++ b/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"] diff --git a/tests/test_fts.py b/tests/test_fts.py index 3f7c5a9c6..64ec645fa 100644 --- a/tests/test_fts.py +++ b/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) diff --git a/tests/test_query.py b/tests/test_query.py index b9822e1c0..0b9f2aef9 100644 --- a/tests/test_query.py +++ b/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 diff --git a/tests/test_rows.py b/tests/test_rows.py index a8a4ca030..f050d5a24 100644 --- a/tests/test_rows.py +++ b/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"}] diff --git a/tests/test_utils.py b/tests/test_utils.py index f728bcd38..3de5e94a8 100644 --- a/tests/test_utils.py +++ b/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