Cleanup scripts for windows in ShellExec (#2902)

## Why
The `cmd.exe` command executor we use to execute scripts in windows
writes the script first to disk before executing it. This is necessary
because `cmd.exe` does not support inlining the script contents as an
argument.

Since ShellExecv needs to exit the current CLI process within the
confines of the `execv` method to mirror the behaviour of the UNIX
version of `execv`, this PR adds a callback that cleans up the temporary
script before exiting the `execv` function call.

## Tests
Unit tests which verify the script is indeed being cleaned up.

Author: shreyas-goenka

internal/testutil/copy.go

@@ -24,24 +24,22 @@ func CopyDirectory(t TestingT, src, dst string) {
 			return os.MkdirAll(filepath.Join(dst, rel), 0o755)
 		}
 
-		// Copy the file to the temporary directory
-		in, err := os.Open(path)
-		if err != nil {
-			return err
-		}
-
-		defer in.Close()
+		CopyFile(t, path, filepath.Join(dst, rel))
+		return nil
+	})
 
-		out, err := os.Create(filepath.Join(dst, rel))
-		if err != nil {
-			return err
-		}
+	require.NoError(t, err)
+}
 
-		defer out.Close()
+func CopyFile(t TestingT, src, dst string) {
+	srcF, err := os.Open(src)
+	require.NoError(t, err)
+	defer srcF.Close()
 
-		_, err = io.Copy(out, in)
-		return err
-	})
+	dstF, err := os.Create(dst)
+	require.NoError(t, err)
+	defer dstF.Close()
 
+	_, err = io.Copy(dstF, srcF)
 	require.NoError(t, err)
 }

libs/exec/execv.go

@@ -1,5 +1,7 @@
 package exec
 
+import "os"
+
 type ExecvOptions struct {
 	// Args is the name of the command to run and its arguments.
 	// Eg: ["echo", "hello"] for "echo hello"
@@ -10,8 +12,21 @@ type ExecvOptions struct {
 
 	// Dir is the working directory of the child process.
 	Dir string
+
+	// It is not possible to execute a cmd.exe script inlined as a argument
+	// to cmd.exe. They have to be serialized to a file and executed.
+	// Thus if [Execv] is used to execution a script though cmd.exe,
+	// the caller should register a cleanup function to clean up the temporary file.
+	cleanup func()
+
+	// Callback to exit the current process in windows. Having this as a function here
+	// helps with testing.
+	windowsExit func(status int)
 }
 
 func Execv(opts ExecvOptions) error {
+	if opts.windowsExit == nil {
+		opts.windowsExit = os.Exit
+	}
 	return execv(opts)
 }

libs/exec/execv_windows.go

@@ -14,11 +14,24 @@ import (
 // and return the exit code.
 // ref: https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/execv-wexecv?view=msvc-170
 func execv(opts ExecvOptions) error {
+	if opts.cleanup != nil {
+		defer opts.cleanup()
+	}
+
+	windowsExit := func(status int) {
+		// First clean up the temporary script if it exists.
+		if opts.cleanup != nil {
+			opts.cleanup()
+		}
+
+		// Then exit the process.
+		opts.windowsExit(status)
+	}
+
 	path, err := exec.LookPath(opts.Args[0])
 	if err != nil {
 		return fmt.Errorf("looking up %q failed: %w", opts.Args[0], err)
 	}
-
 	cmd := exec.Command(path, opts.Args[1:]...)
 
 	cmd.Stdin = os.Stdin
@@ -35,7 +48,8 @@ func execv(opts ExecvOptions) error {
 
 	err = cmd.Wait()
 	if exitErr, ok := err.(*exec.ExitError); ok {
-		os.Exit(exitErr.ExitCode())
+		windowsExit(exitErr.ExitCode())
+		return nil
 	}
 	if err != nil {
 		return fmt.Errorf("running %s failed: %w", strings.Join(opts.Args, " "), err)
@@ -44,6 +58,6 @@ func execv(opts ExecvOptions) error {
 	// Unix implementation of execv never returns control to the CLI process.
 	// To emulate this behavior, we exit early here if the child process exits
 	// successfully.
-	os.Exit(0)
+	windowsExit(0)
 	return nil
 }

libs/exec/shell_execv.go

@@ -18,6 +18,9 @@ func shellExecvOpts(content, dir string, env []string) (ExecvOptions, error) {
 		Args: args,
 		Env:  env,
 		Dir:  dir,
+		cleanup: func() {
+			ec.cleanup()
+		},
 	}, nil
 }
 

libs/exec/shell_execv_test.go

@@ -1,9 +1,13 @@
 package exec
 
 import (
+	"os"
 	"os/exec"
+	"path/filepath"
+	"runtime"
 	"testing"
 
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -21,3 +25,61 @@ func TestShellExecvOpts(t *testing.T) {
 	assert.Equal(t, "-ec", opts.Args[1])
 	assert.Equal(t, "echo hello", opts.Args[2])
 }
+
+func TestShellExecv_Windows(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.Skip("skipping windows test")
+	}
+
+	cmdExePath, err := exec.LookPath("cmd.exe")
+	require.NoError(t, err)
+
+	// Cleanup environment so that other shells like bash and sh are not used.
+	testutil.NullEnvironment(t)
+
+	// Configure PATH so that only cmd.exe shows up.
+	binDir := t.TempDir()
+	testutil.CopyFile(t, cmdExePath, filepath.Join(binDir, "cmd.exe"))
+	os.Setenv("PATH", binDir)
+
+	tests := []struct {
+		name     string
+		content  string
+		exitCode int
+	}{
+		{name: "success", content: "echo hello", exitCode: 0},
+		{name: "non-zero exit", content: "exit 127", exitCode: 127},
+		{name: "command error", content: "not-a-command", exitCode: 1},
+	}
+
+	for _, test := range tests {
+		dir := t.TempDir()
+		t.Setenv("TMP", dir)
+
+		opts, err := shellExecvOpts(test.content, dir, []string{})
+		require.NoError(t, err)
+
+		// Verify that the temporary file is created.
+		files, err := os.ReadDir(dir)
+		require.NoError(t, err)
+		assert.Len(t, files, 1)
+		assert.Regexp(t, "cli-exec.*\\.cmd", files[0].Name())
+
+		exitCode := -1
+		opts.windowsExit = func(status int) {
+			exitCode = status
+		}
+
+		// Execute the script.
+		err = Execv(opts)
+		require.NoError(t, err)
+
+		// Verify that the temporary file is cleaned up after execution.
+		files, err = os.ReadDir(dir)
+		require.NoError(t, err)
+		assert.Len(t, files, 0)
+
+		// Verify that CLI would exit with the correct exit code.
+		assert.Equal(t, test.exitCode, exitCode)
+	}
+}