databricks-cli/libs/dyn/yamlsaver/saver.go

169 lines
4.3 KiB
Go
Raw Normal View History

package yamlsaver
import (
"fmt"
"io"
"os"
"path/filepath"
"sort"
"strconv"
"github.com/databricks/cli/libs/dyn"
"golang.org/x/exp/maps"
"gopkg.in/yaml.v3"
)
type saver struct {
nodesWithStyle map[string]yaml.Style
}
func NewSaver() *saver {
return &saver{}
}
func NewSaverWithStyle(nodesWithStyle map[string]yaml.Style) *saver {
return &saver{
nodesWithStyle: nodesWithStyle,
}
}
func (s *saver) SaveAsYAML(data any, filename string, force bool) error {
err := os.MkdirAll(filepath.Dir(filename), 0755)
if err != nil {
return err
}
// check that file exists
info, err := os.Stat(filename)
if err == nil {
if info.IsDir() {
return fmt.Errorf("%s is a directory", filename)
}
if !force {
return fmt.Errorf("%s already exists. Use --force to overwrite", filename)
}
}
file, err := os.Create(filename)
if err != nil {
return err
}
defer file.Close()
err = s.encode(data, file)
if err != nil {
return err
}
return nil
}
func (s *saver) encode(data any, w io.Writer) error {
yamlNode, err := s.toYamlNode(dyn.V(data))
if err != nil {
return err
}
enc := yaml.NewEncoder(w)
enc.SetIndent(2)
return enc.Encode(yamlNode)
}
func (s *saver) toYamlNode(v dyn.Value) (*yaml.Node, error) {
return s.toYamlNodeWithStyle(v, yaml.Style(0))
}
func (s *saver) toYamlNodeWithStyle(v dyn.Value, style yaml.Style) (*yaml.Node, error) {
switch v.Kind() {
case dyn.KindMap:
m, _ := v.AsMap()
keys := maps.Keys(m)
// We're using location lines to define the order of keys in YAML.
// The location is set when we convert API response struct to config.Value representation
// See convert.convertMap for details
sort.SliceStable(keys, func(i, j int) bool {
return m[keys[i]].Location().Line < m[keys[j]].Location().Line
})
content := make([]*yaml.Node, 0)
for _, k := range keys {
item := m[k]
node := yaml.Node{Kind: yaml.ScalarNode, Value: k, Style: style}
var nestedNodeStyle yaml.Style
if customStyle, ok := s.hasStyle(k); ok {
nestedNodeStyle = customStyle
} else {
nestedNodeStyle = style
}
c, err := s.toYamlNodeWithStyle(item, nestedNodeStyle)
if err != nil {
return nil, err
}
content = append(content, &node)
content = append(content, c)
}
return &yaml.Node{Kind: yaml.MappingNode, Content: content, Style: style}, nil
case dyn.KindSequence:
seq, _ := v.AsSequence()
content := make([]*yaml.Node, 0)
for _, item := range seq {
node, err := s.toYamlNodeWithStyle(item, style)
if err != nil {
return nil, err
}
content = append(content, node)
}
return &yaml.Node{Kind: yaml.SequenceNode, Content: content, Style: style}, nil
case dyn.KindNil:
return &yaml.Node{Kind: yaml.ScalarNode, Value: "null", Style: style}, nil
case dyn.KindString:
// If the string is a scalar value (bool, int, float and etc.), we want to quote it.
if isScalarValueInString(v) {
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustString(), Style: yaml.DoubleQuotedStyle}, nil
}
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustString(), Style: style}, nil
case dyn.KindBool:
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustBool()), Style: style}, nil
case dyn.KindInt:
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustInt()), Style: style}, nil
case dyn.KindFloat:
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustFloat()), Style: style}, nil
case dyn.KindTime:
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustTime().UTC().String(), Style: style}, nil
default:
// Panic because we only want to deal with known types.
panic(fmt.Sprintf("invalid kind: %d", v.Kind()))
}
}
func (s *saver) hasStyle(key string) (yaml.Style, bool) {
style, ok := s.nodesWithStyle[key]
return style, ok
}
func isScalarValueInString(v dyn.Value) bool {
if v.Kind() != dyn.KindString {
return false
}
// Parse value of the string and check if it's a scalar value.
// If it's a scalar value, we want to quote it.
switch v.MustString() {
case "true", "false":
return true
default:
_, err := parseNumber(v.MustString())
return err == nil
}
}
func parseNumber(s string) (any, error) {
if i, err := strconv.ParseInt(s, 0, 64); err == nil {
return i, nil
}
if f, err := strconv.ParseFloat(s, 64); err == nil {
return f, nil
}
return nil, fmt.Errorf("invalid number: %s", s)
}