mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-09-27 20:34:12 +02:00
Fix dishonest progress accounting for multi-byte encoded CSV input (#439)
UpdateWrapper counted decoded characters (len(line)) against a byte-based file length, so utf-16-le input (2 bytes/char) only ever reported ~50% progress even after reading finished, looking like a stall. Track the underlying binary stream's position instead so progress is honest for any encoding, and add regression tests covering a clean utf-16-le insert plus a genuinely corrupt one that must still raise a visible error.
This commit is contained in:
parent
050f11f31f
commit
c989cdf5a2
3 changed files with 108 additions and 2 deletions
|
|
@ -202,15 +202,32 @@ class UpdateWrapper:
|
||||||
def __init__(self, wrapped: io.IOBase, update: Callable[[int], None]) -> None:
|
def __init__(self, wrapped: io.IOBase, update: Callable[[int], None]) -> None:
|
||||||
self._wrapped = wrapped
|
self._wrapped = wrapped
|
||||||
self._update = update
|
self._update = update
|
||||||
|
# If wrapped is a decoded text stream (e.g. io.TextIOWrapper set up
|
||||||
|
# with --encoding), len(line)/len(data) counts characters, not the
|
||||||
|
# bytes actually consumed from the underlying file - for a multi-byte
|
||||||
|
# encoding like UTF-16 that undercounts by ~2x, so progress never
|
||||||
|
# reaches 100% even though reading completed. Track the underlying
|
||||||
|
# binary buffer's position instead so progress stays honest for any
|
||||||
|
# encoding.
|
||||||
|
self._byte_source = getattr(wrapped, "buffer", None)
|
||||||
|
self._last_pos = self._byte_source.tell() if self._byte_source else 0
|
||||||
|
|
||||||
|
def _report(self, data_len: int) -> None:
|
||||||
|
if self._byte_source is not None:
|
||||||
|
pos = self._byte_source.tell()
|
||||||
|
self._update(pos - self._last_pos)
|
||||||
|
self._last_pos = pos
|
||||||
|
else:
|
||||||
|
self._update(data_len)
|
||||||
|
|
||||||
def __iter__(self) -> Iterator[bytes]:
|
def __iter__(self) -> Iterator[bytes]:
|
||||||
for line in self._wrapped:
|
for line in self._wrapped:
|
||||||
self._update(len(line))
|
self._report(len(line))
|
||||||
yield line
|
yield line
|
||||||
|
|
||||||
def read(self, size: int = -1) -> bytes:
|
def read(self, size: int = -1) -> bytes:
|
||||||
data = self._wrapped.read(size)
|
data = self._wrapped.read(size)
|
||||||
self._update(len(data))
|
self._report(len(data))
|
||||||
return data
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -2259,6 +2259,70 @@ def test_insert_encoding(tmpdir):
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_insert_encoding_utf16le(tmpdir):
|
||||||
|
# Regression test for issue #439: against utf-16-le CSV input, progress
|
||||||
|
# accounting used to undercount (it measured decoded characters against
|
||||||
|
# a byte-based file length, so it topped out around 50% - looking like a
|
||||||
|
# stall) even on a fully successful insert. And a genuine decode error
|
||||||
|
# further into the file needs to still surface as a clear exception
|
||||||
|
# rather than being silently swallowed.
|
||||||
|
db_path = str(tmpdir / "test.db")
|
||||||
|
csv_path = str(tmpdir / "test.csv")
|
||||||
|
rows = ["id,name"] + [f"{i},Name {i}" for i in range(500)]
|
||||||
|
with open(csv_path, "wb") as fp:
|
||||||
|
fp.write("\n".join(rows).encode("utf-16-le"))
|
||||||
|
|
||||||
|
result = CliRunner().invoke(
|
||||||
|
cli.cli,
|
||||||
|
[
|
||||||
|
"insert",
|
||||||
|
db_path,
|
||||||
|
"names",
|
||||||
|
csv_path,
|
||||||
|
"--csv",
|
||||||
|
"--encoding",
|
||||||
|
"utf-16-le",
|
||||||
|
"--no-detect-types",
|
||||||
|
],
|
||||||
|
catch_exceptions=False,
|
||||||
|
)
|
||||||
|
assert result.exit_code == 0
|
||||||
|
db = Database(db_path)
|
||||||
|
# All 500 data rows made it in - progress accounting didn't drop any
|
||||||
|
assert db["names"].count == 500
|
||||||
|
assert list(db["names"].rows)[:2] == [
|
||||||
|
{"id": "0", "name": "Name 0"},
|
||||||
|
{"id": "1", "name": "Name 1"},
|
||||||
|
]
|
||||||
|
|
||||||
|
# A file that is genuinely invalid utf-16-le should raise a visible,
|
||||||
|
# descriptive error - not hang or fail silently.
|
||||||
|
bad_csv_path = str(tmpdir / "bad.csv")
|
||||||
|
good_bytes = "\n".join(rows).encode("utf-16-le")
|
||||||
|
# Splice in an unpaired surrogate code unit part-way through the file,
|
||||||
|
# aligned to a 2-byte utf-16 code unit boundary
|
||||||
|
midpoint = (len(good_bytes) // 2) & ~1
|
||||||
|
bad_bytes = good_bytes[:midpoint] + b"\x00\xd8" + good_bytes[midpoint:]
|
||||||
|
with open(bad_csv_path, "wb") as fp:
|
||||||
|
fp.write(bad_bytes)
|
||||||
|
|
||||||
|
bad_result = CliRunner().invoke(
|
||||||
|
cli.cli,
|
||||||
|
[
|
||||||
|
"insert",
|
||||||
|
str(tmpdir / "bad.db"),
|
||||||
|
"names",
|
||||||
|
bad_csv_path,
|
||||||
|
"--csv",
|
||||||
|
"--encoding",
|
||||||
|
"utf-16-le",
|
||||||
|
],
|
||||||
|
catch_exceptions=False,
|
||||||
|
)
|
||||||
|
assert bad_result.exit_code == 1
|
||||||
|
assert "codec can't decode" in bad_result.output
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("fts", ["FTS4", "FTS5"])
|
@pytest.mark.parametrize("fts", ["FTS4", "FTS5"])
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"extra_arg,expected",
|
"extra_arg,expected",
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,6 @@
|
||||||
import csv
|
import csv
|
||||||
import io
|
import io
|
||||||
|
import os
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
|
@ -124,3 +125,27 @@ def test_rows_to_csv_no_headers():
|
||||||
def test_rows_to_csv_empty_rows():
|
def test_rows_to_csv_empty_rows():
|
||||||
assert utils.rows_to_csv([]) == ""
|
assert utils.rows_to_csv([]) == ""
|
||||||
assert utils.rows_to_csv([], headers=["id", "name"]) == "id,name\r\n"
|
assert utils.rows_to_csv([], headers=["id", "name"]) == "id,name\r\n"
|
||||||
|
|
||||||
|
|
||||||
|
def test_update_wrapper_reports_honest_bytes_for_utf16(tmpdir):
|
||||||
|
# Regression test for issue #439: UpdateWrapper used to count decoded
|
||||||
|
# *characters* (len(line)) against a progress length measured in bytes.
|
||||||
|
# For utf-16-le every ASCII character is encoded as 2 bytes, so the
|
||||||
|
# reported progress topped out at ~50% of the file even though reading
|
||||||
|
# had actually finished - looking like a stall. It should now track the
|
||||||
|
# underlying binary stream's position, so total reported progress
|
||||||
|
# matches the file size exactly regardless of encoding.
|
||||||
|
path = str(tmpdir / "utf16.csv")
|
||||||
|
text = "id,name\n" + "".join(f"{i},Alice{i}\n" for i in range(500))
|
||||||
|
with open(path, "wb") as fp:
|
||||||
|
fp.write(text.encode("utf-16-le"))
|
||||||
|
file_length = os.path.getsize(path)
|
||||||
|
|
||||||
|
reported = []
|
||||||
|
with open(path, "rb") as fp:
|
||||||
|
decoded = io.TextIOWrapper(fp, encoding="utf-16-le")
|
||||||
|
wrapper = utils.UpdateWrapper(decoded, reported.append)
|
||||||
|
for _line in wrapper:
|
||||||
|
pass
|
||||||
|
|
||||||
|
assert sum(reported) == file_length
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue