driver/kubernetes: generalize manifest patching

Signed-off-by: abhay1999 <abhaychaurasiya19@gmail.com>

Author: abhay1999

driver/kubernetes/factory.go

@@ -17,8 +17,6 @@ import (
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	corev1 "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/rest"
 )
 
@@ -239,6 +237,8 @@ func (f *factory) processDriverOpts(deploymentName string, namespace string, cfg
 			if err != nil {
 				return nil, "", "", false, 0, errors.Wrap(err, "cannot parse labels")
 			}
+		case k == "manifest-patch":
+			deploymentOpt.ManifestPatch = v
 		case k == "tolerations":
 			ts := strings.Split(v, ";")
 			deploymentOpt.Tolerations = []corev1.Toleration{}
@@ -274,46 +274,6 @@ func (f *factory) processDriverOpts(deploymentName string, namespace string, cfg
 
 				deploymentOpt.Tolerations = append(deploymentOpt.Tolerations, t)
 			}
-		case k == "ownerrefs":
-			refs := strings.Split(v, ";")
-			deploymentOpt.OwnerReferences = []metav1.OwnerReference{}
-			for i := range refs {
-				kvs := strings.Split(refs[i], ",")
-
-				ref := metav1.OwnerReference{}
-
-				for j := range kvs {
-					kv := strings.SplitN(kvs[j], "=", 2)
-					if len(kv) == 2 {
-						switch kv[0] {
-						case "apiVersion":
-							ref.APIVersion = kv[1]
-						case "kind":
-							ref.Kind = kv[1]
-						case "name":
-							ref.Name = kv[1]
-						case "uid":
-							ref.UID = types.UID(kv[1])
-						case "controller":
-							b, err := strconv.ParseBool(kv[1])
-							if err != nil {
-								return nil, "", "", false, 0, errors.Wrap(err, "invalid ownerrefs controller value")
-							}
-							ref.Controller = &b
-						case "blockOwnerDeletion":
-							b, err := strconv.ParseBool(kv[1])
-							if err != nil {
-								return nil, "", "", false, 0, errors.Wrap(err, "invalid ownerrefs blockOwnerDeletion value")
-							}
-							ref.BlockOwnerDeletion = &b
-						default:
-							return nil, "", "", false, 0, errors.Errorf("invalid ownerrefs key %q", kv[0])
-						}
-					}
-				}
-
-				deploymentOpt.OwnerReferences = append(deploymentOpt.OwnerReferences, ref)
-			}
 		case k == "loadbalance":
 			switch v {
 			case LoadbalanceSticky, LoadbalanceRandom:

driver/kubernetes/factory_test.go

@@ -8,8 +8,6 @@ import (
 	"github.com/docker/buildx/driver/bkimage"
 	"github.com/stretchr/testify/require"
 	v1 "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/rest"
 )
 
@@ -271,57 +269,25 @@ func TestFactory_processDriverOpts(t *testing.T) {
 	)
 
 	t.Run(
-		"OwnerRefs", func(t *testing.T) {
-			controller := true
-			blockOwnerDeletion := false
+		"ManifestPatch", func(t *testing.T) {
+			patch := `.metadata.ownerReferences = [{apiVersion: "actions.github.com/v1alpha1", kind: "EphemeralRunner", name: "runner-xyz", uid: "b636330d-26b7-417a-8464-c2641438feed", controller: true, blockOwnerDeletion: false}]`
 			cfg.DriverOpts = map[string]string{
-				"ownerrefs": "apiVersion=actions.github.com/v1alpha1,kind=EphemeralRunner,name=runner-xyz,uid=b636330d-26b7-417a-8464-c2641438feed,controller=true,blockOwnerDeletion=false",
+				"manifest-patch": patch,
 			}
 			r, _, _, _, _, err := f.processDriverOpts(cfg.Name, "test", cfg)
 			require.NoError(t, err)
-			require.Equal(t, []metav1.OwnerReference{
-				{
-					APIVersion:         "actions.github.com/v1alpha1",
-					Kind:               "EphemeralRunner",
-					Name:               "runner-xyz",
-					UID:                types.UID("b636330d-26b7-417a-8464-c2641438feed"),
-					Controller:         &controller,
-					BlockOwnerDeletion: &blockOwnerDeletion,
-				},
-			}, r.OwnerReferences)
+			require.Equal(t, patch, r.ManifestPatch)
 		},
 	)
 
 	t.Run(
-		"MultipleOwnerRefs", func(t *testing.T) {
+		"EmptyManifestPatch", func(t *testing.T) {
 			cfg.DriverOpts = map[string]string{
-				"ownerrefs": "apiVersion=v1,kind=Pod,name=pod1,uid=uid-1;apiVersion=v1,kind=Pod,name=pod2,uid=uid-2",
+				"manifest-patch": "",
 			}
 			r, _, _, _, _, err := f.processDriverOpts(cfg.Name, "test", cfg)
 			require.NoError(t, err)
-			require.Len(t, r.OwnerReferences, 2)
-			require.Equal(t, types.UID("uid-1"), r.OwnerReferences[0].UID)
-			require.Equal(t, types.UID("uid-2"), r.OwnerReferences[1].UID)
-		},
-	)
-
-	t.Run(
-		"InvalidOwnerRefKey", func(t *testing.T) {
-			cfg.DriverOpts = map[string]string{
-				"ownerrefs": "apiVersion=v1,invalid=foo",
-			}
-			_, _, _, _, _, err := f.processDriverOpts(cfg.Name, "test", cfg)
-			require.Error(t, err)
-		},
-	)
-
-	t.Run(
-		"InvalidOwnerRefController", func(t *testing.T) {
-			cfg.DriverOpts = map[string]string{
-				"ownerrefs": "apiVersion=v1,controller=notabool",
-			}
-			_, _, _, _, _, err := f.processDriverOpts(cfg.Name, "test", cfg)
-			require.Error(t, err)
+			require.Empty(t, r.ManifestPatch)
 		},
 	)
 }

driver/kubernetes/manifest/manifest.go

@@ -1,12 +1,15 @@
 package manifest
 
 import (
+	"encoding/json"
 	"fmt"
 	"path"
 	"strings"
 
 	"github.com/docker/buildx/util/platformutil"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/pkg/errors"
+	yaml "go.yaml.in/yaml/v3"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -38,7 +41,7 @@ type DeploymentOpt struct {
 	CustomAnnotations         map[string]string
 	CustomLabels              map[string]string
 	Tolerations               []corev1.Toleration
-	OwnerReferences           []metav1.OwnerReference
+	ManifestPatch             string
 	RequestsCPU               string
 	RequestsMemory            string
 	RequestsEphemeralStorage  string
@@ -155,11 +158,10 @@ func NewDeployment(opt *DeploymentOpt) (d *appsv1.Deployment, s *appsv1.Stateful
 	}
 
 	meta := metav1.ObjectMeta{
-		Namespace:       opt.Namespace,
-		Name:            opt.Name,
-		Labels:          labels,
-		Annotations:     annotations,
-		OwnerReferences: opt.OwnerReferences,
+		Namespace:   opt.Namespace,
+		Name:        opt.Name,
+		Labels:      labels,
+		Annotations: annotations,
 	}
 
 	for _, cfg := range splitConfigFiles(opt.ConfigFiles) {
@@ -331,9 +333,109 @@ func NewDeployment(opt *DeploymentOpt) (d *appsv1.Deployment, s *appsv1.Stateful
 		}
 	}
 
+	if opt.ManifestPatch != "" {
+		if d != nil {
+			if err := applyManifestPatch(d, opt.ManifestPatch); err != nil {
+				return nil, nil, nil, err
+			}
+		}
+		if s != nil {
+			if err := applyManifestPatch(s, opt.ManifestPatch); err != nil {
+				return nil, nil, nil, err
+			}
+		}
+		for _, cfgMap := range c {
+			if err := applyManifestPatch(cfgMap, opt.ManifestPatch); err != nil {
+				return nil, nil, nil, err
+			}
+		}
+	}
+
 	return
 }
 
+func applyManifestPatch(obj any, patch string) error {
+	path, value, err := parseManifestPatch(patch)
+	if err != nil {
+		return err
+	}
+
+	dt, err := json.Marshal(obj)
+	if err != nil {
+		return errors.Wrap(err, "marshal manifest")
+	}
+
+	var doc map[string]any
+	if err := json.Unmarshal(dt, &doc); err != nil {
+		return errors.Wrap(err, "unmarshal manifest")
+	}
+
+	if err := setManifestPatchValue(doc, path, value); err != nil {
+		return err
+	}
+
+	dt, err = json.Marshal(doc)
+	if err != nil {
+		return errors.Wrap(err, "marshal patched manifest")
+	}
+	if err := json.Unmarshal(dt, obj); err != nil {
+		return errors.Wrap(err, "unmarshal patched manifest")
+	}
+	return nil
+}
+
+func parseManifestPatch(patch string) ([]string, any, error) {
+	lhs, rhs, ok := strings.Cut(patch, "=")
+	if !ok {
+		return nil, nil, errors.Errorf("invalid manifest-patch %q: expected assignment expression", patch)
+	}
+
+	lhs = strings.TrimSpace(lhs)
+	rhs = strings.TrimSpace(rhs)
+	if lhs == "" || rhs == "" {
+		return nil, nil, errors.Errorf("invalid manifest-patch %q: expected non-empty assignment", patch)
+	}
+	if !strings.HasPrefix(lhs, ".") {
+		return nil, nil, errors.Errorf("invalid manifest-patch %q: path must start with '.'", patch)
+	}
+
+	path := strings.Split(strings.TrimPrefix(lhs, "."), ".")
+	for _, segment := range path {
+		if segment == "" {
+			return nil, nil, errors.Errorf("invalid manifest-patch %q: empty path segment", patch)
+		}
+	}
+
+	var value any
+	if err := yaml.Unmarshal([]byte(rhs), &value); err != nil {
+		return nil, nil, errors.Wrapf(err, "invalid manifest-patch %q", patch)
+	}
+	return path, value, nil
+}
+
+func setManifestPatchValue(doc map[string]any, path []string, value any) error {
+	current := doc
+	for i, segment := range path {
+		if i == len(path)-1 {
+			current[segment] = value
+			return nil
+		}
+		next, ok := current[segment]
+		if !ok {
+			child := map[string]any{}
+			current[segment] = child
+			current = child
+			continue
+		}
+		child, ok := next.(map[string]any)
+		if !ok {
+			return errors.Errorf("invalid manifest-patch path %q: %q is not an object", strings.Join(path, "."), strings.Join(path[:i+1], "."))
+		}
+		current = child
+	}
+	return nil
+}
+
 func (opt *DeploymentOpt) IsPersistentStorage() bool {
 	return opt.RequestsPersistentStorage != ""
 }

driver/kubernetes/manifest/manifest_test.go

@@ -0,0 +1,47 @@
+package manifest
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+)
+
+func TestNewDeploymentManifestPatch(t *testing.T) {
+	controller := true
+	blockOwnerDeletion := false
+
+	d, s, c, err := NewDeployment(&DeploymentOpt{
+		Namespace:     "test-ns",
+		Name:          "test-builder",
+		Image:         "buildkit:test",
+		Replicas:      1,
+		ManifestPatch: `.metadata.ownerReferences = [{apiVersion: "actions.github.com/v1alpha1", kind: "EphemeralRunner", name: "runner-xyz", uid: "b636330d-26b7-417a-8464-c2641438feed", controller: true, blockOwnerDeletion: false}]`,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, d)
+	require.Nil(t, s)
+	require.Empty(t, c)
+	require.Equal(t, []metav1.OwnerReference{
+		{
+			APIVersion:         "actions.github.com/v1alpha1",
+			Kind:               "EphemeralRunner",
+			Name:               "runner-xyz",
+			UID:                types.UID("b636330d-26b7-417a-8464-c2641438feed"),
+			Controller:         &controller,
+			BlockOwnerDeletion: &blockOwnerDeletion,
+		},
+	}, d.OwnerReferences)
+}
+
+func TestNewDeploymentManifestPatchInvalid(t *testing.T) {
+	_, _, _, err := NewDeployment(&DeploymentOpt{
+		Namespace:     "test-ns",
+		Name:          "test-builder",
+		Image:         "buildkit:test",
+		Replicas:      1,
+		ManifestPatch: `.metadata.ownerReferences = [`,
+	})
+	require.Error(t, err)
+}