Support spooled protocol

Author: mdesmet

README.md

@@ -469,6 +469,19 @@ conn = connect(
 )
 ```
 
+## Spooled protocol
+
+Enable the spooled protocol by specifying a supported encoding in the `encoding` parameter (equires a
+Trino server that supports the spooled protocol).
+
+```python
+from trino.dbapi import connect
+
+conn = connect(
+    encoding="json+zstd"
+)
+```
+
 ## Transactions
 
 The client runs by default in *autocommit* mode. To enable transactions, set

setup.py

@@ -83,11 +83,13 @@
     ],
     python_requires=">=3.9",
     install_requires=[
+        "lz4",
         "python-dateutil",
         "pytz",
         # requests CVE https://github.com/advisories/GHSA-j8r2-6x86-q33q
         "requests>=2.31.0",
         "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
     )
 
 
@@ -1833,8 +1834,22 @@ def test_select_query_spooled_segments(trino_connection):
         stop => 5,
         step => 1)) n""")
     rows = cur.fetchall()
-    # TODO: improve test
-    assert len(rows) > 0
+    assert len(rows) == 300875
+    for row in rows:
+        assert isinstance(row[0], int), f"Expected integer for orderkey, got {type(row[0])}"
+        assert isinstance(row[1], int), f"Expected integer for partkey, got {type(row[1])}"
+        assert isinstance(row[2], int), f"Expected integer for suppkey, got {type(row[2])}"
+        assert isinstance(row[3], int), f"Expected int for linenumber, got {type(row[3])}"
+        assert isinstance(row[4], float), f"Expected float for quantity, got {type(row[4])}"
+        assert isinstance(row[5], float), f"Expected float for extendedprice, got {type(row[5])}"
+        assert isinstance(row[6], float), f"Expected float for discount, got {type(row[6])}"
+        assert isinstance(row[7], float), f"Expected string for tax, got {type(row[7])}"
+        assert isinstance(row[8], str), f"Expected string for returnflag, got {type(row[8])}"
+        assert isinstance(row[9], str), f"Expected string for linestatus, got {type(row[9])}"
+        assert isinstance(row[10], date), f"Expected date for shipdate, got {type(row[10])}"
+        assert isinstance(row[11], date), f"Expected date for commitdate, got {type(row[11])}"
+        assert isinstance(row[12], date), f"Expected date for receiptdate, got {type(row[12])}"
+        assert isinstance(row[13], str), f"Expected string for shipinstruct, got {type(row[13])}"
 
 
 def get_cursor(legacy_prepared_statements, run_trino):

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