Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 5 additions & 7 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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"})
Expand All @@ -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
Expand Down
3 changes: 1 addition & 2 deletions examples/interactive_oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
15 changes: 14 additions & 1 deletion src/databricks/sql/auth/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
2 changes: 1 addition & 1 deletion src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.

Expand Down
10 changes: 9 additions & 1 deletion tests/unit/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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)

Loading