Integrate file updating within Scaffold.Execute

Signed-off-by: Adrian Orive <[email protected]>

pkg/model/file/template.go

@@ -102,3 +102,12 @@ type RequiresValidation interface {
 	// Validate returns true if the template has valid values
 	Validate() error
 }
+
+// UpdatableTemplate is a file that can be inserted code fragments
+type UpdatableTemplate interface {
+	Template
+	// Markers returns the different markers where code fragments will be inserted
+	Markers() []Marker
+	// CodeFragments returns a map that binds markers to code fragments
+	CodeFragments() map[Marker][]string
+}

pkg/scaffold/api.go

@@ -162,19 +162,14 @@ func (s *apiScaffolder) scaffoldV2() error {
 			return fmt.Errorf("error scaffolding APIs: %v", err)
 		}
 
-		kustomizationFile := &crdv2.Kustomization{Resource: s.resource}
 		if err := machinery.NewScaffold().Execute(
 			s.newUniverse(),
-			kustomizationFile,
+			&crdv2.Kustomization{Resource: s.resource},
 			&crdv2.KustomizeConfig{},
 		); err != nil {
 			return fmt.Errorf("error scaffolding kustomization: %v", err)
 		}
 
-		if err := kustomizationFile.Update(); err != nil {
-			return fmt.Errorf("error updating kustomization.yaml: %v", err)
-		}
-
 	} else {
 		// disable generation of example reconcile body if not scaffolding resource
 		// because this could result in a fork-bomb of k8s resources where watching a
@@ -192,26 +187,23 @@ func (s *apiScaffolder) scaffoldV2() error {
 				fmt.Sprintf("%s_controller.go", strings.ToLower(s.resource.Kind))))
 		}
 
-		suiteTestFile := &controllerv2.SuiteTest{Resource: s.resource}
 		if err := machinery.NewScaffold(s.plugins...).Execute(
 			s.newUniverse(),
-			suiteTestFile,
+			&controllerv2.SuiteTest{Resource: s.resource},
 			&controllerv2.Controller{Resource: s.resource},
 		); err != nil {
 			return fmt.Errorf("error scaffolding controller: %v", err)
 		}
-
-		if err := suiteTestFile.Update(); err != nil {
-			return fmt.Errorf("error updating suite_test.go under controllers pkg: %v", err)
-		}
 	}
 
-	main := &templatesv2.Main{
-		WireResource:   s.doResource,
-		WireController: s.doController,
-	}
-	s.newUniverse().InjectInto(main)
-	if err := main.Update(); err != nil {
+	if err := machinery.NewScaffold(s.plugins...).Execute(
+		s.newUniverse(),
+		&templatesv2.Main{
+			Resource:       s.resource,
+			WireResource:   s.doResource,
+			WireController: s.doController,
+		},
+	); err != nil {
 		return fmt.Errorf("error updating main.go: %v", err)
 	}
 

pkg/scaffold/internal/machinery/scaffold.go

@@ -86,6 +86,14 @@ func (s *scaffold) Execute(universe *model.Universe, files ...file.Template) err
 		}
 	}
 
