Preserve YAML comments & style when resolving/applying

go.mod

@@ -4,6 +4,7 @@ go 1.12
 
 require (
 	github.com/docker/docker v1.4.2-0.20180531152204-71cd53e4a197
+	github.com/dprotaso/go-yit v0.0.0-20191028211022-135eb7262960
 	github.com/evanphx/json-patch v4.2.0+incompatible // indirect
 	github.com/fsnotify/fsnotify v1.4.7
 	github.com/google/btree v1.0.0 // indirect
@@ -13,17 +14,15 @@ require (
 	github.com/imdario/mergo v0.3.5 // indirect
 	github.com/mattmoor/dep-notify v0.0.0-20190205035814-a45dec370a17
 	github.com/modern-go/reflect2 v1.0.1 // indirect
-	github.com/petar/GoLLRB v0.0.0-20130427215148-53be0d36a84c // indirect
 	github.com/spf13/afero v1.2.1 // indirect
 	github.com/spf13/cobra v0.0.5
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.3.2
 	golang.org/x/sync v0.0.0-20190423024810-112230192c58
 	golang.org/x/tools v0.0.0-20191001184121-329c8d646ebe
-	gopkg.in/yaml.v2 v2.2.2
+	gopkg.in/yaml.v3 v3.0.0-20191026110619-0b21df46bc1d
 	k8s.io/apimachinery v0.0.0-20180904193909-def12e63c512
 	k8s.io/client-go v8.0.0+incompatible // indirect
 	k8s.io/kubernetes v1.11.10
-	sigs.k8s.io/yaml v1.1.0
 )

pkg/commands/resolver.go

@@ -15,6 +15,7 @@
 package commands
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
@@ -30,6 +31,8 @@ import (
 	"github.com/google/ko/pkg/publish"
 	"github.com/google/ko/pkg/resolve"
 	"github.com/mattmoor/dep-notify/pkg/graph"
+	"gopkg.in/yaml.v3"
+	"k8s.io/apimachinery/pkg/labels"
 )
 
 func gobuildOptions(bo *options.BuildOptions) ([]build.Option, error) {
@@ -116,7 +119,13 @@ func makePublisher(no *options.NameOptions, lo *options.LocalOptions, ta *option
 // resolvedFuture represents a "future" for the bytes of a resolved file.
 type resolvedFuture chan []byte
 
-func resolveFilesToWriter(builder *build.Caching, publisher publish.Interface, fo *options.FilenameOptions, so *options.SelectorOptions, sto *options.StrictOptions, out io.WriteCloser) {
+func resolveFilesToWriter(
+	builder *build.Caching,
+	publisher publish.Interface,
+	fo *options.FilenameOptions,
+	so *options.SelectorOptions,
+	sto *options.StrictOptions,
+	out io.WriteCloser) {
 	defer out.Close()
 
 	// By having this as a channel, we can hook this up to a filesystem
@@ -242,7 +251,23 @@ func resolveFilesToWriter(builder *build.Caching, publisher publish.Interface, f
 	}
 }
 
-func resolveFile(f string, builder build.Interface, pub publish.Interface, so *options.SelectorOptions, sto *options.StrictOptions) (b []byte, err error) {
+func resolveFile(
+	f string,
+	builder build.Interface,
+	pub publish.Interface,
+	so *options.SelectorOptions,
+	sto *options.StrictOptions) (b []byte, err error) {
+
+	var selector labels.Selector
+	if so.Selector != "" {
+		var err error
+		selector, err = labels.Parse(so.Selector)
+
+		if err != nil {
+			return nil, fmt.Errorf("unable to parse selector: %v", err)
+		}
+	}
+
 	if f == "-" {
 		b, err = ioutil.ReadAll(os.Stdin)
 	} else {
@@ -252,12 +277,49 @@ func resolveFile(f string, builder build.Interface, pub publish.Interface, so *o
 		return nil, err
 	}
 
-	if so.Selector != "" {
-		b, err = resolve.FilterBySelector(b, so.Selector)
-		if err != nil {
+	var docNodes []*yaml.Node
+
+	// The loop is to support multi-document yaml files.
+	// This is handled by using a yaml.Decoder and reading objects until io.EOF, see:
+	// https://github.com/go-yaml/yaml/blob/v2.2.1/yaml.go#L124
+	decoder := yaml.NewDecoder(bytes.NewBuffer(b))
+	for {
+		var doc yaml.Node
+		if err := decoder.Decode(&doc); err != nil {
+			if err == io.EOF {
+				break
+			}
 			return nil, err
 		}
+
+		if selector != nil {
+			if match, err := resolve.MatchesSelector(&doc, selector); err != nil {
+				return nil, fmt.Errorf("error evaluating selector: %v", err)
+			} else if !match {
+				continue
+			}
+		}
+
+		docNodes = append(docNodes, &doc)
+
+	}
+
+	if err := resolve.ImageReferences(docNodes, sto.Strict, builder, pub); err != nil {
+		return nil, fmt.Errorf("error resolving image references: %v", err)
 	}
 
-	return resolve.ImageReferences(b, sto.Strict, builder, pub)
+	buf := &bytes.Buffer{}
+	e := yaml.NewEncoder(buf)
+	e.SetIndent(2)
+
+	for _, doc := range docNodes {
+		err := e.Encode(doc)
+		if err != nil {
+			return nil, fmt.Errorf("failed to encode output: %v", err)
+		}
+	}
+	e.Close()
+
+	return buf.Bytes(), nil
+
 }

pkg/commands/resolver_test.go

@@ -0,0 +1,150 @@
+// Copyright 2018 Google LLC All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package commands
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/random"
+	"github.com/google/ko/pkg/commands/options"
+	"github.com/google/ko/pkg/testutil"
+	"gopkg.in/yaml.v3"
+)
+
+var (
+	fooRef      = "github.com/awesomesauce/foo"
+	foo         = mustRandom()
+	fooHash     = mustDigest(foo)
+	barRef      = "github.com/awesomesauce/bar"
+	bar         = mustRandom()
+	barHash     = mustDigest(bar)
+	testBuilder = testutil.NewFixedBuild(map[string]v1.Image{
+		fooRef: foo,
+		barRef: bar,
+	})
+	testHashes = map[string]v1.Hash{
+		fooRef: fooHash,
+		barRef: barHash,
+	}
+)
+
+func TestResolveMultiDocumentYAMLs(t *testing.T) {
+	refs := []string{fooRef, barRef}
+	hashes := []v1.Hash{fooHash, barHash}
+	base := mustRepository("gcr.io/multi-pass")
+
+	buf := bytes.NewBuffer(nil)
+	encoder := yaml.NewEncoder(buf)
+	for _, input := range refs {
+		if err := encoder.Encode(input); err != nil {
+			t.Fatalf("Encode(%v) = %v", input, err)
+		}
+	}
+
+	inputYAML := buf.Bytes()
+
+	outYAML, err := resolveFile(
+		yamlToTmpFile(t, buf.Bytes()),
+		testBuilder,
+		testutil.NewFixedPublish(base, testHashes),
+		&options.SelectorOptions{},
+		&options.StrictOptions{})
+
+	if err != nil {
+		t.Fatalf("ImageReferences(%v) = %v", string(inputYAML), err)
+	}
+
+	buf = bytes.NewBuffer(outYAML)
+	decoder := yaml.NewDecoder(buf)
+	var outStructured []string
+	for {
+		var output string
+		if err := decoder.Decode(&output); err == nil {
+			outStructured = append(outStructured, output)
+		} else if err == io.EOF {
+			outStructured = append(outStructured, output)
+			break
+		} else {
+			t.Errorf("yaml.Unmarshal(%v) = %v", string(outYAML), err)
+		}
+	}
+
+	expectedStructured := []string{
+		testutil.ComputeDigest(base, refs[0], hashes[0]),
+		testutil.ComputeDigest(base, refs[1], hashes[1]),
+	}
+
+	// The multi-document output always seems to leave a trailing --- so we end up with
+	// an extra empty element.
+	expectedStructured = append(expectedStructured, "")
+
+	if want, got := len(expectedStructured), len(outStructured); want != got {
+		t.Errorf("resolveFile(%v) = %v, want %v", string(inputYAML), got, want)
+	}
+
+	if diff := cmp.Diff(expectedStructured, outStructured, cmpopts.EquateEmpty()); diff != "" {
+		t.Errorf("resolveFile(%v); (-want +got) = %v", string(inputYAML), diff)
+	}
+}
+
+func mustRepository(s string) name.Repository {
+	n, err := name.NewRepository(s)
+	if err != nil {
+		panic(err)
+	}
+	return n
+}
+
+func mustDigest(img v1.Image) v1.Hash {
+	d, err := img.Digest()
+	if err != nil {
+		panic(err)
+	}
+	return d
+}
+
+func mustRandom() v1.Image {
+	img, err := random.Image(1024, 5)
+	if err != nil {
+		panic(err)
+	}
+	return img
+}
+
+func yamlToTmpFile(t *testing.T, yaml []byte) string {
+	t.Helper()
+
+	tmpfile, err := ioutil.TempFile("", "doc")
+	if err != nil {
+		t.Fatalf("error creating temp file: %v", err)
+	}
+
+	if _, err := tmpfile.Write(yaml); err != nil {
+		t.Fatalf("error writing temp file: %v", err)
+	}
+
+	if err := tmpfile.Close(); err != nil {
+		t.Fatalf("error closing temp file: %v", err)
+	}
+
+	return tmpfile.Name()
+}