Support custom oauth client id and redirect port (#75)

* 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 <[email protected]>

* Change oauth redirect port range to port

Signed-off-by: Jacky Hu <[email protected]>

* Fix type check issue

Signed-off-by: Jacky Hu <[email protected]>

Signed-off-by: Jacky Hu <[email protected]>

Author: jackyhu-db

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)

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