Repository: trinodb/trino-python-client

Issue:
It's not possible to set role in trino.dbapi.connect method
### Describe the feature

With the recent changes with https://github.com/trinodb/trino-python-client/pull/200 `X-Trino-Role` and `X-Trino-Set-Role` are not allowed to be set. The following code:

```python
connection = trino.dbapi.connect(
    host="trino.example.net",
    port=443,
    http_scheme="https",
    catalog="hive",
    user='abuzer',
    auth=trino.auth.JWTAuthentication('Xvfdas.....'),
    http_headers={'X-Trino-Role': 'hive=ROLE{admin}'}
)
```
gives an error `ValueError: cannot override reserved HTTP header X-Trino-Role`.

We should make the `role` parameter part of `trino.dbapi.connect` method.

### Describe alternatives you've considered

I can to set `role` in `connection._client_session.role` or `connection._http_session.headers`. Since `_client_session` property should be an internal abstraction, `connection._client_session.role='hive=ROLE{admin}'` is ugly and prone to change. We need a proper interface to set such values in `connect` method.

We should keep release notes and publish critical breaking changes in future.

### Are you willing to submit PR?

- [X] Yes I am willing to submit a PR!

Discussion:
hovaesco: Please rebase on master to resolve CI failures. 
mdesmet: @hashar: This would enable us to support grants in dbt-trino.
mdesmet: @hashhar: Please advise on how to get this merged. Note that we had roles already supported (within the session) and they were tested in the same way (now expanded up to http_headers) but not in integration tests as we don't have a connector that supports roles in this project.

As mentioned we need this feature in as we want to support roles security within dbt_trino.
hashhar: I'm fine with current tests for now but as we start to work more on this area we'd eventually need better integration tests - let's create an issue with current ideas and we can follow up their.

@ebyhr WDYT?
hashhar: @mdesmet Can you also update PR description with the template + propose a release note?
mdesmet: > I'm fine with current tests for now but as we start to work more on this area we'd eventually need better integration tests - let's create an issue with current ideas and we can follow up their.

Created #241 241
hashhar: @mdesmet Can you rebase, I see some failures when running tests locally. Possibly something changed on master in the meantime.
mdesmet: @hashhar : Unfortunately something fishy is going on: 351 and 395 pass. However 398 does fail consistently (also locally), while 397 works fine.

I checked the release notes but didn't immediately find something. I'll keep digging.
hashhar: See https://github.com/trinodb/trino/pull/13907/files#r984659618, a "bug" has been fixed - we now don't get result sets for non-SELECT queries like it should've been all along.
mdesmet: @hashhar: I've realigned the implementation on the session_properties and your changes in #242 

It would be nice to get this one in as well.

Patch:
diff --git a/README.md b/README.md
index 0555346..cab2188 100644
--- a/README.md
+++ b/README.md
@@ -107,6 +107,7 @@ engine = create_engine(
       "session_properties": {'query_max_run_time': '1d'},
       "client_tags": ["tag1", "tag2"],
       "experimental_python_types": True,
+      "roles": {"catalog1": "role1"},
     }
 )
 
@@ -115,7 +116,8 @@ engine = create_engine(
     'trino://user@localhost:8080/system?'
     'session_properties={"query_max_run_time": "1d"}'
     '&client_tags=["tag1", "tag2"]'
-    '&experimental_python_types=true',
+    '&experimental_python_types=true'
+    '&roles={"catalog1": "role1"}'
 )
 ```
 
@@ -328,6 +330,20 @@ cur.execute('SELECT * FROM system.runtime.nodes')
 rows = cur.fetchall()
 ```
 
+## Roles
+
+Authorization roles to use for catalogs, specified as a dict with key-value pairs for the catalog and role. For example, `{"catalog1": "roleA", "catalog2": "roleB"}` sets `roleA` for `catalog1` and `roleB` for `catalog2`. See Trino docs.
+
+```python
+import trino
+conn = trino.dbapi.connect(
+    host='localhost',
+    port=443,
+    user='the-user',
+    roles={"catalog1": "roleA", "catalog2": "roleB"},
+)
+```
+
 ## SSL
 
 ### SSL verification
