Allow referencing bundle resources by name (#872)

## Changes
Now we can define variables with values which reference different
Databricks resources by name.
When references like this, DABs automatically looks up the resource by
this name and replaces the reference with ID of the resource referenced.
Thus when the variable is used in the configuration it will contain the
correct resolved ID of resource.

The resolvers are code generated and thus DABs support referencing all
resources which has `GetByName`-like methods in Go SDK.

### Example

```
variables:
  my_cluster_id:
    description: An existing cluster.
    lookup: 
      cluster: "12.2 shared"

resources:
  jobs:
    my_job:
      name: "My Job"
      tasks:
        - task_key: TestTask
          existing_cluster_id: ${var.my_cluster_id}

targets:
  dev:
    variables:
      my_cluster_id:
        lookup: 
           cluster: "dev-cluster"
```

## Tests
Added unit test + manual testing

---------

Co-authored-by: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
This commit is contained in:
Andrew Nester 2024-01-04 22:04:42 +01:00 committed by GitHub
parent 8a1be76910
commit 5fb40f9d07
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 723 additions and 7 deletions

View File

@ -5,7 +5,8 @@
},
"batch": {
".codegen/cmds-workspace.go.tmpl": "cmd/workspace/cmd.go",
".codegen/cmds-account.go.tmpl": "cmd/account/cmd.go"
".codegen/cmds-account.go.tmpl": "cmd/account/cmd.go",
".codegen/lookup.go.tmpl": "bundle/config/variable/lookup.go"
},
"toolchain": {
"required": ["go"],

123
.codegen/lookup.go.tmpl Normal file
View File

@ -0,0 +1,123 @@
// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
package variable
{{ $allowlist :=
list
"alerts"
"clusters"
"cluster-policies"
"clusters"
"dashboards"
"instance-pools"
"jobs"
"metastores"
"pipelines"
"service-principals"
"queries"
"warehouses"
}}
import (
"context"
"fmt"
"github.com/databricks/databricks-sdk-go"
)
type Lookup struct {
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
{{.Singular.PascalName}} string `json:"{{.Singular.SnakeName}},omitempty"`
{{end}}
{{- end}}
}
func LookupFromMap(m map[string]any) *Lookup {
l := &Lookup{}
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if v, ok := m["{{.Singular.KebabName}}"]; ok {
l.{{.Singular.PascalName}} = v.(string)
}
{{end -}}
{{- end}}
return l
}
func (l *Lookup) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
if err := l.validate(); err != nil {
return "", err
}
resolvers := resolvers()
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if l.{{.Singular.PascalName}} != "" {
return resolvers["{{.Singular.KebabName}}"](ctx, w, l.{{.Singular.PascalName}})
}
{{end -}}
{{- end}}
return "", fmt.Errorf("no valid lookup fields provided")
}
func (l *Lookup) String() string {
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if l.{{.Singular.PascalName}} != "" {
return fmt.Sprintf("{{.Singular.KebabName}}: %s", l.{{.Singular.PascalName}})
}
{{end -}}
{{- end}}
return ""
}
func (l *Lookup) validate() error {
// Validate that only one field is set
count := 0
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
if l.{{.Singular.PascalName}} != "" {
count++
}
{{end -}}
{{- end}}
if count != 1 {
return fmt.Errorf("exactly one lookup field must be provided")
}
if strings.Contains(l.String(), "${var") {
return fmt.Errorf("lookup fields cannot contain variable references")
}
return nil
}
type resolverFunc func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error)
func resolvers() map[string](resolverFunc) {
resolvers := make(map[string](resolverFunc), 0)
{{range .Services -}}
{{- if in $allowlist .KebabName -}}
resolvers["{{.Singular.KebabName}}"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.{{.PascalName}}.GetBy{{range .List.NamedIdMap.NamePath}}{{.PascalName}}{{end}}(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity{{ template "field-path" .List.NamedIdMap.IdPath }}), nil
}
{{end -}}
{{- end}}
return resolvers
}
{{- define "field-path" -}}
{{- range .}}.{{.PascalName}}{{end}}
{{- end -}}

1
.gitattributes vendored
View File

@ -1,3 +1,4 @@
bundle/config/variable/lookup.go linguist-generated=true
cmd/account/access-control/access-control.go linguist-generated=true
cmd/account/billable-usage/billable-usage.go linguist-generated=true
cmd/account/budgets/budgets.go linguist-generated=true

