Databricks access_token

databricks-auth-token.py

# https://docs.databricks.com/en/dev-tools/auth/oauth-u2m.html#generate-oauth-pair
# the auth is handled via azure cli
# you need to issue `az login` and set the sub name|id where the ws is attached.

import cmd
import uuid, hashlib, base64


def generate_codes():
    # Generate a UUID.
    uuid1 = uuid.uuid4()

    # Convert the UUID to a string.
    uuid_str1 = str(uuid1).upper()

    # Create the code verifier.
    code_verifier = uuid_str1 + "-" + uuid_str1

    # Create the code challenge based on the code verifier.
    code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode('utf-8')

    # Remove all padding from the code challenge.
    code_challenge = code_challenge.replace('=', '')

    # Print the code verifier and the code challenge.
    # Use these in your calls to manually generate
    # access tokens for OAuth U2M authentication.
    # print(f"code_verifier:  {code_verifier}")
    # print(f"code_challenge: {code_challenge}")

    return code_verifier, code_challenge


def build_authorization_code_url(workspace_url, code_challenge, redirect_url='http://localhost:8020', state_str='databricks-cli'):
    return f"https://{workspace_url}/oidc/v1/authorize?client_id=databricks-cli&redirect_uri={redirect_url}&response_type=code&state={state_str}&code_challenge={code_challenge}&code_challenge_method=S256&scope=all-apis+offline_access"


def build_access_token_url(workspace_url, authorization_code, code_verifier, redirect_url='http://localhost:8020'):
    return f'curl --request POST https://{workspace_url}/oidc/v1/token --data "client_id=databricks-cli" --data "grant_type=authorization_code" --data "scope=all-apis offline_access" --data "redirect_uri={redirect_url}" --data "code_verifier={code_verifier}" --data "code={authorization_code}"'


class DatabricksCmd(cmd.Cmd):

    intro = 'This program obtains an access token from Databricks. Type help or ? to list commands'
    prompt = '(db-cli) '

    workspace_url = None
    redirect_url  = 'http://localhost:8020'
    state_str     = 'databricks-cli'
    code_verifier, code_challenge = generate_codes()

    def do_set_host_url(self, arg):
        """ Set the host url - Can be account or workspace """
        print(arg)
        self.workspace_url = arg

    def do_get_settings(self, arg):
        """ Print current host """
        print(f'workspace_url={self.workspace_url}')
        print(f'redirect_url={self.redirect_url}')
        print(f'state_str={self.state_str}')
        print(f'code_verifier={self.code_verifier}')
        print(f'code_challenge={self.code_challenge}')

    def do_get_authorization_code_url(self, arg):
        """ Constructs an authorization url that
        you can paste in your browser to obtain the authorization code.
        """
        print(build_authorization_code_url(self.workspace_url, self.code_challenge))

    def do_get_authentication_token_url(self, arg):
        """ Constructs url to obtain access token
        Paste in the authorization code obtain from browser.
        The code starts with 'dcod*'.
        """
        authorization_code = arg  # TODO validate it starts with dc
        print(build_access_token_url(self.workspace_url, authorization_code, self.code_verifier))

    def do_bye(self, arg):
        'Close this program.'
        print('Bye')
        return True


if __name__ == '__main__':

    DatabricksCmd().cmdloop()