Encode roles as in JDBC driver

Author: mdesmet

tests/unit/test_client.py

@@ -12,6 +12,7 @@
 import json
 import threading
 import time
+import urllib
 import uuid
 from typing import Dict, Optional
 from unittest import mock
@@ -105,6 +106,13 @@ def test_request_headers(mock_get_and_post):
             headers={
                 accept_encoding_header: accept_encoding_value,
                 client_info_header: client_info_value,
+            },
+            roles={
+                "hive": "ALL",
+                "system": "analyst",
+                "catalog1": "NONE",
+                # ensure backwards compatibility
+                "catalog2": "ROLE{catalog2_role}",
             }
         ),
         http_scheme="http",
@@ -121,7 +129,13 @@ def assert_headers(headers):
         assert headers[constants.HEADER_CLIENT_CAPABILITIES] == "PARAMETRIC_DATETIME"
         assert headers[accept_encoding_header] == accept_encoding_value
         assert headers[client_info_header] == client_info_value
-        assert len(headers.keys()) == 10
+        assert headers[constants.HEADER_ROLE] == (
+            "hive=ALL,"
+            "system=" + urllib.parse.quote("ROLE{analyst}") + ","
+            "catalog1=NONE,"
+            "catalog2=" + urllib.parse.quote("ROLE{catalog2_role}")
+        )
+        assert len(headers.keys()) == 11
 
     req.post("URL")
     _, post_kwargs = post.call_args
@@ -1095,14 +1109,15 @@ def test_request_headers_role_admin(mock_get_and_post):
             roles={"system": "admin"}
         ),
     )
+    roles = "system=" + urllib.parse.quote("ROLE{admin}")
 
     req.post("URL")
     _, post_kwargs = post.call_args
-    assert_headers_with_roles(post_kwargs["headers"], "system=admin")
+    assert_headers_with_roles(post_kwargs["headers"], roles)
 
     req.get("URL")
     _, get_kwargs = get.call_args
-    assert_headers_with_roles(post_kwargs["headers"], "system=admin")
+    assert_headers_with_roles(post_kwargs["headers"], roles)
 
 
 def test_request_headers_role_empty(mock_get_and_post):

trino/client.py

@@ -85,6 +85,8 @@
 MAX_PYTHON_TEMPORAL_PRECISION_POWER = 6
 MAX_PYTHON_TEMPORAL_PRECISION = POWERS_OF_TEN[MAX_PYTHON_TEMPORAL_PRECISION_POWER]
 
+ROLE_PATTERN = re.compile(r"^ROLE\{(.*)\}$")
+
 
 class ClientSession(object):
     """
@@ -143,7 +145,7 @@ def __init__(
         self._transaction_id = transaction_id
         self._extra_credential = extra_credential
         self._client_tags = client_tags.copy() if client_tags is not None else list()
-        self._roles = roles.copy() if roles is not None else {}
+        self._roles = self._format_roles(roles) if roles is not None else {}
         self._prepared_statements: Dict[str, str] = {}
         self._object_lock = threading.Lock()
         self._timezone = timezone or get_localzone_name()
@@ -234,6 +236,15 @@ def timezone(self):
         with self._object_lock:
             return self._timezone
 
+    def _format_roles(self, roles):
+        formatted_roles = {}
+        for catalog, role in roles.items():
+            if role in ("NONE", "ALL") or ROLE_PATTERN.match(role) is not None:
+                formatted_roles[catalog] = role
+            else:
+                formatted_roles[catalog] = f"ROLE{{{role}}}"
+        return formatted_roles
+
     def __getstate__(self):
         state = self.__dict__.copy()
         del state["_object_lock"]