Repository: trinodb/trino-python-client

Issue:
Add integration test for SQLAlchemy

Discussion:

Patch:
diff --git a/trino/client.py b/trino/client.py
index 3dc75b7..19808e6 100644
--- a/trino/client.py
+++ b/trino/client.py
@@ -98,6 +98,8 @@ class ClientSession(object):
     :param extra_credential: extra credentials. as list of ``(key, value)``
                              tuples.
     :param client_tags: Client tags as list of strings.
+    :param role: role for the current session. Some connectors do not
+                 support role management. See connector documentation for more details.
     """
 
     def __init__(
@@ -111,6 +113,7 @@ class ClientSession(object):
         transaction_id: str = None,
         extra_credential: List[Tuple[str, str]] = None,
         client_tags: List[str] = None,
+        role: str = None,
     ):
         self._user = user
         self._catalog = catalog
@@ -121,6 +124,7 @@ class ClientSession(object):
         self._transaction_id = transaction_id
         self._extra_credential = extra_credential
         self._client_tags = client_tags
+        self._role = role
         self._object_lock = threading.Lock()
 
     @property
@@ -192,6 +196,16 @@ class ClientSession(object):
         self.__dict__.update(state)
         self._object_lock = threading.Lock()
 
+    @property
+    def role(self):
+        with self._object_lock:
+            return self._role
+
+    @role.setter
+    def role(self, role):
+        with self._object_lock:
+            self._role = role
+
 
 def get_header_values(headers, header):
     return [val.strip() for val in headers[header].split(",")]
@@ -368,6 +382,7 @@ class TrinoRequest(object):
         headers[constants.HEADER_SCHEMA] = self._client_session.schema
         headers[constants.HEADER_SOURCE] = self._client_session.source
         headers[constants.HEADER_USER] = self._client_session.user
+        headers[constants.HEADER_ROLE] = self._client_session.role
         if self._client_session.client_tags is not None and len(self._client_session.client_tags) > 0:
             headers[constants.HEADER_CLIENT_TAGS] = ",".join(self._client_session.client_tags)
 
@@ -538,6 +553,9 @@ class TrinoRequest(object):
         if constants.HEADER_SET_SCHEMA in http_response.headers:
             self._client_session.schema = http_response.headers[constants.HEADER_SET_SCHEMA]
 
+        if constants.HEADER_SET_ROLE in http_response.headers:
+            self._client_session.role = http_response.headers[constants.HEADER_SET_ROLE]
+
         self._next_uri = response.get("nextUri")
 
         return TrinoStatus(
diff --git a/trino/constants.py b/trino/constants.py
index d225fe3..3004690 100644
--- a/trino/constants.py
+++ b/trino/constants.py
@@ -39,6 +39,9 @@ HEADER_SESSION = "X-Trino-Session"
 HEADER_SET_SESSION = "X-Trino-Set-Session"
 HEADER_CLEAR_SESSION = "X-Trino-Clear-Session"
 
+HEADER_ROLE = "X-Trino-Role"
+HEADER_SET_ROLE = "X-Trino-Set-Role"
+
 HEADER_STARTED_TRANSACTION = "X-Trino-Started-Transaction-Id"
 HEADER_TRANSACTION = "X-Trino-Transaction-Id"
 
diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py
index 50c3a0b..c7056cc 100644
--- a/trino/sqlalchemy/dialect.py
+++ b/trino/sqlalchemy/dialect.py
@@ -22,7 +22,6 @@ from sqlalchemy.engine.url import URL
 from trino import dbapi as trino_dbapi, logging
 from trino.auth import BasicAuthentication, CertificateAuthentication, JWTAuthentication
 from trino.dbapi import Cursor
-from trino.exceptions import TrinoUserError
 from trino.sqlalchemy import compiler, datatype, error
 
 logger = logging.get_logger(__name__)
@@ -320,8 +319,8 @@ class TrinoDialect(DefaultDialect):
             res = connection.execute(sql.text(query))
             version = res.scalar()
             return tuple([version])
-        except TrinoUserError as e:
-            logger.debug(f"Failed to get server version: {e.message}")
+        except exc.ProgrammingError as e:
+            logger.debug(f"Failed to get server version: {e.orig.message}")
             return None
 
     def _get_default_schema_name(self, connection: Connection) -> Optional[str]:


Test patch:
diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py
index ab5b1f4..c5d0dda 100644
--- a/tests/integration/test_dbapi_integration.py
+++ b/tests/integration/test_dbapi_integration.py
@@ -1034,3 +1034,37 @@ def test_use_catalog(run_trino):
     result = cur.fetchall()
     assert result[0][0] == 'tpch'
     assert result[0][1] == 'sf1'
+
+
+@pytest.mark.skipif(trino_version() == '351', reason="Newer Trino versions return the system role")
+def test_set_role_trino_higher_351(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SHOW TABLES FROM information_schema')
+    cur.fetchall()
+    assert cur._request._client_session.role is None
+
+    cur.execute("SET ROLE ALL")
+    cur.fetchall()
+    assert cur._request._client_session.role == "system=ALL"
+
+
+@pytest.mark.skipif(trino_version() != '351', reason="Trino 351 returns the role for the current catalog")
+def test_set_role_trino_351(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SHOW TABLES FROM information_schema')
+    cur.fetchall()
+    assert cur._request._client_session.role is None
+
+    cur.execute("SET ROLE ALL")
+    cur.fetchall()
+    assert cur._request._client_session.role == "tpch=ALL"
diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py
new file mode 100644
index 0000000..1dc8f05
--- /dev/null
+++ b/tests/integration/test_sqlalchemy_integration.py
@@ -0,0 +1,257 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+import pytest
+import sqlalchemy as sqla
+from sqlalchemy.sql import and_, or_, not_
+
+
+@pytest.fixture
+def trino_connection(run_trino, request):
+    _, host, port = run_trino
+    engine = sqla.create_engine(f"trino://test@{host}:{port}/{request.param}",
+                                connect_args={"source": "test", "max_attempts": 1})
+    yield engine, engine.connect()
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_select_query(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    nations = sqla.Table('nation', metadata, schema='tiny', autoload_with=conn)
+    assert_column(nations, "nationkey", sqla.sql.sqltypes.BigInteger)
+    assert_column(nations, "name", sqla.sql.sqltypes.String)
+    assert_column(nations, "regionkey", sqla.sql.sqltypes.BigInteger)
+    assert_column(nations, "comment", sqla.sql.sqltypes.String)
+    query = sqla.select(nations)
+    result = conn.execute(query)
+    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)
+
+
+def assert_column(table, column_name, column_type):
+    assert getattr(table.c, column_name).name == column_name
+    assert isinstance(getattr(table.c, column_name).type, column_type)
+
+
+@pytest.mark.parametrize('trino_connection', ['system'], indirect=True)
+def test_select_specific_columns(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    nodes = sqla.Table('nodes', metadata, schema='runtime', autoload_with=conn)
+    assert_column(nodes, "node_id", sqla.sql.sqltypes.String)
+    assert_column(nodes, "state", sqla.sql.sqltypes.String)
+    query = sqla.select(nodes.c.node_id, nodes.c.state)
+    result = conn.execute(query)
+    rows = result.fetchall()
+    assert len(rows) > 0
+    for row in rows:
+        assert isinstance(row['node_id'], str)
+        assert isinstance(row['state'], str)
+
+
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_define_and_create_table(trino_connection):
+    engine, conn = trino_connection
+    if not engine.dialect.has_schema(engine, "test"):
+        engine.execute(sqla.schema.CreateSchema("test"))
+    metadata = sqla.MetaData()
+    try:
+        sqla.Table('users',
+                   metadata,
+                   sqla.Column('id', sqla.Integer),
+                   sqla.Column('name', sqla.String),
+                   sqla.Column('fullname', sqla.String),
+                   schema="test")
+        metadata.create_all(engine)
+        assert sqla.inspect(engine).has_table('users', schema="test")
+        users = sqla.Table('users', metadata, schema='test', autoload_with=conn)
+        assert_column(users, "id", sqla.sql.sqltypes.Integer)
+        assert_column(users, "name", sqla.sql.sqltypes.String)
+        assert_column(users, "fullname", sqla.sql.sqltypes.String)
+    finally:
+        metadata.drop_all(engine)
+
+
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_insert(trino_connection):
+    engine, conn = trino_connection
+
+    if not engine.dialect.has_schema(engine, "test"):
+        engine.execute(sqla.schema.CreateSchema("test"))
+    metadata = sqla.MetaData()
+    try:
+        users = sqla.Table('users',
+                           metadata,
+                           sqla.Column('id', sqla.Integer),
+                           sqla.Column('name', sqla.String),
+                           sqla.Column('fullname', sqla.String),
+                           schema="test")
+        metadata.create_all(engine)
+        ins = users.insert()
+        conn.execute(ins, {"id": 2, "name": "wendy", "fullname": "Wendy Williams"})
+        query = sqla.select(users)
+        result = conn.execute(query)
+        rows = result.fetchall()
+        assert len(rows) == 1
+        assert rows[0] == (2, "wendy", "Wendy Williams")
+    finally:
+        metadata.drop_all(engine)
+
+
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_insert_multiple_statements(trino_connection):
+    engine, conn = trino_connection
+    if not engine.dialect.has_schema(engine, "test"):
+        engine.execute(sqla.schema.CreateSchema("test"))
+    metadata = sqla.MetaData()
+    users = sqla.Table('users',
+                       metadata,
+                       sqla.Column('id', sqla.Integer),
+                       sqla.Column('name', sqla.String),
+                       sqla.Column('fullname', sqla.String),
+                       schema="test")
+    metadata.create_all(engine)
+    ins = users.insert()
+    conn.execute(ins, [
+        {"id": 2, "name": "wendy", "fullname": "Wendy Williams"},
+        {"id": 3, "name": "john", "fullname": "John Doe"},
+        {"id": 4, "name": "mary", "fullname": "Mary Hopkins"},
+    ])
+    query = sqla.select(users)
+    result = conn.execute(query)
+    rows = result.fetchall()
+    assert len(rows) == 3
+    assert frozenset(rows) == frozenset([
+        (2, "wendy", "Wendy Williams"),
+        (3, "john", "John Doe"),
+        (4, "mary", "Mary Hopkins"),
+    ])
+    metadata.drop_all(engine)
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_operators(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    customers = sqla.Table('nation', metadata, schema='tiny', autoload_with=conn)
+    query = sqla.select(customers).where(customers.c.nationkey == 2)
+    result = conn.execute(query)
+    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)
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_conjunctions(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    customers = sqla.Table('customer', metadata, schema='tiny', autoload_with=conn)
+    query = sqla.select(customers).where(and_(
+        customers.c.name.like('%12%'),
+        customers.c.nationkey == 15,
+        or_(
+            customers.c.mktsegment == 'AUTOMOBILE',
+            customers.c.mktsegment == 'HOUSEHOLD'
+        ),
+        not_(customers.c.acctbal < 0)))
+    result = conn.execute(query)
+    rows = result.fetchall()
+    assert len(rows) == 1
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_textual_sql(trino_connection):
+    _, conn = trino_connection
+    s = sqla.text("SELECT * from tiny.customer where nationkey = :e1 AND acctbal < :e2")
+    result = conn.execute(s, {"e1": 15, "e2": 0})
+    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)
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_alias(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    nations = sqla.Table('nation', metadata, schema='tiny', autoload_with=conn)
+    nations1 = nations.alias("o1")
+    nations2 = nations.alias("o2")
+    s = sqla.select(nations1) \
+        .join(nations2, and_(
+            nations1.c.regionkey == nations2.c.regionkey,
+            nations1.c.nationkey != nations2.c.nationkey,
+            nations1.c.regionkey == 1
+        )) \
+        .distinct()
+    result = conn.execute(s)
+    rows = result.fetchall()
+    assert len(rows) == 5
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_subquery(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    nations = sqla.Table('nation', metadata, schema='tiny', autoload_with=conn)
+    customers = sqla.Table('customer', metadata, schema='tiny', autoload_with=conn)
+    automobile_customers = sqla.select(customers.c.nationkey).where(customers.c.acctbal < -900)
+    automobile_customers_subquery = automobile_customers.subquery()
+    s = sqla.select(nations.c.name).where(nations.c.nationkey.in_(sqla.select(automobile_customers_subquery)))
+    result = conn.execute(s)
+    rows = result.fetchall()
+    assert len(rows) == 15
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_joins(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    nations = sqla.Table('nation', metadata, schema='tiny', autoload_with=conn)
+    customers = sqla.Table('customer', metadata, schema='tiny', autoload_with=conn)
+    s = sqla.select(nations.c.name) \
+        .select_from(nations.join(customers, nations.c.nationkey == customers.c.nationkey)) \
+        .where(customers.c.acctbal < -900) \
+        .distinct()
+    result = conn.execute(s)
+    rows = result.fetchall()
+    assert len(rows) == 15
+
+
+@pytest.mark.parametrize('trino_connection', ['tpch'], indirect=True)
+def test_cte(trino_connection):
+    _, conn = trino_connection
+    metadata = sqla.MetaData()
+    nations = sqla.Table('nation', metadata, schema='tiny', autoload_with=conn)
+    customers = sqla.Table('customer', metadata, schema='tiny', autoload_with=conn)
+    automobile_customers = sqla.select(customers.c.nationkey).where(customers.c.acctbal < -900)
+    automobile_customers_cte = automobile_customers.cte()
+    s = sqla.select(nations).where(nations.c.nationkey.in_(sqla.select(automobile_customers_cte)))
+    result = conn.execute(s)
+    rows = result.fetchall()
+    assert len(rows) == 15
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 2263081..1208930 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -92,9 +92,10 @@ def test_request_headers(mock_get_and_post):
         assert headers[constants.HEADER_SOURCE] == source
         assert headers[constants.HEADER_USER] == user
         assert headers[constants.HEADER_SESSION] == ""
+        assert headers[constants.HEADER_ROLE] is None
         assert headers[accept_encoding_header] == accept_encoding_value
         assert headers[client_info_header] == client_info_value
-        assert len(headers.keys()) == 8
+        assert len(headers.keys()) == 9
 
     req.post("URL")
     _, post_kwargs = post.call_args
@@ -1002,3 +1003,71 @@ def test_retry_with():
     with_retry(FailerUntil(2).__call__)()
     with pytest.raises(SomeException):
         with_retry(FailerUntil(3).__call__)()
+
+
+def assert_headers_with_role(headers, role):
+    assert headers[constants.HEADER_USER] == "test_user"
+    assert headers[constants.HEADER_ROLE] == role
+    assert len(headers.keys()) == 7
+
+
+def test_request_headers_role_hive_all(mock_get_and_post):
+    get, post = mock_get_and_post
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            role="hive=ALL",
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_with_role(post_kwargs["headers"], "hive=ALL")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_with_role(post_kwargs["headers"], "hive=ALL")
+
+
+def test_request_headers_role_admin(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            role="admin",
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_with_role(post_kwargs["headers"], "admin")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_with_role(post_kwargs["headers"], "admin")
+
+
+def test_request_headers_role_empty(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            role="",
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_with_role(post_kwargs["headers"], "")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_with_role(post_kwargs["headers"], "")