fix(oci): walk entire stream of documents

cmd/flipt/bundle.go

@@ -9,9 +9,7 @@ import (
 	"go.flipt.io/flipt/internal/oci"
 )
 
-type bundleCommand struct {
-	rootDir string
-}
+type bundleCommand struct{}
 
 func newBundleCommand() *cobra.Command {
 	bundle := &bundleCommand{}

internal/ext/common.go

@@ -1,77 +1,78 @@
 package ext
 
 import (
+	"encoding/json"
 	"errors"
 )
 
 type Document struct {
-	Version   string     `yaml:"version,omitempty"`
-	Namespace string     `yaml:"namespace,omitempty"`
-	Flags     []*Flag    `yaml:"flags,omitempty"`
-	Segments  []*Segment `yaml:"segments,omitempty"`
+	Version   string     `yaml:"version,omitempty" json:"version,omitempty"`
+	Namespace string     `yaml:"namespace,omitempty" json:"namespace,omitempty"`
+	Flags     []*Flag    `yaml:"flags,omitempty" json:"flags,omitempty"`
+	Segments  []*Segment `yaml:"segments,omitempty" json:"segments,omitempty"`
 }
 
 type Flag struct {
-	Key         string     `yaml:"key,omitempty"`
-	Name        string     `yaml:"name,omitempty"`
-	Type        string     `yaml:"type,omitempty"`
-	Description string     `yaml:"description,omitempty"`
-	Enabled     bool       `yaml:"enabled"`
-	Variants    []*Variant `yaml:"variants,omitempty"`
-	Rules       []*Rule    `yaml:"rules,omitempty"`
-	Rollouts    []*Rollout `yaml:"rollouts,omitempty"`
+	Key         string     `yaml:"key,omitempty" json:"key,omitempty"`
+	Name        string     `yaml:"name,omitempty" json:"name,omitempty"`
+	Type        string     `yaml:"type,omitempty" json:"type,omitempty"`
+	Description string     `yaml:"description,omitempty" json:"description,omitempty"`
+	Enabled     bool       `yaml:"enabled" json:"enabled,omitempty"`
+	Variants    []*Variant `yaml:"variants,omitempty" json:"variants,omitempty"`
+	Rules       []*Rule    `yaml:"rules,omitempty" json:"rules,omitempty"`
+	Rollouts    []*Rollout `yaml:"rollouts,omitempty" json:"rollouts,omitempty"`
 }
 
 type Variant struct {
-	Key         string      `yaml:"key,omitempty"`
-	Name        string      `yaml:"name,omitempty"`
-	Description string      `yaml:"description,omitempty"`
-	Attachment  interface{} `yaml:"attachment,omitempty"`
+	Key         string      `yaml:"key,omitempty" json:"key,omitempty"`
+	Name        string      `yaml:"name,omitempty" json:"name,omitempty"`
+	Description string      `yaml:"description,omitempty" json:"description,omitempty"`
+	Attachment  interface{} `yaml:"attachment,omitempty" json:"attachment,omitempty"`
 }
 
 type Rule struct {
-	Segment       *SegmentEmbed   `yaml:"segment,omitempty"`
-	Rank          uint            `yaml:"rank,omitempty"`
-	Distributions []*Distribution `yaml:"distributions,omitempty"`
+	Segment       *SegmentEmbed   `yaml:"segment,omitempty" json:"segment,omitempty"`
+	Rank          uint            `yaml:"rank,omitempty" json:"rank,omitempty"`
+	Distributions []*Distribution `yaml:"distributions,omitempty" json:"distributions,omitempty"`
 }
 
 type Distribution struct {
-	VariantKey string  `yaml:"variant,omitempty"`
-	Rollout    float32 `yaml:"rollout,omitempty"`
+	VariantKey string  `yaml:"variant,omitempty" json:"variant,omitempty"`
+	Rollout    float32 `yaml:"rollout,omitempty" json:"rollout,omitempty"`
 }
 
 type Rollout struct {
-	Description string         `yaml:"description,omitempty"`
-	Segment     *SegmentRule   `yaml:"segment,omitempty"`
-	Threshold   *ThresholdRule `yaml:"threshold,omitempty"`
+	Description string         `yaml:"description,omitempty" json:"description,omitempty"`
+	Segment     *SegmentRule   `yaml:"segment,omitempty" json:"segment,omitempty"`
+	Threshold   *ThresholdRule `yaml:"threshold,omitempty" json:"threshold,omitempty"`
 }
 
 type SegmentRule struct {
-	Key      string   `yaml:"key,omitempty"`
-	Keys     []string `yaml:"keys,omitempty"`
-	Operator string   `yaml:"operator,omitempty"`
-	Value    bool     `yaml:"value,omitempty"`
+	Key      string   `yaml:"key,omitempty" json:"key,omitempty"`
+	Keys     []string `yaml:"keys,omitempty" json:"keys,omitempty"`
+	Operator string   `yaml:"operator,omitempty" json:"operator,omitempty"`
+	Value    bool     `yaml:"value,omitempty" json:"value,omitempty"`
 }
 
 type ThresholdRule struct {
-	Percentage float32 `yaml:"percentage,omitempty"`
-	Value      bool    `yaml:"value,omitempty"`
+	Percentage float32 `yaml:"percentage,omitempty" json:"percentage,omitempty"`
+	Value      bool    `yaml:"value,omitempty" json:"value,omitempty"`
 }
 
 type Segment struct {
-	Key         string        `yaml:"key,omitempty"`
-	Name        string        `yaml:"name,omitempty"`
-	Description string        `yaml:"description,omitempty"`
-	Constraints []*Constraint `yaml:"constraints,omitempty"`
-	MatchType   string        `yaml:"match_type,omitempty"`
+	Key         string        `yaml:"key,omitempty" json:"key,omitempty"`
+	Name        string        `yaml:"name,omitempty" json:"name,omitempty"`
+	Description string        `yaml:"description,omitempty" json:"description,omitempty"`
+	Constraints []*Constraint `yaml:"constraints,omitempty" json:"constraints,omitempty"`
+	MatchType   string        `yaml:"match_type,omitempty" json:"match_type,omitempty"`
 }
 
 type Constraint struct {
-	Type        string `yaml:"type,omitempty"`
-	Property    string `yaml:"property,omitempty"`
-	Operator    string `yaml:"operator,omitempty"`
-	Value       string `yaml:"value,omitempty"`
-	Description string `yaml:"description,omitempty"`
+	Type        string `yaml:"type,omitempty" json:"type,omitempty"`
+	Property    string `yaml:"property,omitempty" json:"property,omitempty"`
+	Operator    string `yaml:"operator,omitempty" json:"operator,omitempty"`
+	Value       string `yaml:"value,omitempty" json:"value,omitempty"`
+	Description string `yaml:"description,omitempty" json:"description,omitempty"`
 }
 
 type SegmentEmbed struct {
@@ -114,6 +115,42 @@ func (s *SegmentEmbed) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	return errors.New("failed to unmarshal to string or segmentKeys")
 }
 
+// MarshalJSON tries to type assert to either of the following types that implement
+// IsSegment, and returns the marshaled value.
+func (s *SegmentEmbed) MarshalJSON() ([]byte, error) {
+	switch t := s.IsSegment.(type) {
+	case SegmentKey:
+		return json.Marshal(string(t))
+	case *Segments:
+		sk := &Segments{
+			Keys:            t.Keys,
+			SegmentOperator: t.SegmentOperator,
+		}
+		return json.Marshal(sk)
+	}
+
+	return nil, errors.New("failed to marshal to string or segmentKeys")
+}
+
+// UnmarshalJSON attempts to unmarshal a string or `SegmentKeys`, and fails if it can not
+// do so.
+func (s *SegmentEmbed) UnmarshalJSON(v []byte) error {
+	var sk SegmentKey
+
+	if err := json.Unmarshal(v, &sk); err == nil {
+		s.IsSegment = sk
+		return nil
+	}
+
+	var sks *Segments
+	if err := json.Unmarshal(v, &sks); err == nil {
+		s.IsSegment = sks
+		return nil
+	}
+
+	return errors.New("failed to unmarshal to string or segmentKeys")
+}
+
 // IsSegment is used to unify the two types of segments that can come in
 // from the import.
 type IsSegment interface {
@@ -125,8 +162,8 @@ type SegmentKey string
 func (s SegmentKey) IsSegment() {}
 
 type Segments struct {
-	Keys            []string `yaml:"keys,omitempty"`
-	SegmentOperator string   `yaml:"operator,omitempty"`
+	Keys            []string `yaml:"keys,omitempty" json:"keys,omitempty"`
+	SegmentOperator string   `yaml:"operator,omitempty" json:"operator,omitempty"`
 }
 
 func (s *Segments) IsSegment() {}

internal/oci/file.go

@@ -285,25 +285,25 @@ type Bundle struct {
 func (s *Store) Build(ctx context.Context, src fs.FS) (Bundle, error) {
 	store, err := s.store()
 	if err != nil {
-		return Bundle{}, nil
+		return Bundle{}, err
 	}
 
 	layers, err := s.buildLayers(ctx, store, src)
 	if err != nil {
-		return Bundle{}, nil
+		return Bundle{}, err
 	}
 
 	desc, err := oras.PackManifest(ctx, store, oras.PackManifestVersion1_1_RC4, MediaTypeFliptFeatures, oras.PackManifestOptions{
 		ManifestAnnotations: map[string]string{},
 		Layers:              layers,
 	})
 	if err != nil {
-		return Bundle{}, nil
+		return Bundle{}, err
 	}
 
 	if s.reference.Reference != "" {
 		if err := store.Tag(ctx, desc, s.reference.Reference); err != nil {
-			return Bundle{}, nil
+			return Bundle{}, err
 		}
 	}
 
@@ -315,7 +315,7 @@ func (s *Store) Build(ctx context.Context, src fs.FS) (Bundle, error) {
 
 	bundle.CreatedAt, err = parseCreated(desc.Annotations)
 	if err != nil {
-		return Bundle{}, nil
+		return Bundle{}, err
 	}
 
 	return bundle, nil

internal/oci/file_test.go

@@ -19,6 +19,8 @@ import (
 	"oras.land/oras-go/v2/content/oci"
 )
 
+const repo = "testrepo"
+
 func TestNewStore(t *testing.T) {
 	t.Run("unexpected scheme", func(t *testing.T) {
 		_, err := NewStore(zaptest.NewLogger(t), "fake://local/something:latest")
@@ -51,7 +53,7 @@ func TestNewStore(t *testing.T) {
 }
 
 func TestStore_Fetch_InvalidMediaType(t *testing.T) {
-	dir, repo := testRepository(t,
+	dir := testRepository(t,
 		layer("default", `{"namespace":"default"}`, "unexpected.media.type"),
 	)
 
@@ -63,7 +65,7 @@ func TestStore_Fetch_InvalidMediaType(t *testing.T) {
 	_, err = store.Fetch(ctx)
 	require.EqualError(t, err, "layer \"sha256:85ee577ad99c62f314abca9f43ad87c2ee8818513e6383a77690df56d0352748\": type \"unexpected.media.type\": unexpected media type")
 
-	dir, repo = testRepository(t,
+	dir = testRepository(t,
 		layer("default", `{"namespace":"default"}`, MediaTypeFliptNamespace+"+unknown"),
 	)
 
@@ -75,7 +77,7 @@ func TestStore_Fetch_InvalidMediaType(t *testing.T) {
 }
 
 func TestStore_Fetch(t *testing.T) {
-	dir, repo := testRepository(t,
+	dir := testRepository(t,
 		layer("default", `{"namespace":"default"}`, MediaTypeFliptNamespace),
 		layer("other", `namespace: other`, MediaTypeFliptNamespace+"+yaml"),
 	)
@@ -129,7 +131,7 @@ var testdata embed.FS
 
 func TestStore_Build(t *testing.T) {
 	ctx := context.TODO()
-	dir, repo := testRepository(t)
+	dir := testRepository(t)
 
 	store, err := NewStore(zaptest.NewLogger(t), fmt.Sprintf("flipt://local/%s:latest", repo), WithBundleDir(dir))
 	require.NoError(t, err)
@@ -171,15 +173,14 @@ func layer(ns, payload, mediaType string) func(*testing.T, oras.Target) v1.Descr
 	}
 }
 
-func testRepository(t *testing.T, layerFuncs ...func(*testing.T, oras.Target) v1.Descriptor) (dir, repository string) {
+func testRepository(t *testing.T, layerFuncs ...func(*testing.T, oras.Target) v1.Descriptor) (dir string) {
 	t.Helper()
 
-	repository = "testrepo"
 	dir = t.TempDir()
 
-	t.Log("test OCI directory", dir, repository)
+	t.Log("test OCI directory", dir, repo)
 
-	store, err := oci.New(path.Join(dir, repository))
+	store, err := oci.New(path.Join(dir, repo))
 	require.NoError(t, err)
 
 	store.AutoSaveIndex = true

internal/storage/fs/snapshot.go

@@ -122,13 +122,15 @@ func SnapshotFromFiles(logger *zap.Logger, files ...fs.File) (*StoreSnapshot, er
 	for _, fi := range files {
 		defer fi.Close()
 
-		doc, err := documentFromFile(logger, fi)
+		docs, err := documentsFromFile(fi)
 		if err != nil {
 			return nil, err
 		}
 
-		if err := s.addDoc(doc); err != nil {
-			return nil, err
+		for _, doc := range docs {
+			if err := s.addDoc(doc); err != nil {
+				return nil, err
+			}
 		}
 	}
 
@@ -152,21 +154,23 @@ func WalkDocuments(logger *zap.Logger, src fs.FS, fn func(*ext.Document) error)
 		}
 		defer fi.Close()
 
-		doc, err := documentFromFile(logger, fi)
+		docs, err := documentsFromFile(fi)
 		if err != nil {
 			return err
 		}
 
-		if err := fn(doc); err != nil {
-			return err
+		for _, doc := range docs {
+			if err := fn(doc); err != nil {
+				return err
+			}
 		}
 	}
 
 	return nil
 }
 
-// documentFromFile parses and validates a document from a single fs.File instance
-func documentFromFile(logger *zap.Logger, fi fs.File) (*ext.Document, error) {
+// documentsFromFile parses and validates a document from a single fs.File instance
+func documentsFromFile(fi fs.File) ([]*ext.Document, error) {
 	validator, err := cue.NewFeaturesValidator()
 	if err != nil {
 		return nil, err
@@ -184,42 +188,37 @@ func documentFromFile(logger *zap.Logger, fi fs.File) (*ext.Document, error) {
 		return nil, err
 	}
 
+	var docs []*ext.Document
 	extn := filepath.Ext(stat.Name())
+
+	var decode func(any) error
 	switch extn {
 	case ".yaml", ".yml":
-		decoder := yaml.NewDecoder(buf)
 		// Support YAML stream by looping until we reach an EOF.
-		for {
-			doc := &ext.Document{}
-			if err := decoder.Decode(doc); err != nil {
-				if errors.Is(err, io.EOF) {
-					break
-				}
-				return nil, err
-			}
-
-			// set namespace to default if empty in document
-			if doc.Namespace == "" {
-				doc.Namespace = "default"
-			}
-
-			return doc, nil
-		}
+		decode = yaml.NewDecoder(buf).Decode
 	case "", ".json":
+		decode = json.NewDecoder(buf).Decode
+	default:
+		return nil, fmt.Errorf("unexpected extension: %q", extn)
+	}
+
+	for {
 		doc := &ext.Document{}
-		if err := json.NewDecoder(buf).Decode(&doc); err != nil {
+		if err := decode(doc); err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
 			return nil, err
 		}
 
 		// set namespace to default if empty in document
 		if doc.Namespace == "" {
 			doc.Namespace = "default"
 		}
-
-		return doc, nil
+		docs = append(docs, doc)
 	}
 
-	return nil, fmt.Errorf("unexpected extension: %q", extn)
+	return docs, nil
 }
 
 // listStateFiles lists all the file paths in a provided fs.FS containing Flipt feature state