From 0da17f6ec6805778ac04ea8c6f5c5155926018c8 Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Wed, 4 Dec 2024 16:35:54 +0530 Subject: [PATCH] Add default value for `volume_type` for DABs (#1952) ## Changes The Unity Catalog volumes API requires a `volume_type` argument when creating volumes. In the context of DABs, it's unnecessary to require users to specify the volume type every time. We can default to "MANAGED" instead. This PR is similar to https://github.com/databricks/cli/pull/1743 which does the same for dashboards. ## Tests Unit test --- .../mutator/configure_volume_defaults.go | 44 +++++++++++ .../mutator/configure_volume_defaults_test.go | 75 +++++++++++++++++++ bundle/internal/schema/main.go | 19 +++++ .../testdata/fail/incorrect_volume_type.yml | 10 +++ .../internal/schema/testdata/pass/volume.yml | 9 +++ bundle/phases/initialize.go | 1 + bundle/schema/jsonschema.json | 72 ++++++++++++++++++ 7 files changed, 230 insertions(+) create mode 100644 bundle/config/mutator/configure_volume_defaults.go create mode 100644 bundle/config/mutator/configure_volume_defaults_test.go create mode 100644 bundle/internal/schema/testdata/fail/incorrect_volume_type.yml create mode 100644 bundle/internal/schema/testdata/pass/volume.yml diff --git a/bundle/config/mutator/configure_volume_defaults.go b/bundle/config/mutator/configure_volume_defaults.go new file mode 100644 index 000000000..b2bea95d1 --- /dev/null +++ b/bundle/config/mutator/configure_volume_defaults.go @@ -0,0 +1,44 @@ +package mutator + +import ( + "context" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" +) + +type configureVolumeDefaults struct{} + +func ConfigureVolumeDefaults() bundle.Mutator { + return &configureVolumeDefaults{} +} + +func (m *configureVolumeDefaults) Name() string { + return "ConfigureVolumeDefaults" +} + +func (m *configureVolumeDefaults) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + pattern := dyn.NewPattern( + dyn.Key("resources"), + dyn.Key("volumes"), + dyn.AnyKey(), + ) + + // Configure defaults for all volumes. + err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) { + return dyn.MapByPattern(v, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + var err error + v, err = setIfNotExists(v, dyn.NewPath(dyn.Key("volume_type")), dyn.V("MANAGED")) + if err != nil { + return dyn.InvalidValue, err + } + return v, nil + }) + }) + + diags = diags.Extend(diag.FromErr(err)) + return diags +} diff --git a/bundle/config/mutator/configure_volume_defaults_test.go b/bundle/config/mutator/configure_volume_defaults_test.go new file mode 100644 index 000000000..b04006964 --- /dev/null +++ b/bundle/config/mutator/configure_volume_defaults_test.go @@ -0,0 +1,75 @@ +package mutator_test + +import ( + "context" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/mutator" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConfigureVolumeDefaultsVolumeType(t *testing.T) { + b := &bundle.Bundle{ + Config: config.Root{ + Resources: config.Resources{ + Volumes: map[string]*resources.Volume{ + "v1": { + // Empty string is skipped. + // See below for how it is set. + CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{ + VolumeType: "", + }, + }, + "v2": { + // Non-empty string is skipped. + CreateVolumeRequestContent: &catalog.CreateVolumeRequestContent{ + VolumeType: "already-set", + }, + }, + "v3": { + // No volume type set. + }, + "v4": nil, + }, + }, + }, + } + + // We can't set an empty string in the typed configuration. + // Do it on the dyn.Value directly. + bundletest.Mutate(t, b, func(v dyn.Value) (dyn.Value, error) { + return dyn.Set(v, "resources.volumes.v1.volume_type", dyn.V("")) + }) + + diags := bundle.Apply(context.Background(), b, mutator.ConfigureVolumeDefaults()) + require.NoError(t, diags.Error()) + + var v dyn.Value + var err error + + // Set to empty string; unchanged. + v, err = dyn.Get(b.Config.Value(), "resources.volumes.v1.volume_type") + require.NoError(t, err) + assert.Equal(t, "", v.MustString()) + + // Set to non-empty string; unchanged. + v, err = dyn.Get(b.Config.Value(), "resources.volumes.v2.volume_type") + require.NoError(t, err) + assert.Equal(t, "already-set", v.MustString()) + + // Not set; set to default. + v, err = dyn.Get(b.Config.Value(), "resources.volumes.v3.volume_type") + require.NoError(t, err) + assert.Equal(t, "MANAGED", v.MustString()) + + // No valid volume; No change. + _, err = dyn.Get(b.Config.Value(), "resources.volumes.v4.volume_type") + assert.True(t, dyn.IsCannotTraverseNilError(err)) +} diff --git a/bundle/internal/schema/main.go b/bundle/internal/schema/main.go index ddeffe2fd..881ce3496 100644 --- a/bundle/internal/schema/main.go +++ b/bundle/internal/schema/main.go @@ -93,6 +93,24 @@ func removeJobsFields(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema { return s } +// While volume_type is required in the volume create API, DABs automatically sets +// it's value to "MANAGED" if it's not provided. Thus, we make it optional +// in the bundle schema. +func makeVolumeTypeOptional(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema { + if typ != reflect.TypeOf(resources.Volume{}) { + return s + } + + res := []string{} + for _, r := range s.Required { + if r != "volume_type" { + res = append(res, r) + } + } + s.Required = res + return s +} + func main() { if len(os.Args) != 2 { fmt.Println("Usage: go run main.go ") @@ -118,6 +136,7 @@ func main() { p.addDescriptions, p.addEnums, removeJobsFields, + makeVolumeTypeOptional, addInterpolationPatterns, }) if err != nil { diff --git a/bundle/internal/schema/testdata/fail/incorrect_volume_type.yml b/bundle/internal/schema/testdata/fail/incorrect_volume_type.yml new file mode 100644 index 000000000..d540056d4 --- /dev/null +++ b/bundle/internal/schema/testdata/fail/incorrect_volume_type.yml @@ -0,0 +1,10 @@ +bundle: + name: volume with incorrect type + +resources: + volumes: + foo: + catalog_name: main + name: my_volume + schema_name: myschema + volume_type: incorrect_type diff --git a/bundle/internal/schema/testdata/pass/volume.yml b/bundle/internal/schema/testdata/pass/volume.yml new file mode 100644 index 000000000..14684f36e --- /dev/null +++ b/bundle/internal/schema/testdata/pass/volume.yml @@ -0,0 +1,9 @@ +bundle: + name: a volume + +resources: + volumes: + foo: + catalog_name: main + name: my_volume + schema_name: myschema diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index 334587202..6fa0e5fed 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -68,6 +68,7 @@ func Initialize() bundle.Mutator { mutator.SetRunAs(), mutator.OverrideCompute(), mutator.ConfigureDashboardDefaults(), + mutator.ConfigureVolumeDefaults(), mutator.ProcessTargetMode(), mutator.ApplyPresets(), mutator.DefaultQueueing(), diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index cf003f423..f791b8440 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -791,6 +791,51 @@ } ] }, + "resources.Volume": { + "anyOf": [ + { + "type": "object", + "properties": { + "catalog_name": { + "description": "The name of the catalog where the schema and the volume are", + "$ref": "#/$defs/string" + }, + "comment": { + "description": "The comment attached to the volume", + "$ref": "#/$defs/string" + }, + "grants": { + "$ref": "#/$defs/slice/github.com/databricks/cli/bundle/config/resources.Grant" + }, + "name": { + "description": "The name of the volume", + "$ref": "#/$defs/string" + }, + "schema_name": { + "description": "The name of the schema where the volume is", + "$ref": "#/$defs/string" + }, + "storage_location": { + "description": "The storage location on the cloud", + "$ref": "#/$defs/string" + }, + "volume_type": { + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.VolumeType" + } + }, + "additionalProperties": false, + "required": [ + "catalog_name", + "name", + "schema_name" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\.[a-zA-Z]+([-_]?[a-zA-Z0-9]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "variable.Lookup": { "anyOf": [ { @@ -963,6 +1008,9 @@ }, "name": { "$ref": "#/$defs/string" + }, + "uuid": { + "$ref": "#/$defs/string" } }, "additionalProperties": false, @@ -1157,6 +1205,9 @@ }, "schemas": { "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.Schema" + }, + "volumes": { + "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.Volume" } }, "additionalProperties": false @@ -1558,6 +1609,13 @@ } ] }, + "catalog.VolumeType": { + "type": "string", + "enum": [ + "EXTERNAL", + "MANAGED" + ] + }, "compute.Adlsgen2Info": { "anyOf": [ { @@ -5565,6 +5623,20 @@ } ] }, + "resources.Volume": { + "anyOf": [ + { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/github.com/databricks/cli/bundle/config/resources.Volume" + } + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\.[a-zA-Z]+([-_]?[a-zA-Z0-9]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "variable.TargetVariable": { "anyOf": [ {