remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-11-04 01:08:44 +01:00
Code Issues Projects Releases Packages Wiki Activity

fix: validate CSV export cursors before writing

Browse source
This commit is contained in:
Rami Abdelrazzaq 2026-09-13 13:15:08 -05:00
commit 6f064ce6ee
2 changed files with 10 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

17
sqlite_utils/export.py
View file

@ -1,5 +1,5 @@
import csv import csv
from typing import Any, Iterable, Optional, Sequence, TextIO, Union from typing import Any, Sequence, TextIO, Union
CsvDialect = Union[str, csv.Dialect] CsvDialect = Union[str, csv.Dialect]
@ -21,16 +21,11 @@ def rows_to_csv_file(
Additional keyword arguments are forwarded to ``csv.writer``. Additional keyword arguments are forwarded to ``csv.writer``.
""" """
description: Union[Sequence[Sequence[Any]], None] = cursor.description
if description is None:
raise ValueError("Cursor does not have result columns")
writer = csv.writer(file, dialect=dialect, **writer_kwargs) writer = csv.writer(file, dialect=dialect, **writer_kwargs)
if header: if header:
description: Optional[Sequence[Sequence[Any]]] = cursor.description
if description is None:
raise ValueError("Cursor does not have result columns")
writer.writerow([column[0] for column in description]) writer.writerow([column[0] for column in description])
writer.writerows(_rows(cursor)) writer.writerows(cursor)
def _rows(cursor: Iterable[Sequence[Any]]) -> Iterable[Sequence[Any]]:
"""Keep row iteration lazy so large query results are streamed."""
for row in cursor:
yield row

6
tests/test_export.py
View file

@ -40,10 +40,12 @@ def test_rows_to_csv_file_tsv():
db.close() db.close()
def test_rows_to_csv_file_requires_result_columns_for_header(): @pytest.mark.parametrize("header", (True, False))
def test_rows_to_csv_file_requires_result_columns(header):
db = sqlite3.connect(":memory:") db = sqlite3.connect(":memory:")
cursor = db.execute("create table creatures (id integer)") cursor = db.execute("create table creatures (id integer)")
output = io.StringIO(newline="") output = io.StringIO(newline="")
with pytest.raises(ValueError, match="Cursor does not have result columns"): with pytest.raises(ValueError, match="Cursor does not have result columns"):
rows_to_csv_file(cursor, output) rows_to_csv_file(cursor, output, header=header)
assert output.getvalue() == ""
db.close() db.close()