feat: add optional audience parameter to OAuth credential exchange methods and related functions

Author: mconflitti-pbc

src/posit/connect/client.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 from typing_extensions import TYPE_CHECKING, overload
 
 from . import hooks, me
@@ -176,7 +178,11 @@ def __init__(self, *args, **kwargs) -> None:
         self._ctx = Context(self)
 
     @requires("2025.01.0")
-    def with_user_session_token(self, token: str) -> Client:
+    def with_user_session_token(
+        self,
+        token: str,
+        audience: Optional[str] = None,
+    ) -> Client:
         """Create a new Client scoped to the user specified in the user session token.
 
         Create a new Client instance from a user session token exchange for an api key scoped to the
@@ -256,7 +262,9 @@ def user_profile():
             raise ValueError("token must be set to non-empty string.")
 
         visitor_credentials = self.oauth.get_credentials(
-            token, requested_token_type=OAuthTokenType.API_KEY
+            token,
+            requested_token_type=OAuthTokenType.API_KEY,
+            audience=audience,
         )
 
         visitor_api_key = visitor_credentials.get("access_token", "")

src/posit/connect/external/aws.py

@@ -19,7 +19,11 @@ class Credentials(TypedDict):
     expiration: datetime
 
 
-def get_credentials(client: Client, user_session_token: str) -> Credentials:
+def get_credentials(
+    client: Client,
+    user_session_token: str,
+    audience: Optional[str] = None,
+) -> Credentials:
     """
     Get AWS credentials using OAuth token exchange for an AWS Viewer integration.
 
@@ -66,6 +70,7 @@ def get_credentials(client: Client, user_session_token: str) -> Credentials:
     credentials = client.oauth.get_credentials(
         user_session_token=user_session_token,
         requested_token_type=OAuthTokenType.AWS_CREDENTIALS,
+        audience=audience,
     )
 
     # Decode base64 access token
@@ -76,7 +81,9 @@ def get_credentials(client: Client, user_session_token: str) -> Credentials:
 
 
 def get_content_credentials(
-    client: Client, content_session_token: Optional[str] = None
+    client: Client,
+    content_session_token: Optional[str] = None,
+    audience: Optional[str] = None,
 ) -> Credentials:
     """
     Get AWS credentials using OAuth token exchange for an AWS Service Account integration.
@@ -122,6 +129,7 @@ def get_content_credentials(
     credentials = client.oauth.get_content_credentials(
         content_session_token=content_session_token,
         requested_token_type=OAuthTokenType.AWS_CREDENTIALS,
+        audience=audience,
     )
 
     # Decode base64 access token

src/posit/connect/external/databricks.py

@@ -62,11 +62,16 @@ class _PositConnectContentCredentialsProvider:
     * https://github.com/posit-dev/posit-sdk-py/blob/main/src/posit/connect/oauth/oauth.py
     """
 
-    def __init__(self, client: Client):
+    def __init__(
+        self,
+        client: Client,
+        audience: Optional[str] = None,
+    ):
         self._client = client
+        self._audience = audience
 
     def __call__(self) -> Dict[str, str]:
-        credentials = self._client.oauth.get_content_credentials()
+        credentials = self._client.oauth.get_content_credentials(audience=self._audience)
         return _new_bearer_authorization_header(credentials)
 
 
