[PECO-1801] Make OAuth as the default authenticator if no authentication setting is provided (#419)

jackyhu-db · web-flow · commit 2d2b3c1c4422 · 2024-08-01T11:38:15.000-07:00
* [PECO-1801] Make OAuth as the default authenticator if no authentication setting is provided

Signed-off-by: Jacky Hu &lt;jacky.hu@databricks.com&gt;
diff --git a/README.md b/README.md
@@ -24,12 +24,9 @@ For the latest documentation, see
 
 Install the library with `pip install databricks-sql-connector`
 
-Note: Don't hard-code authentication secrets into your Python. Use environment variables
-
 ```bash
 export DATABRICKS_HOST=********.databricks.com
 export DATABRICKS_HTTP_PATH=/sql/1.0/endpoints/****************
-export DATABRICKS_TOKEN=dapi********************************
 ```
 
 Example usage:
@@ -39,12 +36,10 @@ from databricks import sql
 
 host = os.getenv("DATABRICKS_HOST")
 http_path = os.getenv("DATABRICKS_HTTP_PATH")
-access_token = os.getenv("DATABRICKS_TOKEN")
 
 connection = sql.connect(
   server_hostname=host,
-  http_path=http_path,
-  access_token=access_token)
+  http_path=http_path)
 
 cursor = connection.cursor()
 cursor.execute('SELECT :param `p`, * FROM RANGE(10)', {"param": "foo"})
@@ -60,7 +55,10 @@ In the above example:
 - `server-hostname` is the Databricks instance host name.
 - `http-path` is the HTTP Path either to a Databricks SQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef),
 or to a Databricks Runtime interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)
-- `personal-access-token` is the Databricks Personal Access Token for the account that will execute commands and queries
+
+> Note: This example uses [Databricks OAuth U2M](https://docs.databricks.com/en/dev-tools/auth/oauth-u2m.html) 
+> to authenticate the target Databricks user account and needs to open the browser for authentication. So it 
+> can only run on the user's machine.
 
 
 ## Contributing
diff --git a/examples/interactive_oauth.py b/examples/interactive_oauth.py
@@ -14,8 +14,7 @@
 """
 
 with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
-                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
-                 auth_type="databricks-oauth") as connection:
+                 http_path       = os.getenv("DATABRICKS_HTTP_PATH")) as connection:
 
     for x in range(1, 100):
         cursor = connection.cursor()
diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py
@@ -64,7 +64,20 @@ def get_auth_provider(cfg: ClientContext):
         # no op authenticator. authentication is performed using ssl certificate outside of headers
         return AuthProvider()
     else:
-        raise RuntimeError("No valid authentication settings!")
+        if (
+            cfg.oauth_redirect_port_range is not None
+            and cfg.oauth_client_id is not None
+            and cfg.oauth_scopes is not None
+        ):
+            return DatabricksOAuthProvider(
+                cfg.hostname,
+                cfg.oauth_persistence,
+                cfg.oauth_redirect_port_range,
+                cfg.oauth_client_id,
+                cfg.oauth_scopes,
+            )
+        else:
+            raise RuntimeError("No valid authentication settings!")
 
 
 PYSQL_OAUTH_SCOPES = ["sql", "offline_access"]
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
@@ -96,7 +96,7 @@ def __init__(
                 sanitise parameterized inputs to prevent SQL injection.  The inline parameter approach is maintained for
                 legacy purposes and will be deprecated in a future release. When this parameter is `True` you will see
                 a warning log message. To suppress this log message, set `use_inline_params="silent"`.
-            auth_type: `str`, optional
+            auth_type: `str`, optional (default is databricks-oauth if neither `access_token` nor `tls_client_cert_file` is set)
                 `databricks-oauth` : to use Databricks OAuth with fine-grained permission scopes, set to `databricks-oauth`.
                 `azure-oauth` : to use Microsoft Entra ID OAuth flow, set to `azure-oauth`.
 
diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py
@@ -9,7 +9,7 @@
     ExternalAuthProvider,
     AuthType,
 )
-from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
+from databricks.sql.auth.auth import get_python_sql_connector_auth_provider, PYSQL_OAUTH_CLIENT_ID
 from databricks.sql.auth.oauth import OAuthManager
 from databricks.sql.auth.authenticators import DatabricksOAuthProvider
 from databricks.sql.auth.endpoint import (
@@ -178,3 +178,11 @@ def test_get_python_sql_connector_basic_auth(self):
         with self.assertRaises(ValueError) as e:
             get_python_sql_connector_auth_provider("foo.cloud.databricks.com", **kwargs)
         self.assertIn("Username/password authentication is no longer supported", str(e.exception))
+
+    @patch.object(DatabricksOAuthProvider, "_initial_get_token")
+    def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):
+        hostname = "foo.cloud.databricks.com"
+        auth_provider = get_python_sql_connector_auth_provider(hostname)
+        self.assertTrue(type(auth_provider).__name__, "DatabricksOAuthProvider")
+        self.assertTrue(auth_provider._client_id,PYSQL_OAUTH_CLIENT_ID)
+
