Create and apply patch when adding labels to API objects

pkg/skaffold/runner/label/label.go

@@ -17,6 +17,7 @@ limitations under the License.
 package label
 
 import (
+	"encoding/json"
 	"time"
 
 	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
@@ -32,8 +33,174 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	extensionsv1beta1 "k8s.io/api/extensions/v1beta1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	patch "k8s.io/apimachinery/pkg/util/strategicpatch"
 )
 
+type objectType int
+
+// List of API Objects supported by the Skaffold Labeler
+const (
+	_ = iota
+	corev1Pod
+	appsv1Deployment
+	appsv1Beta1Deployment
+	appsv1Beta2Deployment
+	extensionsv1Beta1Deployment
+	corev1Service
+	appv1StatefulSet
+	appsv1Beta1StatefulSet
+	appsv1Beta2StatefulSet
+	extensionsv1Beta1DaemonSet
+	appsv1ReplicaSet
+	appsv1Beta2ReplicaSet
+)
+
+// patcher is responsible for applying a given patch to the provided object
+type patcher func(clientgo.Interface, string, string, []byte) error
+
+// objectMeta is responsible for returning a generic runtime.Object's metadata
+type objectMeta func(runtime.Object) (*metav1.ObjectMeta, bool)
+
+var patchers = map[objectType]patcher{
+	corev1Pod: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.CoreV1().Pods(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1Beta1Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1beta1().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1Beta2Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1beta2().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	extensionsv1Beta1Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.ExtensionsV1beta1().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	corev1Service: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.CoreV1().Services(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appv1StatefulSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1().StatefulSets(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1Beta1StatefulSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1beta1().StatefulSets(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1Beta2StatefulSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1beta2().StatefulSets(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	extensionsv1Beta1DaemonSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.ExtensionsV1beta1().DaemonSets(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1ReplicaSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1().ReplicaSets(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+	appsv1Beta2ReplicaSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
+		_, err := client.AppsV1beta2().ReplicaSets(ns).Patch(name, types.StrategicMergePatchType, p)
+		return err
+	},
+}
+
+var objectMetas = map[objectType]objectMeta{
+	corev1Pod: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*corev1.Pod)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1Deployment: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1.Deployment)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1Beta1Deployment: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1beta1.Deployment)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1Beta2Deployment: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1beta2.Deployment)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	extensionsv1Beta1Deployment: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*extensionsv1beta1.Deployment)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	corev1Service: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*corev1.Service)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appv1StatefulSet: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1.StatefulSet)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1Beta1StatefulSet: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1beta1.StatefulSet)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1Beta2StatefulSet: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1beta1.StatefulSet)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	extensionsv1Beta1DaemonSet: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*extensionsv1beta1.DaemonSet)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1ReplicaSet: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1.ReplicaSet)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+	appsv1Beta2ReplicaSet: func(r runtime.Object) (*metav1.ObjectMeta, bool) {
+		obj, ok := r.(*appsv1beta2.ReplicaSet)
+		if !ok {
+			return nil, ok
+		}
+		return &obj.ObjectMeta, ok
+	},
+}
+
 // retry 3 times to give the object time to propagate to the API server
 const tries int = 3
 const sleeptime time.Duration = 300 * time.Millisecond
@@ -85,71 +252,21 @@ func updateRuntimeObject(client clientgo.Interface, labels map[string]string, re
 		labels[k] = v
 	}
 	var err error