View File

@ -0,0 +1,48 @@
package mutator
import (
"context"
"fmt"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/libs/log"
"golang.org/x/sync/errgroup"
)
type resolveResourceReferences struct{}
func ResolveResourceReferences() bundle.Mutator {
return &resolveResourceReferences{}
}
func (m *resolveResourceReferences) Apply(ctx context.Context, b *bundle.Bundle) error {
errs, errCtx := errgroup.WithContext(ctx)
for k := range b.Config.Variables {
v := b.Config.Variables[k]
if v == nil || v.Lookup == nil {
continue
}
if v.HasValue() {
log.Debugf(ctx, "Ignoring '%s' lookup for the variable '%s' because the value is set", v.Lookup, k)
continue
}
errs.Go(func() error {
id, err := v.Lookup.Resolve(errCtx, b.WorkspaceClient())
if err != nil {
return fmt.Errorf("failed to resolve %s, err: %w", v.Lookup, err)
}
v.Set(id)
return nil
})
}
return errs.Wait()
}
func (*resolveResourceReferences) Name() string {
return "ResolveResourceReferences"
}

View File

@ -0,0 +1,197 @@
package mutator
import (
"context"
"testing"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/bundle/config/variable"
"github.com/databricks/databricks-sdk-go/service/compute"
"github.com/stretchr/testify/require"
)
type MockClusterService struct{}
// ChangeOwner implements compute.ClustersService.
func (MockClusterService) ChangeOwner(ctx context.Context, request compute.ChangeClusterOwner) error {
panic("unimplemented")
}
// Create implements compute.ClustersService.
func (MockClusterService) Create(ctx context.Context, request compute.CreateCluster) (*compute.CreateClusterResponse, error) {
panic("unimplemented")
}
// Delete implements compute.ClustersService.
func (MockClusterService) Delete(ctx context.Context, request compute.DeleteCluster) error {
panic("unimplemented")
}
// Edit implements compute.ClustersService.
func (MockClusterService) Edit(ctx context.Context, request compute.EditCluster) error {
panic("unimplemented")
}
// Events implements compute.ClustersService.
func (MockClusterService) Events(ctx context.Context, request compute.GetEvents) (*compute.GetEventsResponse, error) {
panic("unimplemented")
}
// Get implements compute.ClustersService.
func (MockClusterService) Get(ctx context.Context, request compute.GetClusterRequest) (*compute.ClusterDetails, error) {
panic("unimplemented")
}
// GetPermissionLevels implements compute.ClustersService.
func (MockClusterService) GetPermissionLevels(ctx context.Context, request compute.GetClusterPermissionLevelsRequest) (*compute.GetClusterPermissionLevelsResponse, error) {
panic("unimplemented")
}
// GetPermissions implements compute.ClustersService.
func (MockClusterService) GetPermissions(ctx context.Context, request compute.GetClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
panic("unimplemented")
}
// List implements compute.ClustersService.
func (MockClusterService) List(ctx context.Context, request compute.ListClustersRequest) (*compute.ListClustersResponse, error) {
return &compute.ListClustersResponse{
Clusters: []compute.ClusterDetails{
{ClusterId: "1234-5678-abcd", ClusterName: "Some Custom Cluster"},
{ClusterId: "9876-5432-xywz", ClusterName: "Some Other Name"},
},
}, nil
}
// ListNodeTypes implements compute.ClustersService.
func (MockClusterService) ListNodeTypes(ctx context.Context) (*compute.ListNodeTypesResponse, error) {
panic("unimplemented")
}
// ListZones implements compute.ClustersService.
func (MockClusterService) ListZones(ctx context.Context) (*compute.ListAvailableZonesResponse, error) {
panic("unimplemented")
}
// PermanentDelete implements compute.ClustersService.
func (MockClusterService) PermanentDelete(ctx context.Context, request compute.PermanentDeleteCluster) error {
panic("unimplemented")
}
// Pin implements compute.ClustersService.
func (MockClusterService) Pin(ctx context.Context, request compute.PinCluster) error {
panic("unimplemented")
}
// Resize implements compute.ClustersService.
func (MockClusterService) Resize(ctx context.Context, request compute.ResizeCluster) error {
panic("unimplemented")
}
// Restart implements compute.ClustersService.
func (MockClusterService) Restart(ctx context.Context, request compute.RestartCluster) error {
panic("unimplemented")
}
// SetPermissions implements compute.ClustersService.
func (MockClusterService) SetPermissions(ctx context.Context, request compute.ClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
panic("unimplemented")
}
// SparkVersions implements compute.ClustersService.
func (MockClusterService) SparkVersions(ctx context.Context) (*compute.GetSparkVersionsResponse, error) {
panic("unimplemented")
}
// Start implements compute.ClustersService.
func (MockClusterService) Start(ctx context.Context, request compute.StartCluster) error {
panic("unimplemented")
}
// Unpin implements compute.ClustersService.
func (MockClusterService) Unpin(ctx context.Context, request compute.UnpinCluster) error {
panic("unimplemented")
}
// UpdatePermissions implements compute.ClustersService.
func (MockClusterService) UpdatePermissions(ctx context.Context, request compute.ClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
panic("unimplemented")
}
func TestResolveClusterReference(t *testing.T) {
clusterRef1 := "Some Custom Cluster"
clusterRef2 := "Some Other Name"
justString := "random string"
b := &bundle.Bundle{
Config: config.Root{
Variables: map[string]*variable.Variable{
"my-cluster-id-1": {
Lookup: &variable.Lookup{
Cluster: clusterRef1,
},
},
"my-cluster-id-2": {
Lookup: &variable.Lookup{
Cluster: clusterRef2,
},
},
"some-variable": {
Value: &justString,
},
},
},
}
b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
err := bundle.Apply(context.Background(), b, ResolveResourceReferences())
require.NoError(t, err)
require.Equal(t, "1234-5678-abcd", *b.Config.Variables["my-cluster-id-1"].Value)
require.Equal(t, "9876-5432-xywz", *b.Config.Variables["my-cluster-id-2"].Value)
}
func TestResolveNonExistentClusterReference(t *testing.T) {
clusterRef := "Random"
justString := "random string"
b := &bundle.Bundle{
Config: config.Root{
Variables: map[string]*variable.Variable{
"my-cluster-id": {
Lookup: &variable.Lookup{
Cluster: clusterRef,
},
},
"some-variable": {
Value: &justString,
},
},
},
}
b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
err := bundle.Apply(context.Background(), b, ResolveResourceReferences())
require.ErrorContains(t, err, "failed to resolve cluster: Random, err: ClusterDetails named 'Random' does not exist")
}
func TestNoLookupIfVariableIsSet(t *testing.T) {
clusterRef := "donotexist"
b := &bundle.Bundle{
Config: config.Root{
Variables: map[string]*variable.Variable{
"my-cluster-id": {
Lookup: &variable.Lookup{
Cluster: clusterRef,
},
},
},
},
}
b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
b.Config.Variables["my-cluster-id"].Set("random value")
err := bundle.Apply(context.Background(), b, ResolveResourceReferences())
require.NoError(t, err)
require.Equal(t, "random value", *b.Config.Variables["my-cluster-id"].Value)
}

