Refactor role into roles

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -19,6 +19,7 @@
 
 import trino
 from tests.integration.conftest import trino_version
+from trino import constants
 from trino.exceptions import TrinoQueryError, TrinoUserError, NotSupportedError
 from trino.transaction import IsolationLevel
 
@@ -1045,11 +1046,11 @@ def test_set_role_trino_higher_351(run_trino):
     cur = trino_connection.cursor()
     cur.execute('SHOW TABLES FROM information_schema')
     cur.fetchall()
-    assert cur._request._client_session.role is None
+    assert cur._request._client_session.roles == {}
 
     cur.execute("SET ROLE ALL")
     cur.fetchall()
-    assert cur._request._client_session.role == "system=ALL"
+    assert_role_headers(cur, "system=ALL")
 
 
 @pytest.mark.skipif(trino_version() != '351', reason="Trino 351 returns the role for the current catalog")
@@ -1062,11 +1063,15 @@ def test_set_role_trino_351(run_trino):
     cur = trino_connection.cursor()
     cur.execute('SHOW TABLES FROM information_schema')
     cur.fetchall()
-    assert cur._request._client_session.role is None
+    assert cur._request._client_session.roles == {}
 
     cur.execute("SET ROLE ALL")
     cur.fetchall()
-    assert cur._request._client_session.role == "tpch=ALL"
+    assert_role_headers(cur, "tpch=ALL")
+
+
+def assert_role_headers(cursor, expected_header):
+    assert cursor._request.http_headers[constants.HEADER_ROLE] == expected_header
 
 
 def test_prepared_statements(run_trino):

tests/unit/test_client.py

@@ -13,6 +13,7 @@
 import threading
 import time
 import uuid
+from typing import Optional, Dict
 from unittest import mock
 from urllib.parse import urlparse
 
@@ -92,10 +93,9 @@ def assert_headers(headers):
         assert headers[constants.HEADER_SOURCE] == source
         assert headers[constants.HEADER_USER] == user
         assert headers[constants.HEADER_SESSION] == ""
-        assert headers[constants.HEADER_ROLE] is None
         assert headers[accept_encoding_header] == accept_encoding_value
         assert headers[client_info_header] == client_info_value
-        assert len(headers.keys()) == 9
+        assert len(headers.keys()) == 8
 
     req.post("URL")
     _, post_kwargs = post.call_args
@@ -988,10 +988,12 @@ def __call__(self):
         with_retry(FailerUntil(3).__call__)()
 
 
-def assert_headers_with_role(headers, role):
+def assert_headers_with_roles(headers: Dict[str, str], roles: Optional[str]):
+    if roles is None:
+        assert constants.HEADER_ROLE not in headers
+    else:
+        assert headers[constants.HEADER_ROLE] == roles
     assert headers[constants.HEADER_USER] == "test_user"
-    assert headers[constants.HEADER_ROLE] == role
-    assert len(headers.keys()) == 7
 
 
 def test_request_headers_role_hive_all(mock_get_and_post):
@@ -1001,17 +1003,17 @@ def test_request_headers_role_hive_all(mock_get_and_post):
         port=8080,
         client_session=ClientSession(
             user="test_user",
-            role="hive=ALL",
+            roles={"hive": "ALL"}
         ),
     )
 
     req.post("URL")
     _, post_kwargs = post.call_args
-    assert_headers_with_role(post_kwargs["headers"], "hive=ALL")
+    assert_headers_with_roles(post_kwargs["headers"], "hive=ALL")
 
     req.get("URL")
     _, get_kwargs = get.call_args
-    assert_headers_with_role(post_kwargs["headers"], "hive=ALL")
+    assert_headers_with_roles(post_kwargs["headers"], "hive=ALL")
 
 
 def test_request_headers_role_admin(mock_get_and_post):
@@ -1022,17 +1024,17 @@ def test_request_headers_role_admin(mock_get_and_post):
         port=8080,
         client_session=ClientSession(
             user="test_user",
-            role="admin",
+            roles={"system": "admin"}
         ),
     )
 
     req.post("URL")
     _, post_kwargs = post.call_args
-    assert_headers_with_role(post_kwargs["headers"], "admin")
+    assert_headers_with_roles(post_kwargs["headers"], "system=admin")
 
     req.get("URL")
     _, get_kwargs = get.call_args
