Skip to content

Commit 37bc1dc

Browse files
dbkegleydbkegley
authored
feat!: improve compatibility with databricks sql client (#252)
* feat!: add working databricks-cli fallback credentials strategy * test: add tests for external databricks helpers
1 parent 8648f80 commit 37bc1dc

File tree

13 files changed

+231
-109
lines changed

13 files changed

+231
-109
lines changed

‎examples/connect/.gitignore‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
.posit
1+
**/rsconnect-python/

‎examples/connect/dash/app.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66
import pandas as pd
77
from dash import Dash, Input, Output, dash_table, html
88
from databricks import sql
9-
from databricks.sdk.core import ApiClient, Config
9+
from databricks.sdk.core import ApiClient, Config, databricks_cli
1010
from databricks.sdk.service.iam import CurrentUserAPI
11-
from posit.connect.external.databricks import viewer_credentials_provider
11+
from posit.connect.external.databricks import PositCredentialsStrategy
1212

1313
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1414
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -38,14 +38,16 @@ def update_page(_):
3838
session_token = flask.request.headers.get(
3939
"Posit-Connect-User-Session-Token"
4040
)
41-
credentials_provider = viewer_credentials_provider(
42-
user_session_token=session_token
43-
)
41+
posit_strategy = PositCredentialsStrategy(
42+
local_strategy=databricks_cli,
43+
user_session_token=session_token)
44+
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
48+
credentials_strategy=posit_strategy)
4449

4550
def get_greeting():
46-
cfg = Config(
47-
host=DATABRICKS_HOST_URL, credentials_provider=credentials_provider
48-
)
4951
databricks_user_info = CurrentUserAPI(ApiClient(cfg)).me()
5052
return f"Hello, {databricks_user_info.display_name}!"
5153

@@ -58,8 +60,8 @@ def get_table():
5860
with sql.connect(
5961
server_hostname=DATABRICKS_HOST,
6062
http_path=SQL_HTTP_PATH,
61-
auth_type="databricks-oauth",
62-
credentials_provider=credentials_provider,
63+
# https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
64+
credentials_provider=posit_strategy.sql_credentials_provider(cfg)
6365
) as connection:
6466
with connection.cursor() as cursor:
6567
cursor.execute(query)

‎examples/connect/dash/requirements.txt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
databricks-sql-connector==3.0.1
2-
databricks-sdk==0.20.0
1+
databricks-sql-connector==3.3.0
2+
databricks-sdk==0.29.0
33
dash==2.15.0
4-
git+https://github.com/posit-dev/posit-sdk-py.git
4+
posit-sdk>=0.4.0

‎examples/connect/fastapi/app.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@
44
from typing import Annotated
55

66
from databricks import sql
7+
from databricks.sdk.core import Config, databricks_cli
78
from fastapi import FastAPI, Header
89
from fastapi.responses import JSONResponse
9-
from posit.connect.external.databricks import viewer_credentials_provider
10+
from posit.connect.external.databricks import PositCredentialsStrategy
1011

