| package cq | 
 |  | 
 | import ( | 
 | 	"context" | 
 | 	"io" | 
 | 	"io/ioutil" | 
 | 	osexec "os/exec" | 
 | 	"path/filepath" | 
 | 	"regexp" | 
 |  | 
 | 	"github.com/bazelbuild/buildtools/build" | 
 | 	"go.skia.org/infra/go/exec" | 
 | 	"go.skia.org/infra/go/skerr" | 
 | 	"go.skia.org/infra/go/util" | 
 | ) | 
 |  | 
 | // WithUpdateCQConfig reads the given Starlark config file, runs the given | 
 | // function to modify it, then writes it back to disk and runs lucicfg to | 
 | // generate the proto config files in the given directory. Expects lucicfg to | 
 | // be in PATH and "lucicfg auth-login" to have already been run. generatedDir | 
 | // must be a descendant of the directory which contains starlarkConfigFile. | 
 | func WithUpdateCQConfig(ctx context.Context, starlarkConfigFile, generatedDir string, fn func(*build.File) error) error { | 
 | 	// Make the generatedDir relative to the directory containing | 
 | 	// starlarkConfigFile. | 
 | 	parentDir := filepath.Dir(starlarkConfigFile) | 
 | 	relConfigFile := filepath.Base(starlarkConfigFile) | 
 | 	relGenDir, err := filepath.Rel(parentDir, generatedDir) | 
 | 	if err != nil { | 
 | 		return skerr.Wrapf(err, "could not make %s relative to %s", generatedDir, parentDir) | 
 | 	} | 
 |  | 
 | 	// Check presence and auth status of lucicfg. | 
 | 	if lucicfg, err := osexec.LookPath("lucicfg"); err != nil || lucicfg == "" { | 
 | 		return skerr.Fmt("unable to find lucicfg in PATH; do you have depot tools installed?") | 
 | 	} | 
 | 	if _, err := exec.RunCwd(ctx, ".", "lucicfg", "auth-info"); err != nil { | 
 | 		return skerr.Wrapf(err, "please run `lucicfg auth-login`") | 
 | 	} | 
 |  | 
 | 	// Read the config file. | 
 | 	oldCfgBytes, err := ioutil.ReadFile(starlarkConfigFile) | 
 | 	if err != nil { | 
 | 		return skerr.Wrapf(err, "failed to read %s", starlarkConfigFile) | 
 | 	} | 
 |  | 
 | 	// Update the config. | 
 | 	newCfgBytes, err := WithUpdateCQConfigBytes(starlarkConfigFile, oldCfgBytes, fn) | 
 | 	if err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 |  | 
 | 	// Write the new config. | 
 | 	if err := util.WithWriteFile(starlarkConfigFile, func(w io.Writer) error { | 
 | 		_, err := w.Write(newCfgBytes) | 
 | 		return err | 
 | 	}); err != nil { | 
 | 		return skerr.Wrapf(err, "failed to write config file") | 
 | 	} | 
 |  | 
 | 	// Run lucicfg to generate the proto config files. | 
 | 	cmd := []string{"lucicfg", "generate", "-validate", "-config-dir", relGenDir, relConfigFile} | 
 | 	if _, err := exec.RunCwd(ctx, parentDir, cmd...); err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 | 	return nil | 
 | } | 
 |  | 
 | // WithUpdateCQConfigBytes parses the given bytes as a Config, runs the given | 
 | // function to modify the Config, then returns the updated bytes. | 
 | func WithUpdateCQConfigBytes(filename string, oldCfgBytes []byte, fn func(*build.File) error) ([]byte, error) { | 
 | 	// Parse the Config. | 
 | 	f, err := build.ParseDefault(filename, oldCfgBytes) | 
 | 	if err != nil { | 
 | 		return nil, skerr.Wrapf(err, "failed to parse config file") | 
 | 	} | 
 |  | 
 | 	// Run the passed-in func. | 
 | 	if err := fn(f); err != nil { | 
 | 		return nil, skerr.Wrapf(err, "config modification failed") | 
 | 	} | 
 |  | 
 | 	// Write the new config bytes. | 
 | 	return build.Format(f), nil | 
 | } | 
 |  | 
 | // FindAssignExpr finds the AssignExpr with the given identifier. | 
 | func FindAssignExpr(callExpr *build.CallExpr, identifier string) (int, *build.AssignExpr, error) { | 
 | 	for idx, expr := range callExpr.List { | 
 | 		if assignExpr, ok := expr.(*build.AssignExpr); ok { | 
 | 			if ident, ok := assignExpr.LHS.(*build.Ident); ok { | 
 | 				if ident.Name == identifier { | 
 | 					return idx, assignExpr, nil | 
 | 				} | 
 | 			} | 
 | 		} | 
 | 	} | 
 | 	return 0, nil, skerr.Fmt("no AssignExpr found for %q", identifier) | 
 | } | 
 |  | 
 | // FindExprForBranch finds the CallExpr for the given branch. | 
 | func FindExprForBranch(f *build.File, branch string) (int, *build.CallExpr, error) { | 
 | 	for idx, expr := range f.Stmt { | 
 | 		if callExpr, ok := expr.(*build.CallExpr); ok { | 
 | 			_, assignExpr, err := FindAssignExpr(callExpr, "name") | 
 | 			if err != nil { | 
 | 				continue | 
 | 			} | 
 | 			if stringExpr, ok := assignExpr.RHS.(*build.StringExpr); ok { | 
 | 				if stringExpr.Value == branch { | 
 | 					return idx, callExpr, nil | 
 | 				} | 
 | 			} | 
 | 		} | 
 | 	} | 
 | 	return 0, nil, skerr.Fmt("no config group found for %q", branch) | 
 | } | 
 |  | 
 | // CopyExprSlice returns a shallow copy of the []Expr. | 
 | func CopyExprSlice(slice []build.Expr) []build.Expr { | 
 | 	cp := make([]build.Expr, 0, len(slice)) | 
 | 	for _, expr := range slice { | 
 | 		cp = append(cp, expr) | 
 | 	} | 
 | 	return cp | 
 | } | 
 |  | 
 | // CloneBranch updates the given CQ config to create a config for a new | 
 | // branch based on a given existing branch. Optionally, include experimental | 
 | // tryjobs, include the tree-is-open check, and exclude trybots matching regular | 
 | // expressions. | 
 | func CloneBranch(f *build.File, oldBranch, newBranch string, includeExperimental, includeTreeCheck bool, excludeTrybotRegexp []*regexp.Regexp) error { | 
 | 	// Find the CQ config for the old branch. | 
 | 	_, oldBranchExpr, err := FindExprForBranch(f, oldBranch) | 
 | 	if err != nil { | 
 | 		return err | 
 | 	} | 
 |  | 
 | 	// Copy the old branch and modify it for the new branch. | 
 | 	newBranchExpr, ok := oldBranchExpr.Copy().(*build.CallExpr) | 
 | 	if !ok { | 
 | 		return skerr.Fmt("expected CallExpr for cq_group") | 
 | 	} | 
 | 	newBranchExpr.List = CopyExprSlice(oldBranchExpr.List) | 
 |  | 
 | 	// CQ group name. | 
 | 	nameExprIndex, nameExpr, err := FindAssignExpr(newBranchExpr, "name") | 
 | 	if err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 | 	nameExprCopy := nameExpr.Copy().(*build.AssignExpr) | 
 | 	nameStr, ok := nameExprCopy.RHS.Copy().(*build.StringExpr) | 
 | 	if !ok { | 
 | 		return skerr.Fmt("expected StringExpr for name") | 
 | 	} | 
 | 	nameStr.Value = newBranch | 
 | 	nameExprCopy.RHS = nameStr | 
 | 	newBranchExpr.List[nameExprIndex] = nameExprCopy | 
 |  | 
 | 	// Ref match. | 
 | 	refsetExprIndex, refsetExpr, err := FindAssignExpr(newBranchExpr, "watch") | 
 | 	if err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 | 	refsetExprCopy := refsetExpr.Copy().(*build.AssignExpr) | 
 | 	newBranchExpr.List[refsetExprIndex] = refsetExprCopy | 
 |  | 
 | 	refsetCallExpr, ok := refsetExprCopy.RHS.(*build.CallExpr) | 
 | 	if !ok { | 
 | 		return skerr.Fmt("expected CallExpr for refset") | 
 | 	} | 
 | 	refsetCallExprCopy := refsetCallExpr.Copy().(*build.CallExpr) | 
 | 	refsetExprCopy.RHS = refsetCallExprCopy | 
 | 	refsetCallExprCopy.List = CopyExprSlice(refsetCallExprCopy.List) | 
 |  | 
 | 	refsExprIndex, refsExpr, err := FindAssignExpr(refsetCallExprCopy, "refs") | 
 | 	if err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 | 	refsExprCopy := refsExpr.Copy().(*build.AssignExpr) | 
 | 	refsetCallExprCopy.List[refsExprIndex] = refsExprCopy | 
 |  | 
 | 	refsListExpr, ok := refsExprCopy.RHS.(*build.ListExpr) | 
 | 	if !ok { | 
 | 		return skerr.Fmt("expected ListExpr for refs") | 
 | 	} | 
 | 	refsListExprCopy := refsListExpr.Copy().(*build.ListExpr) | 
 | 	refsExprCopy.RHS = refsListExprCopy | 
 |  | 
 | 	if len(refsListExprCopy.List) != 1 { | 
 | 		return skerr.Fmt("expected a single ref but got %d", len(refsListExprCopy.List)) | 
 | 	} | 
 | 	refExpr, ok := refsListExprCopy.List[0].(*build.StringExpr) | 
 | 	if !ok { | 
 | 		return skerr.Fmt("expected StringExpr for ref") | 
 | 	} | 
 | 	refExprCopy := refExpr.Copy().(*build.StringExpr) | 
 | 	refExprCopy.Value = "refs/heads/" + newBranch | 
 | 	refsListExprCopy.List = []build.Expr{refExprCopy} | 
 |  | 
 | 	// Tryjobs. | 
 | 	verifiersExprIndex, verifiersExpr, err := FindAssignExpr(newBranchExpr, "verifiers") | 
 | 	if err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 | 	verifiersExprCopy := verifiersExpr.Copy().(*build.AssignExpr) | 
 | 	newBranchExpr.List[verifiersExprIndex] = verifiersExprCopy | 
 |  | 
 | 	verifiersListExpr, ok := verifiersExprCopy.RHS.(*build.ListExpr) | 
 | 	if !ok { | 
 | 		return skerr.Fmt("expected ListExpr for verifiers") | 
 | 	} | 
 | 	verifiersListExprCopy := verifiersListExpr.Copy().(*build.ListExpr) | 
 | 	verifiersExprCopy.RHS = verifiersListExprCopy | 
 |  | 
 | 	verifiersListExprCopy.List = make([]build.Expr, 0, len(verifiersListExpr.List)) | 
 | 	for _, expr := range verifiersListExpr.List { | 
 | 		verifierCallExpr, ok := expr.(*build.CallExpr) | 
 | 		if !ok { | 
 | 			return skerr.Fmt("expected CallExpr for verifier") | 
 | 		} | 
 | 		// Include experimental builders? | 
 | 		_, _, err := FindAssignExpr(verifierCallExpr, "experiment_percentage") | 
 | 		if err == nil && !includeExperimental { | 
 | 			continue | 
 | 		} | 
 |  | 
 | 		// Is this builder excluded based on a regex? | 
 | 		_, builder, err := FindAssignExpr(verifierCallExpr, "builder") | 
 | 		if err != nil { | 
 | 			return skerr.Wrap(err) | 
 | 		} | 
 | 		builderStringExpr, ok := builder.RHS.(*build.StringExpr) | 
 | 		if !ok { | 
 | 			return skerr.Fmt("expected StringExpr for builder name") | 
 | 		} | 
 | 		include := true | 
 | 		for _, regex := range excludeTrybotRegexp { | 
 | 			if regex.MatchString(builderStringExpr.Value) { | 
 | 				include = false | 
 | 				break | 
 | 			} | 
 | 		} | 
 | 		if include { | 
 | 			// No need to copy, since we're not modifying the verifier | 
 | 			// expression itself. | 
 | 			verifiersListExprCopy.List = append(verifiersListExprCopy.List, verifierCallExpr) | 
 | 		} | 
 | 	} | 
 |  | 
 | 	// Tree status. | 
 | 	if !includeTreeCheck { | 
 | 		treeCheckIndex, _, err := FindAssignExpr(newBranchExpr, "tree_status_host") | 
 | 		if err == nil { | 
 | 			cp := make([]build.Expr, 0, len(newBranchExpr.List)) | 
 | 			for idx, expr := range newBranchExpr.List { | 
 | 				if idx != treeCheckIndex { | 
 | 					cp = append(cp, expr) | 
 | 				} | 
 | 			} | 
 | 			newBranchExpr.List = cp | 
 | 		} | 
 | 	} | 
 |  | 
 | 	// Add the new branch config. | 
 | 	f.Stmt = append(f.Stmt, newBranchExpr) | 
 | 	return nil | 
 | } | 
 |  | 
 | // DeleteBranch updates the given CQ config to delete the config matching the | 
 | // given branch. | 
 | func DeleteBranch(f *build.File, branch string) error { | 
 | 	branchExprIndex, _, err := FindExprForBranch(f, branch) | 
 | 	if err != nil { | 
 | 		return skerr.Wrap(err) | 
 | 	} | 
 | 	newStmt := make([]build.Expr, 0, len(f.Stmt)-1) | 
 | 	for idx, expr := range f.Stmt { | 
 | 		if idx != branchExprIndex { | 
 | 			newStmt = append(newStmt, expr) | 
 | 		} | 
 | 	} | 
 | 	f.Stmt = newStmt | 
 | 	return nil | 
 | } |