Skip to content

Commit 2d2b3c1

Browse files
authored
[PECO-1801] Make OAuth as the default authenticator if no authentication setting is provided (#419)
* [PECO-1801] Make OAuth as the default authenticator if no authentication setting is provided Signed-off-by: Jacky Hu <[email protected]>
1 parent 8bbcd3e commit 2d2b3c1

File tree

5 files changed

+30
-12
lines changed

5 files changed

+30
-12
lines changed

README.md

+5-7
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,9 @@ For the latest documentation, see
2424

2525
Install the library with `pip install databricks-sql-connector`
2626

27-
Note: Don't hard-code authentication secrets into your Python. Use environment variables
28-
2927
```bash
3028
export DATABRICKS_HOST=********.databricks.com
3129
export DATABRICKS_HTTP_PATH=/sql/1.0/endpoints/****************
32-
export DATABRICKS_TOKEN=dapi********************************
3330
```
3431

3532
Example usage:
@@ -39,12 +36,10 @@ from databricks import sql
3936

4037
host = os.getenv("DATABRICKS_HOST")
4138
http_path = os.getenv("DATABRICKS_HTTP_PATH")
42-
access_token = os.getenv("DATABRICKS_TOKEN")
4339

4440
connection = sql.connect(
4541
server_hostname=host,
46-
http_path=http_path,
47-
access_token=access_token)
42+
http_path=http_path)
4843

4944
cursor = connection.cursor()
5045
cursor.execute('SELECT :param `p`, * FROM RANGE(10)', {"param": "foo"})
@@ -60,7 +55,10 @@ In the above example:
6055
- `server-hostname` is the Databricks instance host name.
6156
- `http-path` is the HTTP Path either to a Databricks SQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef),
6257
or to a Databricks Runtime interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)
63-
- `personal-access-token` is the Databricks Personal Access Token for the account that will execute commands and queries
58+
59+
> Note: This example uses [Databricks OAuth U2M](https://docs.databricks.com/en/dev-tools/auth/oauth-u2m.html)
60+
> to authenticate the target Databricks user account and needs to open the browser for authentication. So it
61+
> can only run on the user's machine.
6462
6563

6664
## Contributing

examples/interactive_oauth.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,7 @@
1414
"""
1515

1616
with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
17-
http_path = os.getenv("DATABRICKS_HTTP_PATH"),
18-
auth_type="databricks-oauth") as connection:
17+
http_path = os.getenv("DATABRICKS_HTTP_PATH")) as connection:
1918

2019
for x in range(1, 100):
2120
cursor = connection.cursor()

src/databricks/sql/auth/auth.py

+14-1
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,20 @@ def get_auth_provider(cfg: ClientContext):
6464
# no op authenticator. authentication is performed using ssl certificate outside of headers
6565
return AuthProvider()
6666
else:
67-
raise RuntimeError("No valid authentication settings!")
67+
if (
68+
cfg.oauth_redirect_port_range is not None
69+
and cfg.oauth_client_id is not None
70+
and cfg.oauth_scopes is not None
71+
):
72+
return DatabricksOAuthProvider(
73+
cfg.hostname,
74+
cfg.oauth_persistence,
75+
cfg.oauth_redirect_port_range,
76+
cfg.oauth_client_id,
77+
cfg.oauth_scopes,
78+
)
79+
else:
80+
raise RuntimeError("No valid authentication settings!")
6881

6982

7083
PYSQL_OAUTH_SCOPES = ["sql", "offline_access"]

src/databricks/sql/client.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ def __init__(
9696
sanitise parameterized inputs to prevent SQL injection. The inline parameter approach is maintained for
9797
legacy purposes and will be deprecated in a future release. When this parameter is `True` you will see
9898
a warning log message. To suppress this log message, set `use_inline_params="silent"`.
99-
auth_type: `str`, optional
99+
auth_type: `str`, optional (default is databricks-oauth if neither `access_token` nor `tls_client_cert_file` is set)
100100
`databricks-oauth` : to use Databricks OAuth with fine-grained permission scopes, set to `databricks-oauth`.
101101
`azure-oauth` : to use Microsoft Entra ID OAuth flow, set to `azure-oauth`.
102102

tests/unit/test_auth.py

+9-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
ExternalAuthProvider,
1010
AuthType,
1111
)
12-
from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
12+
from databricks.sql.auth.auth import get_python_sql_connector_auth_provider, PYSQL_OAUTH_CLIENT_ID
1313
from databricks.sql.auth.oauth import OAuthManager
1414
from databricks.sql.auth.authenticators import DatabricksOAuthProvider
1515
from databricks.sql.auth.endpoint import (
@@ -178,3 +178,11 @@ def test_get_python_sql_connector_basic_auth(self):
178178
with self.assertRaises(ValueError) as e:
179179
get_python_sql_connector_auth_provider("foo.cloud.databricks.com", **kwargs)
180180
self.assertIn("Username/password authentication is no longer supported", str(e.exception))
181+
182+
@patch.object(DatabricksOAuthProvider, "_initial_get_token")
183+
def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):
184+
hostname = "foo.cloud.databricks.com"
185+
auth_provider = get_python_sql_connector_auth_provider(hostname)
186+
self.assertTrue(type(auth_provider).__name__, "DatabricksOAuthProvider")
187+
self.assertTrue(auth_provider._client_id,PYSQL_OAUTH_CLIENT_ID)
188+

0 commit comments

Comments
 (0)