1112
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1213
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -26,18 +27,23 @@ async def get_fares(
2627
"""
2728
global rows
2829

29-
credentials_provider = viewer_credentials_provider(
30-
user_session_token=posit_connect_user_session_token
31-
)
30+
posit_strategy = PositCredentialsStrategy(
31+
local_strategy=databricks_cli,
32+
user_session_token=posit_connect_user_session_token)
33+
cfg = Config(
34+
host=DATABRICKS_HOST_URL,
35+
# uses Posit's custom credential_strategy if running on Connect,
36+
# otherwise falls back to the strategy defined by local_strategy
37+
credentials_strategy=posit_strategy)
3238

3339
if rows is None:
3440
query = "SELECT * FROM samples.nyctaxi.trips LIMIT 10;"
3541

3642
with sql.connect(
3743
server_hostname=DATABRICKS_HOST,
3844
http_path=SQL_HTTP_PATH,
39-
auth_type="databricks-oauth",
40-
credentials_provider=credentials_provider,
45+
# https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
46+
credentials_provider=posit_strategy.sql_credentials_provider(cfg)
4147
) as connection:
4248
with connection.cursor() as cursor:
4349
cursor.execute(query)

‎examples/connect/fastapi/requirements.txt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
databricks-sql-connector==3.0.1
2-
databricks-sdk==0.20.0
1+
databricks-sql-connector==3.3.0
2+
databricks-sdk==0.29.0
33
fastapi==0.110.0
4-
git+https://github.com/posit-dev/posit-sdk-py.git
4+
posit-sdk>=0.4.0

‎examples/connect/flask/app.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33
import os
44

55
from databricks import sql
6+
from databricks.sdk.core import Config, databricks_cli
67
from flask import Flask, request
7-
from posit.connect.external.databricks import viewer_credentials_provider
8+
from posit.connect.external.databricks import PositCredentialsStrategy
89

910
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1011
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -28,18 +29,23 @@ def get_fares():
2829
global rows
2930

3031
session_token = request.headers.get("Posit-Connect-User-Session-Token")
31-
credentials_provider = viewer_credentials_provider(
32-
user_session_token=session_token
33-
)
32+
posit_strategy = PositCredentialsStrategy(
33+
local_strategy=databricks_cli,
34+
user_session_token=session_token)
35+
cfg = Config(
36+
host=DATABRICKS_HOST_URL,
37+
# uses Posit's custom credential_strategy if running on Connect,
38+
# otherwise falls back to the strategy defined by local_strategy
39+
credentials_strategy=posit_strategy)
3440

3541
if rows is None:
3642
query = "SELECT * FROM samples.nyctaxi.trips LIMIT 10;"
3743

3844
with sql.connect(
3945
server_hostname=DATABRICKS_HOST,
4046
http_path=SQL_HTTP_PATH,
41-
auth_type="databricks-oauth",
42-
credentials_provider=credentials_provider,
47+
# https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
48+
credentials_provider=posit_strategy.sql_credentials_provider(cfg)
4349
) as connection:
4450
with connection.cursor() as cursor:
4551
cursor.execute(query)

‎examples/connect/flask/requirements.txt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
databricks-sql-connector==3.0.1
2-
databricks-sdk==0.20.0
1+
databricks-sql-connector==3.3.0
2+
databricks-sdk==0.29.0
33
flask==3.0.2
4-
git+https://github.com/posit-dev/posit-sdk-py.git
4+
posit-sdk>=0.4.0

‎examples/connect/shiny-python/app.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44

55
import pandas as pd
66
from databricks import sql
7-
from databricks.sdk.core import ApiClient, Config
7+
from databricks.sdk.core import ApiClient, Config, databricks_cli
88
from databricks.sdk.service.iam import CurrentUserAPI
9-
from posit.connect.external.databricks import viewer_credentials_provider
9+
from posit.connect.external.databricks import PositCredentialsStrategy
1010
from shiny import App, Inputs, Outputs, Session, render, ui
1111

1212
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
@@ -24,9 +24,14 @@ def server(i: Inputs, o: Outputs, session: Session):
2424
session_token = session.http_conn.headers.get(
2525
"Posit-Connect-User-Session-Token"
2626
)
27-
credentials_provider = viewer_credentials_provider(
28-
user_session_token=session_token
29-
)
27+
posit_strategy = PositCredentialsStrategy(
28+
local_strategy=databricks_cli,
29+
user_session_token=session_token)
30+
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
34+
credentials_strategy=posit_strategy)
3035

3136
@render.data_frame
3237
def result():
@@ -35,8 +40,8 @@ def result():
3540
with sql.connect(
3641
server_hostname=DATABRICKS_HOST,
3742
http_path=SQL_HTTP_PATH,
38-
auth_type="databricks-oauth",
39-
credentials_provider=credentials_provider,
43+
# https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
44+
credentials_provider=posit_strategy.sql_credentials_provider(cfg),
4045
) as connection:
4146
with connection.cursor() as cursor:
4247
cursor.execute(query)
@@ -48,9 +53,6 @@ def result():
4853

4954
@render.text
5055
def text():
51-
cfg = Config(
52-
host=DATABRICKS_HOST_URL, credentials_provider=credentials_provider
53-
)
5456
databricks_user_info = CurrentUserAPI(ApiClient(cfg)).me()
5557
return f"Hello, {databricks_user_info.display_name}!"
5658

‎examples/connect/shiny-python/requirements.txt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
databricks-sql-connector==3.0.1
2-
databricks-sdk==0.20.0
1+
databricks-sql-connector==3.3.0
2+
databricks-sdk==0.29.0
33
shiny==0.7.1
4-
git+https://github.com/posit-dev/posit-sdk-py.git
4+
posit-sdk>=0.4.0

‎examples/connect/streamlit/app.py‎

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,34 +5,32 @@
55
import pandas as pd
66
import streamlit as st
77
from databricks import sql
8-
from databricks.sdk.core import ApiClient, Config
8+
from databricks.sdk.core import ApiClient, Config, databricks_cli
99
from databricks.sdk.service.iam import CurrentUserAPI
10-
from posit.connect.external.databricks import viewer_credentials_provider
11-
from streamlit.web.server.websocket_headers import _get_websocket_headers
10+
from posit.connect.external.databricks import PositCredentialsStrategy
1211

1312
DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
1413
DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
1514
SQL_HTTP_PATH = os.getenv("DATABRICKS_PATH")
1615

17-
session_token = _get_websocket_headers().get(
18-
"Posit-Connect-User-Session-Token"
19-
)
20-
21-
credentials_provider = viewer_credentials_provider(
22-
user_session_token=session_token
23-
)
24-
16+
session_token = st.context.headers.get("Posit-Connect-User-Session-Token")
17+
posit_strategy = PositCredentialsStrategy(
18+
local_strategy=databricks_cli,
19+
user_session_token=session_token)
2520
cfg = Config(
26-
host=DATABRICKS_HOST_URL, credentials_provider=credentials_provider
27-
)
21+
host=DATABRICKS_HOST_URL,
22+
# uses Posit's custom credential_strategy if running on Connect,
23+
# otherwise falls back to the strategy defined by local_strategy
24+
credentials_strategy=posit_strategy)
25+
2826
databricks_user = CurrentUserAPI(ApiClient(cfg)).me()
2927
st.write(f"Hello, {databricks_user.display_name}!")
3028

3129
with sql.connect(
3230
server_hostname=DATABRICKS_HOST,
3331
http_path=SQL_HTTP_PATH,
34-
auth_type="databricks-oauth",
35-
credentials_provider=credentials_provider,
32+
# https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
33+
credentials_provider=posit_strategy.sql_credentials_provider(cfg)
3634
) as connection:
3735
with connection.cursor() as cursor:
3836
cursor.execute("SELECT * FROM samples.nyctaxi.trips LIMIT 10;")

0 commit comments

Comments
 (0)