Add DABs support for Unity Catalog volumes

bundle/artifacts/upload.go

@@ -21,18 +21,13 @@ func (m *cleanUp) Name() string {
 }
 
 func (m *cleanUp) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-	uploadPath, err := libraries.GetUploadBasePath(b)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-
-	client, err := libraries.GetFilerForLibraries(b.WorkspaceClient(), uploadPath)
-	if err != nil {
-		return diag.FromErr(err)
+	client, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b)
+	if diags.HasError() {
+		return diags
 	}
 
 	// We intentionally ignore the error because it is not critical to the deployment
-	err = client.Delete(ctx, ".", filer.DeleteRecursively)
+	err := client.Delete(ctx, ".", filer.DeleteRecursively)
 	if err != nil {
 		log.Errorf(ctx, "failed to delete %s: %v", uploadPath, err)
 	}

bundle/config/mutator/apply_presets.go

@@ -160,6 +160,13 @@ func (m *applyPresets) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnos
 		// the Databricks UI and via the SQL API.
 	}
 
+	// Apply the prefix to volumes
+	for i := range r.Volumes {
+		r.Volumes[i].Name = normalizePrefix(prefix) + r.Volumes[i].Name
+		// HTTP API for volumes doesn't yet support tags. It's only supported in
+		// the Databricks UI and via the SQL API.
+	}
+
 	return nil
 }
 

bundle/config/mutator/apply_presets_test.go

@@ -125,6 +125,62 @@ func TestApplyPresetsPrefixForUcSchema(t *testing.T) {
 	}
 }
 