View File

@ -46,6 +46,12 @@ func setVariable(ctx context.Context, v *variable.Variable, name string) error {
return nil
}
// case: Defined a variable for named lookup for a resource
// It will be resolved later in ResolveResourceReferences mutator
if v.Lookup != nil {
return nil
}
// We should have had a value to set for the variable at this point.
// TODO: use cmdio to request values for unassigned variables if current
// terminal is a tty. Tracked in https://github.com/databricks/cli/issues/379

View File

@ -201,13 +201,25 @@ func (r *Root) MergeTargetOverrides(target *Target) error {
if target.Variables != nil {
for k, v := range target.Variables {
variable, ok := r.Variables[k]
rootVariable, ok := r.Variables[k]
if !ok {
return fmt.Errorf("variable %s is not defined but is assigned a value", k)
}
// we only allow overrides of the default value for a variable
defaultVal := v
variable.Default = &defaultVal
if sv, ok := v.(string); ok {
// we allow overrides of the default value for a variable
defaultVal := sv
rootVariable.Default = &defaultVal
} else if vv, ok := v.(map[string]any); ok {
// we also allow overrides of the lookup value for a variable
lookup, ok := vv["lookup"]
if !ok {
return fmt.Errorf("variable %s is incorrectly defined lookup override, no 'lookup' key defined", k)
}
rootVariable.Lookup = variable.LookupFromMap(lookup.(map[string]any))
} else {
return fmt.Errorf("variable %s is incorrectly defined in target override", k)
}
}
}

View File

@ -30,10 +30,10 @@ type Target struct {
Resources *Resources `json:"resources,omitempty"`
// Override default values for defined variables
// Override default values or lookup name for defined variables
// Does not permit defining new variables or redefining existing ones
// in the scope of an target
Variables map[string]string `json:"variables,omitempty"`
Variables map[string]any `json:"variables,omitempty"`
Git Git `json:"git,omitempty"`

299
bundle/config/variable/lookup.go generated Executable file
View File

@ -0,0 +1,299 @@
// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
package variable
import (
"context"
"fmt"
"strings"
"github.com/databricks/databricks-sdk-go"
)
type Lookup struct {
Alert string `json:"alert,omitempty"`
ClusterPolicy string `json:"cluster_policy,omitempty"`
Cluster string `json:"cluster,omitempty"`
Dashboard string `json:"dashboard,omitempty"`
InstancePool string `json:"instance_pool,omitempty"`
Job string `json:"job,omitempty"`
Metastore string `json:"metastore,omitempty"`
Pipeline string `json:"pipeline,omitempty"`
Query string `json:"query,omitempty"`
ServicePrincipal string `json:"service_principal,omitempty"`
Warehouse string `json:"warehouse,omitempty"`
}
func LookupFromMap(m map[string]any) *Lookup {
l := &Lookup{}
if v, ok := m["alert"]; ok {
l.Alert = v.(string)
}
if v, ok := m["cluster-policy"]; ok {
l.ClusterPolicy = v.(string)
}
if v, ok := m["cluster"]; ok {
l.Cluster = v.(string)
}
if v, ok := m["dashboard"]; ok {
l.Dashboard = v.(string)
}
if v, ok := m["instance-pool"]; ok {
l.InstancePool = v.(string)
}
if v, ok := m["job"]; ok {
l.Job = v.(string)
}
if v, ok := m["metastore"]; ok {
l.Metastore = v.(string)
}
if v, ok := m["pipeline"]; ok {
l.Pipeline = v.(string)
}
if v, ok := m["query"]; ok {
l.Query = v.(string)
}
if v, ok := m["service-principal"]; ok {
l.ServicePrincipal = v.(string)
}
if v, ok := m["warehouse"]; ok {
l.Warehouse = v.(string)
}
return l
}
func (l *Lookup) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
if err := l.validate(); err != nil {
return "", err
}
resolvers := resolvers()
if l.Alert != "" {
return resolvers["alert"](ctx, w, l.Alert)
}
if l.ClusterPolicy != "" {
return resolvers["cluster-policy"](ctx, w, l.ClusterPolicy)
}
if l.Cluster != "" {
return resolvers["cluster"](ctx, w, l.Cluster)
}
if l.Dashboard != "" {
return resolvers["dashboard"](ctx, w, l.Dashboard)
}
if l.InstancePool != "" {
return resolvers["instance-pool"](ctx, w, l.InstancePool)
}
if l.Job != "" {
return resolvers["job"](ctx, w, l.Job)
}
if l.Metastore != "" {
return resolvers["metastore"](ctx, w, l.Metastore)
}
if l.Pipeline != "" {
return resolvers["pipeline"](ctx, w, l.Pipeline)
}
if l.Query != "" {
return resolvers["query"](ctx, w, l.Query)
}
if l.ServicePrincipal != "" {
return resolvers["service-principal"](ctx, w, l.ServicePrincipal)
}
if l.Warehouse != "" {
return resolvers["warehouse"](ctx, w, l.Warehouse)
}
return "", fmt.Errorf("no valid lookup fields provided")
}
func (l *Lookup) String() string {
if l.Alert != "" {
return fmt.Sprintf("alert: %s", l.Alert)
}
if l.ClusterPolicy != "" {
return fmt.Sprintf("cluster-policy: %s", l.ClusterPolicy)
}
if l.Cluster != "" {
return fmt.Sprintf("cluster: %s", l.Cluster)
}
if l.Dashboard != "" {
return fmt.Sprintf("dashboard: %s", l.Dashboard)
}
if l.InstancePool != "" {
return fmt.Sprintf("instance-pool: %s", l.InstancePool)
}
if l.Job != "" {
return fmt.Sprintf("job: %s", l.Job)
}
if l.Metastore != "" {
return fmt.Sprintf("metastore: %s", l.Metastore)
}
if l.Pipeline != "" {
return fmt.Sprintf("pipeline: %s", l.Pipeline)
}
if l.Query != "" {
return fmt.Sprintf("query: %s", l.Query)
}
if l.ServicePrincipal != "" {
return fmt.Sprintf("service-principal: %s", l.ServicePrincipal)
}
if l.Warehouse != "" {
return fmt.Sprintf("warehouse: %s", l.Warehouse)
}
return ""
}
func (l *Lookup) validate() error {
// Validate that only one field is set
count := 0
if l.Alert != "" {
count++
}
if l.ClusterPolicy != "" {
count++
}
if l.Cluster != "" {
count++
}
if l.Dashboard != "" {
count++
}
if l.InstancePool != "" {
count++
}
if l.Job != "" {
count++
}
if l.Metastore != "" {
count++
}
if l.Pipeline != "" {
count++
}
if l.Query != "" {
count++
}
if l.ServicePrincipal != "" {
count++
}
if l.Warehouse != "" {
count++
}
if count != 1 {
return fmt.Errorf("exactly one lookup field must be provided")
}
if strings.Contains(l.String(), "${var") {
return fmt.Errorf("lookup fields cannot contain variable references")
}
return nil
}
type resolverFunc func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error)
func resolvers() map[string](resolverFunc) {
resolvers := make(map[string](resolverFunc), 0)
resolvers["alert"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Alerts.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
resolvers["cluster-policy"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.ClusterPolicies.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.PolicyId), nil
}
resolvers["cluster"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Clusters.GetByClusterName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.ClusterId), nil
}
resolvers["dashboard"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Dashboards.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
resolvers["instance-pool"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.InstancePools.GetByInstancePoolName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.InstancePoolId), nil
}
resolvers["job"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Jobs.GetBySettingsName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.JobId), nil
}
resolvers["metastore"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Metastores.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.MetastoreId), nil
}
resolvers["pipeline"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Pipelines.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.PipelineId), nil
}
resolvers["query"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Queries.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
resolvers["service-principal"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.ServicePrincipals.GetByDisplayName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
resolvers["warehouse"] = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
entity, err := w.Warehouses.GetByName(ctx, name)
if err != nil {
return "", err
}
return fmt.Sprint(entity.Id), nil
}
return resolvers
}

