diff --git a/tests/unit/test_client_spooling.py b/tests/unit/test_client_spooling.py index a56d62df..21912d37 100644 --- a/tests/unit/test_client_spooling.py +++ b/tests/unit/test_client_spooling.py @@ -253,3 +253,77 @@ def test_segment_iterator_retries_failed_segment_without_skipping_it(failing_seg with pytest.raises(StopIteration): next(iterator) assert [seg.acknowledge_count for seg in segs] == [1, 1, 1] + + +def _spooled_segment_with_headers(coordinator_host, custom_headers): + segment_to = { + "type": "spooled", + "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"}, + } + request = TrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_scheme="https", + ) + return SpooledSegment( + segment_to, + request, + coordinator_host=coordinator_host, + custom_headers=custom_headers, + ) + + +def test_send_spooling_request_forwards_custom_headers_to_coordinator(): + custom_headers = {"X-Auth-Gateway-Token": "user-token"} + segment = _spooled_segment_with_headers(coordinator_host="coordinator", custom_headers=custom_headers) + + recorded = {} + + def fake_get(uri, headers=None, **kwargs): + recorded["headers"] = headers + return mock.Mock(ok=True) + + segment._request._get = fake_get + segment._send_spooling_request(segment.uri) + + assert recorded["headers"]["X-Auth-Gateway-Token"] == "user-token" + assert recorded["headers"]["X-Trino-Spooling-Token"] == "token-abc" + + +def test_send_spooling_request_does_not_forward_custom_headers_to_external_storage(): + custom_headers = {"X-Auth-Gateway-Token": "user-token"} + segment = _spooled_segment_with_headers(coordinator_host="coordinator", custom_headers=custom_headers) + + recorded = {} + + def fake_get(uri, headers=None, **kwargs): + recorded["headers"] = headers + return mock.Mock(ok=True) + + segment._request._get = fake_get + external_uri = "https://s3.amazonaws.com/bucket/seg1?X-Amz-Signature=abc" + segment._send_spooling_request(external_uri) + + assert "X-Auth-Gateway-Token" not in recorded["headers"] + assert recorded["headers"]["X-Trino-Spooling-Token"] == "token-abc" + + +def test_send_spooling_request_segment_header_takes_precedence_over_custom_header(): + # Custom header uses the same name as the segment protocol header; the segment header must win. + custom_headers = {"X-Trino-Spooling-Token": "should-not-be-used"} + segment = _spooled_segment_with_headers(coordinator_host="coordinator", custom_headers=custom_headers) + + recorded = {} + + def fake_get(uri, headers=None, **kwargs): + recorded["headers"] = headers + return mock.Mock(ok=True) + + segment._request._get = fake_get + segment._send_spooling_request(segment.uri) + + assert recorded["headers"]["X-Trino-Spooling-Token"] == "token-abc" diff --git a/trino/client.py b/trino/client.py index abf1c2f2..95171803 100644 --- a/trino/client.py +++ b/trino/client.py @@ -1063,7 +1063,12 @@ def _to_segments(self, rows: _SpooledProtocolResponseTO) -> List[DecodableSegmen segments.append(InlineSegment(inline_segment)) elif segment_type == SegmentType.SPOOLED: spooled_segment = cast(_SpooledSegmentTO, segment) - segments.append(SpooledSegment(spooled_segment, self._request.unauthenticated())) + segments.append(SpooledSegment( + spooled_segment, + self._request.unauthenticated(), + coordinator_host=self._request._host, + custom_headers=dict(self._request._client_session.headers), + )) else: raise ValueError(f"Unsupported segment type: {segment_type}") @@ -1238,10 +1243,14 @@ def __init__( self, segment: _SpooledSegmentTO, request: TrinoRequest, + coordinator_host: Optional[str] = None, + custom_headers: Optional[Dict[str, str]] = None, ) -> None: super().__init__(segment) self._segment = cast(_SpooledSegmentTO, segment) self._request = request + self._coordinator_host = coordinator_host + self._custom_headers = custom_headers or {} @property def data(self) -> bytes: @@ -1274,12 +1283,18 @@ def acknowledge_request(): executor.submit(acknowledge_request) def _send_spooling_request(self, uri: str, **kwargs) -> requests.Response: - headers_with_single_value = {} + headers: Dict[str, str] = {} + # Forward user-supplied custom headers (e.g. auth gateway headers) only when the + # request targets the Trino coordinator, never to external storage (e.g. S3 presigned + # URLs) where such headers can break the request. The per-segment protocol headers + # returned by the coordinator always take precedence. + if self._coordinator_host is not None and urllib.parse.urlsplit(uri).hostname == self._coordinator_host: + headers.update(self._custom_headers) for key, values in self.headers.items(): if len(values) > 1: raise ValueError(f"Header '{key}' contains multiple values: {values}") - headers_with_single_value[key] = values[0] - return self._request._get(uri, headers=headers_with_single_value, **kwargs) + headers[key] = values[0] + return self._request._get(uri, headers=headers, **kwargs) def __repr__(self): return (