-    assert_headers_with_role(post_kwargs["headers"], "admin")
+    assert_headers_with_roles(post_kwargs["headers"], "system=admin")
 
 
 def test_request_headers_role_empty(mock_get_and_post):
@@ -1043,14 +1045,14 @@ def test_request_headers_role_empty(mock_get_and_post):
         port=8080,
         client_session=ClientSession(
             user="test_user",
-            role="",
+            roles=None,
         ),
     )
 
     req.post("URL")
     _, post_kwargs = post.call_args
-    assert_headers_with_role(post_kwargs["headers"], "")
+    assert_headers_with_roles(post_kwargs["headers"], None)
 
     req.get("URL")
     _, get_kwargs = get.call_args
-    assert_headers_with_role(post_kwargs["headers"], "")
+    assert_headers_with_roles(post_kwargs["headers"], None)

tests/unit/test_dbapi.py

@@ -248,11 +248,9 @@ def run(self) -> None:
 
 @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.ClientSession.call_args
     assert passed_client_tags["client_tags"] == client_tags

trino/client.py

@@ -98,7 +98,7 @@ class ClientSession(object):
     :param extra_credential: extra credentials. as list of ``(key, value)``
                              tuples.
     :param client_tags: Client tags as list of strings.
-    :param role: role for the current session. Some connectors do not
+    :param roles: roles for the current session. Some connectors do not
                  support role management. See connector documentation for more details.
     """
 
@@ -113,7 +113,7 @@ def __init__(
         transaction_id: str = None,
         extra_credential: List[Tuple[str, str]] = None,
         client_tags: List[str] = None,
-        role: str = None,
+        roles: Dict[str, str] = None,
     ):
         self._user = user
         self._catalog = catalog
@@ -124,7 +124,7 @@ def __init__(
         self._transaction_id = transaction_id
         self._extra_credential = extra_credential
         self._client_tags = client_tags
-        self._role = role
+        self._roles = roles or {}
         self._prepared_statements: Dict[str, str] = {}
         self._object_lock = threading.Lock()
 
@@ -188,24 +188,15 @@ def extra_credential(self):
     def client_tags(self):
         return self._client_tags
 
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        del state["_object_lock"]
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self._object_lock = threading.Lock()
-
     @property
-    def role(self):
+    def roles(self):
         with self._object_lock:
-            return self._role
+            return self._roles
 
-    @role.setter
-    def role(self, role):
+    @roles.setter
+    def roles(self, roles):
         with self._object_lock:
-            self._role = role
+            self._roles = roles
 
     @property
     def prepared_statements(self):
@@ -216,6 +207,15 @@ def prepared_statements(self, prepared_statements):
         with self._object_lock:
             self._prepared_statements = prepared_statements
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state["_object_lock"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._object_lock = threading.Lock()
+
 
 def get_header_values(headers, header):
     return [val.strip() for val in headers[header].split(",")]
@@ -237,6 +237,14 @@ def get_prepared_statement_values(headers, header):
     ]
 
 
+def get_roles_values(headers, header):
+    kvs = get_header_values(headers, header)
+    return [
+        (k.strip(), urllib.parse.unquote_plus(v.strip()))
+        for k, v in (kv.split("=", 1) for kv in kvs)
+    ]
+
+
 class TrinoStatus(object):
     def __init__(self, id, stats, warnings, info_uri, next_uri, update_type, rows, columns=None):
         self.id = id
@@ -400,7 +408,12 @@ 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
-        headers[constants.HEADER_ROLE] = self._client_session.role
+        if len(self._client_session.roles.values()):
+            headers[constants.HEADER_ROLE] = ",".join(
+                # ``name`` must not contain ``=``
+                "{}={}".format(catalog, urllib.parse.quote(str(role)))
+                for catalog, role in self._client_session.roles.items()
+            )
         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)
 
@@ -579,7 +592,10 @@ def process(self, http_response) -> TrinoStatus:
             self._client_session.schema = http_response.headers[constants.HEADER_SET_SCHEMA]
 
         if constants.HEADER_SET_ROLE in http_response.headers:
-            self._client_session.role = http_response.headers[constants.HEADER_SET_ROLE]
+            for key, value in get_roles_values(
+                    http_response.headers, constants.HEADER_SET_ROLE
+            ):
+                self._client_session.roles[key] = value
 
         if constants.HEADER_ADDED_PREPARE in http_response.headers:
             for name, statement in get_prepared_statement_values(