View File

@ -24,6 +24,10 @@ type Variable struct {
// 5. Throw error, since if no default value is defined, then the variable
// is required
Value *string `json:"value,omitempty" bundle:"readonly"`
// The value of this field will be used to lookup the resource by name
// And assign the value of the variable to ID of the resource found.
Lookup *Lookup `json:"lookup,omitempty"`
}
// True if the variable has been assigned a default value. Variables without a

View File

@ -27,6 +27,7 @@ func Initialize() bundle.Mutator {
mutator.ExpandWorkspaceRoot(),
mutator.DefineDefaultWorkspacePaths(),
mutator.SetVariables(),
mutator.ResolveResourceReferences(),
interpolation.Interpolate(
interpolation.IncludeLookupsInPath("bundle"),
interpolation.IncludeLookupsInPath("workspace"),

View File

@ -6,6 +6,11 @@ variables:
b:
description: required variable
d:
description: variable with lookup
lookup:
cluster: some-cluster
bundle:
name: test bundle
@ -30,3 +35,10 @@ targets:
variables:
c: prod-c
b: prod-b
env-overrides-lookup:
variables:
d:
lookup:
cluster: some-test-cluster
b: prod-b

View File

@ -104,3 +104,15 @@ func TestVariablesWithoutDefinition(t *testing.T) {
assert.Equal(t, "foo", *b.Config.Variables["a"].Value)
assert.Equal(t, "bar", *b.Config.Variables["b"].Value)
}
func TestVariablesWithTargetLookupOverrides(t *testing.T) {
b := load(t, "./variables/env_overrides")
err := bundle.Apply(context.Background(), b, bundle.Seq(
mutator.SelectTarget("env-overrides-lookup"),
mutator.SetVariables(),
interpolation.Interpolate(
interpolation.IncludeLookupsInPath(variable.VariableReferencePrefix),
)))
require.NoError(t, err)
assert.Equal(t, "cluster: some-test-cluster", b.Config.Variables["d"].Lookup.String())
}