go/cq/update.go - buildbot - Git at Google

 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
 }
	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
	}