+func TestApplyPresetsPrefixForUcVolumes(t *testing.T) {
+	tests := []struct {
+		name   string
+		prefix string
+		volume *resources.Volume
+		want   string
+	}{
+		{
+			name:   "add prefix to volume",
+			prefix: "[prefix]",
+			volume: &resources.Volume{
+				CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{
+					Name: "volume1",
+				},
+			},
+			want: "prefix_volume1",
+		},
+		{
+			name:   "add empty prefix to volume",
+			prefix: "",
+			volume: &resources.Volume{
+				CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{
+					Name: "volume1",
+				},
+			},
+			want: "volume1",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b := &bundle.Bundle{
+				Config: config.Root{
+					Resources: config.Resources{
+						Volumes: map[string]*resources.Volume{
+							"volume1": tt.volume,
+						},
+					},
+					Presets: config.Presets{
+						NamePrefix: tt.prefix,
+					},
+				},
+			}
+
+			ctx := context.Background()
+			diag := bundle.Apply(ctx, b, mutator.ApplyPresets())
+
+			if diag.HasError() {
+				t.Fatalf("unexpected error: %v", diag)
+			}
+
+			require.Equal(t, tt.want, b.Config.Resources.Volumes["volume1"].Name)
+		})
+	}
+}
+
 func TestApplyPresetsTags(t *testing.T) {
 	tests := []struct {
 		name string

bundle/config/mutator/process_target_mode_test.go

@@ -119,6 +119,9 @@ func mockBundle(mode config.Mode) *bundle.Bundle {
 				Schemas: map[string]*resources.Schema{
 					"schema1": {CreateSchema: &catalog.CreateSchema{Name: "schema1"}},
 				},
+				Volumes: map[string]*resources.Volume{
+					"volume1": {CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{Name: "volume1"}},
+				},
 			},
 		},
 		// Use AWS implementation for testing.
@@ -281,6 +284,8 @@ func TestProcessTargetModeDefault(t *testing.T) {
 	assert.Equal(t, "servingendpoint1", b.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name)
 	assert.Equal(t, "registeredmodel1", b.Config.Resources.RegisteredModels["registeredmodel1"].Name)
 	assert.Equal(t, "qualityMonitor1", b.Config.Resources.QualityMonitors["qualityMonitor1"].TableName)
+	assert.Equal(t, "schema1", b.Config.Resources.Schemas["schema1"].Name)
+	assert.Equal(t, "volume1", b.Config.Resources.Volumes["volume1"].Name)
 }
 
 func TestProcessTargetModeProduction(t *testing.T) {
@@ -322,6 +327,8 @@ func TestProcessTargetModeProduction(t *testing.T) {
 	assert.Equal(t, "servingendpoint1", b.Config.Resources.ModelServingEndpoints["servingendpoint1"].Name)
 	assert.Equal(t, "registeredmodel1", b.Config.Resources.RegisteredModels["registeredmodel1"].Name)
 	assert.Equal(t, "qualityMonitor1", b.Config.Resources.QualityMonitors["qualityMonitor1"].TableName)
+	assert.Equal(t, "schema1", b.Config.Resources.Schemas["schema1"].Name)
+	assert.Equal(t, "volume1", b.Config.Resources.Volumes["volume1"].Name)
 }
 
 func TestProcessTargetModeProductionOkForPrincipal(t *testing.T) {

bundle/config/mutator/run_as_test.go

@@ -40,6 +40,7 @@ func allResourceTypes(t *testing.T) []string {
 		"quality_monitors",
 		"registered_models",
 		"schemas",
+		"volumes",
 	},
 		resourceTypes,
 	)
@@ -138,6 +139,7 @@ func TestRunAsErrorForUnsupportedResources(t *testing.T) {
 		"registered_models",
 		"experiments",
 		"schemas",
+		"volumes",
 	}
 
 	base := config.Root{

bundle/config/resources.go

@@ -19,6 +19,7 @@ type Resources struct {
 	RegisteredModels      map[string]*resources.RegisteredModel      `json:"registered_models,omitempty"`
 	QualityMonitors       map[string]*resources.QualityMonitor       `json:"quality_monitors,omitempty"`
 	Schemas               map[string]*resources.Schema               `json:"schemas,omitempty"`
+	Volumes               map[string]*resources.Volume               `json:"volumes,omitempty"`
 }
 
 type ConfigResource interface {

bundle/config/resources/volume.go

@@ -0,0 +1,27 @@
+package resources
+
+import (
+	"github.com/databricks/databricks-sdk-go/marshal"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
+)
+
+type Volume struct {
+	// List of grants to apply on this schema.
+	Grants []Grant `json:"grants,omitempty"`
+
+	// Full name of the schema (catalog_name.schema_name.volume_name). This value is read from
+	// the terraform state after deployment succeeds.
+	ID string `json:"id,omitempty" bundle:"readonly"`
+
+	*catalog.CreateVolumeRequestContent
+
+	ModifiedStatus ModifiedStatus `json:"modified_status,omitempty" bundle:"internal"`
+}
+
+func (v *Volume) UnmarshalJSON(b []byte) error {
+	return marshal.Unmarshal(b, v)
+}
+
+func (v Volume) MarshalJSON() ([]byte, error) {
+	return marshal.Marshal(v)
+}

bundle/deploy/terraform/convert.go

@@ -394,6 +394,16 @@ func TerraformToBundle(state *resourcesState, config *config.Root) error {
 				}
 				cur.ID = instance.Attributes.ID
 				config.Resources.Schemas[resource.Name] = cur
+			case "databricks_volume":
+				if config.Resources.Volumes == nil {
+					config.Resources.Volumes = make(map[string]*resources.Volume)
+				}
+				cur := config.Resources.Volumes[resource.Name]
+				if cur == nil {
+					cur = &resources.Volume{ModifiedStatus: resources.ModifiedStatusDeleted}
+				}
+				cur.ID = instance.Attributes.ID
+				config.Resources.Volumes[resource.Name] = cur
 			case "databricks_permissions":
 			case "databricks_grants":
 				// Ignore; no need to pull these back into the configuration.
@@ -443,6 +453,11 @@ func TerraformToBundle(state *resourcesState, config *config.Root) error {
 			src.ModifiedStatus = resources.ModifiedStatusCreated
 		}
 	}
+	for _, src := range config.Resources.Volumes {
+		if src.ModifiedStatus == "" && src.ID == "" {
+			src.ModifiedStatus = resources.ModifiedStatusCreated
+		}
+	}
 
 	return nil
 }

bundle/deploy/terraform/convert_test.go

@@ -663,6 +663,14 @@ func TestTerraformToBundleEmptyLocalResources(t *testing.T) {
 					{Attributes: stateInstanceAttributes{ID: "1"}},
 				},
 			},
+			{
+				Type: "databricks_volume",
+				Mode: "managed",
+				Name: "test_volume",
+				Instances: []stateResourceInstance{
+					{Attributes: stateInstanceAttributes{ID: "1"}},
+				},
+			},
 		},
 	}
 	err := TerraformToBundle(&tfState, &config)
