Add support for SET ROLE statement

Author: ulisesojeda

tests/integration/test_dbapi_integration.py

@@ -1034,3 +1034,37 @@ def test_use_catalog(run_trino):
     result = cur.fetchall()
     assert result[0][0] == 'tpch'
     assert result[0][1] == 'sf1'
+
+
+@pytest.mark.skipif(trino_version() == '351', reason="Newer Trino versions return the system role")
+def test_set_role_trino_higher_351(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SHOW TABLES FROM information_schema')
+    cur.fetchall()
+    assert cur._request._client_session.role is None
+
+    cur.execute("SET ROLE ALL")
+    cur.fetchall()
+    assert cur._request._client_session.role == "system=ALL"
+
+
+@pytest.mark.skipif(trino_version() != '351', reason="Trino 351 returns the role for the current catalog")
+def test_set_role_trino_351(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SHOW TABLES FROM information_schema')
+    cur.fetchall()
+    assert cur._request._client_session.role is None
+
+    cur.execute("SET ROLE ALL")
+    cur.fetchall()
+    assert cur._request._client_session.role == "tpch=ALL"

tests/unit/test_client.py

@@ -92,9 +92,10 @@ 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()) == 8
+        assert len(headers.keys()) == 9
 
     req.post("URL")
     _, post_kwargs = post.call_args
@@ -1002,3 +1003,71 @@ def __call__(self):
     with_retry(FailerUntil(2).__call__)()
     with pytest.raises(SomeException):
         with_retry(FailerUntil(3).__call__)()
+
+
+def assert_headers_with_role(headers, role):
+    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):
+    get, post = mock_get_and_post
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            role="hive=ALL",
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_with_role(post_kwargs["headers"], "hive=ALL")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_with_role(post_kwargs["headers"], "hive=ALL")
+
+
+def test_request_headers_role_admin(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            role="admin",
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_with_role(post_kwargs["headers"], "admin")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_with_role(post_kwargs["headers"], "admin")
+
+
+def test_request_headers_role_empty(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            role="",
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_with_role(post_kwargs["headers"], "")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_with_role(post_kwargs["headers"], "")

trino/client.py

@@ -98,6 +98,8 @@ 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
+                 support role management. See connector documentation for more details.
     """
 
     def __init__(
@@ -111,6 +113,7 @@ def __init__(
         transaction_id: str = None,
         extra_credential: List[Tuple[str, str]] = None,
         client_tags: List[str] = None,
+        role: str = None,
     ):
         self._user = user
         self._catalog = catalog
@@ -121,6 +124,7 @@ def __init__(
         self._transaction_id = transaction_id
         self._extra_credential = extra_credential
         self._client_tags = client_tags
+        self._role = role
         self._object_lock = threading.Lock()
 
     @property
@@ -192,6 +196,16 @@ def __setstate__(self, state):
         self.__dict__.update(state)
         self._object_lock = threading.Lock()
 
+    @property
+    def role(self):
+        with self._object_lock:
+            return self._role
+
+    @role.setter
+    def role(self, role):
+        with self._object_lock:
+            self._role = role
+
 
 def get_header_values(headers, header):
     return [val.strip() for val in headers[header].split(",")]
@@ -368,6 +382,7 @@ 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 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)
 
@@ -538,6 +553,9 @@ def process(self, http_response) -> TrinoStatus:
         if constants.HEADER_SET_SCHEMA in http_response.headers:
             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]
+
         self._next_uri = response.get("nextUri")
 
         return TrinoStatus(

trino/constants.py

@@ -39,6 +39,9 @@
 HEADER_SET_SESSION = "X-Trino-Set-Session"
 HEADER_CLEAR_SESSION = "X-Trino-Clear-Session"
 
+HEADER_ROLE = "X-Trino-Role"
+HEADER_SET_ROLE = "X-Trino-Set-Role"
+
 HEADER_STARTED_TRANSACTION = "X-Trino-Started-Transaction-Id"
 HEADER_TRANSACTION = "X-Trino-Transaction-Id"