From 140fa9586276d21fd0cd4b5793e7ef5e82d19836 Mon Sep 17 00:00:00 2001 From: IlyaKhalafi Date: Wed, 14 Oct 2026 11:41:28 +0000 Subject: [PATCH] Count UTF-8 bytes when limiting CSV exports Refs #3004 --- datasette/utils/__init__.py | 7 +++++- tests/test_csv.py | 43 +++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/datasette/utils/__init__.py b/datasette/utils/__init__.py index 0df0a871..747f7327 100644 --- a/datasette/utils/__init__.py +++ b/datasette/utils/__init__.py @@ -1061,7 +1061,12 @@ class LimitedWriter: self.bytes_count = 0 async def write(self, bytes): - self.bytes_count += len(bytes) + content_bytes = ( + bytes.encode("utf-8") + if self.limit_bytes and isinstance(bytes, str) + else bytes + ) + self.bytes_count += len(content_bytes) if self.limit_bytes and (self.bytes_count > self.limit_bytes): raise WriteLimitExceeded(f"CSV contains more than {self.limit_bytes} bytes") await self.writer.write(bytes) diff --git a/tests/test_csv.py b/tests/test_csv.py index adae7e24..4c7fb06b 100644 --- a/tests/test_csv.py +++ b/tests/test_csv.py @@ -1,9 +1,11 @@ import urllib.parse +from unittest.mock import AsyncMock, Mock, call import pytest from bs4 import BeautifulSoup as Soup from datasette.app import Datasette +from datasette.utils import LimitedWriter, WriteLimitExceeded EXPECTED_TABLE_CSV = """id,content 1,hello @@ -18,6 +20,47 @@ hello world """.replace("\n", "\r\n") + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "character,repetitions", [("é", 600000), ("😀", 300000), ("a", 600000)] +) +@pytest.mark.parametrize("limit_mb", [0, 1]) +async def test_csv_limit_counts_utf8_bytes(character, repetitions, limit_mb): + ds = Datasette(settings={"max_csv_mb": limit_mb}) + try: + db = ds.add_memory_database("unicode_csv") + await db.execute_write("CREATE TABLE t (id integer primary key, value text)") + await db.execute_write( + "INSERT INTO t VALUES (?, ?)", [1, character * repetitions] + ) + response = await ds.client.get("/unicode_csv/t.csv") + assert response.status_code == 200 + assert response.content.startswith(b"id,value\r\n") + if limit_mb and character != "a": + assert response.content[-100:].endswith( + b"CSV contains more than 1048576 bytes" + ) + assert len(response.content) < 1048576 + else: + assert response.text.endswith(character * 10 + "\r\n") + assert b"CSV contains more than" not in response.content + finally: + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("content", ["éé", "éé".encode()]) +async def test_csv_writer_utf8_byte_boundary(content): + writer = Mock(write=AsyncMock()) + limited = LimitedWriter(writer, 5 / (1024 * 1024)) + await limited.write("a") + await limited.write(content) + with pytest.raises(WriteLimitExceeded): + await limited.write("!") + assert writer.write.await_args_list == [call("a"), call(content)] + + EXPECTED_TABLE_WITH_LABELS_CSV = """ pk,created,planet_int,on_earth,state,_city_id,_city_id_label,_neighborhood,tags,complex_array,distinct_some_null,n 1,2019-01-14 08:00:00,1,1,CA,1,San Francisco,Mission,"[""tag1"", ""tag2""]","[{""foo"": ""bar""}]",one,n1