@@ -692,6 +700,9 @@ func TestTerraformToBundleEmptyLocalResources(t *testing.T) {
 	assert.Equal(t, "1", config.Resources.Schemas["test_schema"].ID)
 	assert.Equal(t, resources.ModifiedStatusDeleted, config.Resources.Schemas["test_schema"].ModifiedStatus)
 
+	assert.Equal(t, "1", config.Resources.Volumes["test_volume"].ID)
+	assert.Equal(t, resources.ModifiedStatusDeleted, config.Resources.Volumes["test_volume"].ModifiedStatus)
+
 	AssertFullResourceCoverage(t, &config)
 }
 
@@ -754,6 +765,13 @@ func TestTerraformToBundleEmptyRemoteResources(t *testing.T) {
 					},
 				},
 			},
+			Volumes: map[string]*resources.Volume{
+				"test_volume": {
+					CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{
+						Name: "test_volume",
+					},
+				},
+			},
 		},
 	}
 	var tfState = resourcesState{
@@ -786,6 +804,9 @@ func TestTerraformToBundleEmptyRemoteResources(t *testing.T) {
 	assert.Equal(t, "", config.Resources.Schemas["test_schema"].ID)
 	assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.Schemas["test_schema"].ModifiedStatus)
 
+	assert.Equal(t, "", config.Resources.Volumes["test_volume"].ID)
+	assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.Volumes["test_volume"].ModifiedStatus)
+
 	AssertFullResourceCoverage(t, &config)
 }
 
@@ -888,6 +909,18 @@ func TestTerraformToBundleModifiedResources(t *testing.T) {
 					},
 				},
 			},
+			Volumes: map[string]*resources.Volume{
+				"test_volume": {
+					CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{
+						Name: "test_volume",
+					},
+				},
+				"test_volume_new": {
+					CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{
+						Name: "test_volume_new",
+					},
+				},
+			},
 		},
 	}
 	var tfState = resourcesState{
@@ -1020,6 +1053,22 @@ func TestTerraformToBundleModifiedResources(t *testing.T) {
 					{Attributes: stateInstanceAttributes{ID: "2"}},
 				},
 			},
+			{
+				Type: "databricks_volume",
+				Mode: "managed",
+				Name: "test_volume",
+				Instances: []stateResourceInstance{
+					{Attributes: stateInstanceAttributes{ID: "1"}},
+				},
+			},
+			{
+				Type: "databricks_volume",
+				Mode: "managed",
+				Name: "test_volume_old",
+				Instances: []stateResourceInstance{
+					{Attributes: stateInstanceAttributes{ID: "2"}},
+				},
+			},
 		},
 	}
 	err := TerraformToBundle(&tfState, &config)
@@ -1081,6 +1130,13 @@ func TestTerraformToBundleModifiedResources(t *testing.T) {
 	assert.Equal(t, "", config.Resources.Schemas["test_schema_new"].ID)
 	assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.Schemas["test_schema_new"].ModifiedStatus)
 
