[Feature] Add databricks_budget resource

mws/resource_mws_budgets.go

@@ -0,0 +1,61 @@
+package mws
+
+import (
+	"context"
+	"github.com/databricks/databricks-sdk-go/service/billing"
+	"github.com/databricks/terraform-provider-databricks/common"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
+)
+
+func ResourceMwsBudget() common.Resource {
+	s := common.StructToSchema(billing.BudgetConfiguration{}, func(m map[string]*schema.Schema) map[string]*schema.Schema {
+		common.CustomizeSchemaPath(m, "display_name").SetValidateFunc(validation.StringLenBetween(1, 128))
+		for _, p := range []string{"account_id", "budget_configuration_id", "create_time", "update_time"} {
+			common.CustomizeSchemaPath(m, p).SetComputed()
+		}
+		return m
+	})
+	p := common.NewPairSeparatedID("account_id", "budget_configuration_id", "/")
+	return common.Resource{
+		Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			var create billing.CreateBudgetConfigurationBudget
+			common.DataToStructPointer(d, s, &create)
+			acc, err := c.AccountClientWithAccountIdFromConfig(d)
+			if err != nil {
+				return err
+			}
+			budget, err := acc.Budgets.Create(ctx, billing.CreateBudgetConfigurationRequest{Budget: create})
+			if err != nil {
+				return err
+			}
+			d.SetId(budget.Budget.BudgetConfigurationId)
+			d.Set("budget_configuration_id", budget.Budget.BudgetConfigurationId)
+			d.Set("account_id", c.Config.AccountID)
+			common.StructToData(budget, s, d)
+			p.Pack(d)
+			return nil
+		},
+		Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			_, id, err := p.Unpack(d)
+			acc, err := c.AccountClient()
+			if err != nil {
+				return err
+			}
+			res, err := acc.Budgets.GetByBudgetId(ctx, id)
+			if err != nil {
+				return err
+			}
+			return common.StructToData(res.Budget, s, d)
+		},
+		Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			_, id, err := p.Unpack(d)
+			acc, err := c.AccountClient()
+			if err != nil {
+				return err
+			}
+			return acc.Budgets.DeleteByBudgetId(ctx, id)
+		},
+		Schema: s,
+	}
+}

mws/resource_mws_budgets_test.go