diff --git a/trino/client.py b/trino/client.py
index 9b319cf..afb195f 100644
--- a/trino/client.py
+++ b/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 @@ class ClientSession(object):
         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 @@ class ClientSession(object):
         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 @@ class ClientSession(object):
     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 @@ class ClientSession(object):
         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(",")]
@@ -224,7 +224,7 @@ def get_header_values(headers, header):
 def get_session_property_values(headers, header):
     kvs = get_header_values(headers, header)
     return [
-        (k.strip(), urllib.parse.unquote(v.strip()))
+        (k.strip(), urllib.parse.unquote_plus(v.strip()))
         for k, v in (kv.split("=", 1) for kv in kvs)
     ]
 
@@ -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 @@ 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
-        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 @@ class TrinoRequest(object):
             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(
diff --git a/trino/dbapi.py b/trino/dbapi.py
index 8adc4d1..683ced6 100644
--- a/trino/dbapi.py
+++ b/trino/dbapi.py
@@ -110,6 +110,7 @@ class Connection(object):
         http_session=None,
         client_tags=None,
         experimental_python_types=False,
+        roles=None,
     ):
         self.host = host
         self.port = port
@@ -127,7 +128,8 @@ class Connection(object):
             headers=http_headers,
             transaction_id=NO_TRANSACTION,
             extra_credential=extra_credential,
-            client_tags=client_tags
+            client_tags=client_tags,
+            roles=roles,
         )
         # mypy cannot follow module import
         if http_session is None:
diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py
index e967cb6..c5bccf3 100644
--- a/trino/sqlalchemy/dialect.py
+++ b/trino/sqlalchemy/dialect.py
@@ -124,6 +124,9 @@ class TrinoDialect(DefaultDialect):
         if "experimental_python_types" in url.query:
             kwargs["experimental_python_types"] = json.loads(url.query["experimental_python_types"])
 
+        if "roles" in url.query:
+            kwargs["roles"] = json.loads(url.query["roles"])
+
         return args, kwargs
 
     def get_columns(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:


Test patch:
diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py
index 06a869d..276ba98 100644
--- a/tests/integration/test_dbapi_integration.py
+++ b/tests/integration/test_dbapi_integration.py
@@ -19,6 +19,7 @@ import requests
 
 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,28 @@ 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")
+
+
+@pytest.mark.skipif(trino_version() == '351', reason="Newer Trino versions return the system role")
+def test_set_role_in_connection_trino_higher_351(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch", roles={"system": "ALL"}
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SHOW TABLES FROM information_schema')
+    cur.fetchall()
+    assert_role_headers(cur, "system=ALL")
+
+
+def assert_role_headers(cursor, expected_header):
+    assert cursor._request.http_headers[constants.HEADER_ROLE] == expected_header
 
 
 def test_prepared_statements(run_trino):
diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py
index b17f8cf..963b27f 100644
--- a/tests/unit/sqlalchemy/test_dialect.py
+++ b/tests/unit/sqlalchemy/test_dialect.py
@@ -63,6 +63,16 @@ class TestTrinoDialect:
                     experimental_python_types=True,
                 ),
             ),
+            (
+                make_url('trino://user@localhost:8080?roles={"hive":"finance","system":"analyst"}'),
+                list(),
+                dict(host="localhost",
+                     port=8080,
+                     catalog="system",
+                     user="user",
+                     roles={"hive": "finance", "system": "analyst"},
+                     source="trino-sqlalchemy"),
+            ),
         ],
     )
     def test_create_connect_args(self, url: URL, expected_args: List[Any], expected_kwargs: Dict[str, Any]):
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 4100c39..5477a39 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -13,6 +13,7 @@ import json
 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 test_request_headers(mock_get_and_post):
         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 test_retry_with():
         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)
diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py
index 6f2cc50..8138643 100644
--- a/tests/unit/test_dbapi.py
+++ b/tests/unit/test_dbapi.py
@@ -248,11 +248,19 @@ def test_token_retrieved_once_when_multithreaded(sample_post_response_data, samp
 
 @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
+
+
+@patch("trino.dbapi.trino.client")
+def test_role_is_set_when_specified(mock_client):
+    roles = {"system": "finance"}
+    with connect("sample_trino_cluster:443", roles=roles) as conn:
+        conn.cursor().execute("SOME FAKE QUERY")
+
+    _, passed_role = mock_client.ClientSession.call_args
+    assert passed_role["roles"] == roles