Implement roles support in sqlalchemy

trinodb · Sep 23, 2022 · 9b9c848 · 9b9c848
1 parent 75e6556
commit 9b9c848
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -100,6 +100,7 @@ engine = create_engine(
       "session_properties": {'query_max_run_time': '1d'},
       "client_tags": ["tag1", "tag2"],
       "experimental_python_types": True,
+      "roles": {"catalog1": "role1"},
     }
 )
 
@@ -108,7 +109,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"}'
 )
 ```
 

diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py
@@ -124,6 +124,9 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any
         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]]: