Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support setting roles in dbapi and sqlalchemy #212

Merged
merged 4 commits into from
Oct 4, 2022

Conversation

mdesmet
Copy link
Contributor

@mdesmet mdesmet commented Aug 2, 2022

fixes #230

@cla-bot cla-bot bot added the cla-signed label Aug 2, 2022
@mdesmet mdesmet requested a review from hovaesco August 2, 2022 07:27
@mdesmet mdesmet marked this pull request as draft August 12, 2022 11:49
@mdesmet mdesmet marked this pull request as ready for review September 16, 2022 14:08
@mdesmet mdesmet changed the title Support setting role in dbapi and sqlalchemy Support setting roles in dbapi and sqlalchemy Sep 19, 2022
README.md Outdated Show resolved Hide resolved
tests/unit/sqlalchemy/test_dialect.py Outdated Show resolved Hide resolved
tests/integration/test_dbapi_integration.py Outdated Show resolved Hide resolved
trino/client.py Outdated Show resolved Hide resolved
@hovaesco
Copy link
Member

Please rebase on master to resolve CI failures.

@mdesmet mdesmet force-pushed the feature/role branch 2 times, most recently from 760f1a7 to 9b9c848 Compare September 23, 2022 08:50
@mdesmet
Copy link
Contributor Author

mdesmet commented Sep 23, 2022

@hashar: This would enable us to support grants in dbt-trino.

trino/client.py Outdated
def _create_roles(self, roles: Optional[Dict[str, str]]) -> Optional[str]:
if roles is not None:
return ";".join(map(lambda role: f"{role[0]}={role[1]}", roles.items()))
return ""
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not return None? Is the header required even if no role is set?

