github.com/pachyderm/pachyderm@v1.13.4/src/server/pkg/serde/yaml_encoder.go (about)

     1  package serde
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"io"
     7  
     8  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
     9  
    10  	"github.com/gogo/protobuf/jsonpb"
    11  	"github.com/gogo/protobuf/proto"
    12  	"gopkg.in/pachyderm/yaml.v3"
    13  )
    14  
    15  // YAMLEncoder is an implementation of serde.Encoder that operates on YAML data
    16  type YAMLEncoder struct {
    17  	e *yaml.Encoder
    18  
    19  	// OrigName sets whether this YAMLEncoder uses the original (.proto) name of
    20  	// fields when marshalling to protos
    21  	origName bool
    22  }
    23  
    24  // EncodeYAML is a convenience function that encodes yaml data using a
    25  // YAMLEncoder, but can be called inline
    26  func EncodeYAML(v interface{}, options ...EncoderOption) ([]byte, error) {
    27  	var buf bytes.Buffer
    28  	e := NewYAMLEncoder(&buf, options...)
    29  	if err := e.Encode(v); err != nil {
    30  		return nil, err
    31  	}
    32  	return buf.Bytes(), nil
    33  }
    34  
    35  // NewYAMLEncoder returns a new YAMLEncoder that writes to 'w'
    36  func NewYAMLEncoder(w io.Writer, options ...EncoderOption) *YAMLEncoder {
    37  	e := &YAMLEncoder{e: yaml.NewEncoder(w)}
    38  	for _, o := range options {
    39  		o(e)
    40  	}
    41  	return e
    42  }
    43  
    44  // Encode implements the corresponding method of serde.Encoder
    45  func (e *YAMLEncoder) Encode(v interface{}) error {
    46  	return e.EncodeTransform(v, nil)
    47  }
    48  
    49  // EncodeTransform implements the corresponding method of serde.Encoder
    50  func (e *YAMLEncoder) EncodeTransform(v interface{}, f func(map[string]interface{}) error) error {
    51  	// Encode to JSON first
    52  	var buf bytes.Buffer
    53  	j := json.NewEncoder(&buf)
    54  	if err := j.Encode(v); err != nil {
    55  		return errors.Wrapf(err, "serialization error while canonicalizing output")
    56  	}
    57  
    58  	return e.jsonToYAMLTransform(buf.Bytes(), f)
    59  }
    60  
    61  // EncodeProto implements the corresponding method of serde.Encoder
    62  func (e *YAMLEncoder) EncodeProto(v proto.Message) error {
    63  	return e.EncodeProtoTransform(v, nil)
    64  }
    65  
    66  // EncodeProtoTransform implements the corresponding method of serde.Encoder
    67  func (e *YAMLEncoder) EncodeProtoTransform(v proto.Message, f func(map[string]interface{}) error) error {
    68  	// Encode to JSON first
    69  	var buf bytes.Buffer
    70  	m := jsonpb.Marshaler{
    71  		OrigName: e.origName,
    72  	}
    73  	if err := m.Marshal(&buf, v); err != nil {
    74  		return errors.Wrapf(err, "serialization error while canonicalizing output")
    75  	}
    76  
    77  	return e.jsonToYAMLTransform(buf.Bytes(), f)
    78  }
    79  
    80  func (e *YAMLEncoder) jsonToYAMLTransform(intermediateJSON []byte,
    81  	f func(map[string]interface{}) error) error {
    82  	// Unmarshal from JSON to intermediate map ('holder')
    83  	holder := map[string]interface{}{}
    84  	if err := json.Unmarshal(intermediateJSON, &holder); err != nil {
    85  		return errors.Wrapf(err, "deserialization error while canonicalizing output")
    86  	}
    87  
    88  	// transform 'holder' (e.g. de-stringifying TFJob)
    89  	if f != nil {
    90  		if err := f(holder); err != nil {
    91  			return err
    92  		}
    93  	}
    94  
    95  	// Encode 'holder' to YAML
    96  	if err := e.e.Encode(holder); err != nil {
    97  		return errors.Wrapf(err, "serialization error while canonicalizing yaml")
    98  	}
    99  	return nil
   100  }