databricks-cli/cmd/labs/project/proxy.go

147 lines
3.6 KiB
Go
Raw Normal View History

package project
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io/fs"
"path/filepath"
"strings"
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/log"
"github.com/databricks/cli/libs/process"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)
type proxy struct {
Entrypoint `yaml:",inline"`
Name string `yaml:"name"`
Description string `yaml:"description"`
TableTemplate string `yaml:"table_template,omitempty"`
Flags []flag `yaml:"flags,omitempty"`
}
func (cp *proxy) register(parent *cobra.Command) {
cmd := &cobra.Command{
Use: cp.Name,
Short: cp.Description,
RunE: cp.runE,
}
parent.AddCommand(cmd)
flags := cmd.Flags()
for _, flag := range cp.Flags {
flag.register(flags)
}
}
func (cp *proxy) runE(cmd *cobra.Command, _ []string) error {
err := cp.checkUpdates(cmd)
if err != nil {
return err
}
args, err := cp.commandInput(cmd)
if err != nil {
return err
}
envs, err := cp.Prepare(cmd)
if err != nil {
return fmt.Errorf("entrypoint: %w", err)
}
ctx := cmd.Context()
log.Debugf(ctx, "Forwarding subprocess: %s", strings.Join(args, " "))
if cp.TableTemplate != "" {
return cp.renderJsonAsTable(cmd, args, envs)
}
err = process.Forwarded(ctx, args,
cmd.InOrStdin(),
cmd.OutOrStdout(),
cmd.ErrOrStderr(),
process.WithEnvs(envs))
if errors.Is(err, fs.ErrNotExist) && cp.IsPythonProject() {
msg := "cannot find Python %s. Please re-run: databricks labs install %s"
return fmt.Errorf(msg, cp.MinPython, cp.Name)
}
return err
}
// [EXPERIMENTAL] this interface contract may change in the future.
// See https://github.com/databricks/cli/issues/994
func (cp *proxy) renderJsonAsTable(cmd *cobra.Command, args []string, envs map[string]string) error {
var buf bytes.Buffer
ctx := cmd.Context()
err := process.Forwarded(ctx, args,
cmd.InOrStdin(),
&buf,
cmd.ErrOrStderr(),
process.WithEnvs(envs))
if err != nil {
return err
}
var anyVal any
err = json.Unmarshal(buf.Bytes(), &anyVal)
if err != nil {
return err
}
// IntelliJ eagerly replaces tabs with spaces, even though we're not asking for it
fixedTemplate := strings.ReplaceAll(cp.TableTemplate, "\\t", "\t")
return cmdio.RenderWithTemplate(ctx, anyVal, fixedTemplate)
}
func (cp *proxy) commandInput(cmd *cobra.Command) ([]string, error) {
flags := cmd.Flags()
commandInput := struct {
Command string `json:"command"`
Flags map[string]any `json:"flags"`
OutputType string `json:"output_type"`
}{
Command: cp.Name,
Flags: map[string]any{},
}
for _, f := range cp.Flags {
v, err := f.get(flags)
if err != nil {
return nil, fmt.Errorf("get %s flag: %w", f.Name, err)
}
commandInput.Flags[f.Name] = v
}
logLevelFlag := flags.Lookup("log-level")
if logLevelFlag != nil {
commandInput.Flags["log_level"] = logLevelFlag.Value.String()
}
args := []string{}
ctx := cmd.Context()
if cp.IsPythonProject() {
args = append(args, cp.virtualEnvPython(ctx))
libDir := cp.EffectiveLibDir()
entrypoint := filepath.Join(libDir, cp.Main)
args = append(args, entrypoint)
}
raw, err := json.Marshal(commandInput)
if err != nil {
return nil, fmt.Errorf("command input: %w", err)
}
args = append(args, string(raw))
return args, nil
}
type flag struct {
Name string `yaml:"name"`
Description string `yaml:"description"`
Default any `yaml:"default,omitempty"`
}
func (f *flag) register(pf *pflag.FlagSet) {
var dflt string
if f.Default != nil {
dflt = fmt.Sprint(f.Default)
}
pf.String(f.Name, dflt, f.Description)
}
func (f *flag) get(pf *pflag.FlagSet) (any, error) {
return pf.GetString(f.Name)
}