+	for _, f := range files {
+		if updatableFile, canBeUpdated := f.(file.UpdatableTemplate); canBeUpdated {
+			if err := s.updateFile(updatableFile); err != nil {
+				return err
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -182,3 +190,24 @@ func newTemplate(t file.Template) *template.Template {
 		"lower": strings.ToLower,
 	})
 }
+
+func (s *scaffold) updateFile(f file.UpdatableTemplate) error {
+	// TODO(adirio): avoid using GetInput when the path can be obtained directly
+	i, err := f.GetInput()
+	if err != nil {
+		return err
+	}
+
+	codeFragments := f.CodeFragments()
+	mapping := make(map[string][]string, len(codeFragments))
+	for _, marker := range f.Markers() {
+		if value, found := codeFragments[marker]; found {
+			mapping[marker.String()] = value
+		}
+	}
+	if len(mapping) == 0 {
+		return nil
+	}
+
+	return InsertStringsInFile(i.Path, mapping)
+}

pkg/scaffold/internal/templates/v2/controller/controller_suitetest.go

@@ -22,18 +22,9 @@ import (
 
 	"sigs.k8s.io/kubebuilder/pkg/model/file"
 	"sigs.k8s.io/kubebuilder/pkg/model/resource"
-	"sigs.k8s.io/kubebuilder/pkg/scaffold/internal/machinery"
 )
 
-const (
-	importScaffoldMarker    = "// +kubebuilder:scaffold:imports"
-	addSchemeScaffoldMarker = "// +kubebuilder:scaffold:scheme"
-
-	importMarker    = "imports"
-	addSchemeMarker = "scheme"
-)
-
-var _ file.Template = &SuiteTest{}
+var _ file.UpdatableTemplate = &SuiteTest{}
 
 // SuiteTest scaffolds the suite_test.go file to setup the controller test
 type SuiteTest struct {
@@ -65,6 +56,51 @@ func (f *SuiteTest) Validate() error {
 	return f.Resource.Validate()
 }
 
+const (
+	importMarker    = "imports"
+	addSchemeMarker = "scheme"
+)
+
+// Markers implements UpdatableTemplate
+func (f *SuiteTest) Markers() []file.Marker {
+	return []file.Marker{
+		file.NewMarkerFor(f.Path, importMarker),
+		file.NewMarkerFor(f.Path, addSchemeMarker),
+	}
+}
+
+const (
+	apiImportCodeFragment = `%s "%s"
+`
+	addschemeCodeFragment = `err = %s.AddToScheme(scheme.Scheme)
+Expect(err).NotTo(HaveOccurred())
+
+`
+)
+
+// CodeFragments implements UpdatableTemplate
+func (f *SuiteTest) CodeFragments() map[file.Marker][]string {
+	fragments := make(map[file.Marker][]string, 2)
+
+	// Generate import code fragments
+	imports := make([]string, 0)
+	imports = append(imports, fmt.Sprintf(apiImportCodeFragment, f.Resource.ImportAlias, f.Resource.Package))
+
+	// Generate add scheme code fragments
+	addScheme := make([]string, 0)
+	addScheme = append(addScheme, fmt.Sprintf(addschemeCodeFragment, f.Resource.ImportAlias))
+
+	// Only store code fragments in the map if the slices are non-empty
+	if len(imports) != 0 {
+		fragments[file.NewMarkerFor(f.Path, importMarker)] = imports
+	}
+	if len(addScheme) != 0 {
+		fragments[file.NewMarkerFor(f.Path, addSchemeMarker)] = addScheme
+	}
+
+	return fragments
+}
+
 const controllerSuiteTestTemplate = `{{ .Boilerplate }}
 
 package controllers
@@ -126,31 +162,3 @@ var _ = AfterSuite(func() {
 	Expect(err).ToNot(HaveOccurred())
 })
 `
-
-const (
-	apiImportCodeFragment = `%s "%s"
-`
-	addschemeCodeFragment = `err = %s.AddToScheme(scheme.Scheme)
-Expect(err).NotTo(HaveOccurred())
-
-`
-)
-
-// Update updates given file (suite_test.go) with code fragments required for
-// adding import paths and code setup for new types.
-func (f *SuiteTest) Update() error {
-	apiImportCodeFragment := fmt.Sprintf(apiImportCodeFragment, f.Resource.ImportAlias, f.Resource.Package)
-
-	addschemeCodeFragment := fmt.Sprintf(addschemeCodeFragment, f.Resource.ImportAlias)
-
-	err := machinery.InsertStringsInFile(f.Path,
-		map[string][]string{
-			file.NewMarkerFor(f.Path, importMarker).String():    {apiImportCodeFragment},
-			file.NewMarkerFor(f.Path, addSchemeMarker).String(): {addschemeCodeFragment},
-		})
-	if err != nil {
-		return err
-	}
-
-	return nil
-}

pkg/scaffold/internal/templates/v2/crd/kustomization.go

@@ -22,16 +22,9 @@ import (
 
 	"sigs.k8s.io/kubebuilder/pkg/model/file"
 	"sigs.k8s.io/kubebuilder/pkg/model/resource"
-	"sigs.k8s.io/kubebuilder/pkg/scaffold/internal/machinery"
 )
 
-const (
-	resourceMarker         = "crdkustomizeresource"
-	webhookPatchMarker     = "crdkustomizewebhookpatch"
-	caInjectionPatchMarker = "crdkustomizecainjectionpatch"
-)
-
-var _ file.Template = &Kustomization{}
+var _ file.UpdatableTemplate = &Kustomization{}
 
 // Kustomization scaffolds the kustomization file in manager folder.
 type Kustomization struct {
@@ -54,6 +47,21 @@ func (f *Kustomization) GetInput() (file.Input, error) {
 	return f.Input, nil
 }
 
+const (
+	resourceMarker         = "crdkustomizeresource"
+	webhookPatchMarker     = "crdkustomizewebhookpatch"
+	caInjectionPatchMarker = "crdkustomizecainjectionpatch"
+)
+
+// Markers implements UpdatableTemplate
+func (f *Kustomization) Markers() []file.Marker {
+	return []file.Marker{
+		file.NewMarkerFor(f.Path, resourceMarker),
+		file.NewMarkerFor(f.Path, webhookPatchMarker),
+		file.NewMarkerFor(f.Path, caInjectionPatchMarker),
+	}
+}
+
 const (
 	resourceCodeFragment = `- bases/%s_%s.yaml
 `
@@ -63,24 +71,34 @@ const (
 `
 )
 
-func (f *Kustomization) Update() error {
-	if f.Path == "" {
-		f.Path = filepath.Join("config", "crd", "kustomization.yaml")
-	}
+// CodeFragments implements UpdatableTemplate
+func (f *Kustomization) CodeFragments() map[file.Marker][]string {
+	fragments := make(map[file.Marker][]string, 3)
+
+	// Generate resource code fragments
+	res := make([]string, 0)
+	res = append(res, fmt.Sprintf(resourceCodeFragment, f.Resource.Domain, f.Resource.Plural))
+
+	// Generate resource code fragments
+	webhookPatch := make([]string, 0)
+	webhookPatch = append(webhookPatch, fmt.Sprintf(webhookPatchCodeFragment, f.Resource.Plural))
 
-	// TODO(directxman12): not technically valid if something changes from the default
-	// (we'd need to parse the markers)
+	// Generate resource code fragments
+	caInjectionPatch := make([]string, 0)
+	caInjectionPatch = append(caInjectionPatch, fmt.Sprintf(caInjectionPatchCodeFragment, f.Resource.Plural))
 
-	resourceCodeFragment := fmt.Sprintf(resourceCodeFragment, f.Resource.Domain, f.Resource.Plural)
-	webhookPatchCodeFragment := fmt.Sprintf(webhookPatchCodeFragment, f.Resource.Plural)
-	caInjectionPatchCodeFragment := fmt.Sprintf(caInjectionPatchCodeFragment, f.Resource.Plural)
+	// Only store code fragments in the map if the slices are non-empty
+	if len(res) != 0 {
+		fragments[file.NewMarkerFor(f.Path, resourceMarker)] = res
+	}
+	if len(webhookPatch) != 0 {
+		fragments[file.NewMarkerFor(f.Path, webhookPatchMarker)] = webhookPatch
+	}
+	if len(caInjectionPatch) != 0 {
+		fragments[file.NewMarkerFor(f.Path, caInjectionPatchMarker)] = caInjectionPatch
+	}
 
-	return machinery.InsertStringsInFile(f.Path,
-		map[string][]string{
-			file.NewMarkerFor(f.Path, resourceMarker).String():         {resourceCodeFragment},
-			file.NewMarkerFor(f.Path, webhookPatchMarker).String():     {webhookPatchCodeFragment},
-			file.NewMarkerFor(f.Path, caInjectionPatchMarker).String(): {caInjectionPatchCodeFragment},
-		})
+	return fragments
 }
 
 var kustomizationTemplate = `# This kustomization.yaml is not intended to be run by itself,