Skip to content

Commit d695abf

Browse files
dbkegleydbkegley
committed
Update Databricks oauth examples to use the new helpers
1 parent 7f1658d commit d695abf

File tree

5 files changed

+40
-25
lines changed

5 files changed

+40
-25
lines changed

‎examples/connect/databricks/dash/app.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,11 @@
99
from databricks.sdk.core import ApiClient, Config, databricks_cli
1010
from databricks.sdk.service.iam import CurrentUserAPI
1111

12-
from posit.connect.external.databricks import PositCredentialsStrategy
12+
from posit.connect.external.databricks import (
13+
PositCredentialsStrategy,
14+
PositConnectCredentialsStrategy,
15+
PositWorkbenchCredentialsStrategy,
16+
)
1317

1418
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1519
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -39,12 +43,11 @@ def update_page(_):
3943
session_token = flask.request.headers.get("Posit-Connect-User-Session-Token")
4044
posit_strategy = PositCredentialsStrategy(
4145
local_strategy=databricks_cli,
42-
user_session_token=session_token,
46+
workbench_strategy=PositWorkbenchCredentialsStrategy(Config(profile="workbench")),
47+
connect_strategy=PositConnectCredentialsStrategy(user_session_token=session_token),
4348
)
4449
cfg = Config(
45-
host=DATABRICKS_HOST_URL,
46-
# uses Posit's custom credential_strategy if running on Connect,
47-
# otherwise falls back to the strategy defined by local_strategy
50+
host=DATABRICKS_HOST_URL, # required by databricks_cli
4851
credentials_strategy=posit_strategy,
4952
)
5053

‎examples/connect/databricks/fastapi/app.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,11 @@
99
from fastapi import FastAPI, Header
1010
from typing_extensions import TYPE_CHECKING, Annotated
1111

12-
from posit.connect.external.databricks import PositCredentialsStrategy
12+
from posit.connect.external.databricks import (
13+
PositCredentialsStrategy,
14+
PositWorkbenchCredentialsStrategy,
15+
PositConnectCredentialsStrategy,
16+
)
1317

1418
if TYPE_CHECKING:
1519
from fastapi.responses import JSONResponse
@@ -34,12 +38,11 @@ async def get_fares(
3438

3539
posit_strategy = PositCredentialsStrategy(
3640
local_strategy=databricks_cli,
37-
user_session_token=posit_connect_user_session_token,
41+
workbench_strategy=PositWorkbenchCredentialsStrategy(Config(profile="workbench")),
42+
connect_strategy=PositConnectCredentialsStrategy(user_session_token=posit_connect_user_session_token),
3843
)
3944
cfg = Config(
40-
host=DATABRICKS_HOST_URL,
41-
# uses Posit's custom credential_strategy if running on Connect,
42-
# otherwise falls back to the strategy defined by local_strategy
45+
host=DATABRICKS_HOST_URL, # required by databricks_cli
4346
credentials_strategy=posit_strategy,
4447
)
4548

‎examples/connect/databricks/flask/app.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,11 @@
66
from databricks.sdk.core import Config, databricks_cli
77
from flask import Flask, request
88

9-
from posit.connect.external.databricks import PositCredentialsStrategy
9+
from posit.connect.external.databricks import (
10+
PositCredentialsStrategy,
11+
PositConnectCredentialsStrategy,
12+
PositWorkbenchCredentialsStrategy,
13+
)
1014

1115
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1216
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -32,12 +36,11 @@ def get_fares():
3236
session_token = request.headers.get("Posit-Connect-User-Session-Token")
3337
posit_strategy = PositCredentialsStrategy(
3438
local_strategy=databricks_cli,
35-
user_session_token=session_token,
39+
workbench_strategy=PositWorkbenchCredentialsStrategy(Config(profile="workbench")),
40+
connect_strategy=PositConnectCredentialsStrategy(user_session_token=session_token),
3641
)
3742
cfg = Config(
38-
host=DATABRICKS_HOST_URL,
39-
# uses Posit's custom credential_strategy if running on Connect,
40-
# otherwise falls back to the strategy defined by local_strategy
43+
host=DATABRICKS_HOST_URL, # required by databricks_cli
4144
credentials_strategy=posit_strategy,
4245
)
4346

‎examples/connect/databricks/shiny/app.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,11 @@
88
from databricks.sdk.service.iam import CurrentUserAPI
99
from shiny import App, Inputs, Outputs, Session, render, ui
1010

11-
from posit.connect.external.databricks import PositCredentialsStrategy
11+
from posit.connect.external.databricks import (
12+
PositCredentialsStrategy,
13+
PositConnectCredentialsStrategy,
14+
PositWorkbenchCredentialsStrategy,
15+
)
1216

1317
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1418
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -25,12 +29,11 @@ def server(i: Inputs, o: Outputs, session: Session):
2529
session_token = session.http_conn.headers.get("Posit-Connect-User-Session-Token")
2630
posit_strategy = PositCredentialsStrategy(
2731
local_strategy=databricks_cli,
28-
user_session_token=session_token,
32+
workbench_strategy=PositWorkbenchCredentialsStrategy(Config(profile="workbench")),
33+
connect_strategy=PositConnectCredentialsStrategy(user_session_token=session_token),
2934
)
3035
cfg = Config(
31-
host=DATABRICKS_HOST_URL,
32-
# uses Posit's custom credential_strategy if running on Connect,
33-
# otherwise falls back to the strategy defined by local_strategy
36+
host=DATABRICKS_HOST_URL, # required by databricks_cli
3437
credentials_strategy=posit_strategy,
3538
)
3639

‎examples/connect/databricks/streamlit/app.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,11 @@
88
from databricks.sdk.core import ApiClient, Config, databricks_cli
99
from databricks.sdk.service.iam import CurrentUserAPI
1010

11-
from posit.connect.external.databricks import PositCredentialsStrategy
11+
from posit.connect.external.databricks import (
12+
PositCredentialsStrategy,
13+
PositConnectCredentialsStrategy,
14+
PositWorkbenchCredentialsStrategy,
15+
)
1216

1317
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1418
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -17,12 +21,11 @@
1721
session_token = st.context.headers.get("Posit-Connect-User-Session-Token")
1822
posit_strategy = PositCredentialsStrategy(
1923
local_strategy=databricks_cli,
20-
user_session_token=session_token,
24+
workbench_strategy=PositWorkbenchCredentialsStrategy(Config(profile="workbench")),
25+
connect_strategy=PositConnectCredentialsStrategy(user_session_token=session_token),
2126
)
2227
cfg = Config(
23-
host=DATABRICKS_HOST_URL,
24-
# uses Posit's custom credential_strategy if running on Connect,
25-
# otherwise falls back to the strategy defined by local_strategy
28+
host=DATABRICKS_HOST_URL, # required by databricks_cli
2629
credentials_strategy=posit_strategy,
2730
)
2831

0 commit comments

Comments
 (0)