trinodb / trino-python-client

Python client for Trino
Apache License 2.0
311 stars 154 forks source link

It's not possible to set role in trino.dbapi.connect method #230

Closed aksakalli closed 1 year ago

aksakalli commented 1 year ago

Describe the feature

With the recent changes with https://github.com/trinodb/trino-python-client/pull/200 X-Trino-Role and X-Trino-Set-Role are not allowed to be set. The following code:

connection = trino.dbapi.connect(
    host="trino.example.net",
    port=443,
    http_scheme="https",
    catalog="hive",
    user='abuzer',
    auth=trino.auth.JWTAuthentication('Xvfdas.....'),
    http_headers={'X-Trino-Role': 'hive=ROLE{admin}'}
)

gives an error ValueError: cannot override reserved HTTP header X-Trino-Role.

We should make the role parameter part of trino.dbapi.connect method.

Describe alternatives you've considered

I can to set role in connection._client_session.role or connection._http_session.headers. Since _client_session property should be an internal abstraction, connection._client_session.role='hive=ROLE{admin}' is ugly and prone to change. We need a proper interface to set such values in connect method.

We should keep release notes and publish critical breaking changes in future.

Are you willing to submit PR?

mdesmet commented 1 year ago

Hi @aksakalli,

There is already an open PR for this: See #212

Note that the current implementation is not tested with multiple catalogs. I tried adding a test:

@pytest.mark.skipif(trino_version() == '351', reason="Newer Trino versions return the system role")
def test_set_multiple_roles_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 == "system=ALL"

    cur.execute("SET ROLE ALL IN tpch")
    cur.fetchall()
    assert cur._request._client_session.role == "system=ALL;tpch=ALL"

However we need to add the hive connector as tpch or memory doesn't support role management.

You could also try to do cur.execute("SET ROLE admin in hive") as a first query.