Update Databricks sdk helpers to be more explicit about fallback options

Author: dbkegley

examples/connect/databricks/dash/app.py

@@ -6,10 +6,14 @@
 import pandas as pd
 from dash import Dash, Input, Output, dash_table, html
 from databricks import sql
-from databricks.sdk.core import ApiClient, Config, databricks_cli
+from databricks.sdk.core import ApiClient
 from databricks.sdk.service.iam import CurrentUserAPI
 
-from posit.connect.external.databricks import PositCredentialsStrategy
+from posit.connect.external.databricks import (
+    new_config,
+    sql_credentials,
+    ConnectStrategy,
+)
 
 DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
 DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -37,15 +41,10 @@ def update_page(_):
     the first few rows from a table hosted in Databricks.
     """
     session_token = flask.request.headers.get("Posit-Connect-User-Session-Token")
-    posit_strategy = PositCredentialsStrategy(
-        local_strategy=databricks_cli,
-        user_session_token=session_token,
-    )
-    cfg = Config(
-        host=DATABRICKS_HOST_URL,
-        # uses Posit's custom credential_strategy if running on Connect,
-        # otherwise falls back to the strategy defined by local_strategy
-        credentials_strategy=posit_strategy,
+    cfg = new_config(
+        posit_connect_strategy=ConnectStrategy(
+            user_session_token=session_token,
+        ),
     )
 
     def get_greeting():
@@ -61,8 +60,7 @@ def get_table():
             with sql.connect(
                 server_hostname=DATABRICKS_HOST,
                 http_path=SQL_HTTP_PATH,
-                # https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
-                credentials_provider=posit_strategy.sql_credentials_provider(cfg),
+                credentials_provider=sql_credentials(cfg),
             ) as connection:
                 with connection.cursor() as cursor:
                     cursor.execute(query)

examples/connect/databricks/fastapi/app.py

@@ -5,11 +5,14 @@
 import os
 
 from databricks import sql
-from databricks.sdk.core import Config, databricks_cli
 from fastapi import FastAPI, Header
 from typing_extensions import TYPE_CHECKING, Annotated
 
-from posit.connect.external.databricks import PositCredentialsStrategy
+from posit.connect.external.databricks import (
+    new_config,
+    sql_credentials,
+    ConnectStrategy,
+)
 
 if TYPE_CHECKING:
     from fastapi.responses import JSONResponse
@@ -32,15 +35,10 @@ async def get_fares(
     """
     global rows
 
