Add role parameter

Author: damian3031

tests/integration/test_dbapi_integration.py

@@ -1132,7 +1132,7 @@ def test_set_role(run_trino):
         assert_role_headers(cur, "system=ALL")
 
 
-def test_set_role_in_connection(run_trino):
+def test_set_roles_in_connection(run_trino):
     _, host, port = run_trino
 
     trino_connection = trino.dbapi.Connection(
@@ -1144,6 +1144,28 @@ def test_set_role_in_connection(run_trino):
     assert_role_headers(cur, "system=ALL")
 
 
+def test_set_role_in_connection(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch", role="ALL"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SHOW TABLES FROM information_schema')
+    cur.fetchall()
+    assert_role_headers(cur, "system=ALL")
+
+
+def test_set_role_and_roles_in_connection(run_trino):
+    _, host, port = run_trino
+
+    with pytest.raises(ValueError) as e:
+        trino.dbapi.Connection(
+            host=host, port=port, user="test", catalog="tpch", roles={"system": "ALL"}, role="ALL"
+        )
+    assert "specify 'role' or 'roles' parameter, but not both" == str(e.value)
+
+
 def assert_role_headers(cursor, expected_header):
     assert cursor._request.http_headers[constants.HEADER_ROLE] == expected_header
 

trino/client.py

@@ -137,6 +137,7 @@ def __init__(
         client_tags: List[str] = None,
         roles: Dict[str, str] = None,
         timezone: str = None,
+        role: str = None,
     ):
         self._user = user
         self._catalog = catalog
@@ -147,7 +148,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 = self._format_roles(roles) if roles is not None else {}
+        self._roles = self._format_roles(role, roles)
         self._prepared_statements: Dict[str, str] = {}
         self._object_lock = threading.Lock()
         self._timezone = timezone or get_localzone_name()
@@ -238,7 +239,13 @@ def timezone(self):
         with self._object_lock:
             return self._timezone
 
-    def _format_roles(self, roles):
+    def _format_roles(self, role, roles):
+        if role and roles:
+            raise ValueError("specify 'role' or 'roles' parameter, but not both")
+        elif role:
+            roles = {"system": role}
+        elif role is None and roles is None:
+            return {}
         formatted_roles = {}
         for catalog, role in roles.items():
             is_legacy_role_pattern = ROLE_PATTERN.match(role) is not None

trino/dbapi.py

@@ -113,6 +113,7 @@ def __init__(
         legacy_primitive_types=False,
         roles=None,
         timezone=None,
+        role=None,
     ):
         self.host = host
         self.port = port
@@ -133,6 +134,7 @@ def __init__(
             client_tags=client_tags,
             roles=roles,
             timezone=timezone,
+            role=role,
         )
         # mypy cannot follow module import
         if http_session is None: