Processing and completion of positional args to bundle run (#1120)

## Changes

With this change, both job parameters and task parameters can be
specified as positional arguments to bundle run. How the positional
arguments are interpreted depends on the configuration of the job.

### Examples:

For a job that has job parameters configured a user can specify:

```
databricks bundle run my_job -- --param1=value1 --param2=value2
```

And the run is kicked off with job parameters set to:
```json
{
  "param1": "value1",
  "param2": "value2"
}
```

Similarly, for a job that doesn't use job parameters and only has
`notebook_task` tasks, a user can specify:

```
databricks bundle run my_notebook_job -- --param1=value1 --param2=value2
```

And the run is kicked off with task level `notebook_params` configured
as:
```json
{
  "param1": "value1",
  "param2": "value2"
}
```

For a job that doesn't doesn't use job parameters and only has either
`spark_python_task` or `python_wheel_task` tasks, a user can specify:

```
databricks bundle run my_python_file_job -- --flag=value other arguments
```

And the run is kicked off with task level `python_params` configured as:
```json
[
  "--flag=value",
  "other",
  "arguments"
]
```

The same is applied to jobs with only `spark_jar_task` or
`spark_submit_task` tasks.

## Tests

Unit tests. Tested the completions manually.

bundle/run/args.go

@@ -0,0 +1,127 @@
+package run
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+)
+
+// argsHandler defines the (unexported) interface for the runners in this
+// package to implement to handle context-specific positional arguments.
+//
+// For jobs, this means:
+//   - If a job uses job parameters: parse positional arguments into key-value pairs
+//     and pass them as job parameters.
+//   - If a job does not use job parameters AND only has Spark Python tasks:
+//     pass through the positional arguments as a list of Python parameters.
+//   - If a job does not use job parameters AND only has notebook tasks:
+//     parse arguments into key-value pairs and pass them as notebook parameters.
+//   - ...
+//
+// In all cases, we may be able to provide context-aware argument completions.
+type argsHandler interface {
+	// Parse additional positional arguments.
+	ParseArgs(args []string, opts *Options) error
+
+	// Complete additional positional arguments.
+	CompleteArgs(args []string, toComplete string) ([]string, cobra.ShellCompDirective)
+}
+
+// nopArgsHandler is a no-op implementation of [argsHandler].
+// It returns an error if any positional arguments are present and doesn't complete anything.
+type nopArgsHandler struct{}
+
+func (nopArgsHandler) ParseArgs(args []string, opts *Options) error {
+	if len(args) == 0 {
+		return nil
+	}
+
+	return fmt.Errorf("received %d unexpected positional arguments", len(args))
+}
+
+func (nopArgsHandler) CompleteArgs(args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return nil, cobra.ShellCompDirectiveNoFileComp
+}
+
+// argsToKeyValueMap parses key-value pairs from the specified arguments.
+//
+// It accepts these formats:
+//   - `--key=value`
+//   - `--key`, `value`
+//
+// Remaining arguments are returned as-is.
+func argsToKeyValueMap(args []string) (map[string]string, []string) {
+	kv := make(map[string]string)
+	key := ""
+	tail := args
+
+	for i, arg := range args {
+		// If key is set; use the next argument as value.
+		if key != "" {
+			kv[key] = arg
+			key = ""
+			tail = args[i+1:]
+			continue
+		}
+
+		if strings.HasPrefix(arg, "--") {
+			parts := strings.SplitN(arg[2:], "=", 2)
+			if len(parts) == 2 {
+				kv[parts[0]] = parts[1]
+				tail = args[i+1:]
+				continue
+			}
+
+			// Use this argument as key, the next as value.
+			key = parts[0]
+			continue
+		}
+
+		// If we cannot interpret it; return here.
+		break
+	}
+
+	return kv, tail
+}
+
+// genericParseKeyValueArgs parses key-value pairs from the specified arguments.
+// If there are any positional arguments left, it returns an error.
+func genericParseKeyValueArgs(args []string) (map[string]string, error) {
+	kv, args := argsToKeyValueMap(args)
+	if len(args) > 0 {
+		return nil, fmt.Errorf("received %d unexpected positional arguments", len(args))
+	}
+
+	return kv, nil
+}
+
+// genericCompleteKeyValueArgs completes key-value pairs from the specified arguments.
+// Completion options that are already specified are skipped.
+func genericCompleteKeyValueArgs(args []string, toComplete string, options []string) ([]string, cobra.ShellCompDirective) {
+	// If the string to complete contains an equals sign, then we are
+	// completing the value part (which we don't know here).
+	if strings.Contains(toComplete, "=") {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	// Remove already completed key/value pairs.
+	kv, args := argsToKeyValueMap(args)
+
+	// If the list of remaining args is empty, return possible completions.
+	if len(args) == 0 {
+		var completions []string
+		for _, option := range options {
+			// Skip options that have already been specified.
+			if _, ok := kv[option]; ok {
+				continue
+			}
+			completions = append(completions, fmt.Sprintf("--%s=", option))
+		}
+		// Note: we include cobra.ShellCompDirectiveNoSpace to suggest including
+		// the value part right after the equals sign.
+		return completions, cobra.ShellCompDirectiveNoFileComp | cobra.ShellCompDirectiveNoSpace
+	}
+
+	return nil, cobra.ShellCompDirectiveNoFileComp
+}

bundle/run/args_test.go

@@ -0,0 +1,134 @@
+package run
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNopArgsHandler(t *testing.T) {
+	h := nopArgsHandler{}
+	opts := &Options{}
+
+	// No error if no positional arguments are passed.
+	err := h.ParseArgs([]string{}, opts)
+	assert.NoError(t, err)
+
+	// Error if any positional arguments are passed.
+	err = h.ParseArgs([]string{"foo"}, opts)
+	assert.EqualError(t, err, "received 1 unexpected positional arguments")
+
+	// No completions.
+	completions, _ := h.CompleteArgs([]string{}, "")
+	assert.Nil(t, completions)
+}
+
+func TestArgsToKeyValueMap(t *testing.T) {
+	for _, tc := range []struct {
+		input    []string
+		expected map[string]string
+		tail     []string
+		err      error
+	}{
+		{
+			input:    []string{},
+			expected: map[string]string{},
+			tail:     []string{},
+		},
+		{
+			input: []string{"--foo=bar", "--baz", "qux"},
+			expected: map[string]string{
+				"foo": "bar",
+				"baz": "qux",
+			},
+			tail: []string{},
+		},
+		{
+			input: []string{"--foo=bar", "--baz", "qux", "tail"},
+			expected: map[string]string{
+				"foo": "bar",
+				"baz": "qux",
+			},
+			tail: []string{"tail"},
+		},
+		{
+			input: []string{"--foo=bar", "--baz", "qux", "tail", "--foo=bar"},
+			expected: map[string]string{
+				"foo": "bar",
+				"baz": "qux",
+			},
+			tail: []string{"tail", "--foo=bar"},
+		},
+		{
+			input: []string{"--foo=bar", "--baz=qux"},
+			expected: map[string]string{
+				"foo": "bar",
+				"baz": "qux",
+			},
+			tail: []string{},
+		},
+		{
+			input: []string{"--foo=bar", "--baz=--qux"},
+			expected: map[string]string{
+				"foo": "bar",
+				"baz": "--qux",
+			},
+			tail: []string{},
+		},
+		{
+			input: []string{"--foo=bar", "--baz="},
+			expected: map[string]string{
+				"foo": "bar",
+				"baz": "",
+			},
+			tail: []string{},
+		},
+		{
+			input: []string{"--foo=bar", "--baz"},
+			expected: map[string]string{
+				"foo": "bar",
+			},
+			tail: []string{"--baz"},
+		},
+	} {
+		actual, tail := argsToKeyValueMap(tc.input)
+		assert.Equal(t, tc.expected, actual)
+		assert.Equal(t, tc.tail, tail)
+	}
+}
+
+func TestGenericParseKeyValueArgs(t *testing.T) {
+	kv, err := genericParseKeyValueArgs([]string{"--foo=bar", "--baz", "qux"})
+	assert.NoError(t, err)
+	assert.Equal(t, map[string]string{
+		"foo": "bar",
+		"baz": "qux",
+	}, kv)
+
+	_, err = genericParseKeyValueArgs([]string{"--foo=bar", "--baz", "qux", "tail"})
+	assert.EqualError(t, err, "received 1 unexpected positional arguments")
+}
+
+func TestGenericCompleteKeyValueArgs(t *testing.T) {
+	var completions []string
+
+	// Complete nothing if there are no options.
+	completions, _ = genericCompleteKeyValueArgs([]string{}, ``, []string{})
+	assert.Empty(t, completions)
+
+	// Complete nothing if we're in the middle of a key-value pair (as single argument with equals sign).
+	completions, _ = genericCompleteKeyValueArgs([]string{}, `--foo=`, []string{`foo`, `bar`})
+	assert.Empty(t, completions)
+
+	// Complete nothing if we're in the middle of a key-value pair (as two arguments).
+	completions, _ = genericCompleteKeyValueArgs([]string{`--foo`}, ``, []string{`foo`, `bar`})
+	assert.Empty(t, completions)
+
+	// Complete if we're at the beginning.
+	completions, _ = genericCompleteKeyValueArgs([]string{}, ``, []string{`foo`, `bar`})
+	assert.Equal(t, []string{`--foo=`, `--bar=`}, completions)
+
+	// Complete if we have already one key-value pair.
+	completions, _ = genericCompleteKeyValueArgs([]string{`--foo=bar`}, ``, []string{`foo`, `bar`})
+	assert.Equal(t, []string{`--bar=`}, completions)
+}

bundle/run/job.go

@@ -15,6 +15,7 @@ import (
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/fatih/color"
+	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -315,3 +316,11 @@ func (r *jobRunner) Cancel(ctx context.Context) error {
 
 	return errGroup.Wait()
 }
+
+func (r *jobRunner) ParseArgs(args []string, opts *Options) error {
+	return r.posArgsHandler().ParseArgs(args, opts)
+}
+
+func (r *jobRunner) CompleteArgs(args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return r.posArgsHandler().CompleteArgs(args, toComplete)
+}