Problem

When attempting to create a serverless budget policy using the workspace-level API, you receive a 404 error.

ERROR: Failed to create budget policy: 404
Endpoint not found

Cause

The API endpoint used is not correct or accessible for workspace-level policy creation. 

Solution

If you want to create a serverless budget policy at the workspace level, use the UI. For details, refer to the “Create a serverless budget policy” section of the Attribute usage with serverless budget policies (AWS | Azure | GCP) documentation.

To continue using an API to create a serverless budget policy, do so at the account-level instead. For more information, review the Authorizing access to Databricks resources (AWS | Azure | GCP) documentation.

Then, you can use the billing admin role to access the policy in your workspace. 

Account-level API example 

First, define your account console endpoint based on your cloud platform.

• (AWS) cloud.databricks.com 
• (Azure) azuredatabricks.net 
• (GCP) gcp.databricks.com 

Next, create your budget policy. 

POST
https://<account-console-endpoint>/api/2.1/accounts/<account-id>/budget-policies

For more information about the Budget policy account console API, review the Create a budget policy (AWS | Azure | GCP) API documentation. 

Then, update your budget policy rules to provide specific permissions.

PUT
https://<account-console-endpoint>/api/2.0/preview/accounts/<account-id>/access-control/rule-sets

For more information about updating a budget policy rule, review the Update a rule set (AWS | Azure | GCP) API documentation.

Implementation example in Python

The following code example

1. Creates a budget policy.
2. Gets the policy ID and policy etag to perform changes in the budget policy.
3. Defines users to assign permissions for the budget policy.
4. Updates policy rules.

import requests
api_key = "<your-token>"
host = "https://<account-console-endpoint>"
account_id = "<your-account-id>"
head = {
  "Authorization" : "Bearer " + api_key
}

def create_budget_policy(name):
 create = requests.post(f"{host}/api/2.1/accounts/{account_id}/budget-policies", json={
   "policy_name" : name,
   "custom_tags" : [{
     "key" : "tag1", "value" : "val1"
   }, {
     "key" : "tag2", "value" : "val2"
   }]
 }, headers=head)
 return create.json()["policy_id"] #Add error handling as required


def get_policy_id_from_policy_name(name):
 try:
  get_name = requests.get(f"{host}/api/2.1/accounts/{account_id}/budget-policies", headers=head)
  policy_id = None
  for policy in get_name.json()['policies']:
   if policy['policy_name'] == name:
    policy_id = policy['policy_id']
  return policy_id
 except Exception as e:
  return None

def get_etag(policy_id):
 etag = requests.get(f"{host}/api/2.1/preview/accounts/{account_id}/access-control/rule-sets", params={
  "name" : f"accounts/{account_id}/budgetPolicies/{policy_id}/ruleSets/default",
  "etag" : ""
 }, headers=head)
 return etag.json()['etag']

def get_req_body():
 #take below details as parameter in actual script, this is demo script
 users = ["user1@example.com", "user2@example.com"] #Will by default give user role
 service_principal = ["00000-0000-0000-0000-000000", "111111-1111-1111-1111-111111"] #Will give manager role. Please update the script if different service principal need to have different role.
 groups = ["group_example_name"]
  
 data = []
  
 for u in users:
  data.append({
   "role" : "roles/budgetPolicy.user",
   "principals" : [
    f"users/{u}"
   ]
  })
 for s in service_principal:
  data.append({
   "role" : "roles/budgetPolicy.manager",
   "principals" : [
    f"servicePrincipals/{s}"
   ]
  })
   
 for g in groups:
  data.append({
   "role" : "roles/budgetPolicy.manager",
   "principals" : [
    f"groups/{g}"
   ]
  })
   
 return data


def update_budget_rules(policy_id, etag):
  
 print({
  "name" : f"accounts/{account_id}/budgetPolicies/{policy_id}/ruleSets/default",
  "rule_set" : {
  "name" : f"accounts/{account_id}/budgetPolicies/{policy_id}/ruleSets/default",
  "description" : "",
  "grant_rules" : get_req_body(),
  "etag" : etag
 }})
  
 assign = requests.put(f"{host}/api/2.1/preview/accounts/{account_id}/access-control/rule-sets", json={
  "name" : f"accounts/{account_id}/budgetPolicies/{policy_id}/ruleSets/default",
  "rule_set" : {
  "name" : f"accounts/{account_id}/budgetPolicies/{policy_id}/ruleSets/default",
  "description" : "",
  "grant_rules" : get_req_body(),
  "etag" : etag
 }}, headers=head)
 return None

if __name__ == "__main__":
 policy_name = "<your_policy_name>"
 policy_id = get_policy_id_from_policy_name(policy_name)
 policy_id = create_budget_policy(policy_name) if not policy_id else policy_id
 etag = get_etag(policy_id=policy_id)
 update_budget_rules(policy_id=policy_id, etag=etag)