Allow to override built-in databricks_cluster_policy resources (#2941)

* Allow to override built-in `databricks_cluster_policy` resources

Databricks is shipped with a number of built-in cluster policies that are inherited from
the cluster policy families.  For built-in cluster policies we need to provide an
override, not the `definition`.  We also need to handle deletes for them differently -
just remove the overrides because we can't remove built-in policies.  We also need to
distinguish cases when we override built-in policy from cases when we create a new policy
based on the cluster policy family.

This PR makes following changes:

* Discovers a list of cluster policies families dynamically instead of relying on
  hardcoded list of policies
* When creating a new `databricks_cluster_policy`, check if it's built-in policy or not,
  and if it's built-in, then update it with policy override instead of trying to create.
* When deleting a created `databricks_cluster_policy`, check if it's built-in policy or
  not, and if it's built-in, then reset policy overrides to empty string to return back to
  the original state.

This fixes #2935

* Add integration test for policy overrides

Author: alexott

docs/resources/cluster_policy.md

@@ -100,6 +100,38 @@ module "engineering_compute_policy" {
 }
 ```
 
+### Overriding the built-in cluster policies
+
+You can override built-in cluster policies by creating a `databricks_cluster_policy` resource with following attributes:
+
+* `name` - the name of the built-in cluster policy.
+* `policy_family_id` - the ID of the cluster policy family used for built-in cluster policy.
+* `policy_family_definition_overrides` - settings to override in the built-in cluster policy.
+
+You can obtain the list of defined cluster policies families using the `databricks policy-families list` command of the new [Databricks CLI](https://docs.databricks.com/en/dev-tools/cli/index.html), or via [list policy families](https://docs.databricks.com/api/workspace/policyfamilies/list) REST API.
+
+```hcl
+locals {
+  personal_vm_override = {
+    "autotermination_minutes" : {
+      "type" : "fixed",
+      "value" : 220,
+      "hidden" : true
+    },
+    "custom_tags.Team" : {
+      "type" : "fixed",
+      "value" : var.team
+    }
+  }
+}
+
+resource "databricks_cluster_policy" "personal_vm" {
+  policy_family_id                   = "personal-vm"
+  policy_family_definition_overrides = jsonencode(personal_vm_override)
+  name                               = "Personal Compute"
+}
+```
+
 ## Argument Reference
 
 The following arguments are supported:

docs/resources/default_namespace_settings.md

@@ -14,9 +14,9 @@ This setting requires a restart of clusters and SQL warehouses to take effect. A
 
 ```hcl
 resource "databricks_default_namespace_setting" "this" {
-		namespace {
-			value = "namespace_value"
-		}
+  namespace {
+    value = "namespace_value"
+  }
 }
 ```
 

exporter/context.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/compute"
 
 	"github.com/databricks/terraform-provider-databricks/commands"
 	"github.com/databricks/terraform-provider-databricks/common"
@@ -134,6 +135,9 @@ type importContext struct {
 	allDirectories      []workspace.ObjectStatus
 	allWorkspaceObjects []workspace.ObjectStatus
 	wsObjectsMutex      sync.RWMutex
+
+	builtInPolicies      map[string]compute.PolicyFamily
+	builtInPoliciesMutex sync.Mutex
 }
 
 type mount struct {

exporter/exporter_test.go

@@ -253,6 +253,15 @@ var emptyClusterPolicies = qa.HTTPFixture{
 	Response:     compute.ListPoliciesResponse{},
 }
 
+var emptyPolicyFamilies = qa.HTTPFixture{
+	Method:   "GET",
+	Resource: "/api/2.0/policy-families?",
+	Response: compute.ListPolicyFamiliesResponse{
+		PolicyFamilies: []compute.PolicyFamily{},
+	},
+	ReuseRequest: true,
+}
+
 var emptyMlflowWebhooks = qa.HTTPFixture{
 	Method:       "GET",
 	ReuseRequest: true,
@@ -392,6 +401,7 @@ func TestImportingUsersGroupsSecretScopes(t *testing.T) {
 			emptySqlAlerts,
 			emptyPipelines,
 			emptyClusterPolicies,
+			emptyPolicyFamilies,
 			emptyWorkspaceConf,
 			allKnownWorkspaceConfs,
 			dummyWorkspaceConf,
@@ -649,6 +659,7 @@ func TestImportingNoResourcesError(t *testing.T) {
 			emptySqlDashboards,
 			emptySqlAlerts,
 			emptyPipelines,
+			emptyPolicyFamilies,
 			{
 				Method:       "GET",
 				Resource:     "/api/2.0/global-init-scripts",

exporter/importables.go

@@ -12,8 +12,6 @@ import (
 	"strconv"
 	"strings"
 
-	"golang.org/x/exp/slices"
-
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	sdk_jobs "github.com/databricks/databricks-sdk-go/service/jobs"
@@ -37,18 +35,16 @@ import (
 )
 
 var (
-	adlsGen2Regex              = regexp.MustCompile(`^(abfss?)://([^@]+)@([^.]+)\.(?:[^/]+)(/.*)?$`)
-	adlsGen1Regex              = regexp.MustCompile(`^(adls?)://([^.]+)\.(?:[^/]+)(/.*)?$`)
-	wasbsRegex                 = regexp.MustCompile(`^(wasbs?)://([^@]+)@([^.]+)\.(?:[^/]+)(/.*)?$`)
-	s3Regex                    = regexp.MustCompile(`^(s3a?)://([^/]+)(/.*)?$`)
-	gsRegex                    = regexp.MustCompile(`^gs://([^/]+)(/.*)?$`)
-	globalWorkspaceConfName    = "global_workspace_conf"
-	nameNormalizationRegex     = regexp.MustCompile(`\W+`)
-	fileNameNormalizationRegex = regexp.MustCompile(`[^-_\w/.@]`)
-	jobClustersRegex           = regexp.MustCompile(`^((job_cluster|task)\.[0-9]+\.new_cluster\.[0-9]+\.)`)
-	dltClusterRegex            = regexp.MustCompile(`^(cluster\.[0-9]+\.)`)
-	predefinedClusterPolicies  = []string{"personal-vm", "shared-data-science", "shared-compute",
-		"power-user", "job-cluster"}
+	adlsGen2Regex                = regexp.MustCompile(`^(abfss?)://([^@]+)@([^.]+)\.(?:[^/]+)(/.*)?$`)
+	adlsGen1Regex                = regexp.MustCompile(`^(adls?)://([^.]+)\.(?:[^/]+)(/.*)?$`)
+	wasbsRegex                   = regexp.MustCompile(`^(wasbs?)://([^@]+)@([^.]+)\.(?:[^/]+)(/.*)?$`)
+	s3Regex                      = regexp.MustCompile(`^(s3a?)://([^/]+)(/.*)?$`)
+	gsRegex                      = regexp.MustCompile(`^gs://([^/]+)(/.*)?$`)
+	globalWorkspaceConfName      = "global_workspace_conf"
+	nameNormalizationRegex       = regexp.MustCompile(`\W+`)
+	fileNameNormalizationRegex   = regexp.MustCompile(`[^-_\w/.@]`)
+	jobClustersRegex             = regexp.MustCompile(`^((job_cluster|task)\.[0-9]+\.new_cluster\.[0-9]+\.)`)
+	dltClusterRegex              = regexp.MustCompile(`^(cluster\.[0-9]+\.)`)
 	secretPathRegex              = regexp.MustCompile(`^\{\{secrets\/([^\/]+)\/([^}]+)\}\}$`)
 	sqlParentRegexp              = regexp.MustCompile(`^folders/(\d+)$`)
 	dltDefaultStorageRegex       = regexp.MustCompile(`^dbfs:/pipelines/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)
@@ -631,14 +627,18 @@ var resourcesMap map[string]importable = map[string]importable{
 			if err != nil {
 				return err
 			}
-			policies, err := w.ClusterPolicies.ListAll(ic.Context, compute.ListClusterPoliciesRequest{})
+			policiesList, err := w.ClusterPolicies.ListAll(ic.Context, compute.ListClusterPoliciesRequest{})
 			if err != nil {
 				return err
 			}
-			for offset, policy := range policies {
+
+			builtInClusterPolicies := ic.getBuiltinPolicyFamilies()
+			for offset, policy := range policiesList {
 				log.Printf("[TRACE] Scanning %d:  %v", offset+1, policy)
-				if policy.PolicyFamilyId != "" && slices.Contains(predefinedClusterPolicies, policy.PolicyFamilyId) &&
+				family, isBuiltin := builtInClusterPolicies[policy.PolicyFamilyId]
+				if policy.PolicyFamilyId != "" && isBuiltin && family.Name == policy.Name &&
 					policy.PolicyFamilyDefinitionOverrides == "" {
+					log.Printf("[DEBUG] Skipping builtin cluster policy '%s' without overrides", policy.Name)
 					continue
 				}
 				if !ic.MatchesName(policy.Name) {
@@ -650,7 +650,7 @@ var resourcesMap map[string]importable = map[string]importable{
 					ID:       policy.PolicyId,
 				})
 				if offset%10 == 0 {
-					log.Printf("[INFO] Scanned %d of %d cluster policies", offset+1, len(policies))
+					log.Printf("[INFO] Scanned %d of %d cluster policies", offset+1, len(policiesList))
 				}
 			}
 			return nil
@@ -718,7 +718,9 @@ var resourcesMap map[string]importable = map[string]importable{
 				// we need to set definition to empty value because otherwise it will be put into
 				// generated HCL code for data source, and it only supports the `name` attribute
 				r.Data.Set("definition", "")
-				if slices.Contains(predefinedClusterPolicies, policyFamilyId) && policyOverrides == "" {
+				builtInClusterPolicies := ic.getBuiltinPolicyFamilies()
+				_, isBuiltin := builtInClusterPolicies[policyFamilyId]
+				if isBuiltin && policyOverrides == "" {
 					r.Mode = "data"
 				}
 			}

exporter/importables_test.go

@@ -117,6 +117,9 @@ func TestPredefinedClusterPolicy(t *testing.T) {
 	policy, _ := json.Marshal(map[string]map[string]string{})
 	d.Set("definition", string(policy))
 	ic := importContextForTest()
+	ic.builtInPolicies = map[string]compute.PolicyFamily{
+		"job-cluster": {Name: "Job Compute"},
+	}
 	r := resource{ID: "abc", Data: d}
 	err := ic.Importables["databricks_cluster_policy"].Import(ic, &r)
 	assert.NoError(t, err)
@@ -683,10 +686,20 @@ func TestPoliciesListing(t *testing.T) {
 				},
 			},
 		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/policy-families?",
+			Response: compute.ListPolicyFamiliesResponse{
+				PolicyFamilies: []compute.PolicyFamily{
+					{
+						Name:           "Personal Compute",
+						PolicyFamilyId: "personal-vm",
+					},
+				},
+			},
+		},
 	}, func(ctx context.Context, client *common.DatabricksClient) {
-		ic := importContextForTest()
-		ic.Client = client
-		ic.Context = ctx
+		ic := importContextForTestWithClient(ctx, client)
 		err := resourcesMap["databricks_cluster_policy"].List(ic)
 		assert.NoError(t, err)
 		assert.Equal(t, 1, len(ic.testEmits))
@@ -709,10 +722,9 @@ func TestPoliciesListNoNameMatch(t *testing.T) {
 				},
 			},
 		},
+		emptyPolicyFamilies,
 	}, func(ctx context.Context, client *common.DatabricksClient) {
-		ic := importContextForTest()
-		ic.Client = client
-		ic.Context = ctx
+		ic := importContextForTestWithClient(ctx, client)
 		ic.match = "bcd"
 		err := resourcesMap["databricks_cluster_policy"].List(ic)
 		assert.NoError(t, err)

exporter/util.go

@@ -23,6 +23,7 @@ import (
 	"github.com/databricks/terraform-provider-databricks/storage"
 	"github.com/databricks/terraform-provider-databricks/workspace"
 
+	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 
 	"golang.org/x/exp/slices"
@@ -411,6 +412,32 @@ func (ic *importContext) getSpsMapping() {
 	}
 }
 
+func (ic *importContext) getBuiltinPolicyFamilies() map[string]compute.PolicyFamily {
+	ic.builtInPoliciesMutex.Lock()
+	defer ic.builtInPoliciesMutex.Unlock()
+	if ic.builtInPolicies == nil {
+		if !ic.accountLevel {
+			log.Printf("[DEBUG] Going to initialize ic.builtInPolicies. Getting policy families...")
+			families, err := ic.workspaceClient.PolicyFamilies.ListAll(ic.Context, compute.ListPolicyFamiliesRequest{})
+			log.Printf("[DEBUG] Going to initialize ic.builtInPolicies. Getting policy families...")
+			if err == nil {
+				ic.builtInPolicies = make(map[string]compute.PolicyFamily, len(families))
+				for _, f := range families {
+					f2 := f
+					ic.builtInPolicies[f2.PolicyFamilyId] = f2
+				}
+			} else {
+				log.Printf("[ERROR] Can't fetch cluster policy families: %v", err)
+				ic.builtInPolicies = map[string]compute.PolicyFamily{}
+			}
+		} else {
+			log.Print("[WARN] Can't list cluster policy families on account level")
+			ic.builtInPolicies = map[string]compute.PolicyFamily{}
+		}
+	}
+	return ic.builtInPolicies
+}
+
 func (ic *importContext) findSpnByAppID(applicationID string) (u scim.User, err error) {
 	log.Printf("[DEBUG] Looking for SP %s", applicationID)
 	ic.spsMutex.RLocker().Lock()

internal/acceptance/cluster_policy_test.go

@@ -28,3 +28,29 @@ func TestAccClusterPolicyResourceFullLifecycle(t *testing.T) {
 		}`,
 	})
 }
+
+func TestAccClusterPolicyResourceOverrideBuiltIn(t *testing.T) {
+	workspaceLevel(t, step{
+		Template: `resource "databricks_cluster_policy" "personal_vm" {
+			name = "Personal Compute"
+			policy_family_id = "personal-vm"
+			policy_family_definition_overrides = jsonencode({
+				"node_type_id": {
+				  "type": "fixed",
+				  "value": "Standard_DS3_v2"
+				}
+			  })
+		}
+		`,
+	})
+}
+
+func TestAccClusterPolicyResourceOverrideNew(t *testing.T) {
+	workspaceLevel(t, step{
+		Template: `resource "databricks_cluster_policy" "policyoverrideempty" {
+			policy_family_id = "personal-vm"
+			name             = "Policy Override {var.RANDOM}"
+		  }
+		  `,
+	})
+}