Repository: trinodb/trino-python-client

Issue:
SQLAlchemy 2.0 Compatibility
### Describe the feature

SQLAlchemy 2.0 is currently in beta. I tried to use the trino dialect with it, but it errors when trying to access `context.should_autocommit` in this bit of the implementation:
https://github.com/trinodb/trino-python-client/blob/eef6388cd30f93ed1172cb1712b1b7e680276761/trino/sqlalchemy/dialect.py#L366-L370

The attribute `should_autocommit` of `DefaultExecutionContext` was removed in 2.0, as one of the maintainers answered to a question I asked on the SQLAlchemy repository: https://github.com/sqlalchemy/sqlalchemy/discussions/8854#discussioncomment-4195357

They also suggested that the current implementation seems a bit strange and should be re-visited. I don't have much expertise with dialects in general, so I'm creating this issue for now. I'm also willing to submit a PR if someone can explain in more detail what the purpose of that if statement is.

### Describe alternatives you've considered

Alternative: Not using 2.0 - but 2.0 will be out of beta eventually so the dialect should be updated.

### Are you willing to submit PR?

- [X] Yes I am willing to submit a PR!

Discussion:
cla-bot[bot]: Thank you for your pull request and welcome to our community. We could not parse the GitHub identity of the following contributors: **mritter31**.
This is most likely caused by a git client misconfiguration; please make sure to:
1. check if your git client is configured with an email to sign commits `git config --list | grep email`
2. If not, set it up using `git config --global user.email email@example.com`
3. Make sure that the git commit email is configured in your GitHub account settings, see https://github.com/settings/emails
mathiasritter: Thanks @mdesmet, that was what I was looking for.

Unfortunately, the tests for 2.0 are failing because of an issue in `sqlalchemy-utils` library. That issue has been fixed since October, but they haven't released a new version since July...
See: https://github.com/kvesteri/sqlalchemy-utils/issues/666
hashhar: We're using that library only in tests if I remember correctly - can we remove that dependency with some alternative code? It's used in two tests - just to verify whether `sqla.inspect(engine).get_table_names(schema)` returns only tables and whether `sqla.inspect(engine).get_view_names(schema)` only returns views.

Would those tests still work if we directly executed `CREATE VIEW` on the underlying connection instead of going via SQLA?
mathiasritter: @hashhar They have just released a new version of `sqlalchemy-utils` that fixed the issue 😄 got some other issues now though, that I will look into later.
mdesmet: @hashhar : PTAL
cla-bot[bot]: Thank you for your pull request and welcome to our community. We could not parse the GitHub identity of the following contributors: **Mathias Ritter**.
This is most likely caused by a git client misconfiguration; please make sure to:
1. check if your git client is configured with an email to sign commits `git config --list | grep email`
2. If not, set it up using `git config --global user.email email@example.com`
3. Make sure that the git commit email is configured in your GitHub account settings, see https://github.com/settings/emails
hashhar: LGTM. @mathiasritter I assume you've sent the CLA already (and it matches the email you use in your commits)? If not please do. If sent already I'll take care of it.
mathiasritter: @hashhar Yes I’ve already sent it

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 75cc6c4..57e10c3 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -58,6 +58,8 @@ jobs:
           - { python: "3.11", trino: "351", sqlalchemy: "~=1.4.0" }  # first Trino version
           # Test with sqlalchemy 1.3
           - { python: "3.11", trino: "latest", sqlalchemy: "~=1.3.0" }
+          # Test with sqlalchemy 2.0
+          - { python: "3.11", trino: "latest", sqlalchemy: "~=2.0.0rc1" }
     env:
       TRINO_VERSION: "${{ matrix.trino }}"
     steps:
diff --git a/setup.py b/setup.py
index c5aa8a1..9e6003a 100755
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,7 @@ with open("trino/__init__.py", "rb") as f:
     version = str(ast.literal_eval(trino_version.group(1)))
 
 kerberos_require = ["requests_kerberos"]
-sqlalchemy_require = ["sqlalchemy~=1.3"]
+sqlalchemy_require = ["sqlalchemy >= 1.3"]
 external_authentication_token_cache_require = ["keyring"]
 
 # We don't add localstorage_require to all_require as users must explicitly opt in to use keyring.
diff --git a/trino/client.py b/trino/client.py
index 3e72c0b..0f47694 100644
--- a/trino/client.py
+++ b/trino/client.py
@@ -85,6 +85,8 @@ for i in range(0, 13):
 MAX_PYTHON_TEMPORAL_PRECISION_POWER = 6
 MAX_PYTHON_TEMPORAL_PRECISION = POWERS_OF_TEN[MAX_PYTHON_TEMPORAL_PRECISION_POWER]
 
