Merge pull request #32 from twelho/comment-source-access

High-level getter/setter for the YAML tree used for transferring comments

Makefile

@@ -34,7 +34,7 @@ docker-%:
 
 test: docker-test-internal
 test-internal:
-	go test $(addsuffix /...,$(addprefix ./,${SRC_PKGS}))
+	go test -v $(addsuffix /...,$(addprefix ./,${SRC_PKGS}))
 
 tidy: docker-tidy-internal
 tidy-internal: /go/bin/goimports

pkg/serializer/comments.go

@@ -3,6 +3,7 @@ package serializer
 import (
 	"bytes"
 	"encoding/base64"
+	"errors"
 	"fmt"
 
 	"github.com/sirupsen/logrus"
@@ -14,6 +15,13 @@ import (
 
 const preserveCommentsAnnotation = "serializer.libgitops.weave.works/original-data"
 
+var (
+	// TODO: Investigate if we can just depend on `metav1.Object` interface compliance instead of needing to explicitly
+	//  embed the `metav1.ObjectMeta` struct.
+	ErrNoObjectMeta     = errors.New("the given object cannot store comments, it is not metav1.ObjectMeta compliant")
+	ErrNoStoredComments = errors.New("the given object does not have stored comments")
+)
+
 // tryToPreserveComments tries to save the original file data (base64-encoded) into an annotation.
 // This original file data can be used at encoding-time to preserve comments
 func (d *decoder) tryToPreserveComments(doc []byte, obj runtime.Object, ct ContentType) {
@@ -23,66 +31,47 @@ func (d *decoder) tryToPreserveComments(doc []byte, obj runtime.Object, ct Conte
 		return
 	}
 
-	// Convert the object to a metav1.Object (this requires embedding ObjectMeta)
-	metaobj, ok := toMetaObject(obj)
-	if !ok {
-		// If the object doesn't have ObjectMeta embedded, just do nothing
+	// Preserve the original file content in the annotation (this requires embedding ObjectMeta).
+	if !setCommentSourceBytes(obj, doc) {
+		// If the object doesn't have ObjectMeta embedded, just do nothing.
 		logrus.Debugf("Couldn't convert object with GVK %q to metav1.Object, although opts.PreserveComments is enabled", obj.GetObjectKind().GroupVersionKind())
-		return
 	}
-
-	// Preserve the original file content in the annotation
-	setAnnotation(metaobj, preserveCommentsAnnotation, base64.StdEncoding.EncodeToString(doc))
 }
 
 // tryToPreserveComments tries to locate the possibly-saved original file data in the object's annotation
-func (e *encoder) encodeWithCommentSupport(versionEncoder runtime.Encoder, fw FrameWriter, obj runtime.Object, metaobj metav1.Object) error {
+func (e *encoder) encodeWithCommentSupport(versionEncoder runtime.Encoder, fw FrameWriter, obj runtime.Object, metaObj metav1.Object) error {
 	// If the user did not opt into preserving comments, just sanitize ObjectMeta temporarily and and return
 	if !*e.opts.PreserveComments {
 		// Normal encoding without the annotation (so it doesn't leak by accident)
-		return noAnnotationWrapper(metaobj, e.normalEncodeFunc(versionEncoder, fw, obj))
+		return noAnnotationWrapper(metaObj, e.normalEncodeFunc(versionEncoder, fw, obj))
 	}
 
 	// The user requested to preserve comments, but content type is not YAML, so log, sanitize and return
 	if fw.ContentType() != ContentTypeYAML {
 		logrus.Debugf("Asked to preserve comments, but ContentType is not YAML, so ignoring")
 
 		// Normal encoding without the annotation (so it doesn't leak by accident)
-		return noAnnotationWrapper(metaobj, e.normalEncodeFunc(versionEncoder, fw, obj))
+		return noAnnotationWrapper(metaObj, e.normalEncodeFunc(versionEncoder, fw, obj))
 	}
 
-	// Get the encoded previous file data from the annotation or fall back to "normal" encoding
-	encodedPriorData, ok := getAnnotation(metaobj, preserveCommentsAnnotation)
-	if !ok {
-		// no need to delete the annotation as we know it doesn't exist, just do a normal encode
+	priorNode, err := getCommentSourceMeta(metaObj)
+	if errors.Is(err, ErrNoStoredComments) {
+		// No need to delete the annotation as we know it doesn't exist, just do a normal encode
 		return e.normalEncodeFunc(versionEncoder, fw, obj)()
-	}
-
-	// Decode the base64-encoded bytes of the original object (including the comments)
-	priorData, err := base64.StdEncoding.DecodeString(encodedPriorData)
-	if err != nil {
-		// fatal error
+	} else if err != nil {
 		return err
 	}
 
-	// Unmarshal the original YAML document into a *yaml.RNode, including comments
-	priorNode, err := yaml.Parse(string(priorData))
-	if err != nil {
-		// fatal error
-		return err
-	}
-
-	// Encode the new object into a temporary buffer, it should not be written as the "final result" to the fw
+	// Encode the new object into a temporary buffer, it should not be written as the "final result" to the FrameWriter
 	buf := new(bytes.Buffer)
-	if err := noAnnotationWrapper(metaobj, e.normalEncodeFunc(versionEncoder, NewYAMLFrameWriter(buf), obj)); err != nil {
+	if err := noAnnotationWrapper(metaObj, e.normalEncodeFunc(versionEncoder, NewYAMLFrameWriter(buf), obj)); err != nil {
 		// fatal error
 		return err
 	}
-	updatedData := buf.Bytes()
 
 	// Parse the new, upgraded, encoded YAML into *yaml.RNode for addition
 	// of comments from prevNode
-	afterNode, err := yaml.Parse(string(updatedData))
+	afterNode, err := yaml.Parse(buf.String())
 	if err != nil {
 		// fatal error
 		return err
@@ -117,3 +106,70 @@ func noAnnotationWrapper(metaobj metav1.Object, fn func() error) error {
 	// If the annotation isn't present, just run the function
 	return fn()
 }
+
+// GetCommentSource retrieves the YAML tree used as the source for transferring comments for the given runtime.Object.
+// This may be used externally to implement e.g. re-parenting of the comment source tree when moving structs around.
+func GetCommentSource(obj runtime.Object) (*yaml.RNode, error) {
+	// Cast the object to a metav1.Object to get access to annotations.
+	// If this fails, the given object does not support storing comments.
+	metaObj, ok := toMetaObject(obj)
+	if !ok {
+		return nil, ErrNoObjectMeta
+	}
+
+	// Use getCommentSourceMeta to retrieve the comments from the metav1.Object.
+	return getCommentSourceMeta(metaObj)
+}
+
+// getCommentSourceMeta retrieves the YAML tree used as the source for transferring comments for the given metav1.Object.
+func getCommentSourceMeta(metaObj metav1.Object) (*yaml.RNode, error) {
+	// Fetch the source string for the comments. If this fails, the given object does not have any stored comments.
+	sourceStr, ok := getAnnotation(metaObj, preserveCommentsAnnotation)
+	if !ok {
+		return nil, ErrNoStoredComments
+	}
+
+	// Decode the base64-encoded comment source string.
+	sourceBytes, err := base64.StdEncoding.DecodeString(sourceStr)
+	if err != nil {
+		return nil, err
+	}
+
+	// Parse the decoded source data into a *yaml.RNode and return it.
+	return yaml.Parse(string(sourceBytes))
+}
+
+// SetCommentSource sets the given YAML tree as the source for transferring comments for the given runtime.Object.
+// This may be used externally to implement e.g. re-parenting of the comment source tree when moving structs around.
+func SetCommentSource(obj runtime.Object, source *yaml.RNode) error {
+	// Convert the given tree into a string. This also handles the source == nil case.
+	str, err := source.String()
+	if err != nil {
+		return err
+	}
+
+	// Convert the string to bytes and pass it to setCommentSourceBytes to be applied.
+	if !setCommentSourceBytes(obj, []byte(str)) {
+		// If this fails, the passed object is not metav1.ObjectMeta compliant.
+		return ErrNoObjectMeta
+	}
+
+	return nil
+}
+
+// SetCommentSource sets the given bytes as the source for transferring comments for the given runtime.Object.
+func setCommentSourceBytes(obj runtime.Object, source []byte) bool {
+	// Cast the object to a metav1.Object to get access to annotations.
+	// If this fails, the given object does not support storing comments.
+	metaObj, ok := toMetaObject(obj)
+	if !ok {
+		return false
+	}
+
+	// base64-encode the comments string.
+	encodedStr := base64.StdEncoding.EncodeToString(source)
+
+	// Set the value of the comments annotation to the encoded string.
+	setAnnotation(metaObj, preserveCommentsAnnotation, encodedStr)
+	return true
+}

pkg/serializer/comments/lost.go

@@ -2,8 +2,9 @@ package comments
 
 import (
 	"fmt"
-	"sigs.k8s.io/kustomize/kyaml/yaml"
 	"strings"
+
+	"sigs.k8s.io/kustomize/kyaml/yaml"
 )
 
 // lostComment specifies a mapping between a fieldName (in the old structure), which doesn't exist in the

pkg/serializer/comments_test.go

@@ -0,0 +1,202 @@
+package serializer
+
+import (
+	"encoding/base64"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	runtimetest "k8s.io/apimachinery/pkg/runtime/testing"
+	"sigs.k8s.io/kustomize/kyaml/yaml"
+)
+
+const sampleData1 = `# Comment
+
+kind: Test
+spec:
+  # Head comment
+  data:
+  - field # Inline comment
+  - another
+  thing:
+    # Head comment
+    var: true
+`
+
+const sampleData2 = `kind: Test
+spec:
+  # Head comment
+  data:
+  - field # Inline comment
+  - another:
+      subthing: "yes"
+  thing:
+    # Head comment
+    var: true
+status:
+  nested:
+    fields:
+    # Just a comment
+`
+
+type internalSimpleOM struct {
+	runtimetest.InternalSimple
+	metav1.ObjectMeta `json:"metadata,omitempty"`
+}
+
+func withComments(data string) *internalSimpleOM {
+	return &internalSimpleOM{
+		ObjectMeta: metav1.ObjectMeta{
+			Annotations: map[string]string{preserveCommentsAnnotation: data},
+		},
+	}
+}
+
+func parseRNode(t *testing.T, source string) *yaml.RNode {
+	rNode, err := yaml.Parse(source)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return rNode
+}
+
+func TestGetCommentSource(t *testing.T) {
+	testCases := []struct {
+		name        string
+		obj         runtime.Object
+		result      string
+		expectedErr bool
+	}{
+		{
+			name:        "no_ObjectMeta",
+			obj:         &runtimetest.InternalSimple{},
+			expectedErr: true,
+		},
+		{
+			name:        "no_comments",
+			obj:         &internalSimpleOM{},
+			expectedErr: true,
+		},
+		{
+			name:        "invalid_comments",
+			obj:         withComments("ä"),
+			expectedErr: true,
+		},
+		{
+			name:        "successful_parsing",
+			obj:         withComments(base64.StdEncoding.EncodeToString([]byte(sampleData1))),
+			result:      sampleData1,
+			expectedErr: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			source, actualErr := GetCommentSource(tc.obj)
+			if (actualErr != nil) != tc.expectedErr {
+				t.Errorf("expected error %t, but received %t: %v", tc.expectedErr, actualErr != nil, actualErr)
+			}
+
+			if actualErr != nil {
+				// Already handled above.
+				return
+			}
+
+			str, err := source.String()
+			require.NoError(t, err)
+			assert.Equal(t, tc.result, str)
+		})
+	}
+}
+
+func TestSetCommentSource(t *testing.T) {
+	testCases := []struct {
+		name        string
+		obj         runtime.Object
+		source      *yaml.RNode
+		result      string
+		expectedErr bool
+	}{
+		{
+			name:        "no_ObjectMeta",
+			obj:         &runtimetest.InternalSimple{},
+			source:      yaml.NewScalarRNode("test"),
+			expectedErr: true,
+		},
+		{
+			name:        "nil_source",
+			obj:         &internalSimpleOM{},
+			source:      nil,
+			result:      "",
+			expectedErr: false,
+		},
+		{
+			name:        "successful_parsing",
+			obj:         &internalSimpleOM{},
+			source:      parseRNode(t, sampleData1),
+			result:      sampleData1,
+			expectedErr: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			actualErr := SetCommentSource(tc.obj, tc.source)
+			if (actualErr != nil) != tc.expectedErr {
+				t.Errorf("expected error %t, but received %t: %v", tc.expectedErr, actualErr != nil, actualErr)
+			}
+
+			if actualErr != nil {
+				// Already handled above.
+				return
+			}
+
+			meta, ok := toMetaObject(tc.obj)
+			if !ok {
+				t.Fatal("cannot extract metav1.ObjectMeta")
+			}
+
+			annotation, ok := getAnnotation(meta, preserveCommentsAnnotation)
+			if !ok {
+				t.Fatal("expected annotation to be set, but it is not")
+			}
+
+			str, err := base64.StdEncoding.DecodeString(annotation)
+			require.NoError(t, err)
+			assert.Equal(t, tc.result, string(str))
+		})
+	}
+}
+
+func TestCommentSourceSetGet(t *testing.T) {
+	testCases := []struct {
+		name   string
+		source string
+	}{
+		{
+			name:   "encode_decode_1",
+			source: sampleData1,
+		},
+		{
+			name:   "encode_decode_2",
+			source: sampleData2,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			obj := &internalSimpleOM{}
+			assert.NoError(t, SetCommentSource(obj, parseRNode(t, tc.source)))
+
+			rNode, err := GetCommentSource(obj)
+			assert.NoError(t, err)
+
+			str, err := rNode.String()
+			require.NoError(t, err)
+			assert.Equal(t, tc.source, str)
+		})
+	}
+}