From ed84a33b0a70312c04f14908225bb9c119bb84f5 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 24 Oct 2024 15:24:30 +0200 Subject: [PATCH] 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 --- bundle/resources/completion.go | 4 +- bundle/resources/completion_test.go | 26 ++++++++++ bundle/resources/lookup.go | 43 +++++++++++++--- bundle/resources/lookup_test.go | 29 +++++++++++ bundle/run/keys.go | 69 ------------------------- bundle/run/keys_test.go | 25 --------- bundle/run/runner.go | 51 ++++++++---------- bundle/run/runner_test.go | 80 +++-------------------------- cmd/bundle/run.go | 74 +++++++++++++++++++------- 9 files changed, 175 insertions(+), 226 deletions(-) delete mode 100644 bundle/run/keys.go delete mode 100644 bundle/run/keys_test.go diff --git a/bundle/resources/completion.go b/bundle/resources/completion.go index 3ce0510a..c1bcd022 100644 --- a/bundle/resources/completion.go +++ b/bundle/resources/completion.go @@ -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 diff --git a/bundle/resources/completion_test.go b/bundle/resources/completion_test.go index 36ad1a06..2f7942aa 100644 --- a/bundle/resources/completion_test.go +++ b/bundle/resources/completion_test.go @@ -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") + } +} diff --git a/bundle/resources/lookup.go b/bundle/resources/lookup.go index 74aec531..48d86223 100644 --- a/bundle/resources/lookup.go +++ b/bundle/resources/lookup.go @@ -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 config.ConfigResource + + // 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] diff --git a/bundle/resources/lookup_test.go b/bundle/resources/lookup_test.go index d2092c23..b2eaafd1 100644 --- a/bundle/resources/lookup_test.go +++ b/bundle/resources/lookup_test.go @@ -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`) +} diff --git a/bundle/run/keys.go b/bundle/run/keys.go deleted file mode 100644 index 76ec50ac..00000000 --- a/bundle/run/keys.go +++ /dev/null @@ -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)) -} diff --git a/bundle/run/keys_test.go b/bundle/run/keys_test.go deleted file mode 100644 index 5ab73b13..00000000 --- a/bundle/run/keys_test.go +++ /dev/null @@ -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)) -} diff --git a/bundle/run/runner.go b/bundle/run/runner.go index 1cdcc9d8..4c907d06 100644 --- a/bundle/run/runner.go +++ b/bundle/run/runner.go @@ -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 identical to the argument. -// 2. Try to find a resource with . 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 } diff --git a/bundle/run/runner_test.go b/bundle/run/runner_test.go index 85baa192..2fc5fa6f 100644 --- a/bundle/run/runner_test.go +++ b/bundle/run/runner_test.go @@ -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{}})) } diff --git a/cmd/bundle/run.go b/cmd/bundle/run.go index ed5bd2ef..96851d0c 100644 --- a/cmd/bundle/run.go +++ b/cmd/bundle/run.go @@ -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,22 +109,9 @@ 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") - if err != nil { - return err - } - args = append(args, id) - } - - if len(args) < 1 { - return fmt.Errorf("expected a KEY of the resource to run") + arg, err := resolveRunArgument(ctx, b, args) + if err != nil { + return err } diags = bundle.Apply(ctx, b, bundle.Seq( @@ -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 }