+ROLE_PATTERN = re.compile(r"^ROLE\{(.*)\}$")
+
 
 class ClientSession(object):
     """
@@ -143,7 +145,7 @@ class ClientSession(object):
         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 = roles.copy() if roles is not None else {}
+        self._roles = self._format_roles(roles) if roles is not None else {}
         self._prepared_statements: Dict[str, str] = {}
         self._object_lock = threading.Lock()
         self._timezone = timezone or get_localzone_name()
@@ -234,6 +236,15 @@ class ClientSession(object):
         with self._object_lock:
             return self._timezone
 
+    def _format_roles(self, roles):
+        formatted_roles = {}
+        for catalog, role in roles.items():
+            if role in ("NONE", "ALL") or ROLE_PATTERN.match(role) is not None:
+                formatted_roles[catalog] = role
+            else:
+                formatted_roles[catalog] = f"ROLE{{{role}}}"
+        return formatted_roles
+
     def __getstate__(self):
         state = self.__dict__.copy()
         del state["_object_lock"]
diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py
index 2aaac3d..1fa6ed0 100644
--- a/trino/sqlalchemy/dialect.py
+++ b/trino/sqlalchemy/dialect.py
@@ -156,7 +156,7 @@ class TrinoDialect(DefaultDialect):
             ORDER BY "ordinal_position" ASC
         """
         ).strip()
-        res = connection.execute(sql.text(query), schema=schema, table=table_name)
+        res = connection.execute(sql.text(query), {"schema": schema, "table": table_name})
         columns = []
         for record in res:
             column = dict(
@@ -204,7 +204,7 @@ class TrinoDialect(DefaultDialect):
               AND "table_type" = 'BASE TABLE'
         """
         ).strip()
-        res = connection.execute(sql.text(query), schema=schema)
+        res = connection.execute(sql.text(query), {"schema": schema})
         return [row.table_name for row in res]
 
     def get_temp_table_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
@@ -225,7 +225,7 @@ class TrinoDialect(DefaultDialect):
               AND "table_type" = 'VIEW'
         """
         ).strip()
-        res = connection.execute(sql.text(query), schema=schema)
+        res = connection.execute(sql.text(query), {"schema": schema})
         return [row.table_name for row in res]
 
     def get_temp_view_names(self, connection: Connection, schema: str = None, **kw) -> List[str]:
@@ -244,7 +244,7 @@ class TrinoDialect(DefaultDialect):
               AND "table_name" = :view
         """
         ).strip()
-        res = connection.execute(sql.text(query), schema=schema, view=view_name)
+        res = connection.execute(sql.text(query), {"schema": schema, "view": view_name})
         return res.scalar()
 
     def get_indexes(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:
@@ -296,7 +296,7 @@ class TrinoDialect(DefaultDialect):
         try:
             res = connection.execute(
                 sql.text(query),
-                catalog_name=catalog_name, schema_name=schema_name, table_name=table_name
+                {"catalog_name": catalog_name, "schema_name": schema_name, "table_name": table_name}
             )
             return dict(text=res.scalar())
         except error.TrinoQueryError as e:
@@ -314,7 +314,7 @@ class TrinoDialect(DefaultDialect):
             WHERE "schema_name" = :schema
         """
         ).strip()
-        res = connection.execute(sql.text(query), schema=schema)
+        res = connection.execute(sql.text(query), {"schema": schema})
         return res.first() is not None
 
     def has_table(self, connection: Connection, table_name: str, schema: str = None, **kw) -> bool:
@@ -329,7 +329,7 @@ class TrinoDialect(DefaultDialect):
               AND "table_name" = :table
         """
         ).strip()
-        res = connection.execute(sql.text(query), schema=schema, table=table_name)
+        res = connection.execute(sql.text(query), {"schema": schema, "table": table_name})
         return res.first() is not None
 
     def has_sequence(self, connection: Connection, sequence_name: str, schema: str = None, **kw) -> bool:
@@ -363,11 +363,6 @@ class TrinoDialect(DefaultDialect):
         self, cursor: Cursor, statement: str, parameters: Tuple[Any, ...], context: DefaultExecutionContext = None
     ):
         cursor.execute(statement, parameters)
-        if context and context.should_autocommit:
-            # SQL statement only submitted to Trino server when cursor.fetch*() is called.
-            # For DDL (CREATE/ALTER/DROP) and DML (INSERT/UPDATE/DELETE) statement, call cursor.description
-            # to force submit statement immediately.
-            cursor.description  # noqa
 
     def do_rollback(self, dbapi_connection: trino_dbapi.Connection):
         if dbapi_connection.transaction is not None:


Test patch:
diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py
index 9ac2bc0..81f8505 100644
--- a/tests/integration/test_sqlalchemy_integration.py
+++ b/tests/integration/test_sqlalchemy_integration.py
@@ -43,10 +43,10 @@ def test_select_query(trino_connection):
     rows = result.fetchall()
     assert len(rows) == 25
     for row in rows:
-        assert isinstance(row['nationkey'], int)
-        assert isinstance(row['name'], str)
-        assert isinstance(row['regionkey'], int)
-        assert isinstance(row['comment'], str)
+        assert isinstance(row.nationkey, int)
+        assert isinstance(row.name, str)
+        assert isinstance(row.regionkey, int)
+        assert isinstance(row.comment, str)
 
 
 def assert_column(table, column_name, column_type):
@@ -70,8 +70,8 @@ def test_select_specific_columns(trino_connection):
     rows = result.fetchall()
     assert len(rows) > 0
     for row in rows:
-        assert isinstance(row['node_id'], str)
-        assert isinstance(row['state'], str)
+        assert isinstance(row.node_id, str)
+        assert isinstance(row.state, str)
 
 
 @pytest.mark.skipif(
@@ -82,7 +82,8 @@ def test_select_specific_columns(trino_connection):
 def test_define_and_create_table(trino_connection):
     engine, conn = trino_connection
     if not engine.dialect.has_schema(conn, "test"):
-        engine.execute(sqla.schema.CreateSchema("test"))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
     metadata = sqla.MetaData()
     try:
         sqla.Table('users',
@@ -110,7 +111,8 @@ def test_insert(trino_connection):
     engine, conn = trino_connection
 
     if not engine.dialect.has_schema(conn, "test"):
-        engine.execute(sqla.schema.CreateSchema("test"))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
     metadata = sqla.MetaData()
     try:
         users = sqla.Table('users',
@@ -139,7 +141,8 @@ def test_insert(trino_connection):
 def test_insert_multiple_statements(trino_connection):
     engine, conn = trino_connection
     if not engine.dialect.has_schema(conn, "test"):
-        engine.execute(sqla.schema.CreateSchema("test"))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
     metadata = sqla.MetaData()
     users = sqla.Table('users',
                        metadata,
@@ -180,10 +183,10 @@ def test_operators(trino_connection):
     rows = result.fetchall()
     assert len(rows) == 1
     for row in rows:
-        assert isinstance(row['nationkey'], int)
-        assert isinstance(row['name'], str)
-        assert isinstance(row['regionkey'], int)
-        assert isinstance(row['comment'], str)
+        assert isinstance(row.nationkey, int)
+        assert isinstance(row.name, str)
+        assert isinstance(row.regionkey, int)
+        assert isinstance(row.comment, str)
 
 
 @pytest.mark.skipif(
@@ -216,14 +219,14 @@ def test_textual_sql(trino_connection):
     rows = result.fetchall()
     assert len(rows) == 3
     for row in rows:
-        assert isinstance(row['custkey'], int)
-        assert isinstance(row['name'], str)
-        assert isinstance(row['address'], str)
-        assert isinstance(row['nationkey'], int)
-        assert isinstance(row['phone'], str)
-        assert isinstance(row['acctbal'], float)
-        assert isinstance(row['mktsegment'], str)
-        assert isinstance(row['comment'], str)
+        assert isinstance(row.custkey, int)
+        assert isinstance(row.name, str)
+        assert isinstance(row.address, str)
+        assert isinstance(row.nationkey, int)
+        assert isinstance(row.phone, str)
+        assert isinstance(row.acctbal, float)
+        assert isinstance(row.mktsegment, str)
+        assert isinstance(row.comment, str)
 
 
 @pytest.mark.skipif(
@@ -323,7 +326,8 @@ def test_json_column(trino_connection, json_object):
     engine, conn = trino_connection
 
     if not engine.dialect.has_schema(conn, "test"):
-        engine.execute(sqla.schema.CreateSchema("test"))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
     metadata = sqla.MetaData()
 
     try:
@@ -351,7 +355,8 @@ def test_get_table_comment(trino_connection):
     engine, conn = trino_connection
 
     if not engine.dialect.has_schema(conn, "test"):
-        engine.execute(sqla.schema.CreateSchema("test"))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
     metadata = sqla.MetaData()
 
     try:
@@ -378,7 +383,8 @@ def test_get_table_names(trino_connection, schema):
     metadata = sqla.MetaData(schema=schema_name)
 
     if not engine.dialect.has_schema(conn, schema_name):
-        engine.execute(sqla.schema.CreateSchema(schema_name))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema(schema_name))
 
     try:
         sqla.Table(
@@ -388,10 +394,10 @@ def test_get_table_names(trino_connection, schema):
         )
         metadata.create_all(engine)
         view_name = schema_name + ".test_view"
-        conn.execute(f"CREATE VIEW {view_name} AS SELECT * FROM test_get_table_names")
+        conn.execute(sqla.text(f"CREATE VIEW {view_name} AS SELECT * FROM test_get_table_names"))
         assert sqla.inspect(engine).get_table_names(schema_name) == ['test_get_table_names']
     finally:
-        conn.execute(f"DROP VIEW IF EXISTS {view_name}")
+        conn.execute(sqla.text(f"DROP VIEW IF EXISTS {view_name}"))
         metadata.drop_all(engine)
 
 
@@ -411,7 +417,8 @@ def test_get_view_names(trino_connection, schema):
     metadata = sqla.MetaData(schema=schema_name)
 
     if not engine.dialect.has_schema(conn, schema_name):
-        engine.execute(sqla.schema.CreateSchema(schema_name))
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema(schema_name))
 
     try:
         sqla.Table(
@@ -421,10 +428,10 @@ def test_get_view_names(trino_connection, schema):
         )
         metadata.create_all(engine)
         view_name = schema_name + ".test_get_view_names"
-        conn.execute(f"CREATE VIEW {view_name} AS SELECT * FROM test_table")
+        conn.execute(sqla.text(f"CREATE VIEW {view_name} AS SELECT * FROM test_table"))
         assert sqla.inspect(engine).get_view_names(schema_name) == ['test_get_view_names']
     finally:
-        conn.execute(f"DROP VIEW IF EXISTS {view_name}")
+        conn.execute(sqla.text(f"DROP VIEW IF EXISTS {view_name}"))
         metadata.drop_all(engine)
 
 
diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py
index 3316fa6..ff16d6d 100644
--- a/tests/unit/sqlalchemy/test_dialect.py
+++ b/tests/unit/sqlalchemy/test_dialect.py
@@ -16,7 +16,7 @@ from trino.transaction import IsolationLevel
 
 
 class TestTrinoDialect:
-    def setup(self):
+    def setup_method(self):
         self.dialect = TrinoDialect()
 
     @pytest.mark.parametrize(
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 5f1a541..de5f996 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -12,6 +12,7 @@
 import json
 import threading
 import time
+import urllib
 import uuid
 from typing import Dict, Optional
 from unittest import mock
@@ -105,6 +106,13 @@ def test_request_headers(mock_get_and_post):
             headers={
                 accept_encoding_header: accept_encoding_value,
                 client_info_header: client_info_value,
+            },
+            roles={
+                "hive": "ALL",
+                "system": "analyst",
+                "catalog1": "NONE",
+                # ensure backwards compatibility
+                "catalog2": "ROLE{catalog2_role}",
             }
         ),
         http_scheme="http",
@@ -121,7 +129,13 @@ def test_request_headers(mock_get_and_post):
         assert headers[constants.HEADER_CLIENT_CAPABILITIES] == "PARAMETRIC_DATETIME"
         assert headers[accept_encoding_header] == accept_encoding_value
         assert headers[client_info_header] == client_info_value
-        assert len(headers.keys()) == 10
+        assert headers[constants.HEADER_ROLE] == (
+            "hive=ALL,"
+            "system=" + urllib.parse.quote("ROLE{analyst}") + ","
+            "catalog1=NONE,"
+            "catalog2=" + urllib.parse.quote("ROLE{catalog2_role}")
+        )
+        assert len(headers.keys()) == 11
 
     req.post("URL")
     _, post_kwargs = post.call_args
@@ -1095,14 +1109,15 @@ def test_request_headers_role_admin(mock_get_and_post):
             roles={"system": "admin"}
         ),
     )
+    roles = "system=" + urllib.parse.quote("ROLE{admin}")
 
     req.post("URL")
     _, post_kwargs = post.call_args
-    assert_headers_with_roles(post_kwargs["headers"], "system=admin")
+    assert_headers_with_roles(post_kwargs["headers"], roles)
 
     req.get("URL")
     _, get_kwargs = get.call_args
-    assert_headers_with_roles(post_kwargs["headers"], "system=admin")
+    assert_headers_with_roles(post_kwargs["headers"], roles)
 
 
 def test_request_headers_role_empty(mock_get_and_post):