databricks · jackyhu-db · Dec 29, 2022 · Dec 13, 2022 · Dec 29, 2022 · Dec 29, 2022
diff --git 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
@@ -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