From ffb510c56ae922c294837bf88a2c7836a1078854 Mon Sep 17 00:00:00 2001 From: Rami Abdelrazzaq Date: Mon, 14 Sep 2026 12:11:00 -0500 Subject: [PATCH] Reject ambiguous mapping CSV rows --- sqlite_utils/export.py | 22 +++++++++++++++------- tests/test_export.py | 24 ++++++++++++++++++++---- 2 files changed, 35 insertions(+), 11 deletions(-) diff --git a/sqlite_utils/export.py b/sqlite_utils/export.py index fa8eba9..a5a8123 100644 --- a/sqlite_utils/export.py +++ b/sqlite_utils/export.py @@ -1,4 +1,5 @@ import csv +import itertools from collections.abc import Mapping from typing import Any, Sequence, TextIO, Type, Union @@ -26,15 +27,22 @@ def rows_to_csv_file( raise ValueError("Cursor does not have result columns") columns = [column[0] for column in description] + row_iterator = iter(cursor) + sentinel = object() + first_row = next(row_iterator, sentinel) + if isinstance(first_row, Mapping) and len(set(columns)) != len(columns): + raise ValueError("Mapping rows cannot represent duplicate column names") + writer = csv.writer(file, dialect=dialect, **writer_kwargs) if header: writer.writerow(columns) - def rows(): - for row in cursor: - if isinstance(row, Mapping): - yield [row[column] for column in columns] - else: - yield row + def normalize_row(row): + if isinstance(row, Mapping): + return [row[column] for column in columns] + return row - writer.writerows(rows()) + if first_row is not sentinel: + writer.writerows( + normalize_row(row) for row in itertools.chain((first_row,), row_iterator) + ) diff --git a/tests/test_export.py b/tests/test_export.py index 5cfbd51..44c08d4 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -61,10 +61,7 @@ def test_rows_to_csv_file_with_mapping_row_factory(): ) def dict_factory(cursor, row): - return { - column[0]: value - for column, value in zip(cursor.description, row) - } + return {column[0]: value for column, value in zip(cursor.description, row)} db.row_factory = dict_factory cursor = db.execute("select id, name from creatures order by id") @@ -74,6 +71,25 @@ def test_rows_to_csv_file_with_mapping_row_factory(): db.close() +def test_rows_to_csv_file_rejects_duplicate_mapping_columns(): + db = sqlite3.connect(":memory:") + db.execute("create table creatures (id integer, name text)") + db.execute("insert into creatures (id, name) values (1, 'Cleo')") + + def dict_factory(cursor, row): + return {column[0]: value for column, value in zip(cursor.description, row)} + + db.row_factory = dict_factory + cursor = db.execute("select id as value, name as value from creatures") + output = io.StringIO(newline="") + with pytest.raises( + ValueError, match="Mapping rows cannot represent duplicate column names" + ): + rows_to_csv_file(cursor, output, lineterminator="\n") + assert output.getvalue() == "" + db.close() + + @pytest.mark.parametrize("header", (True, False)) def test_rows_to_csv_file_requires_result_columns(header): db = sqlite3.connect(":memory:")