From d229466db20086833c9bca2ecd0a0b680626f3d0 Mon Sep 17 00:00:00 2001 From: RohitBankarRZP Date: Tue, 22 Aug 2023 23:03:01 +0530 Subject: [PATCH 1/3] Add Username + Password Authentication for Trino Client (fixes #829) --- querybook/server/lib/query_executor/clients/trino.py | 6 +++++- .../lib/query_executor/executor_template/templates.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/querybook/server/lib/query_executor/clients/trino.py b/querybook/server/lib/query_executor/clients/trino.py index 24a77d0e5..dd5927785 100644 --- a/querybook/server/lib/query_executor/clients/trino.py +++ b/querybook/server/lib/query_executor/clients/trino.py @@ -52,6 +52,7 @@ def __init__( self, connection_string: str, username: Optional[str] = None, + password: Optional[str] = None, proxy_user: Optional[str] = None, *args: Any, **kwargs: Any, @@ -61,12 +62,15 @@ def __init__( host = trino_conf.host port = 8080 if not trino_conf.port else trino_conf.port + auth = trino.auth.BasicAuthentication(username, password) + connection = trino.dbapi.connect( host=host, port=port, catalog=trino_conf.catalog, schema=trino_conf.schema, - user=proxy_user or username, + auth=auth, + user=proxy_user if proxy_user else username, http_scheme=trino_conf.protocol, ) self._connection = connection diff --git a/querybook/server/lib/query_executor/executor_template/templates.py b/querybook/server/lib/query_executor/executor_template/templates.py index 452d7f6ff..1b57b5964 100644 --- a/querybook/server/lib/query_executor/executor_template/templates.py +++ b/querybook/server/lib/query_executor/executor_template/templates.py @@ -73,7 +73,7 @@ ), ), ("username", FormField(regex="\\w+")), - ("impersonate", FormField(field_type=FormFieldType.Boolean)), + ("password", FormField(hidden=True)), ( "proxy_user_id", FormField( From 2010fe1e5691ffbc7e5fea1b3916785c48c8486b Mon Sep 17 00:00:00 2001 From: RohitBankarRZP Date: Wed, 23 Aug 2023 11:22:26 +0530 Subject: [PATCH 2/3] Making trino authentication backwards compatible --- .../lib/query_executor/clients/trino.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/querybook/server/lib/query_executor/clients/trino.py b/querybook/server/lib/query_executor/clients/trino.py index dd5927785..a6249e2b5 100644 --- a/querybook/server/lib/query_executor/clients/trino.py +++ b/querybook/server/lib/query_executor/clients/trino.py @@ -62,17 +62,27 @@ def __init__( host = trino_conf.host port = 8080 if not trino_conf.port else trino_conf.port - auth = trino.auth.BasicAuthentication(username, password) - - connection = trino.dbapi.connect( - host=host, - port=port, - catalog=trino_conf.catalog, - schema=trino_conf.schema, - auth=auth, - user=proxy_user if proxy_user else username, - http_scheme=trino_conf.protocol, - ) + if username is not None and password is not None: + auth = trino.auth.BasicAuthentication(username, password) + + connection = trino.dbapi.connect( + host=host, + port=port, + catalog=trino_conf.catalog, + schema=trino_conf.schema, + auth=auth, + user=proxy_user if proxy_user else username, + http_scheme=trino_conf.protocol, + ) + else: + connection = trino.dbapi.connect( + host=host, + port=port, + catalog=trino_conf.catalog, + schema=trino_conf.schema, + user=proxy_user or username, + http_scheme=trino_conf.protocol, + ) self._connection = connection super(TrinoClient, self).__init__() From 12fdbcdf7d0f62184134648d3f68e5332bfb2923 Mon Sep 17 00:00:00 2001 From: RohitBankarRZP Date: Wed, 23 Aug 2023 13:58:13 +0530 Subject: [PATCH 3/3] Removing duplicate connection code --- .../lib/query_executor/clients/trino.py | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/querybook/server/lib/query_executor/clients/trino.py b/querybook/server/lib/query_executor/clients/trino.py index a6249e2b5..35e9839d3 100644 --- a/querybook/server/lib/query_executor/clients/trino.py +++ b/querybook/server/lib/query_executor/clients/trino.py @@ -62,27 +62,19 @@ def __init__( host = trino_conf.host port = 8080 if not trino_conf.port else trino_conf.port + auth = trino.constants.DEFAULT_AUTH if username is not None and password is not None: auth = trino.auth.BasicAuthentication(username, password) - connection = trino.dbapi.connect( - host=host, - port=port, - catalog=trino_conf.catalog, - schema=trino_conf.schema, - auth=auth, - user=proxy_user if proxy_user else username, - http_scheme=trino_conf.protocol, - ) - else: - connection = trino.dbapi.connect( - host=host, - port=port, - catalog=trino_conf.catalog, - schema=trino_conf.schema, - user=proxy_user or username, - http_scheme=trino_conf.protocol, - ) + connection = trino.dbapi.connect( + host=host, + port=port, + catalog=trino_conf.catalog, + schema=trino_conf.schema, + auth=auth, + user=proxy_user if proxy_user else username, + http_scheme=trino_conf.protocol, + ) self._connection = connection super(TrinoClient, self).__init__()