Forward custom HTTP headers on coordinator spooling requests

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"

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 (