Custom OAuth app integration - U2M - Consent.exchange_callback_parameters(args).as_dict() response does not contain refresh_token

[milieere]

Discussed in #456

^{Originally posted by milieere November 27, 2023}
Hello,

I registered a custom OAuth app integration, to use Databricks as OAuth provider for my Dash application. I used the SDK to create the app integration:

from databricks.sdk import AccountClient

account_client = AccountClient(host="https://accounts.cloud.databricks.com",
    account_id="XXXXXXXXXXXXXXXXXXXXXXX",
    client_id="XXXXXXXXXXXXXXX",
    client_secret="XXXXXXXXXXXXXX",
)


client = account_client.custom_app_integration.create(
    name='dash-test-conf', redirect_urls=["http://localhost:8030/callback", "https://my-app-url.org/callback"], 
    confidential=True,
    scopes=["all-apis"],
)

Then, I am using the databricks.sdk.oauth classess following example here: https://github.com/databricks/databricks-sdk-py/blob/main/examples/flask_app_with_oauth.py to incorporate the 3-legged OAuth flow to my Dash app. The flow works fine. I am loading credentials from the flask session using a decorator, and I decorate my Dash callback functions with it, to pass connection object to the callback:

def authorize_databricks(func):
    def wrapper(*args, **kwargs):
        if 'databricks_creds' in session:
            credentials_provider = SessionCredentials.from_dict(oauth_client, session["databricks_creds"])
            connection = ConnectionBuilder(credentials_provider, db_config)
            if connection:
                return func(*args, connection=connection, **kwargs)  # Pass the 'token' variable to the function
        else:
            return dcc.Location(pathname=f'/{APP_ABBREVIATION}/databrickslogin', id='redirection')
    
    return wrapper

@callback(
    Output('gaia-data', 'children'),
    Input('gaia-btn', 'n_clicks'),
)
@authorize_databricks
def get_gaia_data(clicks, connection):
    if clicks:
        engine = connection.get_engine()
        result = fetch_data(engine)
        return result

Problem is that I do not obtain refresh_token from Databricks OAuth server. So when the token is expired, the SessionCredentials class is trying to refresh it, using the refresh method:

class SessionCredentials(Refreshable):

    def __init__(self, client: 'OAuthClient', token: Token):
        self._client = client
        super().__init__(token)

    def as_dict(self) -> dict:
        return {'token': self._token.as_dict()}

    @staticmethod
    def from_dict(client: 'OAuthClient', raw: dict) -> 'SessionCredentials':
        return SessionCredentials(client=client, token=Token.from_dict(raw['token']))

    def auth_type(self):
        """Implementing CredentialsProvider protocol"""
        # TODO: distinguish between Databricks IDP and Azure AD
        return 'oauth'

    def __call__(self, *args, **kwargs):
        """Implementing CredentialsProvider protocol"""

        def inner() -> Dict[str, str]:
            return {'Authorization': f"Bearer {self.token().access_token}"}

        return inner

    def refresh(self) -> Token:
        refresh_token = self._token.refresh_token
        if not refresh_token:
            raise ValueError('oauth2: token expired and refresh token is not set')
        params = {'grant_type': 'refresh_token', 'refresh_token': refresh_token}
        headers = {}
        if 'microsoft' in self._client.token_url:
            # Tokens issued for the 'Single-Page Application' client-type may
            # only be redeemed via cross-origin requests
            headers = {'Origin': self._client.redirect_url}
        return retrieve_token(client_id=self._client.client_id,
                              client_secret=self._client.client_secret,
                              token_url=self._client.token_url,
                              params=params,
                              use_params=True,
                              headers=headers)

This leads to raising the ValueError 'oauth2: token expired and refresh token is not set'.

How to make the OAuth server send refresh token? Do I need some special configuration? I haven't seen any fitting options in the API spec here: https://docs.databricks.com/api/account/customappintegration/create , tried setting the app to both confidential=True and confidential=False, but no change. the response from consent.exchange_callback_parameters(request.args).as_dict():

@server.route(app.get_relative_path('/callback'))
def callback():
    consent = Consent.from_dict(oauth_client, session["consent"])
    session["databricks_creds"] = consent.exchange_callback_parameters(request.args).as_dict()
    return redirect(app.get_relative_path('/homepage'))

always lacks refresh_token, the dict only has access_token, and expiry:

{'token': {'access_token': 'eyXXXXXXXX', 'expiry': '2023-11-27T13:09:01.580863', 'token_type': 'Bearer'}}

Thanks for any tips to unblock this.

[milieere]

Update: I developed the following workaround, that checks if the token is expired and if yes, deletes the databricks_cred from session and redirects user to the databrickslogin endpoint, that again initiates the consent, and does the whole flow again:

def authorize_databricks(func):
    def wrapper(*args, **kwargs):
        if 'databricks_creds' in session:
            credentials_provider = SessionCredentials.from_dict(oauth_client, session["databricks_creds"])
            if credentials_provider._token.expired:
                del session['databricks_creds']
                return dcc.Location(pathname=f'/{APP_ABBREVIATION}/databrickslogin', id='redirection')
            connection = ConnectionBuilder(credentials_provider, db_config)
            if connection:
                return func(*args, connection=connection, **kwargs)  # Pass the 'token' variable to the function
        else:
            return dcc.Location(pathname=f'/{APP_ABBREVIATION}/databrickslogin', id='redirection')
    
    return wrapper

But it would be better to make use of the refresh token instead.

Those are my flask OAuth routes:

@server.route(app.get_relative_path('/databrickscallback'))
def callback():
    consent = Consent.from_dict(oauth_client, session["consent"])
    session["databricks_creds"] = consent.exchange_callback_parameters(request.args).as_dict()
    return redirect(app.get_relative_path('/homepage'))

@server.route(app.get_relative_path('/databrickslogin'))
def authorize_databricks():
    if "databricks_creds" not in session:
        consent = oauth_client.initiate_consent()
        session["consent"] = consent.as_dict()
        return redirect(consent.auth_url)
    return redirect(app.get_relative_path('/homepage'))

[tanmay-db]

Hi @milieere, for getting the refresh token, you would need to set offline_access scope. We will update the example for also having this so it is more clear. Thanks for raising the issue.

[milieere]

Hi @tanmay-db thanks for the response. Just to make this clear here if anyone would come to look for help, this is the working solution:

from databricks.sdk import AccountClient
from databricks.sdk.oauth import OAuthClient

# To create custom oauth app integration in account console:
account_client = AccountClient(host="https://accounts.cloud.databricks.com",
    account_id="XXXXXXXXXXXXXXXXXXXXXXX",
    client_id="XXXXXXXXXXXXXXX",
    client_secret="XXXXXXXXXXXXXX",
)

client = account_client.custom_app_integration.create(
    name='dash-test-conf', redirect_urls=["http://localhost:8030/callback", "https://my-app-url.org/callback"], 
    confidential=True,
    scopes=["all-apis", "offline_access"],
)

# To setup oauth client:
oauth_client = OAuthClient(host=databricks_host, 
    client_id=OAUTH_CLIENT_ID,
    client_secret=OAUTH_CLIENT_SECRET,
    scopes = ["all-apis", "offline_access"],
    redirect_url=REDIRECT_URL
)

[mgyucht]

Closing this issue out as it seems the question has been answered.