+	assert.Equal(t, "1", config.Resources.Volumes["test_volume"].ID)
+	assert.Equal(t, "", config.Resources.Volumes["test_volume"].ModifiedStatus)
+	assert.Equal(t, "2", config.Resources.Volumes["test_volume_old"].ID)
+	assert.Equal(t, resources.ModifiedStatusDeleted, config.Resources.Volumes["test_volume_old"].ModifiedStatus)
+	assert.Equal(t, "", config.Resources.Volumes["test_volume_new"].ID)
+	assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.Volumes["test_volume_new"].ModifiedStatus)
+
 	AssertFullResourceCoverage(t, &config)
 }
 

bundle/deploy/terraform/interpolate.go

@@ -58,6 +58,8 @@ func (m *interpolateMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.D
 				path = dyn.NewPath(dyn.Key("databricks_quality_monitor")).Append(path[2:]...)
 			case dyn.Key("schemas"):
 				path = dyn.NewPath(dyn.Key("databricks_schema")).Append(path[2:]...)
+			case dyn.Key("volumes"):
+				path = dyn.NewPath(dyn.Key("databricks_volume")).Append(path[2:]...)
 			default:
 				// Trigger "key not found" for unknown resource types.
 				return dyn.GetByPath(root, path)

bundle/deploy/terraform/interpolate_test.go

@@ -31,6 +31,7 @@ func TestInterpolate(t *testing.T) {
 								"other_model_serving":    "${resources.model_serving_endpoints.other_model_serving.id}",
 								"other_registered_model": "${resources.registered_models.other_registered_model.id}",
 								"other_schema":           "${resources.schemas.other_schema.id}",
+								"other_volume":           "${resources.volumes.other_volume.id}",
 							},
 							Tasks: []jobs.Task{
 								{
@@ -67,6 +68,7 @@ func TestInterpolate(t *testing.T) {
 	assert.Equal(t, "${databricks_model_serving.other_model_serving.id}", j.Tags["other_model_serving"])
 	assert.Equal(t, "${databricks_registered_model.other_registered_model.id}", j.Tags["other_registered_model"])
 	assert.Equal(t, "${databricks_schema.other_schema.id}", j.Tags["other_schema"])
+	assert.Equal(t, "${databricks_volume.other_volume.id}", j.Tags["other_volume"])
 
 	m := b.Config.Resources.Models["my_model"]
 	assert.Equal(t, "my_model", m.Model.Name)

bundle/deploy/terraform/tfdyn/convert_volume.go

@@ -0,0 +1,45 @@
+package tfdyn
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/cli/bundle/internal/tf/schema"
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/dyn/convert"
+	"github.com/databricks/cli/libs/log"
+)
+
+func convertVolumeResource(ctx context.Context, vin dyn.Value) (dyn.Value, error) {
+	// Normalize the output value to the target schema.
+	vout, diags := convert.Normalize(schema.ResourceVolume{}, vin)
+	for _, diag := range diags {
+		log.Debugf(ctx, "volume normalization diagnostic: %s", diag.Summary)
+	}
+
+	return vout, nil
+}
+
+type volumeConverter struct{}
+
+func (volumeConverter) Convert(ctx context.Context, key string, vin dyn.Value, out *schema.Resources) error {
+	vout, err := convertVolumeResource(ctx, vin)
+	if err != nil {
+		return err
+	}
+
+	// Add the converted resource to the output.
+	out.Volume[key] = vout.AsAny()
+
+	// Configure grants for this resource.
+	if grants := convertGrantsResource(ctx, vin); grants != nil {
+		grants.Volume = fmt.Sprintf("${databricks_volume.%s.id}", key)
+		out.Grants["volume_"+key] = grants
+	}
+
+	return nil
+}
+
+func init() {
+	registerConverter("volumes", volumeConverter{})
+}