-	obj := *res.Obj
-	switch obj.(type) {
-	case *corev1.Pod:
-		p := obj.(*corev1.Pod)
-		addSkaffoldLabels(labels, &p.ObjectMeta)
-		pods := client.CoreV1().Pods(retrieveNamespace(res.Namespace, p.ObjectMeta))
-		_, err = pods.UpdateStatus(p)
-	case *appsv1.Deployment:
-		d := obj.(*appsv1.Deployment)
-		addSkaffoldLabels(labels, &d.ObjectMeta)
-		deployments := client.AppsV1().Deployments(retrieveNamespace(res.Namespace, d.ObjectMeta))
-		_, err = deployments.Update(d)
-	case *appsv1beta1.Deployment:
-		d := obj.(*appsv1beta1.Deployment)
-		addSkaffoldLabels(labels, &d.ObjectMeta)
-		deployments := client.AppsV1beta1().Deployments(retrieveNamespace(res.Namespace, d.ObjectMeta))
-		_, err = deployments.Update(d)
-	case *appsv1beta2.Deployment:
-		d := obj.(*appsv1beta2.Deployment)
-		addSkaffoldLabels(labels, &d.ObjectMeta)
-		deployments := client.AppsV1beta2().Deployments(retrieveNamespace(res.Namespace, d.ObjectMeta))
-		_, err = deployments.Update(d)
-	case *extensionsv1beta1.Deployment:
-		d := obj.(*extensionsv1beta1.Deployment)
-		addSkaffoldLabels(labels, &d.ObjectMeta)
-		deployments := client.ExtensionsV1beta1().Deployments(retrieveNamespace(res.Namespace, d.ObjectMeta))
-		_, err = deployments.Update(d)
-	case *corev1.Service:
-		s := obj.(*corev1.Service)
-		addSkaffoldLabels(labels, &s.ObjectMeta)
-		services := client.CoreV1().Services(retrieveNamespace(res.Namespace, s.ObjectMeta))
-		_, err = services.UpdateStatus(s)
-	case *appsv1.StatefulSet:
-		s := obj.(*appsv1.StatefulSet)
-		addSkaffoldLabels(labels, &s.ObjectMeta)
-		sets := client.AppsV1().StatefulSets(retrieveNamespace(res.Namespace, s.ObjectMeta))
-		_, err = sets.UpdateStatus(s)
-	case *appsv1beta1.StatefulSet:
-		s := obj.(*appsv1beta1.StatefulSet)
-		addSkaffoldLabels(labels, &s.ObjectMeta)
-		sets := client.AppsV1beta1().StatefulSets(retrieveNamespace(res.Namespace, s.ObjectMeta))
-		_, err = sets.UpdateStatus(s)
-	case *appsv1beta2.StatefulSet:
-		s := obj.(*appsv1beta2.StatefulSet)
-		addSkaffoldLabels(labels, &s.ObjectMeta)
-		sets := client.AppsV1beta2().StatefulSets(retrieveNamespace(res.Namespace, s.ObjectMeta))
-		_, err = sets.UpdateStatus(s)
-	case *extensionsv1beta1.DaemonSet:
-		s := obj.(*extensionsv1beta1.DaemonSet)
-		addSkaffoldLabels(labels, &s.ObjectMeta)
-		sets := client.ExtensionsV1beta1().DaemonSets(retrieveNamespace(res.Namespace, s.ObjectMeta))
-		_, err = sets.UpdateStatus(s)
-	case *appsv1.ReplicaSet:
-		rs := obj.(*appsv1.ReplicaSet)
-		addSkaffoldLabels(labels, &rs.ObjectMeta)
-		sets := client.AppsV1().ReplicaSets(retrieveNamespace(res.Namespace, rs.ObjectMeta))
-		_, err = sets.UpdateStatus(rs)
-	case *appsv1beta2.ReplicaSet:
-		rs := obj.(*appsv1beta2.ReplicaSet)
-		addSkaffoldLabels(labels, &rs.ObjectMeta)
-		sets := client.AppsV1beta2().ReplicaSets(retrieveNamespace(res.Namespace, rs.ObjectMeta))
-		_, err = sets.UpdateStatus(rs)
-	default:
+	applied := false
+	var metadata *metav1.ObjectMeta
+	originalJSON, _ := json.Marshal(*res.Obj)
+	modifiedObj := (*res.Obj).DeepCopyObject()
+	for typeStr, m := range objectMetas {
+		if metadata, applied = m(modifiedObj); applied {
+			addSkaffoldLabels(labels, metadata)
+			modifiedJSON, _ := json.Marshal(modifiedObj)
+			p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, modifiedObj)
+			err = patchers[typeStr](client, retrieveNamespace(res.Namespace, *metadata), metadata.GetName(), p)
+			break // we should only ever apply one patch, so stop here
+		}
+	}
+	if !applied {
 		logrus.Infof("unknown runtime.Object, skipping label")
-		return nil
 	}
 	return err
 }