Repository: trinodb/trino-python-client

Issue:
Client should expose complex types like ROW in a usable format
Today the STRUCT/ROW fields are returned as plain strings by the Python client.

We should implement the equivalent of https://github.com/dropbox/PyHive/pull/328 here so that clients can properly parse and use the data returned for complex columns.

Discussion:

Patch:
diff --git a/trino/__init__.py b/trino/__init__.py
index 2c1ae6a..dd5f352 100644
--- a/trino/__init__.py
+++ b/trino/__init__.py
@@ -19,4 +19,4 @@ from . import logging
 
 __all__ = ['auth', 'dbapi', 'client', 'constants', 'exceptions', 'logging']
 
-__version__ = "0.312.0"
+__version__ = "0.313.0"
diff --git a/trino/client.py b/trino/client.py
index ed0942c..02616ca 100644
--- a/trino/client.py
+++ b/trino/client.py
@@ -77,6 +77,7 @@ class ClientSession(object):
         headers=None,
         transaction_id=None,
         extra_credential=None,
+        client_tags=None
     ):
         self.catalog = catalog
         self.schema = schema
@@ -88,6 +89,7 @@ class ClientSession(object):
         self._headers = headers or {}
         self.transaction_id = transaction_id
         self.extra_credential = extra_credential
+        self.client_tags = client_tags
 
     @property
     def properties(self):
@@ -226,7 +228,8 @@ class TrinoRequest(object):
         max_attempts: int = MAX_ATTEMPTS,
         request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT,
         handle_retry=exceptions.RetryWithExponentialBackoff(),
-        verify: bool = True
+        verify: bool = True,
+        client_tags: Optional[List[str]] = None
     ) -> None:
         self._client_session = ClientSession(
             catalog,
@@ -237,6 +240,7 @@ class TrinoRequest(object):
             http_headers,
             transaction_id,
             extra_credential,
+            client_tags
         )
 
         self._host = host
@@ -286,6 +290,8 @@ class TrinoRequest(object):
         headers[constants.HEADER_SCHEMA] = self._client_session.schema
         headers[constants.HEADER_SOURCE] = self._client_session.source
         headers[constants.HEADER_USER] = self._client_session.user
+        if self._client_session.client_tags is not None and len(self._client_session.client_tags) > 0:
+            headers[constants.HEADER_CLIENT_TAGS] = ",".join(self._client_session.client_tags)
 
         headers[constants.HEADER_SESSION] = ",".join(
             # ``name`` must not contain ``=``
@@ -524,25 +530,30 @@ class TrinoResult(object):
 
         try:
             if isinstance(value, list):
-                raw_type = {
-                    "typeSignature": data_type["typeSignature"]["arguments"][0]["value"]
-                }
-                return [cls._map_to_python_type((array_item, raw_type)) for array_item in value]
-            if isinstance(value, dict):
-                if raw_type == "map":
-                    raw_key_type = {
+                if raw_type == "array":
+                    raw_type = {
                         "typeSignature": data_type["typeSignature"]["arguments"][0]["value"]
                     }
-                    raw_value_type = {
-                        "typeSignature": data_type["typeSignature"]["arguments"][1]["value"]
-                    }
-                    return {
-                        cls._map_to_python_type((key, raw_key_type)):
-                            cls._map_to_python_type((value[key], raw_value_type))
-                        for key in value
-                    }
-                # TODO: support row type
+                    return [cls._map_to_python_type((array_item, raw_type)) for array_item in value]
+                if raw_type == "row":
+                    raw_types = map(lambda arg: arg["value"], data_type["typeSignature"]["arguments"])
+                    return tuple(
+                        cls._map_to_python_type((array_item, raw_type))
+                        for (array_item, raw_type) in zip(value, raw_types)
+                    )
                 return value
+            if isinstance(value, dict):
+                raw_key_type = {
+                    "typeSignature": data_type["typeSignature"]["arguments"][0]["value"]
+                }
+                raw_value_type = {
+                    "typeSignature": data_type["typeSignature"]["arguments"][1]["value"]
+                }
+                return {
+                    cls._map_to_python_type((key, raw_key_type)):
+                        cls._map_to_python_type((value[key], raw_value_type))
+                    for key in value
+                }
             elif "decimal" in raw_type:
                 return Decimal(value)
             elif raw_type == "double":
diff --git a/trino/constants.py b/trino/constants.py
index 5a31703..5c0e4ab 100644
--- a/trino/constants.py
+++ b/trino/constants.py
@@ -32,6 +32,7 @@ HEADER_SCHEMA = "X-Trino-Schema"
 HEADER_SOURCE = "X-Trino-Source"
 HEADER_USER = "X-Trino-User"
 HEADER_CLIENT_INFO = "X-Trino-Client-Info"
+HEADER_CLIENT_TAGS = "X-Trino-Client-Tags"
 HEADER_EXTRA_CREDENTIAL = "X-Trino-Extra-Credential"
 
 HEADER_SESSION = "X-Trino-Session"
diff --git a/trino/dbapi.py b/trino/dbapi.py
index bb66fa5..bb65539 100644
--- a/trino/dbapi.py
+++ b/trino/dbapi.py
@@ -108,7 +108,8 @@ class Connection(object):
         request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,
         isolation_level=IsolationLevel.AUTOCOMMIT,
         verify=True,
-        http_session=None
+        http_session=None,
+        client_tags=None
     ):
         self.host = host
         self.port = port
@@ -130,6 +131,7 @@ class Connection(object):
         self.redirect_handler = redirect_handler
         self.max_attempts = max_attempts
         self.request_timeout = request_timeout
+        self.client_tags = client_tags
 
         self._isolation_level = isolation_level
         self._request = None
@@ -194,6 +196,7 @@ class Connection(object):
             self.redirect_handler,
             self.max_attempts,
             self.request_timeout,
+            client_tags=self.client_tags
         )
 
     def cursor(self, experimental_python_types=False):
