mirror of https://github.com/databricks/cli.git
128 lines
3.9 KiB
Go
128 lines
3.9 KiB
Go
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
|
|
}
|