Implement roles support in sqlalchemy

trinodb · Sep 16, 2022 · c5355a6 · c5355a6
1 parent 9f54c3c
commit c5355a6
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py
@@ -63,6 +63,16 @@ def setup(self):
                     experimental_python_types=True,
                 ),
             ),
+            (
+                make_url('trino://user@localhost:8080?roles={"hive":"finance"}'),
+                list(),
+                dict(host="localhost",
+                     port=8080,
+                     catalog="system",
+                     user="user",
+                     roles={"hive": "finance"},
+                     source="trino-sqlalchemy"),
+            ),
         ],
     )
     def test_create_connect_args(self, url: URL, expected_args: List[Any], expected_kwargs: Dict[str, Any]):

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]]: