Enable supplying specifying default API headers as optional kwarg

[bjtho08]

The Databricks CLI has built-in support for setting default headers when initializing the ApiClient using the optional default_headers parameter. The default_headers parameter accepts a dictionary, where the keys and values correspond to the header keys and values.

The ability to supply default headers enables programmatically acquiring Azure Service Principal Management tokens, appending them to the API calls with the appropriate header and perform actions that would otherwise require a personal access token or a Databricks Premium subscription.

Example use-case with ApiClient from Databricks-cli SDK:

dbx_host = "https://<dbx_workspace_id>.azuredatabricks.net"
aad_token = subprocess.run(
    "az account get-access-token --scope '2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default' --query 'accessToken' -o tsv",
    shell=True, capture_output=True, text=True).stdout
sp_headers = {
    "X-Databricks-Azure-Workspace-Resource-Id": <Azure_Resource_ID>,
    "X-Databricks-Azure-SP-Management-Token": <Azure_SP_Management_Token>
}
api_client = ApiClient(host=dbx_host, token=aad_token, default_headers=sp_headers)

[copdips]

@bjtho08 just curious, by providing the access-token only but without the sp_headers shown in your example, does the ApiClient still work or not ? The reason why I ask this is because it seems that Databricks API auth only needs an access token.

[bjtho08]

@bjtho08 just curious, by providing the access-token only but without the sp_headers shown in your example, does the ApiClient still work or not ? The reason why I ask this is because it seems that Databricks API auth only needs an access token.

There are two instances where that would work:

1. If the access token is a personal access token generated by the user in the databricks workspace.
2. If the databricks workspace is on the premium price plan.

Using PATs is considered bad practice for large corporations and we currently don't have the premium plan.

[copdips]

Yes, the workspace I'm using is in the premium plan, but I'm not sure the way I grant the access token is not included in the standard plan.

In fact, I don't use the user PAT created from the databricks workspace UI neither, but I use a service principal, by submitting its client_id and client_secret to grant a standard lifetime limited oauth2 access token.

I just read the doc, it says that we need to send the additional headers (as you showed in your example) in case of the SP is not a Databricks workspace user, I'm wondering why not just add the SP to Databricks ?

[bjtho08]

I just read the doc, it says that we need to send the additional headers (as you showed in your example) in case of the SP is not a Databricks workspace user, I'm wondering why not just add the SP to Databricks ?

The way databricks run their business, service principals can only be added to a workspace via the SCIM API, which is only available for the premium plan.

[copdips]

okey, thanks for the confirmation, it's good to know that.

[bjtho08]

Is any work being done on this feature or would it be okay if I take a stab at it myself?

[renardeinside]

hi @bjtho08 , unfortunately development capacity is limited as of the moment.

If you have a concept/idea/implementation of this feature, I'm more than happy to merge your PR - any help would be much appreciated.