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

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