@@ -81,12 +86,21 @@ class _PositConnectViewerCredentialsProvider:
     * https://github.com/posit-dev/posit-sdk-py/blob/main/src/posit/connect/oauth/oauth.py
     """
 
-    def __init__(self, client: Client, user_session_token: str):
+    def __init__(
+        self,
+        client: Client,
+        user_session_token: str,
+        audience: Optional[str] = None,
+    ):
         self._client = client
         self._user_session_token = user_session_token
+        self._audience = audience
 
     def __call__(self) -> Dict[str, str]:
-        credentials = self._client.oauth.get_credentials(self._user_session_token)
+        credentials = self._client.oauth.get_credentials(
+            self._user_session_token,
+            audience=self._audience,
+        )
         return _new_bearer_authorization_header(credentials)
 
 
@@ -174,10 +188,12 @@ def __init__(
         self,
         client: Optional[Client] = None,
         user_session_token: Optional[str] = None,
+        audience: Optional[str] = None,
     ):
         self._cp: Optional[CredentialsProvider] = None
         self._client = client
         self._user_session_token = user_session_token
+        self._audience = audience
 
     def auth_type(self) -> str:
         return POSIT_OAUTH_INTEGRATION_AUTH_TYPE
@@ -194,13 +210,18 @@ def __call__(self, *args, **kwargs) -> CredentialsProvider:  # noqa: ARG002
         if self._cp is None:
             if self._user_session_token:
                 self._cp = _PositConnectViewerCredentialsProvider(
-                    self._client, self._user_session_token
+                    self._client,
+                    self._user_session_token,
+                    audience=self._audience,
                 )
             else:
                 logger.info(
                     "ConnectStrategy will attempt to use OAuth Service Account credentials because user_session_token is not set"
                 )
-                self._cp = _PositConnectContentCredentialsProvider(self._client)
+                self._cp = _PositConnectContentCredentialsProvider(
+                    self._client,
+                    audience=self._audience,
+                )
         return self._cp
 
 

src/posit/connect/external/snowflake.py

@@ -69,10 +69,12 @@ def __init__(
         local_authenticator: Optional[str] = None,
         client: Optional[Client] = None,
         user_session_token: Optional[str] = None,
+        audience: Optional[str] = None,
     ):
         self._local_authenticator = local_authenticator
         self._client = client
         self._user_session_token = user_session_token
+        self._audience = audience
 
     @property
     def authenticator(self) -> Optional[str]:
@@ -93,5 +95,8 @@ def token(self) -> Optional[str]:
         if self._client is None:
             self._client = Client()
 
-        credentials = self._client.oauth.get_credentials(self._user_session_token)
+        credentials = self._client.oauth.get_credentials(
+            self._user_session_token,
+            audience=self._audience,
+        )
         return credentials.get("access_token")

src/posit/connect/oauth/associations.py

@@ -1,10 +1,16 @@
 """OAuth association resources."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from typing_extensions import List
 
-from ..context import Context
 from ..resources import BaseResource, Resources
 
+if TYPE_CHECKING:
+    from ..context import Context
+
 
 class Association(BaseResource):
     pass
@@ -66,9 +72,11 @@ def delete(self) -> None:
         path = f"v1/content/{self.content_guid}/oauth/integrations/associations"
         self._ctx.client.put(path, json=data)
 
-    def update(self, integration_guid: str) -> None:
+    def update(self, integration_guid: str | list[str]) -> None:
         """Set integration associations."""
-        data = [{"oauth_integration_guid": integration_guid}]
+        if isinstance(integration_guid, str):
+            integration_guid = [integration_guid]
+        data = [{"oauth_integration_guid": guid} for guid in integration_guid]
 
         path = f"v1/content/{self.content_guid}/oauth/integrations/associations"
         self._ctx.client.put(path, json=data)

src/posit/connect/oauth/oauth.py

@@ -62,6 +62,7 @@ def get_credentials(
         self,
         user_session_token: Optional[str] = None,
         requested_token_type: Optional[str | OAuthTokenType] = None,
+        audience: Optional[str] = None,
     ) -> Credentials:
         """Perform an oauth credential exchange with a user-session-token."""
         # craft a credential exchange request
@@ -72,6 +73,8 @@ def get_credentials(
             data["subject_token"] = user_session_token
         if requested_token_type:
             data["requested_token_type"] = requested_token_type
+        if audience:
+            data["audience"] = audience
 
         response = self._ctx.client.post(self._path, data=data)
         return Credentials(**response.json())
@@ -80,6 +83,7 @@ def get_content_credentials(
         self,
         content_session_token: Optional[str] = None,
         requested_token_type: Optional[str | OAuthTokenType] = None,
+        audience: Optional[str] = None,
     ) -> Credentials:
         """Perform an oauth credential exchange with a content-session-token."""
         # craft a credential exchange request
@@ -89,6 +93,8 @@ def get_content_credentials(
         data["subject_token"] = content_session_token or _get_content_session_token()
         if requested_token_type:
             data["requested_token_type"] = requested_token_type
+        if audience:
+            data["audience"] = audience
 
         response = self._ctx.client.post(self._path, data=data)
         return Credentials(**response.json())