steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Reuse resource resolution code for the run command (#1858) 

## Changes

As of #1846 we have a generalized package for doing resource lookups and
completion.

This change updates the run command to use this instead of more specific
code under `bundle/run`.

## Tests

* Unit tests pass
* Manually confirmed that completion and prompting works
This commit is contained in:
Pieter Noordhuis 2024-10-24 15:24:30 +02:00 committed by GitHub
parent eaea308254
commit ed84a33b0a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 175 additions and 226 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
bundle/resources/completion.go
View File

@ -4,9 +4,9 @@ import "github.com/databricks/cli/bundle"
// Completions returns the same as [References] except
// that every key maps directly to a single reference.
func Completions(b *bundle.Bundle) map[string]Reference {
func Completions(b *bundle.Bundle, filters ...Filter) map[string]Reference {
out := make(map[string]Reference)
keyOnlyRefs, _ := References(b)
keyOnlyRefs, _ := References(b, filters...)
for k, refs := range keyOnlyRefs {
if len(refs) != 1 {
continue

26
bundle/resources/completion_test.go
View File

@ -30,3 +30,29 @@ func TestCompletions_SkipDuplicates(t *testing.T) {
assert.Contains(t, out, "bar")
}
}
func TestCompletions_Filter(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"foo": {},
},
Pipelines: map[string]*resources.Pipeline{
"bar": {},
},
},
},
}
includeJobs := func(ref Reference) bool {
_, ok := ref.Resource.(*resources.Job)
return ok
}
// Test that this does not include the pipeline.
out := Completions(b, includeJobs)
if assert.Len(t, out, 1) {
assert.Contains(t, out, "foo")
}
}

41
bundle/resources/lookup.go
View File

@ -10,13 +10,36 @@ import (
// Reference is a reference to a resource.
// It includes the resource type description, and a reference to the resource itself.
type Reference struct {
// Key is the unique key of the resource, e.g. "my_job".
Key string
// KeyWithType is the unique key of the resource, including the resource type, e.g. "jobs.my_job".
KeyWithType string
// Description is the resource type description.
Description config.ResourceDescription
// Resource is the resource itself.
Resource config.ConfigResource
}
// Map is the core type for resource lookup and completion.
type Map map[string][]Reference
// Filter defines the function signature for filtering resources.
type Filter func(Reference) bool
// includeReference checks if the specified reference passes all filters.
// If the list of filters is empty, the reference is always included.
func includeReference(filters []Filter, ref Reference) bool {
for _, filter := range filters {
if !filter(ref) {
return false
}
}
return true
}
// References returns maps of resource keys to a slice of [Reference].
//
// The first map is indexed by the resource key only.
@ -24,7 +47,7 @@ type Map map[string][]Reference
//
// While the return types allows for multiple resources to share the same key,
// this is confirmed not to happen in the [validate.UniqueResourceKeys] mutator.
func References(b *bundle.Bundle) (Map, Map) {
func References(b *bundle.Bundle, filters ...Filter) (Map, Map) {
keyOnly := make(Map)
keyWithType := make(Map)
@ -32,13 +55,19 @@ func References(b *bundle.Bundle) (Map, Map) {
for _, group := range b.Config.Resources.AllResources() {
for k, v := range group.Resources {
ref := Reference{
Key: k,
KeyWithType: fmt.Sprintf("%s.%s", group.Description.PluralName, k),
Description: group.Description,
Resource: v,
}
kt := fmt.Sprintf("%s.%s", group.Description.PluralName, k)
keyOnly[k] = append(keyOnly[k], ref)
keyWithType[kt] = append(keyWithType[kt], ref)
// Skip resources that do not pass all filters.
if !includeReference(filters, ref) {
continue
}
keyOnly[ref.Key] = append(keyOnly[ref.Key], ref)
keyWithType[ref.KeyWithType] = append(keyWithType[ref.KeyWithType], ref)
}
}
@ -48,8 +77,8 @@ func References(b *bundle.Bundle) (Map, Map) {
// Lookup returns the resource with the specified key.
// If the key maps to more than one resource, an error is returned.
// If the key does not map to any resource, an error is returned.
func Lookup(b *bundle.Bundle, key string) (Reference, error) {
keyOnlyRefs, keyWithTypeRefs := References(b)
func Lookup(b *bundle.Bundle, key string, filters ...Filter) (Reference, error) {
keyOnlyRefs, keyWithTypeRefs := References(b, filters...)
refs, ok := keyOnlyRefs[key]
if !ok {
refs, ok = keyWithTypeRefs[key]

29
bundle/resources/lookup_test.go
View File

@ -86,3 +86,32 @@ func TestLookup_Nominal(t *testing.T) {
assert.Equal(t, "Foo job", out.Resource.GetName())
}
}
func TestLookup_NominalWithFilters(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"foo": {},
},
Pipelines: map[string]*resources.Pipeline{
"bar": {},
},
},
},
}
includeJobs := func(ref Reference) bool {
_, ok := ref.Resource.(*resources.Job)
return ok
}
// This should succeed because the filter includes jobs.
_, err := Lookup(b, "foo", includeJobs)
require.NoError(t, err)
// This should fail because the filter excludes pipelines.
_, err = Lookup(b, "bar", includeJobs)
require.Error(t, err)
assert.ErrorContains(t, err, `resource with key "bar" not found`)
}

69
bundle/run/keys.go
View File

@ -1,69 +0,0 @@
package run
import (
"fmt"
"github.com/databricks/cli/bundle"
"golang.org/x/exp/maps"
)
// RunnerLookup maps identifiers to a list of workloads that match that identifier.
// The list can have more than 1 entry if resources of different types use the
// same key. When this happens, the user should disambiguate between them.
type RunnerLookup map[string][]Runner
// ResourceKeys computes a map with
func ResourceKeys(b *bundle.Bundle) (keyOnly RunnerLookup, keyWithType RunnerLookup) {
keyOnly = make(RunnerLookup)
keyWithType = make(RunnerLookup)
r := b.Config.Resources
for k, v := range r.Jobs {
kt := fmt.Sprintf("jobs.%s", k)
w := jobRunner{key: key(kt), bundle: b, job: v}
keyOnly[k] = append(keyOnly[k], &w)
keyWithType[kt] = append(keyWithType[kt], &w)
}
for k, v := range r.Pipelines {
kt := fmt.Sprintf("pipelines.%s", k)
w := pipelineRunner{key: key(kt), bundle: b, pipeline: v}
keyOnly[k] = append(keyOnly[k], &w)
keyWithType[kt] = append(keyWithType[kt], &w)
}
return
}
// ResourceCompletionMap returns a map of resource keys to their respective names.
func ResourceCompletionMap(b *bundle.Bundle) map[string]string {
out := make(map[string]string)
keyOnly, keyWithType := ResourceKeys(b)
// Keep track of resources we have seen by their fully qualified key.
seen := make(map[string]bool)
// First add resources that can be identified by key alone.
for k, v := range keyOnly {
// Invariant: len(v) >= 1. See [ResourceKeys].
if len(v) == 1 {
seen[v[0].Key()] = true
out[k] = v[0].Name()
}
}
// Then add resources that can only be identified by their type and key.
for k, v := range keyWithType {
// Invariant: len(v) == 1. See [ResourceKeys].
_, ok := seen[v[0].Key()]
if ok {
continue
}
out[k] = v[0].Name()
}
return out
}
// ResourceCompletions returns a list of keys that unambiguously reference resources in the bundle.
func ResourceCompletions(b *bundle.Bundle) []string {
return maps.Keys(ResourceCompletionMap(b))
}

25
bundle/run/keys_test.go
View File

@ -1,25 +0,0 @@
package run
import (
"testing"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/bundle/config/resources"
"github.com/stretchr/testify/assert"
)
func TestResourceCompletionsUnique(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"foo": {},
"bar": {},
},
},
},
}
assert.ElementsMatch(t, []string{"foo", "bar"}, ResourceCompletions(b))
}