@@ -0,0 +1,143 @@
+package mws
+
+import (
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/billing"
+	"github.com/stretchr/testify/mock"
+	"testing"
+
+	"github.com/databricks/terraform-provider-databricks/qa"
+)
+
+func getTestBudget() *billing.BudgetConfiguration {
+	return &billing.BudgetConfiguration{
+		AccountId: "account_id",
+		AlertConfigurations: []billing.AlertConfiguration{
+			{
+				ActionConfigurations: []billing.ActionConfiguration{
+					{
+						ActionType: billing.ActionConfigurationTypeEmailNotification,
+						Target:     "me@databricks.com",
+					},
+				},
+				AlertConfigurationId: "alert_configuration_id",
+				QuantityThreshold:    "840",
+				QuantityType:         billing.AlertConfigurationQuantityTypeListPriceDollarsUsd,
+				TimePeriod:           billing.AlertConfigurationTimePeriodMonth,
+				TriggerType:          billing.AlertConfigurationTriggerTypeCumulativeSpendingExceeded,
+			},
+		},
+		Filter: &billing.BudgetConfigurationFilter{
+			Tags: []billing.BudgetConfigurationFilterTagClause{
+				{
+					Key: "Environment",
+					Value: &billing.BudgetConfigurationFilterClause{
+						Operator: billing.BudgetConfigurationFilterOperatorIn,
+						Values:   []string{"Testing"},
+					},
+				},
+			},
+			WorkspaceId: &billing.BudgetConfigurationFilterWorkspaceIdClause{
+				Operator: billing.BudgetConfigurationFilterOperatorIn,
+				Values: []int64{
+					1234567890098765,
+				},
+			},
+		},
+		BudgetConfigurationId: "budget_configuration_id",
+		DisplayName:           "budget_name",
+	}
+}
+
+func TestResourceBudgetCreate(t *testing.T) {
+	qa.ResourceFixture{
+		MockAccountClientFunc: func(a *mocks.MockAccountClient) {
+			api := a.GetMockbudgetsAPI().EXPECT()
+			api.Create(mock.Anything, billing.CreateBudgetConfigurationRequest{
+				Budget: billing.CreateBudgetConfigurationBudget{
+					AlertConfigurations: []billing.CreateBudgetConfigurationBudgetAlertConfigurations{
+						{
+							ActionConfigurations: []billing.CreateBudgetConfigurationBudgetActionConfigurations{
+								{
+									ActionType: billing.ActionConfigurationTypeEmailNotification,
+									Target:     "me@databricks.com",
+								},
+							},
+							QuantityThreshold: "840",
+							QuantityType:      billing.AlertConfigurationQuantityTypeListPriceDollarsUsd,
+							TimePeriod:        billing.AlertConfigurationTimePeriodMonth,
+							TriggerType:       billing.AlertConfigurationTriggerTypeCumulativeSpendingExceeded,
+						},
+					},
+					DisplayName: getTestBudget().DisplayName,
+					Filter:      getTestBudget().Filter,
+				},
+			}).Return(&billing.CreateBudgetConfigurationResponse{Budget: getTestBudget()}, nil)
+			api.GetByBudgetId(mock.Anything, "budget_configuration_id").Return(
+				&billing.GetBudgetConfigurationResponse{Budget: getTestBudget()}, nil,
+			)
+		},
+		Create:    true,
+		AccountID: "account_id",
+		HCL: `
+		display_name = "budget_name"
+
+		alert_configurations {
+			time_period         = "MONTH"
+			trigger_type        = "CUMULATIVE_SPENDING_EXCEEDED" 
+			quantity_type       = "LIST_PRICE_DOLLARS_USD"
+			quantity_threshold  = "840"
+
+			action_configurations {
+				action_type = "EMAIL_NOTIFICATION"
+				target      = "me@databricks.com"
+			}
+		}
+
+		filter {
+			tags {
+				key   = "Environment"
+				value {
+					operator = "IN"
+					values = ["Testing"]
+				}
+			}
+
+			workspace_id {
+				operator = "IN"
+				values   = [
+					1234567890098765
+				]                              
+			}
+		}
+		`,
+		Resource: ResourceMwsBudget(),
+	}.ApplyAndExpectData(t, nil) // ???
+}
+
+func TestResourceMwsBudgetRead(t *testing.T) {
+	qa.ResourceFixture{
+		MockAccountClientFunc: func(a *mocks.MockAccountClient) {
+			a.GetMockbudgetsAPI().EXPECT().
+				GetByBudgetId(mock.Anything, "budget_configuration_id").
+				Return(&billing.GetBudgetConfigurationResponse{Budget: getTestBudget()}, nil)
+		},
+		Resource:  ResourceMwsBudget(),
+		Read:      true,
+		New:       true,
+		AccountID: "account_id",
+		ID:        "account_id/budget_configuration_id",
+	}.ApplyAndExpectData(t, nil) // ???
+}
+
+func TestResourceMwsBudgetDelete(t *testing.T) {
+	qa.ResourceFixture{
+		MockAccountClientFunc: func(a *mocks.MockAccountClient) {
+			a.GetMockbudgetsAPI().EXPECT().DeleteByBudgetId(mock.Anything, "budget_configuration_id").Return(nil)
+		},
+		Resource:  ResourceMwsBudget(),
+		AccountID: "account_id",
+		Delete:    true,
+		ID:        "account_id/budget_configuration_id",
+	}.ApplyAndExpectData(t, nil)
+}

provider/provider.go

@@ -161,6 +161,7 @@ func DatabricksProvider() *schema.Provider {
 			"databricks_mlflow_webhook":                  mlflow.ResourceMlflowWebhook().ToResource(),
 			"databricks_model_serving":                   serving.ResourceModelServing().ToResource(),
 			"databricks_mount":                           storage.ResourceMount().ToResource(),
+			"databricks_mws_budget":                      mws.ResourceMwsBudget().ToResource(),
 			"databricks_mws_customer_managed_keys":       mws.ResourceMwsCustomerManagedKeys().ToResource(),
 			"databricks_mws_credentials":                 mws.ResourceMwsCredentials().ToResource(),
 			"databricks_mws_log_delivery":                mws.ResourceMwsLogDelivery().ToResource(),