update

Author: leon-inf

pkg/controller/kubebuilderx/plan_builder.go

@@ -165,18 +165,29 @@ func buildOrderedVertices(transCtx *transformContext, currentTree *ObjectTree, d
 			workloadVertices = append(workloadVertices, vertex)
 		}
 	}
+	graphOptions := func(options ObjectOptions) []model.GraphOption {
+		opts := []model.GraphOption{
+			inDataContext4G(),
+		}
+		for i := range options.Hooks {
+			opts = append(opts, model.WithHook(options.Hooks[i]))
+		}
+		return opts
+	}
 	createNewObjects := func() {
 		for name := range createSet {
-			if desiredTree.childrenOptions[name].SkipToReconcile {
+			options := desiredTree.childrenOptions[name]
+			if options.SkipToReconcile {
 				continue
 			}
-			v := model.NewObjectVertex(nil, assign(ctx, newSnapshot[name]), model.ActionCreatePtr(), inDataContext4G())
+			v := model.NewObjectVertex(nil, assign(ctx, newSnapshot[name]), model.ActionCreatePtr(), graphOptions(options)...)
 			findAndAppend(v)
 		}
 	}
 	updateObjects := func() {
 		for name := range updateSet {
-			if desiredTree.childrenOptions[name].SkipToReconcile {
+			options := desiredTree.childrenOptions[name]
+			if options.SkipToReconcile {
 				continue
 			}
 			oldObj := oldSnapshot[name]
@@ -190,21 +201,23 @@ func buildOrderedVertices(transCtx *transformContext, currentTree *ObjectTree, d
 				var v *model.ObjectVertex
 				subResource := desiredTree.childrenOptions[*name].SubResource
 				if subResource != "" {
-					v = model.NewObjectVertex(oldObj, newObj, model.ActionUpdatePtr(), inDataContext4G(), model.WithSubResource(subResource))
+					opts := append(graphOptions(options), model.WithSubResource(subResource))
+					v = model.NewObjectVertex(oldObj, newObj, model.ActionUpdatePtr(), opts...)
 				} else {
-					v = model.NewObjectVertex(oldObj, newObj, model.ActionUpdatePtr(), inDataContext4G())
+					v = model.NewObjectVertex(oldObj, newObj, model.ActionUpdatePtr(), graphOptions(options)...)
 				}
 				findAndAppend(v)
 			}
 		}
 	}
-	deleteOrphanObjects := func() {
+	deleteObjects := func() {
 		for name := range deleteSet {
-			if desiredTree.childrenOptions[name].SkipToReconcile {
+			options := desiredTree.childrenOptions[name]
+			if options.SkipToReconcile {
 				continue
 			}
 			object := oldSnapshot[name]
-			v := model.NewObjectVertex(nil, object, model.ActionDeletePtr(), inDataContext4G())
+			v := model.NewObjectVertex(nil, object, model.ActionDeletePtr(), graphOptions(options)...)
 			findAndAppend(v)
 		}
 	}
@@ -218,7 +231,7 @@ func buildOrderedVertices(transCtx *transformContext, currentTree *ObjectTree, d
 	// objects to be updated
 	updateObjects()
 	// objects to be deleted
-	deleteOrphanObjects()
+	deleteObjects()
 	// handle object dependencies
 	handleDependencies()
 	return vertices
@@ -246,6 +259,9 @@ func (b *PlanBuilder) defaultWalkFunc(v graph.Vertex) error {
 	if vertex.Action == nil {
 		return errors.New("vertex action can't be nil")
 	}
+	if err := b.callHooks(vertex); err != nil {
+		return fmt.Errorf("vertex call hooks failed: %v, obj: %s, action: %v", err, vertex.Obj.GetName(), *vertex.Action)
+	}
 	ctx := b.transCtx.ctx
 	switch *vertex.Action {
 	case model.CREATE:
@@ -262,6 +278,15 @@ func (b *PlanBuilder) defaultWalkFunc(v graph.Vertex) error {
 	return nil
 }
 
+func (b *PlanBuilder) callHooks(vertex *model.ObjectVertex) error {
+	for i := range vertex.Hooks {
+		if err := vertex.Hooks[i](vertex.Obj); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (b *PlanBuilder) createObject(ctx context.Context, vertex *model.ObjectVertex) error {
 	err := b.cli.Create(ctx, vertex.Obj, clientOption(vertex))
 	if err != nil && !apierrors.IsAlreadyExists(err) {

pkg/controller/kubebuilderx/reconciler.go

@@ -43,6 +43,9 @@ type ObjectOptions struct {
 
 	// if true, the object should not be reconciled
 	SkipToReconcile bool
+
+	// hooks are called before the object is manipulated
+	Hooks []func(client.Object) error
 }
 
 type WithSubResource string
@@ -57,6 +60,12 @@ func (o SkipToReconcile) ApplyToObject(opts *ObjectOptions) {
 	opts.SkipToReconcile = bool(o)
 }
 
+type WithHook func(client.Object) error
+
+func (o WithHook) ApplyToObject(opts *ObjectOptions) {
+	opts.Hooks = append(opts.Hooks, o)
+}
+
 type ObjectTree struct {
 	// TODO(free6om): should find a better place to hold these two params?
 	context.Context
@@ -252,8 +261,20 @@ func (t *ObjectTree) Delete(objects ...client.Object) error {
 	return nil
 }
 
-func (t *ObjectTree) DeleteSecondaryObjects() {
-	t.children = make(model.ObjectSnapshot)
+func (t *ObjectTree) DeleteWithOption(object client.Object, options ...ObjectOption) error {
+	name, err := model.GetGVKName(object)
+	if err != nil {
+		return err
+	}
+	delete(t.children, *name)
+	if len(options) > 0 {
+		option := ObjectOptions{}
+		for _, opt := range options {
+			opt.ApplyToObject(&option)
+		}
+		t.childrenOptions[*name] = option
+	}
+	return nil
 }
 
 func (t *ObjectTree) SetFinalizer(finalizer string) {

pkg/controller/kubebuilderx/reconciler_test.go

@@ -67,8 +67,6 @@ var _ = Describe("reconciler test", func() {
 			Expect(tree.Update(obj0Update)).Should(Succeed())
 			Expect(tree.List(&corev1.Pod{})[0]).Should(Equal(obj0Update))
 			Expect(tree.GetSecondaryObjects()).Should(HaveLen(1))
-			tree.DeleteSecondaryObjects()
-			Expect(tree.GetSecondaryObjects()).Should(HaveLen(0))
 
 			By("DeepCopy")
 			tree.SetRoot(root)

pkg/controller/model/graph_options.go

@@ -19,14 +19,17 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 package model
 
-import "sigs.k8s.io/controller-runtime/pkg/client"
+import (
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
 
 type GraphOptions struct {
 	replaceIfExisting     bool
 	haveDifferentTypeWith bool
 	clientOpt             any
 	propagationPolicy     client.PropagationPolicy
 	subResource           string
+	hooks                 []func(client.Object) error
 }
 
 type GraphOption interface {
@@ -99,3 +102,19 @@ func WithSubResource(subResource string) GraphOption {
 		subResource: subResource,
 	}
 }
+
+type hookOption struct {
+	hook func(object client.Object) error
+}
+
+var _ GraphOption = &hookOption{}
+
+func (o *hookOption) ApplyTo(opts *GraphOptions) {
+	opts.hooks = append(opts.hooks, o.hook)
+}
+
+func WithHook(hook func(client.Object) error) GraphOption {
+	return &hookOption{
+		hook: hook,
+	}
+}

pkg/controller/model/transform_types.go

@@ -72,6 +72,7 @@ type ObjectVertex struct {
 	SubResource       string
 	ClientOpt         any
 	PropagationPolicy client.PropagationPolicy
+	Hooks             []func(client.Object) error
 }
 
 func (v *ObjectVertex) String() string {
@@ -92,6 +93,7 @@ func NewObjectVertex(oldObj, newObj client.Object, action *Action, opts ...Graph
 		Action:      action,
 		SubResource: graphOpts.subResource,
 		ClientOpt:   graphOpts.clientOpt,
+		Hooks:       graphOpts.hooks,
 	}
 }