feat: find tree root via a user-specified command

Signed-off-by: Brian McGee <[email protected]>
Co-authored-by: Matt Sturgeon <[email protected]>

Author: brianmcgee

cmd/init/init.toml

@@ -36,6 +36,11 @@
 # Env $TREEFMT_TREE_ROOT_FILE
 # tree-root-file = ".git/config"
 
+# Command to run to find the tree root. It is parsed using shlex, to allow quoting arguments that contain whitespace.
+# If you wish to pass arguments containing quotes, you should use nested quotes e.g. "'" or '"'
+# Env $TREEFMT_TREE_ROOT_CMD
+# tree-root-cmd = "git rev-parse --show-toplevel"
+
 # Set the verbosity of logs
 # 0 = warn, 1 = info, 2 = debug
 # Env $TREEFMT_VERBOSE

cmd/root_test.go

@@ -1531,6 +1531,116 @@ func TestGit(t *testing.T) {
 	)
 }
 
+func TestTreeRootCmd(t *testing.T) {
+	as := require.New(t)
+
+	tempDir := test.TempExamples(t)
+	configPath := filepath.Join(tempDir, "/treefmt.toml")
+
+	test.ChangeWorkDir(t, tempDir)
+
+	// basic config
+	cfg := &config.Config{
+		FormatterConfigs: map[string]*config.Formatter{
+			"echo": {
+				Command:  "echo", // will not generate any underlying changes in the file
+				Includes: []string{"*"},
+			},
+		},
+	}
+
+	test.WriteConfig(t, configPath, cfg)
+
+	// construct a tree root command with some error logging and dumping output on stdout
+	treeRootCmd := func(output string) string {
+		return fmt.Sprintf("bash -c '>&2 echo -e \"some error text\nsome more error text\" && echo %s'", output)
+	}
+
+	// helper for checking the contents of stderr matches our expected debug output
+	checkStderr := func(buf []byte) {
+		output := string(buf)
+		as.Contains(output, "DEBU tree-root-cmd | stderr: some error text\n")
+		as.Contains(output, "DEBU tree-root-cmd | stderr: some more error text\n")
+	}
+
+	// run treefmt with DEBUG logging enabled and with tree root cmd being the root of the temp directory
+	treefmt(t,
+		withArgs("-vv", "--tree-root-cmd", treeRootCmd(tempDir)),
+		withNoError(t),
+		withStderr(checkStderr),
+		withConfig(configPath, cfg),
+		withStats(t, map[stats.Type]int{
+			stats.Traversed: 32,
+			stats.Matched:   32,
+			stats.Formatted: 32,
+			stats.Changed:   0,
+		}),
+	)
+
+	// run from a subdirectory, mixing things up by specifying the command via an env variable
+	treefmt(t,
+		withArgs("-vv"),
+		withEnv(map[string]string{
+			"TREEFMT_TREE_ROOT_CMD": treeRootCmd(filepath.Join(tempDir, "go")),
+		}),
+		withNoError(t),
+		withStderr(checkStderr),
+		withConfig(configPath, cfg),
+		withStats(t, map[stats.Type]int{
+			stats.Traversed: 2,
+			stats.Matched:   2,
+			stats.Formatted: 2,
+			stats.Changed:   0,
+		}),
+	)
+
+	// run from a subdirectory, mixing things up by specifying the command via config
+	cfg.TreeRootCmd = treeRootCmd(filepath.Join(tempDir, "haskell"))
+
+	treefmt(t,
+		withArgs("-vv"),
+		withNoError(t),
+		withStderr(checkStderr),
+		withConfig(configPath, cfg),
+		withStats(t, map[stats.Type]int{
+			stats.Traversed: 7,
+			stats.Matched:   7,
+			stats.Formatted: 7,
+			stats.Changed:   0,
+		}),
+	)
+
+	// run with a long-running command (2 seconds or more)
+	treefmt(t,
+		withArgs(
+			"-vv",
+			"--tree-root-cmd", fmt.Sprintf(
+				"bash -c 'sleep 2 && echo %s'",
+				tempDir,
+			),
+		),
+		withError(func(as *require.Assertions, err error) {
+			as.ErrorContains(err, "tree-root-cmd was killed after taking more than 2s to execute")
+		}),
+		withConfig(configPath, cfg),
+	)
+
+	// run with a command that outputs multiple lines
+	treefmt(t,
+		withArgs(
+			"-vv",
+			"--tree-root-cmd", fmt.Sprintf(
+				"bash -c 'echo %s && echo %s'",
+				tempDir, tempDir,
+			),
+		),
+		withError(func(as *require.Assertions, err error) {
+			as.ErrorContains(err, "tree-root-cmd cannot output multiple lines")
+		}),
+		withConfig(configPath, cfg),
+	)
+}
+
 func TestTreeRootExclusivity(t *testing.T) {
 	tempDir := test.TempExamples(t)
 	configPath := filepath.Join(tempDir, "/treefmt.toml")
@@ -1580,7 +1690,7 @@ func TestTreeRootExclusivity(t *testing.T) {
 		},
 	}
 
-	permutations := [][]string{
+	invalidCombinations := [][]string{
 		{"tree-root", "tree-root-cmd"},
 		{"tree-root", "tree-root-file"},
 		{"tree-root-cmd", "tree-root-file"},
@@ -1591,11 +1701,11 @@ func TestTreeRootExclusivity(t *testing.T) {
 	// Given that ultimately everything is being reduced into the config object after parsing from viper, I'm fairly
 	// confident if these tests all pass then the mixed methods should yield the same result.
 
-	// test permutations of the same type
-	for _, perm := range permutations {
+	// for each set of invalid args, test them with flags, environment variables, and config entries.
+	for _, combination := range invalidCombinations {
 		// test flags
 		var args []string
-		for _, key := range perm {
+		for _, key := range combination {
 			args = append(args, flagValues[key]...)
 		}
 
@@ -1607,7 +1717,7 @@ func TestTreeRootExclusivity(t *testing.T) {
 		// test env variables
 		env := make(map[string]string)
 
-		for _, key := range perm {
+		for _, key := range combination {
 			entry := envValues[key]
 			env[entry[0]] = entry[1]
 		}
@@ -1622,7 +1732,7 @@ func TestTreeRootExclusivity(t *testing.T) {
 			FormatterConfigs: formatterConfigs,
 		}
 
-		for _, key := range perm {
+		for _, key := range combination {
 			entry := configValues[key]
 			entry(cfg)
 		}

config/config.go

@@ -1,9 +1,11 @@
 package config
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -108,12 +110,15 @@ func SetFlags(fs *pflag.FlagSet) {
 	)
 	fs.String(
 		"tree-root", "",
-		"The root directory from which treefmt will start walking the filesystem (defaults to the directory "+
-			"containing the config file). (env $TREEFMT_TREE_ROOT)",
+		"The root directory from which treefmt will start walking the filesystem. "+
+			"Defaults to the root of the current git worktree. If not in a git repo, defaults to the directory "+
+			"containing the config file. (env $TREEFMT_TREE_ROOT)",
 	)
 	fs.String(
 		"tree-root-cmd", "",
-		"Command to run to find the tree root. (env $TREEFMT_TREE_ROOT_CMD)",
+		"Command to run to find the tree root. It is parsed using shlex, to allow quoting arguments that "+
+			"contain whitespace. If you wish to pass arguments containing quotes, you should use nested quotes "+
+			"e.g. \"'\" or '\"'. (env $TREEFMT_TREE_ROOT_CMD)",
 	)
 	fs.String(
 		"tree-root-file", "",
@@ -164,6 +169,8 @@ func NewViper() (*viper.Viper, error) {
 
 // FromViper takes a viper instance and produces a Config instance.
 func FromViper(v *viper.Viper) (*Config, error) {
+	logger := log.WithPrefix("config")
+
 	configReset := map[string]any{
 		"ci":          false,
 		"clear-cache": false,
@@ -198,7 +205,7 @@ func FromViper(v *viper.Viper) (*Config, error) {
 	}
 
 	// determine tree root
-	if err = determineTreeRoot(v, cfg); err != nil {
+	if err = determineTreeRoot(v, cfg, logger); err != nil {
 		return nil, fmt.Errorf("failed to determine tree root: %w", err)
 	}
 
@@ -258,7 +265,7 @@ func FromViper(v *viper.Viper) (*Config, error) {
 	return cfg, nil
 }
 
-func determineTreeRoot(v *viper.Viper, cfg *Config) error {
+func determineTreeRoot(v *viper.Viper, cfg *Config, logger *log.Logger) error {
 	var err error
 
 	// enforce the various tree root options are mutually exclusive
@@ -279,41 +286,44 @@ func determineTreeRoot(v *viper.Viper, cfg *Config) error {
 	}
 
 	if count > 1 {
-		return errors.New("only one of tree-root, tree-root-cmd or tree-root-file can be specified")
+		return errors.New("at most one of tree-root, tree-root-cmd or tree-root-file can be specified")
 	}
 
 	// set git-based tree root command if the walker is git and no tree root has been specified
 	if cfg.Walk == walk.Git.String() && count == 0 {
 		cfg.TreeRootCmd = "git rev-parse --show-toplevel"
 
-		log.Infof(
+		logger.Infof(
 			"git walker enabled and tree root has not been specified: defaulting tree-root-cmd to '%s'",
 			cfg.TreeRootCmd,
 		)
 	}
 
 	switch {
 	case cfg.TreeRoot != "":
-		log.Debugf("tree root specified explicitly: %s", cfg.TreeRoot)
+		logger.Debugf("tree root specified explicitly: %s", cfg.TreeRoot)
 
 	case cfg.TreeRootFile != "":
-		log.Debugf("searching for tree root using --tree-root-file: %s", cfg.TreeRootFile)
+		logger.Debugf("searching for tree root using tree-root-file: %s", cfg.TreeRootFile)
 
 		_, cfg.TreeRoot, err = FindUp(cfg.WorkingDirectory, cfg.TreeRootFile)
 		if err != nil {
 			return fmt.Errorf("failed to find tree-root based on tree-root-file: %w", err)
 		}
 
 	case cfg.TreeRootCmd != "":
-		log.Debugf("searching for tree root using --tree-root-cmd: %s", cfg.TreeRootCmd)
+		logger.Debugf("searching for tree root using tree-root-cmd: %s", cfg.TreeRootCmd)
 
 		if cfg.TreeRoot, err = execTreeRootCmd(cfg); err != nil {
 			return err
 		}
 
 	default:
 		// no tree root was specified
-		log.Debugf("no tree root specified, defaulting to the directory containing the config file: %s", v.ConfigFileUsed())
+		logger.Debugf(
+			"no tree root specified, defaulting to the directory containing the config file: %s",
+			v.ConfigFileUsed(),
+		)
 
 		cfg.TreeRoot = filepath.Dir(v.ConfigFileUsed())
 	}
@@ -323,7 +333,7 @@ func determineTreeRoot(v *viper.Viper, cfg *Config) error {
 		return fmt.Errorf("failed to get absolute path for tree root: %w", err)
 	}
 
-	log.Debugf("tree root: %s", cfg.TreeRoot)
+	logger.Debugf("tree root: %s", cfg.TreeRoot)
 
 	return nil
 }
@@ -337,24 +347,83 @@ func execTreeRootCmd(cfg *Config) (string, error) {
 
 	// set a reasonable timeout of 2 seconds to wait for the command to return
 	// it shouldn't take anywhere near this amount of time unless there's a problem
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	executionTimeout := 2 * time.Second
+
+	ctx, cancel := context.WithTimeout(context.Background(), executionTimeout)
 	defer cancel()
 
 	// construct the command, setting the correct working directory
 	//nolint:gosec
 	cmd := exec.CommandContext(ctx, parts[0], parts[1:]...)
 	cmd.Dir = cfg.WorkingDirectory
 
-	// execute
-	out, cmdErr := cmd.CombinedOutput()
-	if cmdErr != nil {
-		log.Errorf("tree-root-cmd output: \n%s", out)
+	// setup some pipes to capture stdout and stderr
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return "", fmt.Errorf("failed to create stdout pipe for tree-root-cmd: %w", err)
+	}
+
+	stderr, err := cmd.StderrPipe()
+	if err != nil {
+		return "", fmt.Errorf("failed to create stderr pipe for tree-root-cmd: %w", err)
+	}
+
+	// start processing stderr before we begin executing the command
+	go func() {
+		// capture stderr line by line and log
+		l := log.WithPrefix("tree-root-cmd | stderr")
+
+		scanner := bufio.NewScanner(stderr)
+		for scanner.Scan() {
+			l.Debugf("%s", scanner.Text())
+		}
+	}()
+
+	// start executing without waiting
+	if cmdErr := cmd.Start(); cmdErr != nil {
+		return "", fmt.Errorf("failed to start tree-root-cmd: %w", cmdErr)
+	}
+
+	// read stdout until it is closed (command exits)
+	output, err := io.ReadAll(stdout)
+	if err != nil {
+		return "", fmt.Errorf("failed to read stdout from tree-root-cmd: %w", err)
+	}
+
+	log.WithPrefix("tree-root-cmd | stdout").Debugf("%s", output)
+
+	// check execution error
+	if cmdErr := cmd.Wait(); cmdErr != nil {
+		var exitErr *exec.ExitError
+
+		// by experimenting, I noticed that sometimes we received the deadline exceeded error first, other times
+		// the exit error indicating the process was killed, therefore, we look for both
+		tookTooLong := errors.Is(cmdErr, context.DeadlineExceeded)
+		tookTooLong = tookTooLong || (errors.As(cmdErr, &exitErr) && exitErr.ProcessState.String() == "signal: killed")
+
+		if tookTooLong {
+			return "", fmt.Errorf(
+				"tree-root-cmd was killed after taking more than %v to execute",
+				executionTimeout,
+			)
+		}
+
+		// otherwise, some other kind of error occurred
+		return "", fmt.Errorf("failed to execute tree-root-cmd: %w", cmdErr)
+	}
+
+	// trim the output and check it's not empty
+	treeRoot := strings.TrimSpace(string(output))
+
+	if treeRoot == "" {
+		return "", fmt.Errorf("empty output received after executing tree-root-cmd: %s", cfg.TreeRootCmd)
+	}
 
-		return "", fmt.Errorf("failed to run tree-root-cmd: %w", cmdErr)
+	if strings.Contains(treeRoot, "\n") {
+		return "", fmt.Errorf("tree-root-cmd cannot output multiple lines: %s", cfg.TreeRootCmd)
 	}
 
-	// trim the output and return
-	return strings.TrimSpace(string(out)), nil
+	return treeRoot, nil
 }
 
 func Find(searchDir string, fileNames ...string) (path string, err error) {