diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c5d33dbd..39271590 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1716,7 +1716,7 @@ def statement_uri(token): "encoding": "json", "segments": [{ "type": "inline", - "metadata": {"uncompressedSize": "5", "segmentSize": "5"}, + "metadata": {"uncompressedSize": 5, "segmentSize": 5}, "data": encoded_rows, }], }, diff --git a/tests/unit/test_client_spooling.py b/tests/unit/test_client_spooling.py index 21912d37..0f5071de 100644 --- a/tests/unit/test_client_spooling.py +++ b/tests/unit/test_client_spooling.py @@ -14,15 +14,19 @@ from unittest import mock import pytest +import zstandard from trino.client import _RequestHeartbeat from trino.client import ClientSession from trino.client import DecodableSegment from trino.client import InlineSegment +from trino.client import JsonQueryDataDecoder from trino.client import SegmentIterator from trino.client import SpooledSegment from trino.client import TrinoQuery from trino.client import TrinoRequest +from trino.client import ZStdQueryDataDecoder +from trino.mapper import NoOpRowMapper def _mock_trino_request(): @@ -176,7 +180,7 @@ def _spooled_fetch_response(): "segments": [ { "type": "inline", - "metadata": {"uncompressedSize": "10", "segmentSize": "10"}, + "metadata": {"uncompressedSize": 10, "segmentSize": 10}, "data": "", } ], @@ -210,7 +214,7 @@ def __init__(self, name): "type": "spooled", "uri": f"http://storage/{name}", "ackUri": f"http://storage/{name}/ack", - "metadata": {"uncompressedSize": "10", "segmentSize": "10"}, + "metadata": {"uncompressedSize": 10, "segmentSize": 10}, }, request=None, ) @@ -261,7 +265,7 @@ def _spooled_segment_with_headers(coordinator_host, custom_headers): "uri": "https://coordinator/v1/spooled/download/seg1", "ackUri": "https://coordinator/v1/spooled/ack/seg1", "headers": {"X-Trino-Spooling-Token": ["token-abc"]}, - "metadata": {"segmentSize": "1", "uncompressedSize": "1"}, + "metadata": {"segmentSize": 1, "uncompressedSize": 1}, } request = TrinoRequest( host="coordinator", @@ -327,3 +331,32 @@ def fake_get(uri, headers=None, **kwargs): segment._send_spooling_request(segment.uri) assert recorded["headers"]["X-Trino-Spooling-Token"] == "token-abc" + + +def test_zstd_decoder_decodes_compressed_segment(): + rows = [[1, "a"], [2, "b"]] + uncompressed = json.dumps(rows).encode("utf8") + compressed = zstandard.ZstdCompressor().compress(uncompressed) + + metadata = { + "segmentSize": len(compressed), + "uncompressedSize": len(uncompressed), + } + decoder = ZStdQueryDataDecoder(JsonQueryDataDecoder(NoOpRowMapper())) + + assert decoder.decode(compressed, metadata) == rows + + +def test_zstd_decoder_rejects_a_segment_whose_size_does_not_match(): + rows = [[1, "a"]] + uncompressed = json.dumps(rows).encode("utf8") + compressed = zstandard.ZstdCompressor().compress(uncompressed) + + metadata = { + "segmentSize": len(compressed) + 1, + "uncompressedSize": len(uncompressed), + } + decoder = ZStdQueryDataDecoder(JsonQueryDataDecoder(NoOpRowMapper())) + + with pytest.raises(RuntimeError, match="Expected to read"): + decoder.decode(compressed, metadata) diff --git a/trino/client.py b/trino/client.py index 9d4956bf..f2e5fd6e 100644 --- a/trino/client.py +++ b/trino/client.py @@ -1196,8 +1196,8 @@ class _SpooledProtocolResponseTO(TypedDict): class _SegmentMetadataTO(TypedDict): - uncompressedSize: str - segmentSize: str + uncompressedSize: int + segmentSize: int class _SegmentTO(_SegmentMetadataTO): @@ -1583,5 +1583,5 @@ def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: class Lz4QueryDataDecoder(CompressedQueryDataDecoder): def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: expected_uncompressed_size = metadata["uncompressedSize"] - decoded_bytes = lz4.block.decompress(data, uncompressed_size=int(expected_uncompressed_size)) + decoded_bytes = lz4.block.decompress(data, uncompressed_size=expected_uncompressed_size) return decoded_bytes