Skip to content

Commit

Permalink
[PECO-1801] Make OAuth as the default authenticator if no authenticat…
Browse files Browse the repository at this point in the history
…ion 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]>
  • Loading branch information
jackyhu-db authored Aug 1, 2024
1 parent 8bbcd3e commit 2d2b3c1
Show file tree
Hide file tree
Showing 5 changed files with 30 additions and 12 deletions.
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)

0 comments on commit 2d2b3c1

Please sign in to comment.