Accept shorthand format for system roles

README.md

@@ -359,6 +359,18 @@ conn = trino.dbapi.connect(
 )
 ```
 
+You could also pass `system` role without explicitly specifing "system" catalog:
+
+```python
+import trino
+conn = trino.dbapi.connect(
+ host='localhost',
+ port=443,
+ user='the-user',
+ roles="role1" # equivalent to {"system": "role1"}
+)
+```
+
 ## Timezone
 
 The time zone for the session can be explicitly set using the IANA time zone

tests/integration/test_dbapi_integration.py

@@ -1144,6 +1144,18 @@ def test_set_role_in_connection(run_trino):
  assert_role_headers(cur, "system=ALL")
 
 
+def test_set_system_role_in_connection(run_trino):
+ _, host, port = run_trino
+
+ trino_connection = trino.dbapi.Connection(
+ host=host, port=port, user="test", catalog="tpch", roles="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
 

trino/client.py

@@ -135,7 +135,7 @@ def __init__(
  transaction_id: str = None,
  extra_credential: List[Tuple[str, str]] = None,
  client_tags: List[str] = None,
- roles: Dict[str, str] = None,
+ roles: Union[Dict[str, str], str] = None,
  timezone: str = None,
  ):
  self._user = user
@@ -239,6 +239,8 @@ def timezone(self):
  return self._timezone
 
  def _format_roles(self, roles):
+ if isinstance(roles, str):
+ roles = {"system": roles}
  formatted_roles = {}
  for catalog, role in roles.items():
  is_legacy_role_pattern = ROLE_PATTERN.match(role) is not None