feat(pkg-drv): allow setting custom aspect and kinds (#3328)

* feat(pkg-drv): allow setting custom aspect and kinds

* fixup! use kinds array

* fixup! format

go/tools/gopackagesdriver/bazel_json_builder.go

@@ -30,12 +30,15 @@ const (
 	RulesGoStdlibLabel = "@io_bazel_rules_go//:stdlib"
 )
 
+var _defaultKinds = []string{"go_library", "go_test", "go_binary"}
+
 func (b *BazelJSONBuilder) fileQuery(filename string) string {
 	if filepath.IsAbs(filename) {
 		fp, _ := filepath.Rel(b.bazel.WorkspaceRoot(), filename)
 		filename = fp
 	}
-	return fmt.Sprintf(`kind("go_library|go_test|go_binary", same_pkg_direct_rdeps("%s"))`, filename)
+	kinds := append(_defaultKinds, additionalKinds...)
+	return fmt.Sprintf(`kind("%s", same_pkg_direct_rdeps("%s"))`, strings.Join(kinds, "|"), filename)
 }
 
 func (b *BazelJSONBuilder) packageQuery(importPath string) string {
@@ -120,7 +123,7 @@ func (b *BazelJSONBuilder) Build(ctx context.Context, mode LoadMode) ([]string,
 		"--experimental_convenience_symlinks=ignore",
 		"--ui_event_filters=-info,-stderr",
 		"--noshow_progress",
-		"--aspects=" + rulesGoRepositoryName + "//go/tools/gopackagesdriver:aspect.bzl%go_pkg_info_aspect",
+		"--aspects=" + customAspect,
 		"--output_groups=" + b.outputGroupsForMode(mode),
 		"--keep_going", // Build all possible packages
 	}, bazelFlags, bazelBuildFlags, labels)

go/tools/gopackagesdriver/main.go

@@ -58,6 +58,8 @@ var (
 	bazelQueryScope       = getenvDefault("GOPACKAGESDRIVER_BAZEL_QUERY_SCOPE", "")
 	bazelBuildFlags       = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_BUILD_FLAGS"))
 	workspaceRoot         = os.Getenv("BUILD_WORKSPACE_DIRECTORY")
+	customAspect          = getenvDefault("GOPACKAGESDRIVER_BAZEL_ASPECT", rulesGoRepositoryName+"//go/tools/gopackagesdriver:aspect.bzl%go_pkg_info_aspect")
+	additionalKinds       = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_KINDS"))
 	emptyResponse         = &driverResponse{
 		NotHandled: false,
 		Sizes:      types.SizesFor("gc", "amd64").(*types.StdSizes),