51
bundle/run/runner.go
View File

@ -3,9 +3,10 @@ package run
import (
"context"
"fmt"
"strings"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/config/resources"
refs "github.com/databricks/cli/bundle/resources"
"github.com/databricks/cli/bundle/run/output"
)
@ -38,34 +39,24 @@ type Runner interface {
argsHandler
}
// Find locates a runner matching the specified argument.
//
// Its behavior is as follows:
// 1. Try to find a resource with <key> identical to the argument.
// 2. Try to find a resource with <type>.<key> identical to the argument.
//
// If an argument resolves to multiple resources, it returns an error.
func Find(b *bundle.Bundle, arg string) (Runner, error) {
keyOnly, keyWithType := ResourceKeys(b)
if len(keyWithType) == 0 {
return nil, fmt.Errorf("bundle defines no resources")
// IsRunnable returns a filter that only allows runnable resources.
func IsRunnable(ref refs.Reference) bool {
switch ref.Resource.(type) {
case *resources.Job, *resources.Pipeline:
return true
default:
return false
}
}
// ToRunner converts a resource reference to a runnable resource.
func ToRunner(b *bundle.Bundle, ref refs.Reference) (Runner, error) {
switch resource := ref.Resource.(type) {
case *resources.Job:
return &jobRunner{key: key(ref.KeyWithType), bundle: b, job: resource}, nil
case *resources.Pipeline:
return &pipelineRunner{key: key(ref.KeyWithType), bundle: b, pipeline: resource}, nil
default:
return nil, fmt.Errorf("unsupported resource type: %T", resource)
}
runners, ok := keyOnly[arg]
if !ok {
runners, ok = keyWithType[arg]
if !ok {
return nil, fmt.Errorf("no such resource: %s", arg)
}
}
if len(runners) != 1 {
var keys []string
for _, runner := range runners {
keys = append(keys, runner.Key())
}
return nil, fmt.Errorf("ambiguous: %s (can resolve to all of %s)", arg, strings.Join(keys, ", "))
}
return runners[0], nil
}

80
bundle/run/runner_test.go
View File

@ -3,82 +3,14 @@ package run
import (
"testing"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/bundle/config/resources"
refs "github.com/databricks/cli/bundle/resources"
"github.com/stretchr/testify/assert"
)
func TestFindNoResources(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{},
},
}
_, err := Find(b, "foo")
assert.ErrorContains(t, err, "bundle defines no resources")
}
func TestFindSingleArg(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"foo": {},
},
},
},
}
_, err := Find(b, "foo")
assert.NoError(t, err)
}
func TestFindSingleArgNotFound(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"foo": {},
},
},
},
}
_, err := Find(b, "bar")
assert.ErrorContains(t, err, "no such resource: bar")
}
func TestFindSingleArgAmbiguous(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"key": {},
},
Pipelines: map[string]*resources.Pipeline{
"key": {},
},
},
},
}
_, err := Find(b, "key")
assert.ErrorContains(t, err, "ambiguous: ")
}
func TestFindSingleArgWithType(t *testing.T) {
b := &bundle.Bundle{
Config: config.Root{
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"key": {},
},
},
},
}
_, err := Find(b, "jobs.key")
assert.NoError(t, err)
func TestRunner_IsRunnable(t *testing.T) {
assert.True(t, IsRunnable(refs.Reference{Resource: &resources.Job{}}))
assert.True(t, IsRunnable(refs.Reference{Resource: &resources.Pipeline{}}))
assert.False(t, IsRunnable(refs.Reference{Resource: &resources.MlflowModel{}}))
assert.False(t, IsRunnable(refs.Reference{Resource: &resources.MlflowExperiment{}}))
}