README.md Outdated
@@ -321,6 +321,20 @@ cur.execute('SELECT * FROM system.runtime.nodes')
rows = cur.fetchall()
```

## Roles

Authorization roles to use for catalogs, specified as a dict with key-value pairs for the catalog and role. For example, `{"catalog1": "roleA", "catalog2": "roleB"}` sets roleA for catalog1 and roleB for catalog2.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Authorization roles to use for catalogs, specified as a dict with key-value pairs for the catalog and role. For example, `{"catalog1": "roleA", "catalog2": "roleB"}` sets roleA for catalog1 and roleB for catalog2.
Authorization roles to use for catalogs, specified as a dict with key-value pairs for the catalog and role. For example, `{"catalog1": "roleA", "catalog2": "roleB"}` sets `roleA` for `catalog1` and `roleB` for `catalog2`.

Also might be useful to link to Trino docs about roles since AFAIK there is also something called a system role.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There seems no specific page about roles, only about the SQL syntax. I think a proper page about how roles work in Trino will be definitely beneficial.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc: @mosabua @dain regarding some discussion of what roles mean - we can just explain the concept because AFAIK there is no plugin providing roles in Trino - just the SPI.

tests/unit/test_dbapi.py Outdated Show resolved Hide resolved
@@ -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"}'
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no test for sqlalchemy integration?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mdesmet
Copy link
Contributor Author

mdesmet commented Sep 27, 2022

@hashhar: Please advise on how to get this merged. Note that we had roles already supported (within the session) and they were tested in the same way (now expanded up to http_headers) but not in integration tests as we don't have a connector that supports roles in this project.

As mentioned we need this feature in as we want to support roles security within dbt_trino.

@hashhar
Copy link
Member

hashhar commented Sep 28, 2022

I'm fine with current tests for now but as we start to work more on this area we'd eventually need better integration tests - let's create an issue with current ideas and we can follow up their.

@ebyhr WDYT?

@hashhar
Copy link
Member

hashhar commented Sep 28, 2022

@mdesmet Can you also update PR description with the template + propose a release note?

@@ -19,6 +19,7 @@

import trino
from tests.integration.conftest import trino_version
from trino import constants
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Refactor role into roles -> Add support for system and connector roles?

Does it make sense (and is it accurate?)?

Copy link
Contributor Author

@mdesmet mdesmet Sep 28, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@hashhar: It is really just a rename of role into roles. There is no new feature being introduced in that commit actually. The role field actually already contained multiple roles. I thought this is more consistent with the roles being passed on from dbapi.connect, introduced in the next commit.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's the reason why the first commit by itself fails tests. See

======================================================================================================================================= FAILURES ========================================================================================================================================
_____________________________________________________________________________________________________________________ test_set_role_in_connection_trino_higher_351 ______________________________________________________________________________________________________________________

run_trino = (<Popen: returncode: None args: ['docker', 'run', '--rm', '-p', '56257:8080',...>, 'localhost', 56257)

    @pytest.mark.skipif(trino_version() == '351', reason="Newer Trino versions return the system role")
    def test_set_role_in_connection_trino_higher_351(run_trino):
        _, host, port = run_trino

>       trino_connection = trino.dbapi.Connection(
            host=host, port=port, user="test", catalog="tpch", roles={"system": "ALL"}
        )
E       TypeError: Connection.__init__() got an unexpected keyword argument 'roles'

tests/integration/test_dbapi_integration.py:1077: TypeError
------------------------------------------------------------------------------------------------------------------------------- Captured stdout teardown --------------------------------------------------------------------------------------------------------------------------------
trino-python-client-tests-637e3b1
____________________________________________________________________________________________________________________________ test_role_is_set_when_specified ____________________________________________________________________________________________________________________________

mock_client = <MagicMock name='client' id='4532039360'>

    @patch("trino.dbapi.trino.client")
    def test_role_is_set_when_specified(mock_client):
        roles = {"system": "finance"}
>       with connect("sample_trino_cluster:443", roles=roles) as conn:

tests/unit/test_dbapi.py:262:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

args = ('sample_trino_cluster:443',), kwargs = {'roles': {'system': 'finance'}}

    def connect(*args, **kwargs):
        """Constructor for creating a connection to the database.

        See class :py:class:`Connection` for arguments.

        :returns: a :py:class:`Connection` object.
        """
>       return Connection(*args, **kwargs)
E       TypeError: Connection.__init__() got an unexpected keyword argument 'roles'

trino/dbapi.py:82: TypeError
================================================================================================================================ short test summary info ================================================================================================================================
FAILED tests/integration/test_dbapi_integration.py::test_set_role_in_connection_trino_higher_351 - TypeError: Connection.__init__() got an unexpected keyword argument 'roles'
FAILED tests/unit/test_dbapi.py::test_role_is_set_when_specified - TypeError: Connection.__init__() got an unexpected keyword argument 'roles'

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I probably messed up something in my fixup comments locally. All good now.

Copy link
Member

@hashhar hashhar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM % comments.

@mdesmet
Copy link
Contributor Author

mdesmet commented Sep 29, 2022

I'm fine with current tests for now but as we start to work more on this area we'd eventually need better integration tests - let's create an issue with current ideas and we can follow up their.

Created #241

hashhar
hashhar previously approved these changes Sep 30, 2022
@hashhar
Copy link
Member

hashhar commented Sep 30, 2022

@mdesmet Can you rebase, I see some failures when running tests locally. Possibly something changed on master in the meantime.

@hashhar
Copy link
Member

hashhar commented Sep 30, 2022

Seems like some changes in post-397 affect query cancellations causing "most" of the test failures. The last successful run was against 397.

@mdesmet
Copy link
Contributor Author

mdesmet commented Sep 30, 2022

@hashhar : Unfortunately something fishy is going on: 351 and 395 pass. However 398 does fail consistently (also locally), while 397 works fine.

I checked the release notes but didn't immediately find something. I'll keep digging.

@hashhar
Copy link
Member

hashhar commented Sep 30, 2022

See https://github.com/trinodb/trino/pull/13907/files#r984659618, a "bug" has been fixed - we now don't get result sets for non-SELECT queries like it should've been all along.

@mdesmet
Copy link
Contributor Author

mdesmet commented Oct 3, 2022

@hashhar: I've realigned the implementation on the session_properties and your changes in #242

It would be nice to get this one in as well.

trino/client.py Outdated
def get_roles_values(headers, header):
kvs = get_header_values(headers, header)
return [
(k.strip(), urllib.parse.unquote(v.strip()))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unquote_plus

@hashhar hashhar merged commit b4bd746 into trinodb:master Oct 4, 2022
@hashhar hashhar mentioned this pull request Oct 4, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Development

Successfully merging this pull request may close these issues.

It's not possible to set role in trino.dbapi.connect method
3 participants