-    posit_strategy = PositCredentialsStrategy(
-        local_strategy=databricks_cli,
-        user_session_token=posit_connect_user_session_token,
-    )
-    cfg = Config(
-        host=DATABRICKS_HOST_URL,
-        # uses Posit's custom credential_strategy if running on Connect,
-        # otherwise falls back to the strategy defined by local_strategy
-        credentials_strategy=posit_strategy,
+    cfg = new_config(
+        posit_connect_strategy=ConnectStrategy(
+            user_session_token=posit_connect_user_session_token,
+        ),
     )
 
     if rows is None:
@@ -49,8 +47,7 @@ async def get_fares(
         with sql.connect(
             server_hostname=DATABRICKS_HOST,
             http_path=SQL_HTTP_PATH,
-            # https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
-            credentials_provider=posit_strategy.sql_credentials_provider(cfg),
+            credentials_provider=sql_credentials(cfg),
         ) as connection:
             with connection.cursor() as cursor:
                 cursor.execute(query)

examples/connect/databricks/flask/app.py

@@ -3,10 +3,13 @@
 import os
 
 from databricks import sql
-from databricks.sdk.core import Config, databricks_cli
 from flask import Flask, request
 
-from posit.connect.external.databricks import PositCredentialsStrategy
+from posit.connect.external.databricks import (
+    new_config,
+    sql_credentials,
+    ConnectStrategy,
+)
 
 DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
 DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -30,15 +33,10 @@ def get_fares():
     global rows
 
     session_token = request.headers.get("Posit-Connect-User-Session-Token")
-    posit_strategy = PositCredentialsStrategy(
-        local_strategy=databricks_cli,
-        user_session_token=session_token,
-    )
-    cfg = Config(
-        host=DATABRICKS_HOST_URL,
-        # uses Posit's custom credential_strategy if running on Connect,
-        # otherwise falls back to the strategy defined by local_strategy
-        credentials_strategy=posit_strategy,
+    cfg = new_config(
+        posit_connect_strategy=ConnectStrategy(
+            user_session_token=session_token
+        ),
     )
 
     if rows is None:
@@ -47,8 +45,7 @@ def get_fares():
         with sql.connect(
             server_hostname=DATABRICKS_HOST,
             http_path=SQL_HTTP_PATH,
-            # https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
-            credentials_provider=posit_strategy.sql_credentials_provider(cfg),
+            credentials_provider=sql_credentials(cfg),
         ) as connection:
             with connection.cursor() as cursor:
                 cursor.execute(query)

examples/connect/databricks/shiny/app.py

@@ -4,11 +4,15 @@
 
 import pandas as pd
 from databricks import sql
-from databricks.sdk.core import ApiClient, Config, databricks_cli
+from databricks.sdk.core import ApiClient
 from databricks.sdk.service.iam import CurrentUserAPI
 from shiny import App, Inputs, Outputs, Session, render, ui
 
-from posit.connect.external.databricks import PositCredentialsStrategy
+from posit.connect.external.databricks import (
+    new_config,
+    sql_credentials,
+    ConnectStrategy,
+)
 
 DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
 DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
@@ -23,15 +27,10 @@ def server(i: Inputs, o: Outputs, session: Session):
     the first few rows from a table hosted in Databricks.
     """
     session_token = session.http_conn.headers.get("Posit-Connect-User-Session-Token")
-    posit_strategy = PositCredentialsStrategy(
-        local_strategy=databricks_cli,
-        user_session_token=session_token,
-    )
-    cfg = Config(
-        host=DATABRICKS_HOST_URL,
-        # uses Posit's custom credential_strategy if running on Connect,
-        # otherwise falls back to the strategy defined by local_strategy
-        credentials_strategy=posit_strategy,
+    cfg = new_config(
+        posit_connect_strategy=ConnectStrategy(
+            user_session_token=session_token
+        ),
     )
 
     @render.data_frame
@@ -41,8 +40,7 @@ def result():
         with sql.connect(
             server_hostname=DATABRICKS_HOST,
             http_path=SQL_HTTP_PATH,
-            # https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
-            credentials_provider=posit_strategy.sql_credentials_provider(cfg),
+            credentials_provider=sql_credentials(cfg),
         ) as connection:
             with connection.cursor() as cursor:
                 cursor.execute(query)

examples/connect/databricks/streamlit/app.py

@@ -5,25 +5,24 @@
 import pandas as pd
 import streamlit as st
 from databricks import sql
-from databricks.sdk.core import ApiClient, Config, databricks_cli
+from databricks.sdk.core import ApiClient
 from databricks.sdk.service.iam import CurrentUserAPI
 
-from posit.connect.external.databricks import PositCredentialsStrategy
+from posit.connect.external.databricks import (
+    new_config,
+    sql_credentials,
+    ConnectStrategy,
+)
 
 DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
 DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
 SQL_HTTP_PATH = os.getenv("DATABRICKS_PATH")
 
 session_token = st.context.headers.get("Posit-Connect-User-Session-Token")
-posit_strategy = PositCredentialsStrategy(
-    local_strategy=databricks_cli,
-    user_session_token=session_token,
-)
-cfg = Config(
-    host=DATABRICKS_HOST_URL,
-    # uses Posit's custom credential_strategy if running on Connect,
-    # otherwise falls back to the strategy defined by local_strategy
-    credentials_strategy=posit_strategy,
+cfg = new_config(
+    posit_connect_strategy=ConnectStrategy(
+        user_session_token=session_token
+    ),
 )
 
 databricks_user = CurrentUserAPI(ApiClient(cfg)).me()
@@ -32,8 +31,7 @@
 with sql.connect(
     server_hostname=DATABRICKS_HOST,
     http_path=SQL_HTTP_PATH,
-    # https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
-    credentials_provider=posit_strategy.sql_credentials_provider(cfg),
+    credentials_provider=sql_credentials(cfg),
 ) as connection:
     with connection.cursor() as cursor:
         cursor.execute("SELECT * FROM samples.nyctaxi.trips LIMIT 10;")

pyproject.toml

@@ -131,7 +131,15 @@ exclude = ".*"
 [dependency-groups]
 build = ["build"]
 coverage = ["coverage"]
-examples = ["rsconnect-python", "pandas", "databricks", "shiny"]
+examples = [
+    "databricks-sdk>=0.29.0",
+    "databricks-sql-connector>=3.3.0",
+    "pandas",
+    "rsconnect-python",
+    "shiny",
+    "streamlit>=1.37.0"
+]
+external = ["databricks-sdk"]
 git = ["pre-commit"]
 lint = ["ruff", "pyright"]
 test = ["rsconnect-python", "responses", "pytest", "pyjson5"]
@@ -143,6 +151,7 @@ dev = [
     # # `numpy==1.24.4` requires `distutils`, which is not available in python3.12
     # # `numpy>1.24.4` does not require `distutils`
     # { include-group = "examples" },
+    { include-group = "external" },
     { include-group = "git" },
     { include-group = "lint" },
     { include-group = "test" },

src/posit/connect/_utils.py

@@ -32,11 +32,25 @@ def update_dict_values(obj: dict[str, Any], /, **kwargs: Any) -> None:
     dict.update(obj, **kwargs)
 
 
-def is_local() -> bool:
-    """Returns true if called from a piece of content running on a Connect server.
+def is_workbench() -> bool:
+    """Attempts to return true if called from a piece of content running on Posit Workbench.
+
+    There is not yet a definitive way to determine if the content is running on Workbench. This method is best-effort.
+    """
+    return "RSW_LAUNCHER" in os.environ or \
+        "RSTUDIO_MULTI_SESSION" in os.environ or \
+        "RS_SERVER_ADDRESS" in os.environ or \
+        "RS_SERVER_URL" in os.environ
+
 
-    The connect server will always set the environment variable `RSTUDIO_PRODUCT=CONNECT`.
-    We can use this environment variable to determine if the content is running locally
-    or on a Connect server.
+def is_connect() -> bool:
+    """Returns true if called from a piece of content running on Posit Connect.
+
+    The Connect content will always set the environment variable `RSTUDIO_PRODUCT=CONNECT`.
     """
-    return os.getenv("RSTUDIO_PRODUCT") != "CONNECT"
+    return os.getenv("RSTUDIO_PRODUCT") == "CONNECT"
+
+
+def is_local() -> bool:
+    """Returns true if called from a piece of content running locally."""
+    return not is_connect() and not is_workbench()