diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index 31198a617..d0a213aa4 100644 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -89,8 +89,10 @@ 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["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"), ) return get_auth_provider(cfg) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 53b0c9715..b339ba8d7 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -60,6 +60,13 @@ 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 built-in client_id of databricks-sql-python. + + 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`. When `auth_type` is set to `databricks-oauth` without persisting the oauth token in a persistence storage