From 4f43fb9acf8892cc8a95905e057e19ca42deb10d Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Sun, 2 Mar 2025 18:11:46 +0100 Subject: [PATCH] [WIP] Add bundle exec --- acceptance/acceptance_test.go | 2 +- acceptance/bundle/exec/basic/output.txt | 5 ++ acceptance/bundle/exec/basic/script | 1 + acceptance/bundle/help/bundle/output.txt | 1 + cmd/bundle/bundle.go | 1 + cmd/bundle/exec.go | 109 +++++++++++++++++++++++ 6 files changed, 118 insertions(+), 1 deletion(-) create mode 100644 acceptance/bundle/exec/basic/output.txt create mode 100644 acceptance/bundle/exec/basic/script create mode 100644 cmd/bundle/exec.go diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 066a84299..f8b259643 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -41,7 +41,7 @@ var ( // In order to debug CLI running under acceptance test, set this to full subtest name, e.g. "bundle/variables/empty" // Then install your breakpoints and click "debug test" near TestAccept in VSCODE. // example: var SingleTest = "bundle/variables/empty" -var SingleTest = "" +var SingleTest = "bundle/exec/basic" // If enabled, instead of compiling and running CLI externally, we'll start in-process server that accepts and runs // CLI commands. The $CLI in test scripts is a helper that just forwards command-line arguments to this server (see bin/callserver.py). diff --git a/acceptance/bundle/exec/basic/output.txt b/acceptance/bundle/exec/basic/output.txt new file mode 100644 index 000000000..bc6165d26 --- /dev/null +++ b/acceptance/bundle/exec/basic/output.txt @@ -0,0 +1,5 @@ + +>>> errcode [CLI] bundle exec -- echo Hello, World! +Error: Please add a '--' separator. Usage: 'databricks bundle exec -- arg1 arg2 ...' + +Exit code: 1 diff --git a/acceptance/bundle/exec/basic/script b/acceptance/bundle/exec/basic/script new file mode 100644 index 000000000..f0d714b9a --- /dev/null +++ b/acceptance/bundle/exec/basic/script @@ -0,0 +1 @@ +trace errcode $CLI bundle exec echo "Hello,\ World" diff --git a/acceptance/bundle/help/bundle/output.txt b/acceptance/bundle/help/bundle/output.txt index fc6dd623d..bb885c80e 100644 --- a/acceptance/bundle/help/bundle/output.txt +++ b/acceptance/bundle/help/bundle/output.txt @@ -11,6 +11,7 @@ Available Commands: deploy Deploy bundle deployment Deployment related commands destroy Destroy deployed bundle resources + exec Execute a command using the same authentication context as the bundle generate Generate bundle configuration init Initialize using a bundle template open Open a resource in the browser diff --git a/cmd/bundle/bundle.go b/cmd/bundle/bundle.go index fb88cd7d0..e0818c2f9 100644 --- a/cmd/bundle/bundle.go +++ b/cmd/bundle/bundle.go @@ -28,5 +28,6 @@ func New() *cobra.Command { cmd.AddCommand(newDebugCommand()) cmd.AddCommand(deployment.NewDeploymentCommand()) cmd.AddCommand(newOpenCommand()) + cmd.AddCommand(newExecCommand()) return cmd } diff --git a/cmd/bundle/exec.go b/cmd/bundle/exec.go new file mode 100644 index 000000000..efc8164f5 --- /dev/null +++ b/cmd/bundle/exec.go @@ -0,0 +1,109 @@ +package bundle + +import ( + "bufio" + "fmt" + "os/exec" + "strings" + "sync" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/auth" + "github.com/spf13/cobra" +) + +// TODO: Confirm that quoted strings are parsed as a single argument. +// TODO: test that -- works with flags as well. +// TODO CONTINUE: Making the bundle exec function work. +// TODO CONTINUE: Adding the scripts section to DABs. + +func newExecCommand() *cobra.Command { + variadicArgs := []string{} + + execCmd := &cobra.Command{ + Use: "exec", + Short: "Execute a command using the same authentication context as the bundle", + Args: cobra.MinimumNArgs(1), + Long: `Examples: + // TODO: Ensure that these multi work strings work with the exec command. +1. databricks bundle exec -- echo "Hello, world!" +2. databricks bundle exec -- /bin/bash -c "echo 'Hello, world!'" +3. databricks bundle exec -- uv run pytest"`, + RunE: func(cmd *cobra.Command, args []string) error { + if cmd.ArgsLenAtDash() != 0 { + return fmt.Errorf("Please add a '--' separator. Usage: 'databricks bundle exec -- %s'", strings.Join(args, " ")) + } + + // Load the bundle configuration to get the authentication credentials + // set in the context. + // TODO: What happens when no bundle is configured? + _, diags := root.MustConfigureBundle(cmd) + if diags.HasError() { + return diags.Error() + } + + childCmd := exec.Command(args[1], args[2:]...) + childCmd.Env = auth.ProcessEnv(root.ConfigUsed(cmd.Context())) + + // Create pipes for stdout and stderr + stdout, err := childCmd.StdoutPipe() + if err != nil { + return fmt.Errorf("Error creating stdout pipe: %w", err) + } + + stderr, err := childCmd.StderrPipe() + if err != nil { + return fmt.Errorf("Error creating stderr pipe: %w", err) + } + + // Start the command + if err := childCmd.Start(); err != nil { + return fmt.Errorf("Error starting command: %s\n", err) + } + + // Stream both stdout and stderr to the user. + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + scanner := bufio.NewScanner(stdout) + for scanner.Scan() { + fmt.Println(scanner.Text()) + } + }() + + go func() { + defer wg.Done() + scanner := bufio.NewScanner(stderr) + for scanner.Scan() { + fmt.Println(scanner.Text()) + } + }() + + // Wait for the command to finish. + // TODO: Pretty exit codes? + // TODO: Make CLI return the same exit codes? + err = childCmd.Wait() + if exitErr, ok := err.(*exec.ExitError); ok { + return fmt.Errorf("Command exited with code: %d", exitErr.ExitCode()) + } + + if err := childCmd.Wait(); err != nil { + return fmt.Errorf("Error waiting for command: %w", err) + } + + // Wait for the goroutines to finish printing to stdout and stderr. + wg.Wait() + + return nil + }, + } + + // TODO: Is this needed to make -- work with flags? + // execCmd.Flags().SetInterspersed(false) + + // TODO: func (c *Command) ArgsLenAtDash() int solves my problems here. + + return execCmd +}