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

feat(sqlalchemy): Add Support for externalAuthentication #344

Merged
merged 1 commit into from
Sep 22, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion tests/unit/sqlalchemy/test_dialect.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import pytest
from sqlalchemy.engine.url import URL, make_url

from trino.auth import BasicAuthentication
from trino.auth import BasicAuthentication, OAuth2Authentication
from trino.dbapi import Connection
from trino.sqlalchemy import URL as trino_url
from trino.sqlalchemy.dialect import (
Expand Down Expand Up @@ -296,3 +296,12 @@ def test_trino_connection_certificate_auth():
assert isinstance(cparams['auth'], CertificateAuthentication)
assert cparams['auth']._cert == cert
assert cparams['auth']._key == key


def test_trino_connection_oauth2_auth():
dialect = TrinoDialect()
url = make_url('trino://host/?externalAuthentication=true')
_, cparams = dialect.create_connect_args(url)

assert cparams['http_scheme'] == "https"
assert isinstance(cparams['auth'], OAuth2Authentication)
11 changes: 10 additions & 1 deletion trino/sqlalchemy/dialect.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,12 @@

from trino import dbapi as trino_dbapi
from trino import logging
from trino.auth import BasicAuthentication, CertificateAuthentication, JWTAuthentication
from trino.auth import (
BasicAuthentication,
CertificateAuthentication,
JWTAuthentication,
OAuth2Authentication,
)
from trino.dbapi import Cursor
from trino.sqlalchemy import compiler, datatype, error

Expand Down Expand Up @@ -113,6 +118,10 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any
kwargs["http_scheme"] = "https"
kwargs["auth"] = CertificateAuthentication(unquote_plus(url.query['cert']), unquote_plus(url.query['key']))

if "externalAuthentication" in url.query:
Copy link
Member

Choose a reason for hiding this comment

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

I understand that it's aligning with JDBC driver, but you can simply pass "auth": OAuth2Authentication() in connect_args.

Copy link
Author

Choose a reason for hiding this comment

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

That's what I've been doing when using the client directly. But I've been running into an issue passing it as a class when trying to use it from some other projects such as ipython-sql.
Eventually I found out that ipython-sql allows you to actually pass the -creator parameter to bypass the configuration by sending a connection straight away but it'd be great if it would work without this workaround.

Copy link
Member

Choose a reason for hiding this comment

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

Do we have good behaviour when someone has externalAuthentication in URI and has explicitly set a different auth mechanism? What takes precedence? I'd like to avoid this situation of having multiple ways of configuring something since it leads to more combination of things to test and edge cases to think about.

Copy link
Author

Choose a reason for hiding this comment

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

Hey @hashhar, I have similar feelings to be honest. I just extended the way it's being done for SQLAlchemy URIs (It's already done this way for Basic Authentication, JWT Authentication and Cert Authentication).

I think it makes sense to align with the JDBC driver and have this for SQLAlchemy... If the user passes multiple Auth Methods I think it's on the user side to be honest, but I guess it could be improved by validating if they did this somehow (checking the query args and connect_args)

Copy link
Member

Choose a reason for hiding this comment

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

This validation seems useful but out of scope of this PR because it needs to apply to other auth methods as outlined above.

Copy link
Author

Choose a reason for hiding this comment

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

You mean validating if the user passes multiple auth methods?

So you think it'd be fine to keep that responsibility on the user side so far? I'd be fine with it (=

kwargs["http_scheme"] = "https"
hashhar marked this conversation as resolved.
Show resolved Hide resolved
kwargs["auth"] = OAuth2Authentication()

if "source" in url.query:
kwargs["source"] = unquote_plus(url.query["source"])
else:
Expand Down