From ad3dba8d217f791ef672a5723d71a6f04b99715e Mon Sep 17 00:00:00 2001 From: David Kegley Date: Fri, 23 Feb 2024 17:32:34 -0500 Subject: [PATCH 01/12] Add Databricks OAuthIntegration credentials provider --- src/posit/connect/client.py | 16 +++++ src/posit/connect/client_test.py | 1 - src/posit/connect/external/__init__.py | 0 src/posit/connect/external/databricks.py | 60 +++++++++++++++++++ src/posit/connect/external/databricks_test.py | 0 src/posit/connect/oauth.py | 35 +++++++++++ src/posit/connect/oauth_test.py | 0 7 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 src/posit/connect/external/__init__.py create mode 100644 src/posit/connect/external/databricks.py create mode 100644 src/posit/connect/external/databricks_test.py create mode 100644 src/posit/connect/oauth.py create mode 100644 src/posit/connect/oauth_test.py diff --git a/src/posit/connect/client.py b/src/posit/connect/client.py index 06142572..1a6c83a3 100644 --- a/src/posit/connect/client.py +++ b/src/posit/connect/client.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from requests import Response, Session from typing import Optional @@ -7,9 +8,17 @@ from .auth import Auth from .config import Config +from .oauth import OAuthIntegration from .users import Users +# Connect sets the value of the environment variable RSTUDIO_PRODUCT = CONNECT +# when content is running on a Connect server. Use this var to determine if the +# client SDK was initialized from a piece of content running on a Connect server. +def is_local() -> bool: + return not os.getenv("RSTUDIO_PRODUCT") == "CONNECT" + + class Client: def __init__( self, @@ -37,6 +46,7 @@ def __init__( # Place to cache the server settings self.server_settings = None + self._oauth = None @property def connect_version(self): @@ -44,6 +54,12 @@ def connect_version(self): self.server_settings = self.get("server_settings").json() return self.server_settings["version"] + @property + def oauth(self) -> OAuthIntegration: + if self._oauth is None: + self._oauth = OAuthIntegration(config=self.config, session=self.session) + return self._oauth + @property def users(self) -> Users: return Users(config=self.config, session=self.session) diff --git a/src/posit/connect/client_test.py b/src/posit/connect/client_test.py index 349fb3af..6046109c 100644 --- a/src/posit/connect/client_test.py +++ b/src/posit/connect/client_test.py @@ -3,7 +3,6 @@ from unittest.mock import MagicMock, patch - from .client import Client diff --git a/src/posit/connect/external/__init__.py b/src/posit/connect/external/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/posit/connect/external/databricks.py b/src/posit/connect/external/databricks.py new file mode 100644 index 00000000..e7337f2d --- /dev/null +++ b/src/posit/connect/external/databricks.py @@ -0,0 +1,60 @@ +import abc +from typing import Callable, Dict, Optional + +from ..client import Client, is_local +from ..oauth import OAuthIntegration + +HeaderFactory = Callable[[], Dict[str, str]] + +# https://github.com/databricks/databricks-sdk-py/blob/v0.20.0/databricks/sdk/credentials_provider.py +# https://github.com/databricks/databricks-sql-python/blob/v3.1.0/src/databricks/sql/auth/authenticators.py +# In order to keep compatibility with the Databricks SDK +class CredentialsProvider(abc.ABC): + """CredentialsProvider is the protocol (call-side interface) + for authenticating requests to Databricks REST APIs""" + + @abc.abstractmethod + def auth_type(self) -> str: + ... + + @abc.abstractmethod + def __call__(self, *args, **kwargs) -> HeaderFactory: + ... + + +class PositOAuthIntegrationCredentialsProvider(CredentialsProvider): + def __init__(self, posit_oauth: OAuthIntegration, user_identity: str): + self.posit_oauth = posit_oauth + self.user_identity = user_identity + + def auth_type(self) -> str: + return "posit-oauth-integration" + + def __call__(self, *args, **kwargs) -> HeaderFactory: + def inner() -> Dict[str, str]: + access_token = self.posit_oauth.get_credentials(self.user_identity).json()['access_token'] + return {"Authorization": f"Bearer {access_token}"} + return inner + + +def viewer_credentials_provider(client: Optional[Client], user_identity: Optional[str]) -> Optional[CredentialsProvider]: + + # If the content is not running on Connect then viewer auth should + # fall back to the locally configured credentials hierarchy + if is_local(): + return None + + if client is None: + client = Client() + + + # If the user-identity-token wasn't provided and we're running on Connect then we raise an exception. + # user_identity is required to impersonate the viewer. + if user_identity is None: + raise Exception("The user-identity-token is required for viewer authentication.") + + return PositOAuthIntegrationCredentialsProvider(client.oauth, user_identity) + + +def service_account_credentials_provider(client: Optional[Client]): + raise NotImplemented diff --git a/src/posit/connect/external/databricks_test.py b/src/posit/connect/external/databricks_test.py new file mode 100644 index 00000000..e69de29b diff --git a/src/posit/connect/oauth.py b/src/posit/connect/oauth.py new file mode 100644 index 00000000..ce8c11df --- /dev/null +++ b/src/posit/connect/oauth.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from requests import Response, Session +from typing import Optional + +from . import urls +from .config import Config + + +class OAuthIntegration: + + def __init__( + self, config: Config, session: Session + ) -> None: + self.url = urls.append_path(config.url, "v1/oauth/integrations/credentials") + self.config = config + self.session = session + + + def get_credentials(self, user_identity: Optional[str]) -> Response: + + # craft a basic credential exchange request where the self.config.api_key owner + # is requesting their own credentials + data = dict() + data["grant_type"] = "urn:ietf:params:oauth:grant-type:token-exchange" + data["subject_token_type"] = "urn:posit:connect:api-key" + data["subject_token"] = self.config.api_key + + # if this content is running on Connect, then it is allowed to request + # the content viewer's credentials + if user_identity: + data["subject_token_type"] = "urn:posit:connect:user-identity-token" + data["subject_token"] = user_identity + + return self.session.post(self.url, json=data) diff --git a/src/posit/connect/oauth_test.py b/src/posit/connect/oauth_test.py new file mode 100644 index 00000000..e69de29b From 325ba3037e7db2df1adaa7effe07b0be3f85e154 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 09:32:27 -0500 Subject: [PATCH 02/12] Return Credentials response. Dont cache oauthIntegration resource in client --- src/posit/connect/client.py | 5 +---- src/posit/connect/oauth.py | 15 +++++++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/src/posit/connect/client.py b/src/posit/connect/client.py index 1a6c83a3..7dd2e1b0 100644 --- a/src/posit/connect/client.py +++ b/src/posit/connect/client.py @@ -46,7 +46,6 @@ def __init__( # Place to cache the server settings self.server_settings = None - self._oauth = None @property def connect_version(self): @@ -56,9 +55,7 @@ def connect_version(self): @property def oauth(self) -> OAuthIntegration: - if self._oauth is None: - self._oauth = OAuthIntegration(config=self.config, session=self.session) - return self._oauth + return OAuthIntegration(config=self.config, session=self.session) @property def users(self) -> Users: diff --git a/src/posit/connect/oauth.py b/src/posit/connect/oauth.py index ce8c11df..3ff61e09 100644 --- a/src/posit/connect/oauth.py +++ b/src/posit/connect/oauth.py @@ -1,12 +1,18 @@ from __future__ import annotations -from requests import Response, Session -from typing import Optional +from requests import Session +from typing import Optional, TypedDict from . import urls from .config import Config +class Credentials(TypedDict, total=False): + access_token: str + issued_token_type: str + token_type: str + + class OAuthIntegration: def __init__( @@ -17,7 +23,7 @@ def __init__( self.session = session - def get_credentials(self, user_identity: Optional[str]) -> Response: + def get_credentials(self, user_identity: Optional[str]) -> Credentials: # craft a basic credential exchange request where the self.config.api_key owner # is requesting their own credentials @@ -32,4 +38,5 @@ def get_credentials(self, user_identity: Optional[str]) -> Response: data["subject_token_type"] = "urn:posit:connect:user-identity-token" data["subject_token"] = user_identity - return self.session.post(self.url, json=data) + response = self.session.post(self.url, json=data) + return Credentials(**response.json()) From 29dfd5449b6199735e53f9b8150d8c62d4c7e842 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 10:23:44 -0500 Subject: [PATCH 03/12] Add oauth integration tests --- src/posit/connect/oauth.py | 4 +-- src/posit/connect/oauth_test.py | 51 +++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/src/posit/connect/oauth.py b/src/posit/connect/oauth.py index 3ff61e09..f6ddfc54 100644 --- a/src/posit/connect/oauth.py +++ b/src/posit/connect/oauth.py @@ -23,7 +23,7 @@ def __init__( self.session = session - def get_credentials(self, user_identity: Optional[str]) -> Credentials: + def get_credentials(self, user_identity: Optional[str]=None) -> Credentials: # craft a basic credential exchange request where the self.config.api_key owner # is requesting their own credentials @@ -38,5 +38,5 @@ def get_credentials(self, user_identity: Optional[str]) -> Credentials: data["subject_token_type"] = "urn:posit:connect:user-identity-token" data["subject_token"] = user_identity - response = self.session.post(self.url, json=data) + response = self.session.post(self.url, data=data) return Credentials(**response.json()) diff --git a/src/posit/connect/oauth_test.py b/src/posit/connect/oauth_test.py index e69de29b..ca721351 100644 --- a/src/posit/connect/oauth_test.py +++ b/src/posit/connect/oauth_test.py @@ -0,0 +1,51 @@ +import responses + +from .client import Client + +class TestOAuthIntegrations: + + @responses.activate + def test_get_credentials(self): + responses.post( + "https://connect.example/__api__/v1/oauth/integrations/credentials", + match=[ + responses.matchers.urlencoded_params_matcher( + { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:posit:connect:user-identity-token", + "subject_token": "cit", + } + ) + ], + json={ + "access_token": "viewer-token", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "token_type": "Bearer", + }, + ) + responses.post( + "https://connect.example/__api__/v1/oauth/integrations/credentials", + match=[ + responses.matchers.urlencoded_params_matcher( + { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:posit:connect:api-key", + "subject_token": "12345", + } + ) + ], + json={ + "access_token": "sdk-user-token", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "token_type": "Bearer", + }, + ) + con = Client(api_key="12345", url="https://connect.example/") + assert ( + con.oauth.get_credentials()["access_token"] + == "sdk-user-token" + ) + assert ( + con.oauth.get_credentials("cit")["access_token"] + == "viewer-token" + ) From d74f2216aebe365e3fa99162d63a9213526aa368 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 11:33:02 -0500 Subject: [PATCH 04/12] Test databricks content on Connect --- src/posit/connect/external/databricks.py | 47 ++++++++++++------------ 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/src/posit/connect/external/databricks.py b/src/posit/connect/external/databricks.py index e7337f2d..4d8f4bcd 100644 --- a/src/posit/connect/external/databricks.py +++ b/src/posit/connect/external/databricks.py @@ -4,25 +4,25 @@ from ..client import Client, is_local from ..oauth import OAuthIntegration -HeaderFactory = Callable[[], Dict[str, str]] - -# https://github.com/databricks/databricks-sdk-py/blob/v0.20.0/databricks/sdk/credentials_provider.py -# https://github.com/databricks/databricks-sql-python/blob/v3.1.0/src/databricks/sql/auth/authenticators.py -# In order to keep compatibility with the Databricks SDK -class CredentialsProvider(abc.ABC): - """CredentialsProvider is the protocol (call-side interface) - for authenticating requests to Databricks REST APIs""" - - @abc.abstractmethod - def auth_type(self) -> str: - ... - - @abc.abstractmethod - def __call__(self, *args, **kwargs) -> HeaderFactory: - ... - - -class PositOAuthIntegrationCredentialsProvider(CredentialsProvider): +#HeaderFactory = Callable[[], Dict[str, str]] +# +## https://github.com/databricks/databricks-sdk-py/blob/v0.20.0/databricks/sdk/credentials_provider.py +## https://github.com/databricks/databricks-sql-python/blob/v3.1.0/src/databricks/sql/auth/authenticators.py +## In order to keep compatibility with the Databricks SDK +#class CredentialsProvider(abc.ABC): +# """CredentialsProvider is the protocol (call-side interface) +# for authenticating requests to Databricks REST APIs""" +# +# @abc.abstractmethod +# def auth_type(self) -> str: +# ... +# +# @abc.abstractmethod +# def __call__(self, *args, **kwargs) -> HeaderFactory: +# ... + + +class PositOAuthIntegrationCredentialsProvider: def __init__(self, posit_oauth: OAuthIntegration, user_identity: str): self.posit_oauth = posit_oauth self.user_identity = user_identity @@ -32,12 +32,12 @@ def auth_type(self) -> str: def __call__(self, *args, **kwargs) -> HeaderFactory: def inner() -> Dict[str, str]: - access_token = self.posit_oauth.get_credentials(self.user_identity).json()['access_token'] + access_token = self.posit_oauth.get_credentials(self.user_identity)['access_token'] return {"Authorization": f"Bearer {access_token}"} return inner -def viewer_credentials_provider(client: Optional[Client], user_identity: Optional[str]) -> Optional[CredentialsProvider]: +def viewer_credentials_provider(client: Optional[Client] = None, user_identity: Optional[str] = None) -> Optional[CredentialsProvider]: # If the content is not running on Connect then viewer auth should # fall back to the locally configured credentials hierarchy @@ -47,14 +47,13 @@ def viewer_credentials_provider(client: Optional[Client], user_identity: Optiona if client is None: client = Client() - # If the user-identity-token wasn't provided and we're running on Connect then we raise an exception. # user_identity is required to impersonate the viewer. if user_identity is None: - raise Exception("The user-identity-token is required for viewer authentication.") + raise ValueError("The user-identity-token is required for viewer authentication.") return PositOAuthIntegrationCredentialsProvider(client.oauth, user_identity) -def service_account_credentials_provider(client: Optional[Client]): +def service_account_credentials_provider(client: Optional[Client] = None): raise NotImplemented From 0b1360039670b90ea27980939db4e7457a144e37 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 16:10:02 -0500 Subject: [PATCH 05/12] Move is_local; check CONNECT_SERVER and CONNECT_CONTENT_GUID --- src/posit/connect/client.py | 8 ----- src/posit/connect/external/databricks.py | 41 ++++++++++++++---------- 2 files changed, 24 insertions(+), 25 deletions(-) diff --git a/src/posit/connect/client.py b/src/posit/connect/client.py index c5f7093d..fe568580 100644 --- a/src/posit/connect/client.py +++ b/src/posit/connect/client.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os from requests import Response, Session from typing import Optional @@ -12,13 +11,6 @@ from .users import User, Users -# Connect sets the value of the environment variable RSTUDIO_PRODUCT = CONNECT -# when content is running on a Connect server. Use this var to determine if the -# client SDK was initialized from a piece of content running on a Connect server. -def is_local() -> bool: - return not os.getenv("RSTUDIO_PRODUCT") == "CONNECT" - - class Client: def __init__( self, diff --git a/src/posit/connect/external/databricks.py b/src/posit/connect/external/databricks.py index 4d8f4bcd..63e5d241 100644 --- a/src/posit/connect/external/databricks.py +++ b/src/posit/connect/external/databricks.py @@ -1,25 +1,26 @@ import abc +import os from typing import Callable, Dict, Optional -from ..client import Client, is_local +from ..client import Client from ..oauth import OAuthIntegration -#HeaderFactory = Callable[[], Dict[str, str]] -# -## https://github.com/databricks/databricks-sdk-py/blob/v0.20.0/databricks/sdk/credentials_provider.py -## https://github.com/databricks/databricks-sql-python/blob/v3.1.0/src/databricks/sql/auth/authenticators.py -## In order to keep compatibility with the Databricks SDK -#class CredentialsProvider(abc.ABC): -# """CredentialsProvider is the protocol (call-side interface) -# for authenticating requests to Databricks REST APIs""" -# -# @abc.abstractmethod -# def auth_type(self) -> str: -# ... -# -# @abc.abstractmethod -# def __call__(self, *args, **kwargs) -> HeaderFactory: -# ... +HeaderFactory = Callable[[], Dict[str, str]] + +# https://github.com/databricks/databricks-sdk-py/blob/v0.20.0/databricks/sdk/credentials_provider.py +# https://github.com/databricks/databricks-sql-python/blob/v3.1.0/src/databricks/sql/auth/authenticators.py +# In order to keep compatibility with the Databricks SDK +class CredentialsProvider(abc.ABC): + """CredentialsProvider is the protocol (call-side interface) + for authenticating requests to Databricks REST APIs""" + + @abc.abstractmethod + def auth_type(self) -> str: + raise NotImplemented + + @abc.abstractmethod + def __call__(self, *args, **kwargs) -> HeaderFactory: + raise NotImplemented class PositOAuthIntegrationCredentialsProvider: @@ -37,6 +38,12 @@ def inner() -> Dict[str, str]: return inner +# Use these environment varariables to determine if the +# client SDK was initialized from a piece of content running on a Connect server. +def is_local() -> bool: + return not os.getenv("CONNECT_SERVER") and not os.getenv("CONNECT_CONTENT_GUID") + + def viewer_credentials_provider(client: Optional[Client] = None, user_identity: Optional[str] = None) -> Optional[CredentialsProvider]: # If the content is not running on Connect then viewer auth should From 76903480a7e28e895ce7eda6001f8f4c10af21b4 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 16:16:43 -0500 Subject: [PATCH 06/12] Use post body temporarily for oauth/integration/credentials endpoints --- src/posit/connect/oauth.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/posit/connect/oauth.py b/src/posit/connect/oauth.py index f6ddfc54..89f116c0 100644 --- a/src/posit/connect/oauth.py +++ b/src/posit/connect/oauth.py @@ -38,5 +38,6 @@ def get_credentials(self, user_identity: Optional[str]=None) -> Credentials: data["subject_token_type"] = "urn:posit:connect:user-identity-token" data["subject_token"] = user_identity - response = self.session.post(self.url, data=data) + # TODO: use data= when the endpoint is updated to use a form-post + response = self.session.post(self.url, json=data) return Credentials(**response.json()) From ea91a374b9fab936d0deae41e9faad364c1ed0e2 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 16:28:35 -0500 Subject: [PATCH 07/12] use form xml post in credentials request --- src/posit/connect/oauth.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/posit/connect/oauth.py b/src/posit/connect/oauth.py index 89f116c0..f6ddfc54 100644 --- a/src/posit/connect/oauth.py +++ b/src/posit/connect/oauth.py @@ -38,6 +38,5 @@ def get_credentials(self, user_identity: Optional[str]=None) -> Credentials: data["subject_token_type"] = "urn:posit:connect:user-identity-token" data["subject_token"] = user_identity - # TODO: use data= when the endpoint is updated to use a form-post - response = self.session.post(self.url, json=data) + response = self.session.post(self.url, data=data) return Credentials(**response.json()) From 57508a6c2976a0ce554471e7ad5a68e598df11f1 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 16:32:50 -0500 Subject: [PATCH 08/12] Add sample databricks content --- examples/connect/.gitignore | 1 + examples/connect/databricks/README.md | 19 +++++++ examples/connect/databricks/requirements.txt | 56 +++++++++++++++++++ examples/connect/databricks/sample-content.py | 43 ++++++++++++++ 4 files changed, 119 insertions(+) create mode 100644 examples/connect/.gitignore create mode 100644 examples/connect/databricks/README.md create mode 100644 examples/connect/databricks/requirements.txt create mode 100644 examples/connect/databricks/sample-content.py diff --git a/examples/connect/.gitignore b/examples/connect/.gitignore new file mode 100644 index 00000000..afe36a39 --- /dev/null +++ b/examples/connect/.gitignore @@ -0,0 +1 @@ +.posit diff --git a/examples/connect/databricks/README.md b/examples/connect/databricks/README.md new file mode 100644 index 00000000..fb1c223f --- /dev/null +++ b/examples/connect/databricks/README.md @@ -0,0 +1,19 @@ +```bash +# start streamlit locally +DATABRICKS_TOKEN= \ +streamlit run ./sample-content.py + +# deploy the app the first time +publisher deploy -a localhost:3939 -n databricks ./ + +# re-deploy the databricks app +publisher redeploy databricks +``` + +TODO: Test this content with databricks-connect + + +``` +# install the sdk from this branch +pip install git+https://github.com/posit-dev/posit-sdk-py.git@kegs/databricks-oauth-2 +``` diff --git a/examples/connect/databricks/requirements.txt b/examples/connect/databricks/requirements.txt new file mode 100644 index 00000000..4833a049 --- /dev/null +++ b/examples/connect/databricks/requirements.txt @@ -0,0 +1,56 @@ +altair==5.2.0 +attrs==23.2.0 +blinker==1.7.0 +cachetools==5.3.2 +certifi==2024.2.2 +charset-normalizer==3.3.2 +click==8.1.7 +databricks-sdk==0.20.0 +databricks-sql-connector==3.1.0 +et-xmlfile==1.1.0 +gitdb==4.0.11 +GitPython==3.1.42 +google-auth==2.28.0 +idna==3.6 +importlib-metadata==7.0.1 +Jinja2==3.1.3 +jsonschema==4.21.1 +jsonschema-specifications==2023.12.1 +lz4==4.3.3 +markdown-it-py==3.0.0 +MarkupSafe==2.1.5 +mdurl==0.1.2 +numpy==1.26.4 +oauthlib==3.2.2 +openpyxl==3.1.2 +packaging==23.2 +pandas==2.1.4 +pillow==10.2.0 +posit-sdk @ git+https://github.com/posit-dev/posit-sdk-py.git@ea91a374b9fab936d0deae41e9faad364c1ed0e2 +protobuf==4.25.3 +pyarrow==14.0.2 +pyasn1==0.5.1 +pyasn1-modules==0.3.0 +pydeck==0.8.1b0 +Pygments==2.17.2 +python-dateutil==2.8.2 +pytz==2024.1 +referencing==0.33.0 +requests==2.31.0 +rich==13.7.0 +rpds-py==0.18.0 +rsa==4.9 +six==1.16.0 +smmap==5.0.1 +streamlit==1.31.1 +tenacity==8.2.3 +thrift==0.16.0 +toml==0.10.2 +toolz==0.12.1 +tornado==6.4 +typing_extensions==4.9.0 +tzdata==2024.1 +tzlocal==5.2 +urllib3==2.2.1 +validators==0.22.0 +zipp==3.17.0 diff --git a/examples/connect/databricks/sample-content.py b/examples/connect/databricks/sample-content.py new file mode 100644 index 00000000..1fc1c5a7 --- /dev/null +++ b/examples/connect/databricks/sample-content.py @@ -0,0 +1,43 @@ +import os + +from posit.connect.external.databricks import viewer_credentials_provider + +from databricks import sql +from databricks.sdk.service.iam import CurrentUserAPI +from databricks.sdk.core import ApiClient, Config + +import pandas as pd +import streamlit as st +from streamlit.web.server.websocket_headers import _get_websocket_headers + +DB_PAT=os.getenv("DATABRICKS_TOKEN") + +DB_HOST=os.getenv("DB_HOST") +DB_HOST_URL = f"https://{DB_HOST}" +SQL_HTTP_PATH=os.getenv("SQL_HTTP_PATH") + +USER_IDENTITY = None + +# Read the viewer's individual content identity token from the streamlit ws header. +headers = _get_websocket_headers() +if headers: + USER_IDENTITY = headers.get('Posit-Connect-User-Identity') + +credentials_provider = viewer_credentials_provider(user_identity=USER_IDENTITY) +cfg = Config(host=DB_HOST_URL, credentials_provider=credentials_provider) +#cfg = Config(host=DB_HOST_URL, token=DB_PAT) + +databricks_user = CurrentUserAPI(ApiClient(cfg)).me() +st.write(f"Hello, {databricks_user.display_name}!") + +with sql.connect( + server_hostname=DB_HOST, + http_path=SQL_HTTP_PATH, + #access_token=DB_PAT) as connection: + auth_type='databricks-oauth', + credentials_provider=credentials_provider) as connection: + with connection.cursor() as cursor: + cursor.execute("SELECT * FROM data") + result = cursor.fetchall() + st.table(pd.DataFrame(result)) + From 24ad71458de56b01b8168e80441ea860236f1933 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 16:40:42 -0500 Subject: [PATCH 09/12] Fix linter --- examples/connect/databricks/sample-content.py | 2 ++ src/posit/connect/external/databricks.py | 8 ++++---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/connect/databricks/sample-content.py b/examples/connect/databricks/sample-content.py index 1fc1c5a7..1a440dcb 100644 --- a/examples/connect/databricks/sample-content.py +++ b/examples/connect/databricks/sample-content.py @@ -1,3 +1,5 @@ +# -*- coding: utf-8 -*- +# mypy: ignore-errors import os from posit.connect.external.databricks import viewer_credentials_provider diff --git a/src/posit/connect/external/databricks.py b/src/posit/connect/external/databricks.py index 63e5d241..5e8ccda8 100644 --- a/src/posit/connect/external/databricks.py +++ b/src/posit/connect/external/databricks.py @@ -16,14 +16,14 @@ class CredentialsProvider(abc.ABC): @abc.abstractmethod def auth_type(self) -> str: - raise NotImplemented + raise NotImplementedError @abc.abstractmethod def __call__(self, *args, **kwargs) -> HeaderFactory: - raise NotImplemented + raise NotImplementedError -class PositOAuthIntegrationCredentialsProvider: +class PositOAuthIntegrationCredentialsProvider(CredentialsProvider): def __init__(self, posit_oauth: OAuthIntegration, user_identity: str): self.posit_oauth = posit_oauth self.user_identity = user_identity @@ -63,4 +63,4 @@ def viewer_credentials_provider(client: Optional[Client] = None, user_identity: def service_account_credentials_provider(client: Optional[Client] = None): - raise NotImplemented + raise NotImplementedError From 34f74491dedc4659689368dbc83d1b7a1ba441a0 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Mon, 26 Feb 2024 16:41:29 -0500 Subject: [PATCH 10/12] update sample content requirements --- examples/connect/databricks/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/connect/databricks/requirements.txt b/examples/connect/databricks/requirements.txt index 4833a049..99719b8d 100644 --- a/examples/connect/databricks/requirements.txt +++ b/examples/connect/databricks/requirements.txt @@ -26,7 +26,7 @@ openpyxl==3.1.2 packaging==23.2 pandas==2.1.4 pillow==10.2.0 -posit-sdk @ git+https://github.com/posit-dev/posit-sdk-py.git@ea91a374b9fab936d0deae41e9faad364c1ed0e2 +posit-sdk @ git+https://github.com/posit-dev/posit-sdk-py.git@24ad71458de56b01b8168e80441ea860236f1933 protobuf==4.25.3 pyarrow==14.0.2 pyasn1==0.5.1 From 8091ddc00ef8ef08e65ec32da373a1efed00e05c Mon Sep 17 00:00:00 2001 From: David Kegley Date: Tue, 27 Feb 2024 10:06:50 -0500 Subject: [PATCH 11/12] Revert env var change --- src/posit/connect/external/databricks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/posit/connect/external/databricks.py b/src/posit/connect/external/databricks.py index 5e8ccda8..ab634fc6 100644 --- a/src/posit/connect/external/databricks.py +++ b/src/posit/connect/external/databricks.py @@ -38,10 +38,10 @@ def inner() -> Dict[str, str]: return inner -# Use these environment varariables to determine if the +# Use this environment variable to determine if the # client SDK was initialized from a piece of content running on a Connect server. def is_local() -> bool: - return not os.getenv("CONNECT_SERVER") and not os.getenv("CONNECT_CONTENT_GUID") + return not os.getenv("RSTUDIO_PRODUCT") == "CONNECT" def viewer_credentials_provider(client: Optional[Client] = None, user_identity: Optional[str] = None) -> Optional[CredentialsProvider]: From 8e3a86542a805b210113d0aacf2f98d5e1731899 Mon Sep 17 00:00:00 2001 From: David Kegley Date: Tue, 27 Feb 2024 12:04:34 -0500 Subject: [PATCH 12/12] Use user-session-token instead of content-identity-token --- examples/connect/databricks/sample-content.py | 8 ++++---- src/posit/connect/external/databricks.py | 18 +++++++++--------- src/posit/connect/oauth.py | 8 ++++---- src/posit/connect/oauth_test.py | 2 +- 4 files changed, 18 insertions(+), 18 deletions(-) diff --git a/examples/connect/databricks/sample-content.py b/examples/connect/databricks/sample-content.py index 1a440dcb..3620bbc7 100644 --- a/examples/connect/databricks/sample-content.py +++ b/examples/connect/databricks/sample-content.py @@ -18,14 +18,14 @@ DB_HOST_URL = f"https://{DB_HOST}" SQL_HTTP_PATH=os.getenv("SQL_HTTP_PATH") -USER_IDENTITY = None +USER_SESSION_TOKEN = None -# Read the viewer's individual content identity token from the streamlit ws header. +# Read the viewer's user session token from the streamlit ws header. headers = _get_websocket_headers() if headers: - USER_IDENTITY = headers.get('Posit-Connect-User-Identity') + USER_SESSION_TOKEN = headers.get('Posit-Connect-User-Session') -credentials_provider = viewer_credentials_provider(user_identity=USER_IDENTITY) +credentials_provider = viewer_credentials_provider(user_session_token=USER_SESSION_TOKEN) cfg = Config(host=DB_HOST_URL, credentials_provider=credentials_provider) #cfg = Config(host=DB_HOST_URL, token=DB_PAT) diff --git a/src/posit/connect/external/databricks.py b/src/posit/connect/external/databricks.py index ab634fc6..cce1a81d 100644 --- a/src/posit/connect/external/databricks.py +++ b/src/posit/connect/external/databricks.py @@ -24,16 +24,16 @@ def __call__(self, *args, **kwargs) -> HeaderFactory: class PositOAuthIntegrationCredentialsProvider(CredentialsProvider): - def __init__(self, posit_oauth: OAuthIntegration, user_identity: str): + def __init__(self, posit_oauth: OAuthIntegration, user_session_token: str): self.posit_oauth = posit_oauth - self.user_identity = user_identity + self.user_session_token = user_session_token def auth_type(self) -> str: return "posit-oauth-integration" def __call__(self, *args, **kwargs) -> HeaderFactory: def inner() -> Dict[str, str]: - access_token = self.posit_oauth.get_credentials(self.user_identity)['access_token'] + access_token = self.posit_oauth.get_credentials(self.user_session_token)['access_token'] return {"Authorization": f"Bearer {access_token}"} return inner @@ -44,7 +44,7 @@ def is_local() -> bool: return not os.getenv("RSTUDIO_PRODUCT") == "CONNECT" -def viewer_credentials_provider(client: Optional[Client] = None, user_identity: Optional[str] = None) -> Optional[CredentialsProvider]: +def viewer_credentials_provider(client: Optional[Client] = None, user_session_token: Optional[str] = None) -> Optional[CredentialsProvider]: # If the content is not running on Connect then viewer auth should # fall back to the locally configured credentials hierarchy @@ -54,12 +54,12 @@ def viewer_credentials_provider(client: Optional[Client] = None, user_identity: if client is None: client = Client() - # If the user-identity-token wasn't provided and we're running on Connect then we raise an exception. - # user_identity is required to impersonate the viewer. - if user_identity is None: - raise ValueError("The user-identity-token is required for viewer authentication.") + # If the user-session-token wasn't provided and we're running on Connect then we raise an exception. + # user_session_token is required to impersonate the viewer. + if user_session_token is None: + raise ValueError("The user-session-token is required for viewer authentication.") - return PositOAuthIntegrationCredentialsProvider(client.oauth, user_identity) + return PositOAuthIntegrationCredentialsProvider(client.oauth, user_session_token) def service_account_credentials_provider(client: Optional[Client] = None): diff --git a/src/posit/connect/oauth.py b/src/posit/connect/oauth.py index f6ddfc54..9fecdb40 100644 --- a/src/posit/connect/oauth.py +++ b/src/posit/connect/oauth.py @@ -23,7 +23,7 @@ def __init__( self.session = session - def get_credentials(self, user_identity: Optional[str]=None) -> Credentials: + def get_credentials(self, user_session_token: Optional[str]=None) -> Credentials: # craft a basic credential exchange request where the self.config.api_key owner # is requesting their own credentials @@ -34,9 +34,9 @@ def get_credentials(self, user_identity: Optional[str]=None) -> Credentials: # if this content is running on Connect, then it is allowed to request # the content viewer's credentials - if user_identity: - data["subject_token_type"] = "urn:posit:connect:user-identity-token" - data["subject_token"] = user_identity + if user_session_token: + data["subject_token_type"] = "urn:posit:connect:user-session-token" + data["subject_token"] = user_session_token response = self.session.post(self.url, data=data) return Credentials(**response.json()) diff --git a/src/posit/connect/oauth_test.py b/src/posit/connect/oauth_test.py index ca721351..f81389b3 100644 --- a/src/posit/connect/oauth_test.py +++ b/src/posit/connect/oauth_test.py @@ -12,7 +12,7 @@ def test_get_credentials(self): responses.matchers.urlencoded_params_matcher( { "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", - "subject_token_type": "urn:posit:connect:user-identity-token", + "subject_token_type": "urn:posit:connect:user-session-token", "subject_token": "cit", } )