@@ -391,8 +394,10 @@ class Cursor(object):
         if isinstance(param, list):
             return "ARRAY[%s]" % ','.join(map(self._format_prepared_param, param))
 
+        if isinstance(param, tuple):
+            return "ROW(%s)" % ','.join(map(self._format_prepared_param, param))
+
         if isinstance(param, dict):
-            # TODO: support mixed types in dicts and convert to ROW
             keys = list(param.keys())
             values = [param[key] for key in keys]
             return "MAP({}, {})".format(


Test patch:
diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py
index 0d3a758..6febe84 100644
--- a/tests/integration/test_dbapi_integration.py
+++ b/tests/integration/test_dbapi_integration.py
@@ -15,6 +15,7 @@ from decimal import Decimal
 
 import pytest
 import pytz
+import requests
 
 import trino
 from tests.integration.conftest import trino_version
@@ -615,6 +616,32 @@ def test_boolean_query_param(trino_connection):
     assert rows[0][0] is False
 
 
+def test_row(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+    params = (1, Decimal("2.0"), datetime(2020, 1, 1, 0, 0, 0))
+    cur.execute("SELECT ?", (params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
+def test_nested_row(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+    params = ((1, "test", Decimal("3.1")), Decimal("2.0"), datetime(2020, 1, 1, 0, 0, 0))
+    cur.execute("SELECT ?", (params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
+def test_named_row(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+    cur.execute("SELECT CAST(ROW(1, 2e0) AS ROW(x BIGINT, y DOUBLE))")
+    rows = cur.fetchall()
+
+    assert rows[0][0] == (1, 2.0)
+
+
 def test_float_query_param(trino_connection):
     cur = trino_connection.cursor()
     cur.execute("SELECT ?", params=(1.1,))
@@ -867,3 +894,46 @@ def test_info_uri(trino_connection):
     cur.fetchall()
     assert cur.info_uri is not None
     assert cur._query.query_id in cur.info_uri
+
+
+def test_client_tags_single_tag(run_trino):
+    client_tags = ["foo"]
+    query_client_tags = retrieve_client_tags_from_query(run_trino, client_tags)
+    assert query_client_tags == client_tags
+
+
+def test_client_tags_multiple_tags(run_trino):
+    client_tags = ["foo", "bar"]
+    query_client_tags = retrieve_client_tags_from_query(run_trino, client_tags)
+    assert query_client_tags == client_tags
+
+
+def test_client_tags_special_characters(run_trino):
+    client_tags = ["foo %20", "bar=test"]
+    query_client_tags = retrieve_client_tags_from_query(run_trino, client_tags)
+    assert query_client_tags == client_tags
+
+
+def retrieve_client_tags_from_query(run_trino, client_tags):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host,
+        port=port,
+        user="test",
+        client_tags=client_tags,
+    )
+
+    cur = trino_connection.cursor()
+    cur.execute('SELECT 1')
+    cur.fetchall()
+
+    api_url = "http://" + trino_connection.host + ":" + str(trino_connection.port)
+    query_info = requests.post(api_url + "/ui/login", data={
+        "username": "admin",
+        "password": "",
+        "redirectPath": api_url + '/ui/api/query/' + cur._query.query_id
+    }).json()
+
+    query_client_tags = query_info['session']['clientTags']
+    return query_client_tags
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index f03b5f7..6a70812 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -172,6 +172,49 @@ def test_request_invalid_http_headers():
     assert str(value_error.value).startswith("cannot override reserved HTTP header")
 
 
+def test_request_client_tags_headers(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        user="test_user",
+        client_tags=["tag1", "tag2"]
+    )
+
+    def assert_headers(headers):
+        assert headers[constants.HEADER_CLIENT_TAGS] == "tag1,tag2"
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers(post_kwargs["headers"])
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers(get_kwargs["headers"])
+
+
+def test_request_client_tags_headers_no_client_tags(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        user="test_user"
+    )
+
+    def assert_headers(headers):
+        assert constants.HEADER_CLIENT_TAGS not in headers
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers(post_kwargs["headers"])
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers(get_kwargs["headers"])
+
+
 def test_enabling_https_automatically_when_using_port_443(mock_get_and_post):
     _, post = mock_get_and_post
 
diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py
index 690aa50..f847007 100644
--- a/tests/unit/test_dbapi.py
+++ b/tests/unit/test_dbapi.py
@@ -222,3 +222,15 @@ def test_token_retrieved_once_when_multithreaded(sample_post_response_data):
         thread.join()
 
     assert len(_get_token_requests(challenge_id)) == 1
+
+
+@patch("trino.dbapi.trino.client")
+def test_tags_are_set_when_specified(mock_client):
+    # WHEN
+    client_tags = ["TAG1", "TAG2"]
+    with connect("sample_trino_cluster:443", client_tags=client_tags) as conn:
+        conn.cursor().execute("SOME FAKE QUERY")
+
+    # THEN
+    _, passed_client_tags = mock_client.TrinoRequest.call_args
+    assert passed_client_tags["client_tags"] == client_tags