From 8fa4415475cee0ef1ac6f792a7d82bcfecdaab84 Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Mon, 12 Dec 2022 22:25:53 -0800 Subject: [PATCH 1/3] Support custom oauth client id and rediret port range PySQL is used by other tools/CLIs which have own oauth client id, we need to expose oauth_client_id and oauth_redirect_port_range as the connection parameters to support this customization. Signed-off-by: Jacky Hu --- src/databricks/sql/auth/auth.py | 5 +++-- src/databricks/sql/client.py | 6 ++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index 31198a617..dcc5cb11f 100644 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -89,8 +89,9 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs): use_cert_as_auth=kwargs.get("_use_cert_as_auth"), tls_client_cert_file=kwargs.get("_tls_client_cert_file"), oauth_scopes=PYSQL_OAUTH_SCOPES, - oauth_client_id=PYSQL_OAUTH_CLIENT_ID, - oauth_redirect_port_range=PYSQL_OAUTH_REDIRECT_PORT_RANGE, + oauth_client_id=kwargs.get("oauth_client_id") or PYSQL_OAUTH_CLIENT_ID, + oauth_redirect_port_range=kwargs.get("oauth_redirect_port_range") if kwargs.get( + "oauth_client_id") else PYSQL_OAUTH_REDIRECT_PORT_RANGE, oauth_persistence=kwargs.get("experimental_oauth_persistence"), ) return get_auth_provider(cfg) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 53b0c9715..4c578ad84 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -60,6 +60,12 @@ def __init__( any IDP configured. This is only for interactive python applications and open a browser window. Note this is beta (private preview) + oauth_client_id: `str`, optional + custom oauth client_id. If not specified, it will use the client_id of databricks-sql-python + + oauth_redirect_port_range: `List[str]`, optional + oauth redirect port range. This is required when custom oauth client_id `oauth_client_id` is set + experimental_oauth_persistence: configures preferred storage for persisting oauth tokens. This has to be a class implementing `OAuthPersistence`. When `auth_type` is set to `databricks-oauth` without persisting the oauth token in a persistence storage From 472209c6f886cab1e9ff15cfcae6413a7da89822 Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Thu, 29 Dec 2022 12:23:31 -0800 Subject: [PATCH 2/3] Change oauth redirect port range to port Signed-off-by: Jacky Hu --- src/databricks/sql/auth/auth.py | 5 +++-- src/databricks/sql/client.py | 7 ++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index dcc5cb11f..29fe1a41d 100644 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -90,8 +90,9 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs): tls_client_cert_file=kwargs.get("_tls_client_cert_file"), oauth_scopes=PYSQL_OAUTH_SCOPES, oauth_client_id=kwargs.get("oauth_client_id") or PYSQL_OAUTH_CLIENT_ID, - oauth_redirect_port_range=kwargs.get("oauth_redirect_port_range") if kwargs.get( - "oauth_client_id") else PYSQL_OAUTH_REDIRECT_PORT_RANGE, + oauth_redirect_port_range=[kwargs.get("oauth_redirect_port")] + if kwargs.get("oauth_client_id") + else PYSQL_OAUTH_REDIRECT_PORT_RANGE, oauth_persistence=kwargs.get("experimental_oauth_persistence"), ) return get_auth_provider(cfg) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 4c578ad84..b339ba8d7 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -61,10 +61,11 @@ def __init__( Note this is beta (private preview) oauth_client_id: `str`, optional - custom oauth client_id. If not specified, it will use the client_id of databricks-sql-python + custom oauth client_id. If not specified, it will use the built-in client_id of databricks-sql-python. - oauth_redirect_port_range: `List[str]`, optional - oauth redirect port range. This is required when custom oauth client_id `oauth_client_id` is set + oauth_redirect_port: `int`, optional + port of the oauth redirect uri (localhost). This is required when custom oauth client_id + `oauth_client_id` is set experimental_oauth_persistence: configures preferred storage for persisting oauth tokens. This has to be a class implementing `OAuthPersistence`. From 55446652470c357676f4c0ce85850128d8096fbd Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Thu, 29 Dec 2022 12:47:45 -0800 Subject: [PATCH 3/3] Fix type check issue Signed-off-by: Jacky Hu --- src/databricks/sql/auth/auth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index 29fe1a41d..d0a213aa4 100644 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -90,8 +90,8 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs): tls_client_cert_file=kwargs.get("_tls_client_cert_file"), oauth_scopes=PYSQL_OAUTH_SCOPES, oauth_client_id=kwargs.get("oauth_client_id") or PYSQL_OAUTH_CLIENT_ID, - oauth_redirect_port_range=[kwargs.get("oauth_redirect_port")] - if kwargs.get("oauth_client_id") + oauth_redirect_port_range=[kwargs["oauth_redirect_port"]] + if kwargs.get("oauth_client_id") and kwargs.get("oauth_redirect_port") else PYSQL_OAUTH_REDIRECT_PORT_RANGE, oauth_persistence=kwargs.get("experimental_oauth_persistence"), )