cmd.go

Documentation: edge-infra.dev/hack/tools/tp-container-dep-migrator

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  
    11  	"github.com/bazelbuild/buildtools/build"
    12  
    13  	container "edge-infra.dev/hack/build/rules/container/gazelle/language"
    14  	"edge-infra.dev/pkg/lib/build/bazel"
    15  	"edge-infra.dev/pkg/lib/cli/sink"
    16  )
    17  
    18  func New() *sink.Command {
    19  	return &sink.Command{
    20  		Use:   "tp-migrate [space separated list of files to migrate]",
    21  		Short: "Migrate .bzl files containing dictionaries to use third_party_container_dep macros",
    22  		Exec: func(_ context.Context, r sink.Run) error {
    23  			filesToMigrate := r.Args()
    24  			wd := bazel.ResolveWdOrDie()
    25  
    26  			for _, fPath := range filesToMigrate {
    27  				if !strings.Contains(fPath, ".bzl") {
    28  					r.Log.Error(errors.New("skipping non-bzl file"), "filePath", fPath)
    29  					continue
    30  				}
    31  				r.Log.Info("found", "bzl file", fPath)
    32  				migratedPath, migrationRequired, err := migrateBzlFile(wd, fPath, r)
    33  				if err != nil {
    34  					return err
    35  				}
    36  
    37  				if migrationRequired {
    38  					r.Log.Info("migrated", "file", migratedPath)
    39  				} else {
    40  					r.Log.Info("no image dictionaries present to migrate", "file", fPath)
    41  				}
    42  			}
    43  
    44  			return nil
    45  		},
    46  	}
    47  }
    48  
    49  // Main function body to migrate a .bzl file's image dictionaries to a wrapper function
    50  // containing third_party_container_dep rules instead
    51  func migrateBzlFile(wdPath string, fPath string, r sink.Run) (string, bool, error) {
    52  	migratedFileName := fmt.Sprintf("%s_migrated%s", strings.TrimSuffix(filepath.Base(fPath), filepath.Ext(fPath)), filepath.Ext(fPath))
    53  	migratedFilePath := fmt.Sprintf("%s/%s/%s", wdPath, filepath.Dir(fPath), migratedFileName)
    54  
    55  	originalPath := filepath.Join(wdPath, fPath)
    56  	bzlBytes, err := os.ReadFile(originalPath)
    57  	if err != nil {
    58  		r.Log.Error(err, "error reading original file", "file path", originalPath)
    59  		return "", false, err
    60  	}
    61  
    62  	bzlFile, err := build.ParseBzl(originalPath, bzlBytes)
    63  	if err != nil {
    64  		r.Log.Error(err, "error parsing bzl file", "bzl file path", originalPath)
    65  		return "", false, err
    66  	}
    67  
    68  	imageDictReplaceMap := map[int]build.Expr{}
    69  	for i, s := range bzlFile.Stmt {
    70  		if assignExpr, isAssign := s.(*build.AssignExpr); isAssign {
    71  			if dictStmt, isDict := assignExpr.RHS.(*build.DictExpr); isDict {
    72  				wrapperDefName := strings.ReplaceAll(assignExpr.LHS.(*build.Ident).Name, "-", "_")
    73  				wrapperDefName = strings.ToLower(wrapperDefName)
    74  
    75  				wrapperDef := imageDictToTpDepWrapperDef(wrapperDefName, dictStmt)
    76  
    77  				imageDictReplaceMap[i] = wrapperDef
    78  			}
    79  		}
    80  	}
    81  	if len(imageDictReplaceMap) == 0 {
    82  		return "", false, nil
    83  	}
    84  
    85  	migratedFile := bzlFile.Copy().(*build.File)
    86  
    87  	for replaceIndex, wrapperDef := range imageDictReplaceMap {
    88  		migratedFile.Stmt[replaceIndex] = wrapperDef
    89  	}
    90  
    91  	tpDepLoad := &build.LoadStmt{
    92  		Module:       &build.StringExpr{Value: "//hack/build/rules/container:third_party_images.bzl"},
    93  		From:         []*build.Ident{&build.Ident{Name: container.ThirdPartyContainerDepRuleName}},
    94  		To:           []*build.Ident{&build.Ident{Name: container.ThirdPartyContainerDepRuleName}},
    95  		ForceCompact: true,
    96  	}
    97  	migratedFile.Stmt = append(migratedFile.Stmt, tpDepLoad)
    98  	err = os.WriteFile(migratedFilePath, build.Format(migratedFile), 0644)
    99  
   100  	build.Format(migratedFile)
   101  	if err != nil {
   102  		r.Log.Error(err, "error writing file", "file", migratedFilePath)
   103  		return "", false, err
   104  	}
   105  
   106  	return migratedFilePath, true, nil
   107  }
   108  
   109  // imageDictToTpDepWrappeDef returns a *build.DefExpr that represents a wrapper function
   110  // containing all third_party_container_dep calls previously represented as dictionary
   111  // items
   112  func imageDictToTpDepWrapperDef(wrapperDefName string, dict *build.DictExpr) *build.DefStmt {
   113  	defExpr := &build.DefStmt{Name: wrapperDefName}
   114  
   115  	for _, baseKVExpr := range dict.List {
   116  		callExpr := &build.CallExpr{
   117  			X:              &build.Ident{Name: container.ThirdPartyContainerDepRuleName},
   118  			ForceMultiLine: true,
   119  		}
   120  
   121  		ruleName := strings.ReplaceAll(baseKVExpr.Key.(*build.StringExpr).Value, "-", "_")
   122  
   123  		assignExprs := dictValsToAssignExprs(ruleName, baseKVExpr.Value.(*build.DictExpr))
   124  		callExpr.List = append(callExpr.List, assignExprs...)
   125  
   126  		defExpr.Body = append(defExpr.Body, callExpr)
   127  	}
   128  	return defExpr
   129  }
   130  
   131  // Convert dictionary values to AssignExprs
   132  func dictValsToAssignExprs(name string, dictExpr *build.DictExpr) []build.Expr {
   133  	assignExprs := []build.Expr{
   134  		&build.AssignExpr{
   135  			LHS: &build.Ident{Name: "name"},
   136  			Op:  "=",
   137  			RHS: &build.StringExpr{Value: name},
   138  		},
   139  	}
   140  
   141  	for _, e := range dictExpr.List {
   142  		attr := strings.ToLower(e.Key.(*build.StringExpr).Value)
   143  		attr = attrRenames(attr)
   144  		assignExpr := &build.AssignExpr{
   145  			LHS: &build.Ident{Name: attr},
   146  			Op:  "=",
   147  			RHS: e.Value,
   148  		}
   149  
   150  		assignExprs = append(assignExprs, assignExpr)
   151  	}
   152  	return assignExprs
   153  }
   154  
   155  // Any attributes that might need to be renamed go here
   156  func attrRenames(attr string) string {
   157  	renameMap := map[string]string{
   158  		"repo": "repository",
   159  	}
   160  
   161  	if rename, found := renameMap[attr]; found {
   162  		return rename
   163  	}
   164  	return attr
   165  }
   166
View as plain text