Fix compressed spooling: compare decoded length against integer segment sizes

[arpitjain099]

Description

CompressedQueryDataDecoder.decode guards the compressed spooling path by checking the segment length against the sizes the coordinator reports in the segment metadata. Those sizes arrive as strings (the _SegmentMetadataTO TypedDict types segmentSize and uncompressedSize as str, and the metadata built across the spooling tests uses string values like "10"), but the code compared them directly to len(data), which is an int. An int is never equal to a decimal string, so not len(data) == metadata["segmentSize"] is always true and every compressed segment raised RuntimeError before it was ever decompressed. The error text gives it away, it reads "Expected to read 29 bytes but got 29" and still raises.

This makes json+zstd and json+lz4 spooling unusable from the client even though both encodings are advertised in the X-Trino-Encoding header. The lz4 path already coerces the size with int(...) when calling lz4.block.decompress, so the two size checks here just needed the same treatment.

The fix wraps both metadata sizes in int() before comparing, matching the existing lz4 code. I also added a unit test that decodes a real zstd-compressed segment with the string-typed metadata the protocol actually sends. It fails against the old code with the "expected N, got N" RuntimeError and passes with the fix. The existing spooling tests missed this because they all decode with encoding="json" or mock the decoder, so the compressed decode path had no coverage.

Non-technical explanation

Compressed results returned through the spooling protocol were always rejected with an error. This fixes the size check so compressed segments decode correctly.

Release notes

(x) Release notes are required, with the following suggested text:

* Fix decoding of compressed segments in the spooling protocol, which previously failed with a size-mismatch error for every `json+zstd` and `json+lz4` segment.

[hashhar]

Thanks for digging into this. I tried to reproduce the failure before merging and couldn't get it to happen against a real coordinator. Were you able to reproduce this?

I checked out this branch reverted just the int(...) casts to get back to the old behavior and ran it against a coordinator with spooling and json+zstd/json+lz4 enabled. I didn't see any failures. I also printed type(metadata["segmentSize"]) directly in the decoder and it came back as int, not str.

I think the existing int(...) cast in Lz4QueryDataDecoder.decompress is defensive/leftover code.

There might be a bug in _SegmentMetadataTO since it types segmentSize/uncompressedSize as str and the mocks in test_client_spooling.py follow that wrong type which is probably why this looked like a str-vs-int issue.

trino/client.py has ignore_errors = true in mypy config so nothing catches this mismatch. We should instead fix the annotation and the test fixtures directly instead of adding casts around it.

[arpitjain099]

No, I could not reproduce it against a real coordinator. I do not have one with spooling enabled, and I should have said that in the PR rather than leaving it implied.

My test built its metadata from _SegmentMetadataTO and the existing fixtures in test_client_spooling.py, which use string sizes. So it failed before and passed after, but all that demonstrated was that the fixtures disagree with the wire format, not that any coordinator sends strings. You checked the actual type and I did not. That is the part I got wrong.

Your diagnosis fits better: the annotation is wrong, the fixtures inherited the error, and ignore_errors = true on trino/client.py in mypy means nothing flagged it.

Happy to redo this as the fix you described: correct segmentSize and uncompressedSize to int in _SegmentMetadataTO, update the fixtures in test_client_spooling.py to match, and drop the leftover int(...) in Lz4QueryDataDecoder.decompress. Say the word and I will force-push that over this branch, or close this and open a fresh one, whichever you prefer.

[hashhar]

yes, please adjust the PR to fix the incorrect fixtures and see if mypy can be enabled to catch this in the future.

[arpitjain099]

Done in deba88b, force-pushed over the branch and rebased onto current master.

The change is now what you described:

• _SegmentMetadataTO.segmentSize and .uncompressedSize are int
• fixtures in test_client_spooling.py and test_client.py use integers
• both int(...) casts I had added are gone, and so is the leftover one in Lz4QueryDataDecoder.decompress

The test I had added asserted the wrong premise, so I replaced it rather than adjusting it. There are now two: one that a zstd segment round-trips with integer sizes, and one that a mismatched segmentSize still raises, so the size check cannot be quietly lost later.

On mypy, the short answer is that enabling it would not have caught this, and I would rather show you the numbers than assert that. Everything below is mypy 1.13.0, the version pinned in .pre-commit-config.yaml.

Dropping trino.client from the ignore_errors list gives 85 errors, 42 of them no-untyped-def:

  42 [no-untyped-def]
  11 [assignment]
   9 [no-untyped-call]
   8 [arg-type]
   3 [union-attr]
   3 [typeddict-item]

That is its own cleanup, not something to fold in here. But more to the point, none of those 85 is this bug. With the old str annotation restored, the size comparison passes a plain check, because lz4.block.decompress is untyped so nothing constrains what it is handed, and len(data) == metadata["segmentSize"] is not an error by default.

What does catch it is strict_equality. With the wrong annotation:

trino/client.py:1560: error: Non-overlapping equality check (left operand type: "int", right operand type: "str")  [comparison-overlap]
trino/client.py:1564: error: Non-overlapping equality check (left operand type: "int", right operand type: "str")  [comparison-overlap]

which is exactly the two lines this PR started from, and both disappear with the corrected annotation.

It is not free, though. Turning it on for this module also reports three things that are not bugs: 1099 and 1102 compare segment["type"] against SegmentType.INLINE/SPOOLED, and since SegmentType is a str, Enum those are True at runtime and mypy is simply wrong about them, plus a var-annotated on segments at 1096. So the smallest config that gets the check is a [mypy-trino.client] section with strict_equality = true, a disable_error_code list covering the noisy categories, and three suppressions.

I have that working locally and it comes out clean, but it is a long disable list that will drift, and it is a different change from this one. Happy to push it here, send it as a separate PR, or leave it. Your call.

Unit tests: 376 passed. The 5 failures in test_auth_gssapi.py and the trino/auth.py:503 mypy error are both present on a clean master checkout here, so neither is from this change.