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
This commit is contained in:
shreyas-goenka 2024-12-04 16:35:54 +05:30 committed by GitHub
parent 0e088eb9f8
commit 0da17f6ec6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 230 additions and 0 deletions

View File

@ -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
}

View File

@ -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))
}

View File

@ -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 <output-file>")
@ -118,6 +136,7 @@ func main() {
p.addDescriptions,
p.addEnums,
removeJobsFields,
makeVolumeTypeOptional,
addInterpolationPatterns,
})
if err != nil {

View File

@ -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

View File

@ -0,0 +1,9 @@
bundle:
name: a volume
resources:
volumes:
foo:
catalog_name: main
name: my_volume
schema_name: myschema

View File

@ -68,6 +68,7 @@ func Initialize() bundle.Mutator {
mutator.SetRunAs(),
mutator.OverrideCompute(),
mutator.ConfigureDashboardDefaults(),
mutator.ConfigureVolumeDefaults(),
mutator.ProcessTargetMode(),
mutator.ApplyPresets(),
mutator.DefaultQueueing(),

View File

@ -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": [
{