Implement roles support in dbapi.connect()

Author: mdesmet

README.md

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

tests/integration/test_dbapi_integration.py

@@ -1070,6 +1070,19 @@ def test_set_role_trino_351(run_trino):
     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
 

tests/unit/test_dbapi.py

@@ -254,3 +254,13 @@ def test_tags_are_set_when_specified(mock_client):
 
     _, 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

trino/dbapi.py

@@ -110,6 +110,7 @@ def __init__(
         http_session=None,
         client_tags=None,
         experimental_python_types=False,
+        roles=None,
     ):
         self.host = host
         self.port = port
@@ -127,7 +128,8 @@ def __init__(
             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: