feat: Add client credential oauth integration support + related databricks helpers to SDK #348

zackverham · 2024-12-02T20:54:47Z

This PR adds a new function to the oauth tooling. It knows how to craft a request against the credential exchange endpoint that can interface with client credential flows. This supports service-account-like oauth interactions.

Connect sets an environment variable called CONNECT_CONTENT_SESSION_TOKEN which can be used by the executing content when making the request. This means that the syntax for using this new credential exchange function is slightly different.

This PR also adds the databricks extensions needed to use this new oauth exchange flow when interfacing with databricks resources.

This feature won't land in Connect until the upcoming release gets published - I don't know how we typically queue up work like this in sync with product releases.

github-actions · 2024-12-02T20:55:09Z

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
1950	1808	93%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
src/posit/connect/content.py	96%	🟢
src/posit/connect/external/databricks.py	91%	🟢
src/posit/connect/external/snowflake.py	91%	🟢
src/posit/connect/oauth/_init_.py	100%	🟢
src/posit/connect/oauth/oauth.py	100%	🟢
TOTAL	96%	🟢

updated for commit: 59862c4 by action🐍

zackverham · 2024-12-02T20:59:53Z

Here's a code snippet that illustrates the ergonomics of using this code in published content (this is pulled from a Quarto doc I have running locally):

from posit.connect.external.databricks import PositContentCredentialsStrategy

import requests

import pandas as pd
from databricks import sql
from databricks.sdk.core import ApiClient, Config, databricks_cli
from databricks.sdk.service.iam import CurrentUserAPI

# env vars
DATABRICKS_HOST = "<REDACTED>"
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
SQL_HTTP_PATH = "<REDACTED>"

posit_strategy = PositContentCredentialsStrategy(local_strategy=databricks_cli)

cfg = Config(host=DATABRICKS_HOST_URL, credentials_strategy=posit_strategy)

databricks_user_info = CurrentUserAPI(ApiClient(cfg)).me()
print(f"Hello, {databricks_user_info.display_name}!")

query = "SELECT * FROM samples.nyctaxi.trips LIMIT 10;"
with sql.connect(
    server_hostname=DATABRICKS_HOST,
    http_path=SQL_HTTP_PATH,
    credentials_provider=posit_strategy.sql_credentials_provider(cfg),
) as connection:
    with connection.cursor() as cursor:
        cursor.execute(query)
        rows = cursor.fetchall()
        print(pd.DataFrame([row.asDict() for row in rows]))

src/posit/connect/external/databricks.py

mconflitti-pbc · 2024-12-03T18:21:33Z

Odd ergonomics with the passing around of that config object! But read through the convo on the databricks repo and it seems to be the best option for now! LGTM

Co-authored-by: Barret Schloerke <[email protected]>

dbkegley · 2024-12-04T15:23:23Z

src/posit/connect/external/databricks.py

+    DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
+    SQL_HTTP_PATH = "<REDACTED>"
+
+    posit_strategy = PositContentCredentialsStrategy(local_strategy=databricks_cli)


The other credentialsStrategy docstring uses a sample py-shiny app so that the user can (hopefully) see that the location that the initialization happens is really important. That context is lost in this example. I think it would be worthwhile to show both and call out explicitly why the initialization is different for service accounts than for viewer auth

yeah, I guess the question ultimately is how much of this documentation should exist in the docstrings vs in our docs pages (both the cookbook examples themselves, and prose discussions explaining things).

I think this might be a question for @schloerke - I don't know exactly what expectations would be for SDK consumers.

I think docs strings should show usage examples.

I think the current Connect recipes are turning into small usage examples and should be moved to the doc strings. (And the recipes should be larger/longer goal-oriented stories combining many small usage examples.)

tl/dr; More usage examples in doc strings, please!

This example needs to be updated to use the PositLocalContentCredentialsStrategy

dbkegley · 2024-12-04T15:25:39Z

src/posit/connect/external/databricks.py

+    import requests
+
+    from databricks import sql
+    from databricks.sdk.core import ApiClient, Config, databricks_cli


@zackverham did you ever figure out a workaround for getting local service accounts tokens via the cli fallback?

I have not. Databricks labeled my ticket in the CLI repo as a bug, but that's where I'm stuck currently:

databricks/cli#1939

I'm not totally sure what we want to do with that. Options seem to be:

remove the local option from our implementation and re-introduce it when databricks supports it

leave the option to provide a local strategy + code comments explaining the databricks cli doesn't support it. A particularly motivated user could probably fork the databricks cli and implement a token command that follows the pattern established in other cli operations to get it to work. This isn't super likely and obviously pretty rough, but the local strategy here doesn't explicitly state it has to be provided by an official vendor.

… into zack-client-creds

schloerke · 2024-12-04T18:15:44Z

Happy to have more usage examples if you'd like. Either way, feel free to merge when ready.

mconflitti-pbc · 2024-12-06T14:10:02Z

src/posit/connect/external/databricks.py

+        # HTTP session headers are available in this context.
+        session_token = session.http_conn.headers.get("Posit-Connect-User-Session-Token")
+        posit_strategy = PositCredentialsStrategy(
+            local_strategy=databricks_cli, user_session_token=session_token


should this now be the patched local strategy?

its sort of ambiguous imo. the user could still use the databricks_cli as the local strategy, if say for example their databricks admin doesn't want to hand out service principal credentials. They can still use their own identity for local development with the databricks cli - that was my thinking for keeping this as-is.

Added some comments to the docstring to clarify the various use cases.

ah good point!

ah, I myself got confused - to clarify, this strategy is supporting viewer auth, which has no issues. The local strategy here works fine. Its just the content strategies that need the patch. (I still added docstring updates in the right places)

mconflitti-pbc · 2024-12-06T14:19:55Z

src/posit/connect/external/databricks.py

+
+    See Also
+    --------
+    * https://docs.databricks.com/en/dev-tools/auth/oauth-m2m.html#manually-generate-a-workspace-level-access-token


Seeing this example as well: https://docs.databricks.com/en/dev-tools/python-sql-connector.html#oauth-machine-to-machine-m2m-authentication

Is the problem using their version of this that they have omitted client_credentials scope?
https://github.com/databricks/databricks-sdk-py/blob/980166030c688e08a52d0f55389abd766f1d928d/databricks/sdk/credentials_provider.py#L164-L174

Good question - this is somewhat confusing, because the development of their CLI/SDK tooling and the sql connector are happening on independent threads and don't always sync up.

Here's the credentials provider we care about - it should fall back to whatever the CLI does to acquire a token:
https://github.com/databricks/databricks-sdk-py/blob/980166030c688e08a52d0f55389abd766f1d928d/databricks/sdk/credentials_provider.py#L630

Which calls this:
https://github.com/databricks/databricks-sdk-py/blob/980166030c688e08a52d0f55389abd766f1d928d/databricks/sdk/credentials_provider.py#L578

Which executes the databricks auth token command to acquire a token.

The real crux of the problem is in the auth token command - their token command isn't written to support anything aside from auth code flows:

https://github.com/databricks/cli/blob/main/cmd/auth/token.go#L45

So you can test this even outside of our code - just running databricks auth token with an oauth-m2m profile doesn't work. It ends up returning your own token and ignoring the profile.

Thanks for clarifying!

mconflitti-pbc · 2024-12-06T14:22:57Z

src/posit/connect/external/databricks.py

+    local_strategy = PositLocalContentCredentialsStrategy(
+        token_endpoint_url=TOKEN_ENDPOINT_URL,
+        client_id=CLIENT_ID,
+        client_secret=CLIENT_SECRET,
+    ) 


will it be on the user to make sure they have the proper env vars wired up? Curious if it makes sense to pull the values from their .databrickscfg file for them?

https://docs.databricks.com/en/dev-tools/cli/authentication.html#oauth-machine-to-machine-m2m-authentication

Maybe just a nice to have for now.

its a good thought, for sure - I'm hesitant to add that, since (a) this is really just supposed to be a band-aid to get people unstuck until the databricks CLI is working again, and (b) we'd then have to implement a parser that understands the databrickscfg file (profile settings, etc).

thats fair!

zackverham added 3 commits November 27, 2024 06:20

initial proposal for client credentials support in python sdk

67c7af3

a bit of refactoring + some more documentation

67edfbb

test coverage, linting

08796e2

zackverham requested a review from tdstein as a code owner December 2, 2024 20:54

zackverham requested review from schloerke and dbkegley December 2, 2024 20:55

zackverham self-assigned this Dec 2, 2024

zackverham changed the title ~~Add client credential oauth integration support + related databricks helpers to SDK~~ feat: Add client credential oauth integration support + related databricks helpers to SDK Dec 2, 2024

zackverham requested a review from christierney December 2, 2024 21:00

schloerke reviewed Dec 3, 2024

View reviewed changes

christierney reviewed Dec 3, 2024

View reviewed changes

src/posit/connect/external/databricks.py Outdated Show resolved Hide resolved

zackverham and others added 6 commits December 3, 2024 16:45

Update src/posit/connect/external/databricks.py

e654b0b

Co-authored-by: Barret Schloerke <[email protected]>

Update src/posit/connect/external/databricks.py

dd02020

Co-authored-by: Barret Schloerke <[email protected]>

responding to PR comments

36634c7

fix comment typo

d7ba5e6

Wrap examples in markdown blocks

e146ddd

Expose external on docs

d58a743

dbkegley reviewed Dec 4, 2024

View reviewed changes

zackverham added 4 commits December 4, 2024 12:05

updating docstrings in response to PR comments

b229b4b

Merge branch 'zack-client-creds' of github.com:posit-dev/posit-sdk-py…

3ec143d

… into zack-client-creds

docstring polish - backticks

89497bd

fix linting issues

3f42531

schloerke approved these changes Dec 4, 2024

View reviewed changes

schloerke and others added 3 commits December 4, 2024 14:46

Copy in existing example app and add title to each file

fb80f1d

add patched-in local strategy for client credential access tokens

73569a3

add implementing strategy that can use local credentials provider

91877ee

zackverham added 4 commits December 5, 2024 16:59

missing params

c38fecd

commit to force develop dependency to update

4a1a512

add raise_for_status

54a47c9

adding docstrings, tests

144e895

christierney approved these changes Dec 5, 2024

View reviewed changes

mconflitti-pbc reviewed Dec 6, 2024

View reviewed changes

sweep through docstrings in response to PR comments

2042938

mconflitti-pbc reviewed Dec 6, 2024

View reviewed changes

mconflitti-pbc approved these changes Dec 6, 2024

View reviewed changes

linting

59862c4

schloerke approved these changes Dec 6, 2024

View reviewed changes

zackverham merged commit 6b82ba3 into main Dec 6, 2024
35 checks passed

zackverham deleted the zack-client-creds branch December 6, 2024 20:51

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add client credential oauth integration support + related databricks helpers to SDK #348

feat: Add client credential oauth integration support + related databricks helpers to SDK #348

zackverham commented Dec 2, 2024 •

edited

Loading

github-actions bot commented Dec 2, 2024 •

edited

Loading

zackverham commented Dec 2, 2024 •

edited

Loading

mconflitti-pbc commented Dec 3, 2024

dbkegley Dec 4, 2024

zackverham Dec 4, 2024

schloerke Dec 4, 2024

mconflitti-pbc Dec 6, 2024

dbkegley Dec 4, 2024

zackverham Dec 4, 2024

schloerke commented Dec 4, 2024

mconflitti-pbc Dec 6, 2024

zackverham Dec 6, 2024

zackverham Dec 6, 2024

mconflitti-pbc Dec 6, 2024

zackverham Dec 6, 2024

mconflitti-pbc Dec 6, 2024

zackverham Dec 6, 2024

mconflitti-pbc Dec 6, 2024

mconflitti-pbc Dec 6, 2024

mconflitti-pbc Dec 6, 2024

zackverham Dec 6, 2024

mconflitti-pbc Dec 6, 2024

feat: Add client credential oauth integration support + related databricks helpers to SDK #348

feat: Add client credential oauth integration support + related databricks helpers to SDK #348

Conversation

zackverham commented Dec 2, 2024 • edited Loading

github-actions bot commented Dec 2, 2024 • edited Loading

☂️ Python Coverage

Overall Coverage

New Files

Modified Files

zackverham commented Dec 2, 2024 • edited Loading

mconflitti-pbc commented Dec 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schloerke commented Dec 4, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zackverham commented Dec 2, 2024 •

edited

Loading

github-actions bot commented Dec 2, 2024 •

edited

Loading

zackverham commented Dec 2, 2024 •

edited

Loading