Support spooled protocol

Author: mdesmet

setup.py

@@ -83,11 +83,14 @@
     ],
     python_requires=">=3.9",
     install_requires=[
+        "lz4",
         "python-dateutil",
         "pytz",
         # requests CVE https://github.com/advisories/GHSA-j8r2-6x86-q33q
         "requests>=2.31.0",
+        "typing_extensions",
         "tzlocal",
+        "zstandard",
     ],
     extras_require={
         "all": all_require,

tests/integration/test_dbapi_integration.py

@@ -29,12 +29,13 @@
 from trino.transaction import IsolationLevel
 
 
-@pytest.fixture
-def trino_connection(run_trino):
+@pytest.fixture(params=[None, "json+zstd", "json+lz4", "json"])
+def trino_connection(request, run_trino):
     host, port = run_trino
+    encoding = request.param
 
     yield trino.dbapi.Connection(
-        host=host, port=port, user="test", source="test", max_attempts=1
+        host=host, port=port, user="test", source="test", max_attempts=1, encoding=encoding
     )
 
 

tests/integration/test_types_integration.py

@@ -12,12 +12,18 @@
 from tests.integration.conftest import trino_version
 
 
-@pytest.fixture
-def trino_connection(run_trino):
+@pytest.fixture(params=[None, "json+zstd", "json+lz4", "json"])
+def trino_connection(request, run_trino):
     host, port = run_trino
+    encoding = request.param
 
     yield trino.dbapi.Connection(
-        host=host, port=port, user="test", source="test", max_attempts=1
+        host=host,
+        port=port,
+        user="test",
+        source="test",
+        max_attempts=1,
+        encoding=encoding
     )
 
 

tests/unit/test_client.py

@@ -99,6 +99,7 @@ def test_request_headers(mock_get_and_post):
     accept_encoding_value = "identity,deflate,gzip"
     client_info_header = constants.HEADER_CLIENT_INFO
     client_info_value = "some_client_info"
+    encoding = "json+zstd"
 
     with pytest.deprecated_call():
         req = TrinoRequest(
@@ -111,6 +112,7 @@ def test_request_headers(mock_get_and_post):
                 catalog=catalog,
                 schema=schema,
                 timezone=timezone,
+                encoding=encoding,
                 headers={
                     accept_encoding_header: accept_encoding_value,
                     client_info_header: client_info_value,
@@ -145,7 +147,8 @@ def assert_headers(headers):
             "catalog2=" + urllib.parse.quote("ROLE{catalog2_role}")
         )
         assert headers["User-Agent"] == f"{constants.CLIENT_NAME}/{__version__}"
-        assert len(headers.keys()) == 13
+        assert headers[constants.HEADER_ENCODING] == encoding
+        assert len(headers.keys()) == 14
 
     req.post("URL")
     _, post_kwargs = post.call_args

trino/client.py

@@ -34,22 +34,28 @@
 """
 from __future__ import annotations
 
+import base64
 import copy
 import functools
+import json
 import os
 import random
 import re
 import threading
 import urllib.parse
 import warnings
+from abc import abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
 from email.utils import parsedate_to_datetime
 from time import sleep
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union, cast
 from zoneinfo import ZoneInfo
 
+import lz4.block
 import requests
+import zstandard
+from typing_extensions import TypedDict
 from tzlocal import get_localzone_name  # type: ignore
 
 import trino.logging
@@ -107,6 +113,7 @@ class ClientSession:
     :param roles: roles for the current session. Some connectors do not
                  support role management. See connector documentation for more details.
     :param timezone: The timezone for query processing. Defaults to the system's local timezone.
+    :param encoding: The encoding for the spooled protocol. Defaults to None.
     """
 
     def __init__(
@@ -123,6 +130,7 @@ def __init__(
         client_tags: List[str] = None,
         roles: Union[Dict[str, str], str] = None,
         timezone: str = None,
+        encoding: str = None,
     ):
         self._user = user
         self._authorization_user = authorization_user
@@ -140,6 +148,7 @@ def __init__(
         self._timezone = timezone or get_localzone_name()
         if timezone:  # Check timezone validity
             ZoneInfo(timezone)
+        self._encoding = encoding
 
     @property
     def user(self):
@@ -235,6 +244,11 @@ def timezone(self):
         with self._object_lock:
             return self._timezone
 
+    @property
+    def encoding(self):
+        with self._object_lock:
+            return self._encoding
+
     def _format_roles(self, roles):
         if isinstance(roles, str):
             roles = {"system": roles}
@@ -299,7 +313,7 @@ class TrinoStatus:
     next_uri: Optional[str]
     update_type: Optional[str]
     update_count: Optional[int]
-    rows: List[Any]
+    rows: Union[List[Any], Dict[str, Any]]
     columns: List[Any]
 
     def __repr__(self):
@@ -462,6 +476,7 @@ def http_headers(self) -> Dict[str, str]:
         headers[constants.HEADER_USER] = self._client_session.user
         headers[constants.HEADER_AUTHORIZATION_USER] = self._client_session.authorization_user
         headers[constants.HEADER_TIMEZONE] = self._client_session.timezone
+        headers[constants.HEADER_ENCODING] = self._client_session.encoding
         headers[constants.HEADER_CLIENT_CAPABILITIES] = 'PARAMETRIC_DATETIME'
         headers["user-agent"] = f"{constants.CLIENT_NAME}/{__version__}"
         if len(self._client_session.roles.values()):
@@ -522,6 +537,8 @@ def max_attempts(self, value) -> None:
             self._get = self._http_session.get
             self._post = self._http_session.post
             self._delete = self._http_session.delete
+            self._prepare_request = self._http_session.prepare_request
+            self._send = self._http_session.send
             return
 
         with_retry = _retry_with(
@@ -537,6 +554,8 @@ def max_attempts(self, value) -> None:
         self._get = with_retry(self._http_session.get)
         self._post = with_retry(self._http_session.post)
         self._delete = with_retry(self._http_session.delete)
+        self._prepare_request = self._http_session.prepare_request
+        self._send = with_retry(self._http_session.send)
 
     def get_url(self, path) -> str:
         return "{protocol}://{host}:{port}{path}".format(
@@ -835,7 +854,7 @@ def _update_state(self, status):
         if status.columns:
             self._columns = status.columns
 
-    def fetch(self) -> List[List[Any]]:
+    def fetch(self) -> List[Union[List[Any]], Any]:
         """Continue fetching data for the current query_id"""
         try:
             response = self._request.get(self._request.next_uri)
@@ -849,7 +868,30 @@ def fetch(self) -> List[List[Any]]:
         if not self._row_mapper:
             return []
 
-        return self._row_mapper.map(status.rows)
+        rows = status.rows
+        if isinstance(rows, dict):
+            # spooled protocol
+            rows = cast(_SpooledProtocolResponseTO, rows)
+            segments = self._to_segments(rows)
+            return list(SegmentIterator(segments))
+        else:
+            return self._row_mapper.map(rows)
+
+    def _to_segments(self, rows: _SpooledProtocolResponseTO) -> SpooledData:
+        encoding = rows["encoding"]
+        segments = []
+        for segment in rows["segments"]:
+            segment_type = segment["type"]
+            if segment_type == "inline":
+                inline_segment = cast(_InlineSegmentTO, segment)
+                segments.append(InlineSegment(inline_segment, self._row_mapper))
+            elif segment_type == "spooled":
+                spooled_segment = cast(_SpooledSegmentTO, segment)
+                segments.append(SpooledSegment(rows, spooled_segment, self._row_mapper, self._request))
+            else:
+                raise ValueError(f"Unsupported segment type: {segment_type}")
+
+        return SpooledData(encoding, segments)
 
     def cancel(self) -> None:
         """Cancel the current query"""
@@ -925,3 +967,149 @@ def _parse_retry_after_header(retry_after):
         retry_date = parsedate_to_datetime(retry_after)
         now = datetime.utcnow()
         return (retry_date - now).total_seconds()
+
+
+# Trino Spooled protocol transfer objects
+class _SpooledProtocolResponseTO(TypedDict):
+    encoding: Literal["json", "json+std", "json+lz4"]
+    segments: List[_SegmentTO]
+
+
+class _SegmentMetadataTO(TypedDict):
+    uncompressedSize: str
+
+
+class _SegmentTO(_SegmentMetadataTO):
+    type: Literal["spooled", "inline"]
+    metadata: _SegmentMetadataTO
+
+
+class _SpooledSegmentTO(_SegmentTO):
+    uri: str
+    ackUri: str
+    headers: Dict[str, List[str]]
+
+
+class _InlineSegmentTO(_SegmentTO):
+    data: str
+
+
+class Segment:
+    def __init__(self, segment: _SegmentTO) -> None:
+        self._segment = segment
+        self.metadata = segment["metadata"]
+
+    @property
+    @abstractmethod
+    def rows(self):
+        pass
+
+
+class InlineSegment(Segment):
+    def __init__(self, segment: _InlineSegmentTO, row_mapper: RowMapper) -> None:
+        super().__init__(segment)
+        self._segment = cast(_InlineSegmentTO, segment)
+        self._row_mapper = row_mapper
+
+    @property
+    def rows(self) -> List[List[Any]]:
+        inline_segment = self._segment
+        data = inline_segment["data"]
+        decoded_string = base64.b64decode(data)
+        rows = self._row_mapper.map(json.loads(decoded_string))
+        return list(iter(rows))
+
+
+class SpooledSegment(Segment):
+    def __init__(
+        self,
+        response: _SpooledProtocolResponseTO,
+        segment: _SpooledSegmentTO,
+        row_mapper: RowMapper,
+        request: TrinoRequest
+    ) -> None:
+        super().__init__(segment)
+        self._segment = cast(_SpooledSegmentTO, segment)
+        self._row_mapper = row_mapper
+        self._request = request
+        self._encoding = response["encoding"]
+        self._headers = segment.get("headers")
+
+    @property
+    def rows(self) -> List[List[Any]]:
+        return self._row_mapper.map(json.loads(self._load_spooled_segment()))
+
+    def acknowledge(self):
+        http_response = self._send_spooling_request(self._segment["ackUri"], self._segment)
+        if not http_response.ok:
+            self._request.raise_response_error(http_response)
+
+    def _load_spooled_segment(self) -> str:
+        segment = self._segment
+        http_response = self._send_spooling_request(segment["uri"], segment)
+        if not http_response.ok:
+            self._request.raise_response_error(http_response)
+        content = http_response.content
+
+        encoding = self._encoding
+        if encoding == "json+zstd":
+            zstd_decompressor = zstandard.ZstdDecompressor()
+            return zstd_decompressor.decompress(content).decode('utf-8')
+        elif encoding == "json+lz4":
+            expected_size = segment["metadata"]["uncompressedSize"]
+            return lz4.block.decompress(content, uncompressed_size=int(expected_size)).decode('utf-8')
+        elif encoding == "json":
+            return content.decode('utf-8')
+        else:
+            raise ValueError(f"Unsupported encoding: {encoding}")
+
+    def _send_spooling_request(self, uri: str, segment: _SpooledSegmentTO) -> requests.Response:
+        req = requests.Request(
+            "GET",
+            uri,
+            headers=list(map(lambda item: [(item[0], value) for value in item[1]], segment.get("headers", {}).items()))
+        )
+        prepared = self._request._prepare_request(req)
+        return self._request._send(prepared)
+
+
+class SpooledData:
+    def __init__(self, encoding: str, segments: List[Segment]) -> None:
+        self._encoding = encoding
+        self._segments = iter(segments)
+
+    def __iter__(self) -> Iterator[Tuple["SpooledData", "Segment"]]:
+        return self
+
+    def __next__(self) -> Tuple["SpooledData", "Segment"]:
+        return self, next(self._segments)
+
+
+class SegmentIterator:
+    def __init__(self, spooled_data: SpooledData):
+        self._segments = iter(spooled_data._segments)
+        self._rows: Iterator[List[List[Any]]] = iter([])
+        self._finished = False
+        self._current_segment: Optional[Segment] = None
+
+    def __iter__(self) -> Iterator[List[Any]]:
+        return self
+
+    def __next__(self) -> List[Any]:
+        # If rows are exhausted, fetch the next segment
+        while True:
+            try:
+                return next(self._rows)
+            except StopIteration:
+                if self._current_segment and isinstance(self._current_segment, SpooledSegment):
+                    self._current_segment.acknowledge()
+                if self._finished:
+                    raise StopIteration
+                self._load_next_row_set()
+
+    def _load_next_row_set(self):
+        try:
+            self._current_segment = segment = next(self._segments)
+            self._rows = iter(segment.rows)
+        except StopIteration:
+            self._finished = True

trino/constants.py

@@ -37,6 +37,7 @@
 HEADER_CLIENT_TAGS = "X-Trino-Client-Tags"
 HEADER_EXTRA_CREDENTIAL = "X-Trino-Extra-Credential"
 HEADER_TIMEZONE = "X-Trino-Time-Zone"
+HEADER_ENCODING = "X-Trino-Query-Data-Encoding"
 
 HEADER_SESSION = "X-Trino-Session"
 HEADER_SET_SESSION = "X-Trino-Set-Session"

trino/dbapi.py

@@ -153,6 +153,7 @@ def __init__(
         legacy_prepared_statements=None,
         roles=None,
         timezone=None,
+        encoding=None,
     ):
         # Automatically assign http_schema, port based on hostname
         parsed_host = urlparse(host, allow_fragments=False)
@@ -176,6 +177,7 @@ def __init__(
             client_tags=client_tags,
             roles=roles,
             timezone=timezone,
+            encoding=encoding,
         )
         # mypy cannot follow module import
         if http_session is None: