Mark deps explicitly added with //go get as direct (#4027)

**What type of PR is this?** Feature **What does this PR do? Why is it needed?** When running `go get example.com/[email protected]`, the module `example.com/foo` will be added to `go.mod` with an `// indirect` comment, just like its transitive dependencies. This results in `go_deps` not being able to distinguish this explicitly requested new dep from other indirect deps that shouldn't be imported by the root module. Running `go mod tidy` or `go get` without arguments is required to have the comment removed after adding a reference to the new module in code. This change makes it so that `bazel run @rules_go//go get example.com/[email protected]` marks the requested module as a direct dependency. This realizes golang/go#68593 in our custom wrapper, thus making this command the only one needed to add a new module dependency and have it work with a subsequent Bazel build (assuming that also updates the BUILD files with Gazelle). This [matches the behavior of `gopls` when adding a new dependency](https://github.com/golang/tools/blob/ec1a81bfec7cc6563474b7aa160db30df9bfce6b/gopls/internal/server/command.go#L805-L809). --------- Co-authored-by: Son Luong Ngoc <[email protected]>
bazelbuild · Sep 15, 2024 · 2f062f8 · 2f062f8
1 parent 7eba57f
commit 2f062f8
Showing 1 changed file with 114 additions and 23 deletions.
diff --git a/go/tools/go_bin_runner/main.go b/go/tools/go_bin_runner/main.go
@@ -20,61 +20,89 @@ var GoBinRlocationPath = "not set"
 var ConfigRlocationPath = "not set"
 var HasBazelModTidy = "not set"
 
+type bazelEnvVars struct {
+	workspaceDir string
+	workingDir   string
+	binary       string
+}
+
+var bazelEnv = bazelEnvVars{
+	workspaceDir: os.Getenv("BUILD_WORKSPACE_DIRECTORY"),
+	workingDir:   os.Getenv("BUILD_WORKING_DIRECTORY"),
+	binary:       os.Getenv("BAZEL"),
+}
+
 // Produced by gazelle's go_deps extension.
 type Config struct {
 	GoEnv     map[string]string `json:"go_env"`
 	DepsFiles []string          `json:"dep_files"`
 }
 
 func main() {
+	if err := run(os.Args, os.Stdout, os.Stderr); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func run(args []string, stdout, stderr io.Writer) error {
+	// Force usage of the Bazel-configured Go SDK.
+	err := os.Setenv("GOTOOLCHAIN", "local")
+	if err != nil {
+		return err
+	}
+
 	goBin, err := runfiles.Rlocation(GoBinRlocationPath)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
 	cfg, err := parseConfig()
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
 	env, err := getGoEnv(goBin, cfg)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
 	hashesBefore, err := hashWorkspaceRelativeFiles(cfg.DepsFiles)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
-	args := append([]string{goBin}, os.Args[1:]...)
-	cwd := os.Getenv("BUILD_WORKING_DIRECTORY")
-	err = runProcess(args, env, cwd)
-	if err != nil {
-		log.Fatal(err)
+	goArgs := append([]string{goBin}, args[1:]...)
+	if err = runProcess(goArgs, env, stdout, stderr); err != nil {
+		return err
+	}
+
+	if len(args) > 1 && args[1] == "get" {
+		if err = markRequiresAsDirect(goBin, args[2:], stderr); err != nil {
+			return err
+		}
 	}
 
 	hashesAfter, err := hashWorkspaceRelativeFiles(cfg.DepsFiles)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
 	diff := diffMaps(hashesBefore, hashesAfter)
 	if len(diff) > 0 {
 		if HasBazelModTidy == "True" {
-			bazel := os.Getenv("BAZEL")
+			bazel := bazelEnv.binary
 			if bazel == "" {
 				bazel = "bazel"
 			}
-			_, _ = fmt.Fprintf(os.Stderr, "\nrules_go: Running '%s mod tidy' since %s changed...\n", bazel, strings.Join(diff, ", "))
-			err = runProcess([]string{bazel, "mod", "tidy"}, os.Environ(), cwd)
-			if err != nil {
-				log.Fatal(err)
+			_, _ = fmt.Fprintf(stderr, "rules_go: Running '%s mod tidy' since %s changed...\n", bazel, strings.Join(diff, ", "))
+			if err = runProcess([]string{bazel, "mod", "tidy"}, nil, stdout, stderr); err != nil {
+				return err
 			}
 		} else {
-			_, _ = fmt.Fprintf(os.Stderr, "\nrules_go: %s changed, please apply any buildozer fixes suggested by Bazel\n", strings.Join(diff, ", "))
+			_, _ = fmt.Fprintf(stderr, "rules_go: %s changed, please apply any buildozer fixes suggested by Bazel\n", strings.Join(diff, ", "))
 		}
 	}
+	return nil
 }
 
 func parseConfig() (Config, error) {
@@ -115,12 +143,75 @@ func getGoEnv(goBin string, cfg Config) ([]string, error) {
 	return append(env, "GOROOT="+goRoot), nil
 }
 
-func hashWorkspaceRelativeFiles(relativePaths []string) (map[string]string, error) {
-	workspace := os.Getenv("BUILD_WORKSPACE_DIRECTORY")
+// Make every explicitly specified module a direct dep by removing the
+// "// indirect" comment. This results in the go_deps extension treating the
+// new module as visible to the main module, without the user having to first
+// add a reference in code and run `go mod tidy`.
+// https://github.com/golang/go/issues/68593
+func markRequiresAsDirect(goBin string, getArgs []string, stderr io.Writer) error {
+	var pkgs []string
+	for _, arg := range getArgs {
+		if strings.HasPrefix(arg, "-") {
+			// Skip flags.
+			continue
+		}
+		pkgs = append(pkgs, strings.Split(arg, "@")[0])
+	}
+
+	// Run 'go mod edit -json' to get the list of requires.
+	cmd := exec.Command(goBin, "mod", "edit", "-json")
+	cmd.Dir = bazelEnv.workingDir
+	out, err := cmd.Output()
+	if err != nil {
+		return err
+	}
 
+	var modJson struct {
+		Require []struct{
+			Path     string
+			Version  string
+			Indirect bool
+		}
+	}
+	if err = json.Unmarshal(out, &modJson); err != nil {
+		return err
+	}
+
+	// Make every explicitly specified module a direct dep by dropping and
+	// re-adding the require directive - this is the only way to remove the
+	// indirect comment with go mod edit.
+	// Note that we do not use golang.org/x/mod/modfile to edit the go.mod file
+	// as this would cause @rules_go//go to fail if there is an issue with this
+	// module dep such as a missing sum.
+	var editArgs []string
+	for _, require := range modJson.Require {
+		if !require.Indirect {
+			continue
+		}
+		for _, pkg := range pkgs {
+			if strings.HasPrefix(pkg, require.Path) && (len(pkg) == len(require.Path) || pkg[len(require.Path)] == '/') {
+				editArgs = append(editArgs, "-droprequire", require.Path, "-require", require.Path+"@"+require.Version)
+				break
+			}
+		}
+	}
+
+	if len(editArgs) > 0 {
+		_, _ = fmt.Fprintln(stderr, "rules_go: Marking requested modules as direct dependencies...")
+		cmd = exec.Command(goBin, append([]string{"mod", "edit"}, editArgs...)...)
+		cmd.Dir = bazelEnv.workingDir
+		if err = cmd.Run(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func hashWorkspaceRelativeFiles(relativePaths []string) (map[string]string, error) {
 	hashes := make(map[string]string)
 	for _, p := range relativePaths {
-		h, err := hashFile(filepath.Join(workspace, p))
+		h, err := hashFile(filepath.Join(bazelEnv.workspaceDir, p))
 		if err != nil {
 			return nil, err
 		}
@@ -154,11 +245,11 @@ func hashFile(path string) (string, error) {
 	return hex.EncodeToString(h.Sum(nil)), nil
 }
 
-func runProcess(args, env []string, dir string) error {
+func runProcess(args, env []string, stdout, stderr io.Writer) error {
 	cmd := exec.Command(args[0], args[1:]...)
-	cmd.Dir = dir
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
+	cmd.Dir = bazelEnv.workingDir
+	cmd.Stdout = stdout
+	cmd.Stderr = stderr
 	cmd.Env = env
 	err := cmd.Run()
 	if exitErr, ok := err.(*exec.ExitError); ok {