Add support for client tags in Python API

Co-authored-by: abrisan <[email protected]>

Author: 2 people

tests/integration/test_dbapi_integration.py

@@ -15,6 +15,7 @@
 
 import pytest
 import pytz
+import requests
 
 import trino
 from tests.integration.conftest import trino_version
@@ -867,3 +868,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

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
 

tests/unit/test_dbapi.py

@@ -222,3 +222,15 @@ def run(self) -> None:
         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

trino/client.py

@@ -77,6 +77,7 @@ def __init__(
         headers=None,
         transaction_id=None,
         extra_credential=None,
+        client_tags=None
     ):
         self.catalog = catalog
         self.schema = schema
@@ -88,6 +89,7 @@ def __init__(
         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 @@ def __init__(
         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 @@ def __init__(
             http_headers,
             transaction_id,
             extra_credential,
+            client_tags
         )
 
         self._host = host
@@ -286,6 +290,8 @@ def http_headers(self) -> Dict[str, str]:
         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 ``=``

trino/constants.py

@@ -32,6 +32,7 @@
 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"

trino/dbapi.py

@@ -108,7 +108,8 @@ def __init__(
         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 @@ def __init__(
         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 @@ def _create_request(self):
             self.redirect_handler,
             self.max_attempts,
             self.request_timeout,
+            client_tags=self.client_tags
         )
 
     def cursor(self, experimental_python_types=False):