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

Author: pietern

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,35 +10,64 @@ 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.
 // The second map is indexed by the resource type name and its key.
 //
 // 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)
 
 	// Collect map of resource references indexed by their keys.
 	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]

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

@@ -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")
-	}
-
-	runners, ok := keyOnly[arg]
-	if !ok {
-		runners, ok = keyWithType[arg]
-		if !ok {
-			return nil, fmt.Errorf("no such resource: %s", arg)
-		}
+// 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
 	}
+}
 
-	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, ", "))
+// 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)
 	}
-
-	return runners[0], nil
 }