70
cmd/bundle/run.go
View File

@ -1,12 +1,14 @@
package bundle
import (
"context"
"encoding/json"
"fmt"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/deploy/terraform"
"github.com/databricks/cli/bundle/phases"
"github.com/databricks/cli/bundle/resources"
"github.com/databricks/cli/bundle/run"
"github.com/databricks/cli/bundle/run/output"
"github.com/databricks/cli/cmd/bundle/utils"
@ -14,8 +16,54 @@ import (
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/flags"
"github.com/spf13/cobra"
"golang.org/x/exp/maps"
)
func promptRunArgument(ctx context.Context, b *bundle.Bundle) (string, error) {
// Compute map of "Human readable name of resource" -> "resource key".
inv := make(map[string]string)
for k, ref := range resources.Completions(b, run.IsRunnable) {
title := fmt.Sprintf("%s: %s", ref.Description.SingularTitle, ref.Resource.GetName())
inv[title] = k
}
key, err := cmdio.Select(ctx, inv, "Resource to run")
if err != nil {
return "", err
}
return key, nil
}
func resolveRunArgument(ctx context.Context, b *bundle.Bundle, args []string) (string, error) {
// If no arguments are specified, prompt the user to select something to run.
if len(args) == 0 && cmdio.IsPromptSupported(ctx) {
return promptRunArgument(ctx, b)
}
if len(args) < 1 {
return "", fmt.Errorf("expected a KEY of the resource to run")
}
return args[0], nil
}
func keyToRunner(b *bundle.Bundle, arg string) (run.Runner, error) {
// Locate the resource to run.
ref, err := resources.Lookup(b, arg, run.IsRunnable)
if err != nil {
return nil, err
}
// Convert the resource to a runnable resource.
runner, err := run.ToRunner(b, ref)
if err != nil {
return nil, err
}
return runner, nil
}
func newRunCommand() *cobra.Command {
cmd := &cobra.Command{
Use: "run [flags] KEY",
@ -61,23 +109,10 @@ task or a Python wheel task, the second example applies.
return err
}
// If no arguments are specified, prompt the user to select something to run.
if len(args) == 0 && cmdio.IsPromptSupported(ctx) {
// Invert completions from KEY -> NAME, to NAME -> KEY.
inv := make(map[string]string)
for k, v := range run.ResourceCompletionMap(b) {
inv[v] = k
}
id, err := cmdio.Select(ctx, inv, "Resource to run")
arg, err := resolveRunArgument(ctx, b, args)
if err != nil {
return err
}
args = append(args, id)
}
if len(args) < 1 {
return fmt.Errorf("expected a KEY of the resource to run")
}
diags = bundle.Apply(ctx, b, bundle.Seq(
terraform.Interpolate(),
@ -89,7 +124,7 @@ task or a Python wheel task, the second example applies.
return err
}
runner, err := run.Find(b, args[0])
runner, err := keyToRunner(b, arg)
if err != nil {
return err
}
@ -146,10 +181,11 @@ task or a Python wheel task, the second example applies.
}
if len(args) == 0 {
return run.ResourceCompletions(b), cobra.ShellCompDirectiveNoFileComp
completions := resources.Completions(b, run.IsRunnable)
return maps.Keys(completions), cobra.ShellCompDirectiveNoFileComp
} else {
// If we know the resource to run, we can complete additional positional arguments.
runner, err := run.Find(b, args[0])
runner, err := keyToRunner(b, args[0])
if err != nil {
return nil, cobra.ShellCompDirectiveError
}