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

docs/_quarto.yml

@@ -107,3 +107,7 @@ quartodoc:
       contents:
         - connect.metrics
         - connect.metrics.usage
+    - title: External Integrations
+      contents:
+        - connect.external.databricks
+        - connect.external.snowflake

src/posit/connect/content.py

@@ -755,8 +755,8 @@ def find_by(
         -------
         Optional[ContentItem]
 
-        Example
-        -------
+        Examples
+        --------
         >>> find_by(name="example-content-name")
         """
         attr_items = attrs.items()

src/posit/connect/external/databricks.py

@@ -2,22 +2,27 @@
 from typing import Callable, Dict, Optional
 
 from ..client import Client
+from ..oauth import Credentials
 from .external import is_local
 
 """
 NOTE: These APIs are provided as a convenience and are subject to breaking changes:
 https://github.com/databricks/databricks-sdk-py#interface-stability
 """
 
+POSIT_OAUTH_INTEGRATION_AUTH_TYPE = "posit-oauth-integration"
+
 # The Databricks SDK CredentialsProvider == Databricks SQL HeaderFactory
 CredentialsProvider = Callable[[], Dict[str, str]]
 
 
 class CredentialsStrategy(abc.ABC):
     """Maintain compatibility with the Databricks SQL/SDK client libraries.
 
-    https://github.com/databricks/databricks-sql-python/blob/v3.3.0/src/databricks/sql/auth/authenticators.py#L19-L33
-    https://github.com/databricks/databricks-sdk-py/blob/v0.29.0/databricks/sdk/credentials_provider.py#L44-L54
+    See Also
+    --------
+    * https://github.com/databricks/databricks-sql-python/blob/v3.3.0/src/databricks/sql/auth/authenticators.py#L19-L33
+    * https://github.com/databricks/databricks-sdk-py/blob/v0.29.0/databricks/sdk/credentials_provider.py#L44-L54
     """
 
     @abc.abstractmethod
@@ -29,29 +34,122 @@ def __call__(self, *args, **kwargs) -> CredentialsProvider:
         raise NotImplementedError
 
 
+def _new_bearer_authorization_header(credentials: Credentials) -> Dict[str, str]:
+    """Helper to transform an Credentials object into the Bearer auth header consumed by databricks.
+
+    Raises
+    ------
+        ValueError: If provided Credentials object does not contain an access token
+
+    Returns
+    -------
+        Dict[str, str]
+    """
+    access_token = credentials.get("access_token")
+    if access_token is None:
+        raise ValueError("Missing value for field 'access_token' in credentials.")
+    return {"Authorization": f"Bearer {access_token}"}
+
+
+def _get_auth_type(local_auth_type: str) -> str:
+    """Returns the auth type currently in use.
+
+    The databricks-sdk client uses the configured auth_type to create
+    a user-agent string which is used for attribution. We should only
+    overwrite the auth_type if we are using the PositCredentialsStrategy (non-local),
+    otherwise, we should return the auth_type of the configured local_strategy instead
+    to avoid breaking someone elses attribution.
+
+    NOTE: The databricks-sql client does not use auth_type to set the user-agent.
+    https://github.com/databricks/databricks-sql-python/blob/v3.3.0/src/databricks/sql/client.py#L214-L219
+
+    See Also
+    --------
+    * https://github.com/databricks/databricks-sdk-py/blob/v0.29.0/databricks/sdk/config.py#L261-L269
+
+    Returns
+    -------
+        str
+    """
+    if is_local():
+        return local_auth_type
+
+    return POSIT_OAUTH_INTEGRATION_AUTH_TYPE
+
+
+class PositContentCredentialsProvider:
+    """CredentialsProvider implementation which initiates a credential exchange using a content-session-token."""
+
+    def __init__(self, client: Client):
+        self._client = client
+
+    def __call__(self) -> Dict[str, str]:
+        credentials = self._client.oauth.get_content_credentials()
+        return _new_bearer_authorization_header(credentials)
+
+
 class PositCredentialsProvider:
+    """CredentialsProvider implementation which initiates a credential exchange using a user-session-token."""
+
     def __init__(self, client: Client, user_session_token: str):
         self._client = client
         self._user_session_token = user_session_token
 
     def __call__(self) -> Dict[str, str]:
         credentials = self._client.oauth.get_credentials(self._user_session_token)
-        access_token = credentials.get("access_token")
-        if access_token is None:
-            raise ValueError("Missing value for field 'access_token' in credentials.")
-        return {"Authorization": f"Bearer {access_token}"}
+        return _new_bearer_authorization_header(credentials)
 
 
-class PositCredentialsStrategy(CredentialsStrategy):
+class PositContentCredentialsStrategy(CredentialsStrategy):
+    """`CredentialsStrategy` implementation which supports interacting with Service Account OAuth integrations on Connect.
+
+    This strategy callable class returns a `PositContentCredentialsProvider` when hosted on Connect, and
+    its `local_strategy` strategy otherwise.
+
+    Examples
+    --------
+    ```python
+    from posit.connect.external.databricks import PositContentCredentialsStrategy
+
+    import pandas as pd
+    import requests
+
+    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]))
+    ```
+    """
+
     def __init__(
         self,
         local_strategy: CredentialsStrategy,
         client: Optional[Client] = None,
-        user_session_token: Optional[str] = None,
     ):
         self._local_strategy = local_strategy
         self._client = client
-        self._user_session_token = user_session_token
 
     def sql_credentials_provider(self, *args, **kwargs):
         """The sql connector attempts to call the credentials provider w/o any args.
@@ -66,23 +164,103 @@ def sql_credentials_provider(self, *args, **kwargs):
         return lambda: self.__call__(*args, **kwargs)
 
     def auth_type(self) -> str:
-        """Returns the auth type currently in use.
+        return _get_auth_type(self._local_strategy.auth_type())
+
+    def __call__(self, *args, **kwargs) -> CredentialsProvider:
+        # If the content is not running on Connect then fall back to local_strategy
+        if is_local():
+            return self._local_strategy(*args, **kwargs)
 
-        The databricks-sdk client uses the configurated auth_type to create
-        a user-agent string which is used for attribution. We should only
-        overwrite the auth_type if we are using the PositCredentialsStrategy (non-local),
-        otherwise, we should return the auth_type of the configured local_strategy instead
-        to avoid breaking someone elses attribution.
+        if self._client is None:
+            self._client = Client()
+
+        return PositContentCredentialsProvider(self._client)
+
+
+class PositCredentialsStrategy(CredentialsStrategy):
+    """`CredentialsStrategy` implementation which supports interacting with Viewer OAuth integrations on Connect.
+
+    This strategy callable class returns a `PositCredentialsProvider` when hosted on Connect, and
+    its `local_strategy` strategy otherwise.
+
+    Examples
+    --------
+    ```python
+    import os
+
+    import pandas as pd
+    from databricks import sql
+    from databricks.sdk.core import ApiClient, Config, databricks_cli
+    from databricks.sdk.service.iam import CurrentUserAPI
+    from posit.connect.external.databricks import PositCredentialsStrategy
+    from shiny import App, Inputs, Outputs, Session, render, ui
+
+    # env vars
+    DATABRICKS_HOST = "<REDACTED>"
+    DATABRICKS_HOST_URL = f"https://{DATABRICKS_HOST}"
+    SQL_HTTP_PATH = "<REDACTED>"
+
+    app_ui = ui.page_fluid(ui.output_text("text"), ui.output_data_frame("result"))
+
+
+    def server(i: Inputs, o: Outputs, session: Session):
+        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, credentials_strategy=posit_strategy)
+
+        @render.data_frame
+        def result():
+            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()
+                    df = pd.DataFrame(rows, columns=[col[0] for col in cursor.description])
+                    return df
+
+        @render.text
+        def text():
+            databricks_user_info = CurrentUserAPI(ApiClient(cfg)).me()
+            return f"Hello, {databricks_user_info.display_name}!"
+
+
+    app = App(app_ui, server)
+    ```
+    """
+
+    def __init__(
+        self,
+        local_strategy: CredentialsStrategy,
+        client: Optional[Client] = None,
+        user_session_token: Optional[str] = None,
+    ):
+        self._local_strategy = local_strategy
+        self._client = client
+        self._user_session_token = user_session_token
+
+    def sql_credentials_provider(self, *args, **kwargs):
+        """The sql connector attempts to call the credentials provider w/o any args.
 
-        https://github.com/databricks/databricks-sdk-py/blob/v0.29.0/databricks/sdk/config.py#L261-L269
+        The SQL client's `ExternalAuthProvider` is not compatible w/ the SDK's implementation of
+        `CredentialsProvider`, so create a no-arg lambda that wraps the args defined by the real caller.
+        This way we can pass in a databricks `Config` object required by most of the SDK's `CredentialsProvider`
+        implementations from where `sql.connect` is called.
 
-        NOTE: The databricks-sql client does not use auth_type to set the user-agent.
-        https://github.com/databricks/databricks-sql-python/blob/v3.3.0/src/databricks/sql/client.py#L214-L219
+        See Also
+        --------
+        * https://github.com/databricks/databricks-sql-python/issues/148#issuecomment-2271561365
         """
-        if is_local():
-            return self._local_strategy.auth_type()
-        else:
-            return "posit-oauth-integration"
+        return lambda: self.__call__(*args, **kwargs)
+
+    def auth_type(self) -> str:
+        return _get_auth_type(self._local_strategy.auth_type())
 
     def __call__(self, *args, **kwargs) -> CredentialsProvider:
         # If the content is not running on Connect then fall back to local_strategy

src/posit/connect/oauth/__init__.py

@@ -1 +1,2 @@
+from .oauth import Credentials as Credentials
 from .oauth import OAuth as OAuth

src/posit/connect/oauth/oauth.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from typing import Optional
 
 from typing_extensions import TypedDict
@@ -8,12 +9,36 @@
 from .integrations import Integrations
 from .sessions import Sessions
 
+GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
+USER_SESSION_TOKEN_TYPE = "urn:posit:connect:user-session-token"
+CONTENT_SESSION_TOKEN_TYPE = "urn:posit:connect:content-session-token"
+
+def _get_content_session_token() -> str:
+    """Return the content session token.
+
+    Reads the environment variable 'CONNECT_CONTENT_SESSION_TOKEN'.
+
+    Raises
+    ------
+        ValueError: If CONNECT_CONTENT_SESSION_TOKEN is not set or invalid
+
+    Returns
+    -------
+        str
+    """
+    value = os.environ.get("CONNECT_CONTENT_SESSION_TOKEN")
+    if not value:
+        raise ValueError("Invalid value for 'CONNECT_CONTENT_SESSION_TOKEN': Must be a non-empty string.")
+    return value
 
 class OAuth(Resources):
     def __init__(self, params: ResourceParameters, api_key: str) -> None:
         super().__init__(params)
         self.api_key = api_key
 
+    def _get_credentials_url(self) -> str: 
+        return self.params.url + "v1/oauth/integrations/credentials"
+
     @property
     def integrations(self):
         return Integrations(self.params)
@@ -23,18 +48,27 @@ def sessions(self):
         return Sessions(self.params)
 
     def get_credentials(self, user_session_token: Optional[str] = None) -> Credentials:
-        url = self.params.url + "v1/oauth/integrations/credentials"
-
+        """Perform an oauth credential exchange with a user-session-token."""
         # craft a credential exchange request
         data = {}
-        data["grant_type"] = "urn:ietf:params:oauth:grant-type:token-exchange"
-        data["subject_token_type"] = "urn:posit:connect:user-session-token"
+        data["grant_type"] = GRANT_TYPE
+        data["subject_token_type"] = USER_SESSION_TOKEN_TYPE
         if user_session_token:
             data["subject_token"] = user_session_token
 
-        response = self.params.session.post(url, data=data)
+        response = self.params.session.post(self._get_credentials_url(), data=data)
         return Credentials(**response.json())
 
+    def get_content_credentials(self, content_session_token: Optional[str] = None) -> Credentials:
+        """Perform an oauth credential exchange with a content-session-token."""
+        # craft a credential exchange request
+        data = {}
+        data["grant_type"] = GRANT_TYPE
+        data["subject_token_type"] = CONTENT_SESSION_TOKEN_TYPE
+        data["subject_token"] = content_session_token or _get_content_session_token() 
+
+        response = self.params.session.post(self._get_credentials_url(), data=data)
+        return Credentials(**response.json())
 
 class Credentials(TypedDict, total=False):
     access_token: str