[Exporter] Added support for databricks_budget_policy resource (#5217)

## Changes
<!-- Summary of your changes that are easy to understand -->

## Tests
<!--
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [x] `make test` run locally
- [x] relevant change in `docs/` folder
- [ ] covered with integration tests in `internal/acceptance`
- [x] using Go SDK
- [x] using TF Plugin Framework
- [x] has entry in `NEXT_CHANGELOG.md` file

Co-authored-by: Omer Lachish <289488+rauchy@users.noreply.github.com>

Author: alexott

NEXT_CHANGELOG.md

@@ -26,6 +26,7 @@
 ### Exporter
 
 * Added support for `databricks_data_quality_monitor` resource ([#5193](https://github.com/databricks/terraform-provider-databricks/pull/5193)).
+* Added support for `databricks_budget_policy` resource ([#5217](https://github.com/databricks/terraform-provider-databricks/pull/5217)).
 * Fix typo in the name of environment variable ([#5158](https://github.com/databricks/terraform-provider-databricks/pull/5158)).
 * Export permission assignments on workspace level ([#5169](https://github.com/databricks/terraform-provider-databricks/pull/5169)).
 * Added support for UC Tag policies ([#5213](https://github.com/databricks/terraform-provider-databricks/pull/5213)).

docs/guides/experimental-exporter.md

@@ -175,7 +175,7 @@ Services could be specified in combination with predefined aliases (`all` - for
 * `access` -  **listing** [databricks_permissions](../resources/permissions.md), [databricks_instance_profile](../resources/instance_profile.md), [databricks_ip_access_list](../resources/ip_access_list.md), and [databricks_access_control_rule_set](../resources/access_control_rule_set.md).   *Please note that for `databricks_permissions` we list only `authorization = "tokens"`, the permissions for other objects (notebooks, ...) will be emitted when corresponding objects are processed!*
 * `alerts` - **listing** [databricks_alert](../resources/alert.md) and [databricks_alert_v2](../resources/alert_v2.md).
 * `apps` - **listing** [databricks_app](../resources/app.md) and [databricks_apps_settings_custom_template](../resources/apps_settings_custom_template.md).
-* `billing` - **listing** [databricks_budget](../resources/budget.md).
+* `billing` - **listing** [databricks_budget](../resources/budget.md) and [databricks_budget_policy](../resources/budget_policy.md).
 * `compute` - **listing** [databricks_cluster](../resources/cluster.md).
 * `dashboards` - **listing** [databricks_dashboard](../resources/dashboard.md).
 * `directories` - **listing** [databricks_directory](../resources/directory.md).  *Please note that directories aren't listed when running in the incremental mode! Only directories with updated notebooks will be emitted.*
@@ -249,6 +249,7 @@ Exporter aims to generate HCL code for most of the resources within the Databric
 | [databricks_apps_settings_custom_template](../resources/apps_settings_custom_template.md) | Yes | No | Yes | No |
 | [databricks_artifact_allowlist](../resources/artifact_allowlist.md) | Yes | No | Yes | No |
 | [databricks_budget](../resources/budget.md) | Yes | Yes | No | Yes |
+| [databricks_budget_policy](../resources/budget_policy.md) | Yes | Yes | No | Yes |
 | [databricks_catalog](../resources/catalog.md) | Yes | Yes | Yes | No |
 | [databricks_cluster](../resources/cluster.md) | Yes | No | Yes | No |
 | [databricks_cluster_policy](../resources/cluster_policy.md) | Yes | No | Yes | No |

exporter/exporter_test.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/service/apps"
+	"github.com/databricks/databricks-sdk-go/service/billing"
 	sdk_uc "github.com/databricks/databricks-sdk-go/service/catalog"
 	sdk_compute "github.com/databricks/databricks-sdk-go/service/compute"
 	sdk_dashboards "github.com/databricks/databricks-sdk-go/service/dashboards"
@@ -398,6 +399,16 @@ var emptyDatabaseInstances = qa.HTTPFixture{
 	ReuseRequest: true,
 }
 
+var emptyBudgetPolicies = qa.HTTPFixture{
+	Method:   "GET",
+	Resource: "/api/2.0/accounts/[^/]+/budget/policies?",
+	Response: billing.ListBudgetPoliciesResponse{
+		Policies:      []billing.BudgetPolicy{},
+		NextPageToken: "",
+	},
+	ReuseRequest: true,
+}
+
 var emptyIpAccessLIst = qa.HTTPFixture{
 	Method:   http.MethodGet,
 	Resource: "/api/2.0/ip-access-lists",
@@ -575,6 +586,7 @@ func TestImportingUsersGroupsSecretScopes(t *testing.T) {
 			noCurrentMetastoreAttached,
 			emptyApps,
 			emptyAppsSettingsCustomTemplates,
+			emptyBudgetPolicies,
 			emptyLakeviewList,
 			emptyMetastoreList,
 			meAdminFixture,
@@ -850,6 +862,7 @@ func TestImportingNoResourcesError(t *testing.T) {
 			},
 			emptyApps,
 			emptyAppsSettingsCustomTemplates,
+			emptyBudgetPolicies,
 			emptyDataQualityMonitors,
 			emptyDatabaseInstances,
 			emptyUsersList,

exporter/impl_apps.go

@@ -105,6 +105,14 @@ func importApp(ic *importContext, r *resource) error {
 		}
 	}
 
+	// Budget Policy
+	if app.BudgetPolicyId != "" {
+		ic.Emit(&resource{
+			Resource: "databricks_budget_policy",
+			ID:       app.BudgetPolicyId,
+		})
+	}
+
 	// Emit permissions
 	ic.emitPermissionsIfNotIgnored(r, fmt.Sprintf("/apps/%s", app.Name), "app_"+r.Name)
 	return nil

exporter/impl_billing.go

@@ -0,0 +1,114 @@
+package exporter
+
+import (
+	"log"
+	"strconv"
+
+	"github.com/databricks/databricks-sdk-go/service/billing"
+	"github.com/databricks/terraform-provider-databricks/common"
+)
+
+func listBudgetPolicies(ic *importContext) error {
+	if ic.accountClient == nil {
+		return nil
+	}
+	policies, err := ic.accountClient.BudgetPolicy.ListAll(ic.Context, billing.ListBudgetPoliciesRequest{})
+	if err != nil {
+		return err
+	}
+	for _, policy := range policies {
+		if policy.PolicyId == "" {
+			continue
+		}
+		if !ic.MatchesName(policy.PolicyName) {
+			continue
+		}
+		ic.Emit(&resource{
+			Resource: "databricks_budget_policy",
+			ID:       policy.PolicyId,
+		})
+	}
+	return nil
+}
+
+func importBudgetPolicy(ic *importContext, r *resource) error {
+	// Get binding_workspace_ids directly from DataWrapper
+	if r.DataWrapper == nil {
+		log.Printf("[WARN] DataWrapper is nil for budget policy %s", r.ID)
+		return nil
+	}
+
+	accountID := ic.Client.Config.AccountID
+	// Emit access control rule set for the budget policy
+	ic.Emit(&resource{
+		Resource: "databricks_access_control_rule_set",
+		ID:       "accounts/" + accountID + "/budgetPolicies/" + r.ID + "/ruleSets/default",
+	})
+
+	bindingWorkspaceIdsRaw := r.DataWrapper.Get("binding_workspace_ids")
+	if bindingWorkspaceIdsRaw != nil {
+		// Convert to slice of int64
+		var bindingWorkspaceIds []int64
+		if workspaceIdsList, ok := bindingWorkspaceIdsRaw.([]int64); ok {
+			bindingWorkspaceIds = workspaceIdsList
+		}
+		// Emit workspace resources for each binding_workspace_id
+		if !ic.Client.Config.IsAzure() {
+			for _, workspaceId := range bindingWorkspaceIds {
+				ic.Emit(&resource{
+					Resource: "databricks_mws_workspaces",
+					ID:       accountID + "/" + strconv.FormatInt(workspaceId, 10),
+				})
+			}
+		}
+	}
+
+	return nil
+}
+
+func listBudgets(ic *importContext) error {
+	updatedSinceMs := ic.getUpdatedSinceMs()
+	budgets, err := ic.accountClient.Budgets.ListAll(ic.Context, billing.ListBudgetConfigurationsRequest{})
+	if err != nil {
+		return err
+	}
+	for _, budget := range budgets {
+		if ic.incremental && budget.CreateTime < updatedSinceMs {
+			log.Printf("[DEBUG] skipping budget '%s' that was updated at %d (last active=%d)",
+				budget.DisplayName, budget.UpdateTime, updatedSinceMs)
+			continue
+		}
+		ic.Emit(&resource{
+			Resource: "databricks_budget",
+			ID:       ic.accountClient.Config.AccountID + "|" + budget.BudgetConfigurationId,
+			Name:     budget.DisplayName,
+		})
+	}
+	return nil
+}
+
+func importBudget(ic *importContext, r *resource) error {
+	var budget billing.BudgetConfiguration
+	s := ic.Resources["databricks_budget"].Schema
+	common.DataToStructPointer(r.Data, s, &budget)
+	if budget.Filter != nil && budget.Filter.WorkspaceId != nil && !ic.accountClient.Config.IsAzure() {
+		for _, workspaceId := range budget.Filter.WorkspaceId.Values {
+			ic.Emit(&resource{
+				Resource: "databricks_mws_workspaces",
+				ID:       ic.accountClient.Config.AccountID + "/" + strconv.FormatInt(workspaceId, 10),
+			})
+		}
+	}
+	for _, alert := range budget.AlertConfigurations {
+		for _, action := range alert.ActionConfigurations {
+			if action.ActionType == billing.ActionConfigurationTypeEmailNotification {
+				ic.Emit(&resource{
+					Resource:  "databricks_user",
+					Attribute: "user_name",
+					Value:     action.Target,
+				})
+			}
+		}
+	}
+	return nil
+}

exporter/importables.go

@@ -12,7 +12,6 @@ import (
 	"strings"
 
 	"github.com/databricks/databricks-sdk-go/apierr"
-	"github.com/databricks/databricks-sdk-go/service/billing"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/ml"
@@ -1351,7 +1350,7 @@ var resourcesMap map[string]importable = map[string]importable{
 				IsValidApproximation: createIsMatchingScopeAndKey("scope", "key")},
 			{Path: "resources.uc_securable.securable_full_name", Resource: "databricks_volume"},
 			{Path: "resources.database.instance_name", Resource: "databricks_database_instance", Match: "name"},
-			// {Path: "budget_policy_id", Resource: "databricks_budget"},
+			{Path: "budget_policy_id", Resource: "databricks_budget_policy", Match: "policy_id"},
 		},
 	},
 	"databricks_pipeline": {
@@ -2065,6 +2064,8 @@ var resourcesMap map[string]importable = map[string]importable{
 				Regexp: regexp.MustCompile("^accounts/[^/]+/servicePrincipals/([^/]+)/ruleSets/default$")},
 			{Path: "name", Resource: "databricks_group", MatchType: MatchRegexp,
 				Regexp: regexp.MustCompile("^accounts/[^/]+/groups/([^/]+)/ruleSets/default$")},
+			{Path: "name", Resource: "databricks_budget_policy", Match: "policy_id", MatchType: MatchRegexp,
+				Regexp: regexp.MustCompile(`^accounts/[^/]+/budgetPolicies/([^/]+)/ruleSets/default$`)},
 		},
 		Ignore: func(ic *importContext, r *resource) bool {
 			// We're ignoring ACLs without grant rules because we don't know about that at time of emitting from groups/service principals
@@ -3328,54 +3329,23 @@ var resourcesMap map[string]importable = map[string]importable{
 			{Path: "credentials_id", Resource: "databricks_mws_credentials", Match: "credentials_id"},
 		},
 	},
+	"databricks_budget_policy": {
+		AccountLevel:    true,
+		PluginFramework: true,
+		Service:         "billing",
+		Name:            func(ic *importContext, d *schema.ResourceData) string { return d.Id() },
+		List:            listBudgetPolicies,
+		Import:          importBudgetPolicy,
+		Ignore:          generateIgnoreObjectWithEmptyAttributeValue("databricks_budget_policy", "policy_id"),
+		Depends: []reference{
+			{Path: "binding_workspace_ids", Resource: "databricks_mws_workspaces", Match: "workspace_id"},
+		},
+	},
 	"databricks_budget": {
 		AccountLevel: true,
 		Service:      "billing",
-		List: func(ic *importContext) error {
-			updatedSinceMs := ic.getUpdatedSinceMs()
-			budgets, err := ic.accountClient.Budgets.ListAll(ic.Context, billing.ListBudgetConfigurationsRequest{})
-			if err != nil {
-				return err
-			}
-			for _, budget := range budgets {
-				if ic.incremental && budget.CreateTime < updatedSinceMs {
-					log.Printf("[DEBUG] skipping budget '%s' that was updated at %d (last active=%d)",
-						budget.DisplayName, budget.UpdateTime, updatedSinceMs)
-					continue
-				}
-				ic.Emit(&resource{
-					Resource: "databricks_budget",
-					ID:       ic.accountClient.Config.AccountID + "|" + budget.BudgetConfigurationId,
-					Name:     budget.DisplayName,
-				})
-			}
-			return nil
-		},
-		Import: func(ic *importContext, r *resource) error {
-			var budget billing.BudgetConfiguration
-			s := ic.Resources["databricks_budget"].Schema
-			common.DataToStructPointer(r.Data, s, &budget)
-			if budget.Filter != nil && budget.Filter.WorkspaceId != nil && !ic.accountClient.Config.IsAzure() {
-				for _, workspaceId := range budget.Filter.WorkspaceId.Values {
-					ic.Emit(&resource{
-						Resource: "databricks_mws_workspaces",
-						ID:       ic.accountClient.Config.AccountID + "/" + strconv.FormatInt(workspaceId, 10),
-					})
-				}
-			}
-			for _, alert := range budget.AlertConfigurations {
-				for _, action := range alert.ActionConfigurations {
-					if action.ActionType == billing.ActionConfigurationTypeEmailNotification {
-						ic.Emit(&resource{
-							Resource:  "databricks_user",
-							Attribute: "user_name",
-							Value:     action.Target,
-						})
-					}
-				}
-			}
-			return nil
-		},
+		List:         listBudgets,
+		Import:       importBudget,
 		Depends: []reference{
 			{Path: "filter.workspace_id.values", Resource: "databricks_mws_workspaces", Match: "workspace_id"},
 			{Path: "alert_configurations.action_configurations.target", Resource: "databricks_user", Match: "user_name"},