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

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")
+	}
+}

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 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)
+			// Skip resources that do not pass all filters.
-			keyOnly[k] = append(keyOnly[k], ref)
+			if !includeReference(filters, ref) {
-			keyWithType[kt] = append(keyWithType[kt], 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) {
+func Lookup(b *bundle.Bundle, key string, filters ...Filter) (Reference, error) {
-	keyOnlyRefs, keyWithTypeRefs := References(b)
+	keyOnlyRefs, keyWithTypeRefs := References(b, filters...)
 	refs, ok := keyOnlyRefs[key]
 	if !ok {
 		refs, ok = keyWithTypeRefs[key]

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`)
+}

bundle/run/keys.go

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

bundle/run/keys_test.go

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

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.
+// IsRunnable returns a filter that only allows runnable resources.
-//
+func IsRunnable(ref refs.Reference) bool {
-// Its behavior is as follows:
+	switch ref.Resource.(type) {
-//  1. Try to find a resource with <key> identical to the argument.
+	case *resources.Job, *resources.Pipeline:
-//  2. Try to find a resource with <type>.<key> identical to the argument.
+		return true
-//
+	default:
-// If an argument resolves to multiple resources, it returns an error.
+		return false
-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")
+// 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
 }

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) {
+func TestRunner_IsRunnable(t *testing.T) {
-	b := &bundle.Bundle{
+	assert.True(t, IsRunnable(refs.Reference{Resource: &resources.Job{}}))
-		Config: config.Root{
+	assert.True(t, IsRunnable(refs.Reference{Resource: &resources.Pipeline{}}))
-			Resources: config.Resources{},
+	assert.False(t, IsRunnable(refs.Reference{Resource: &resources.MlflowModel{}}))
-		},
+	assert.False(t, IsRunnable(refs.Reference{Resource: &resources.MlflowExperiment{}}))
-	}
-
-	_, 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)
 }

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,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.
+		arg, err := resolveRunArgument(ctx, b, args)
-		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")
-		}
 
 		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
 			}