remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-10-23 17:57:03 +02: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
from typing import Any, Iterable, Optional, Sequence, TextIO, Union
from typing import Any, Sequence, TextIO, Union
CsvDialect = Union[str, csv.Dialect]
@ -21,16 +21,11 @@ def rows_to_csv_file(
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)
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.writerows(_rows(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
writer.writerows(cursor)

6
tests/test_export.py
View file

@ -40,10 +40,12 @@ def test_rows_to_csv_file_tsv():
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:")
cursor = db.execute("create table creatures (id integer)")
output = io.StringIO(newline="")
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()