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) }