storage/yaml.go

package storage

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"local/storage/resolve"
	"os"
	"path"
	"path/filepath"
	"strings"
	"unicode"

	yaml "gopkg.in/yaml.v2"
)

const (
	yamlExt = ".yaml"
)

type Yaml struct {
	path string
}

func NewYaml(p string) (*Yaml, error) {
	_, err := os.Stat(path.Dir(p))
	if err != nil {
		return nil, err
	}
	p, err = filepath.Abs(p)
	if err != nil {
		return nil, err
	}
	return &Yaml{
		path: p,
	}, os.MkdirAll(path.Dir(p), os.ModePerm)
}

func (y *Yaml) Namespaces() ([][]string, error) {
	m, err := y.getMap()
	if err != nil {
		return nil, err
	}
	return keysDFS(m)
}

func (y *Yaml) List(ns []string, limits ...string) ([]string, error) {
	namespace := resolve.Namespace(ns)
	m, err := y.getMap(namespace)
	if err != nil {
		return nil, err
	}
	limits = resolve.Limits(limits)
	ks := make([]string, 0, len(m))
	for k := range m {
		if k >= limits[0] && k <= limits[1] {
			ks = append(ks, k)
		}
	}
	return ks, nil
}

func (y *Yaml) Get(key string, ns ...string) ([]byte, error) {
	r, err := y.GetStream(key, ns...)
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(r)
}

func (y *Yaml) GetStream(key string, ns ...string) (io.Reader, error) {
	namespace := resolve.Namespace(ns)
	m, err := y.getMap(namespace)
	if err != nil {
		return nil, err
	}
	v, ok := m[key]
	if !ok {
		return nil, ErrNotFound
	}
	s, ok := v.(string)
	if !ok {
		return nil, ErrNotFound
	}
	if strings.HasPrefix(s, "b64://") {
		b, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(s, "b64://"))
		return bytes.NewReader(b), err
	}
	return strings.NewReader(s), nil
}

func (y *Yaml) Set(key string, value []byte, ns ...string) error {
	r := bytes.NewReader(value)
	if value == nil {
		return y.Del(key, ns...)
	}
	return y.SetStream(key, r, ns...)
}

func (y *Yaml) Del(key string, ns ...string) error {
	return y.SetStream(key, nil, ns...)
}

func isASCII(s string) bool {
	for _, c := range s {
		if c > unicode.MaxASCII {
			return false
		}
	}
	return true
}

func (y *Yaml) SetStream(key string, r io.Reader, ns ...string) error {
	namespace := resolve.Namespace(ns)
	var v interface{} = nil
	if r != nil {
		b, err := ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		if isASCII(string(b)) {
			v = string(b)
		} else {
			v = "b64://" + base64.StdEncoding.EncodeToString(b)
		}
	}
	m, err := y.getMap()
	if err != nil {
		return err
	}
	if err := setInMap(m, []string{namespace}, key, v); err != nil {
		return err
	}
	return y.setMap(m)
}

func (y *Yaml) Close() error {
	return nil
}

func (y *Yaml) getMap(keys ...string) (map[string]interface{}, error) {
	b, err := y.get()
	if err != nil {
		return nil, err
	}
	var mBad map[interface{}]interface{}
	if err := yaml.Unmarshal(b, &mBad); err != nil {
		return nil, err
	}
	m, err := mbadToM(mBad)
	if err != nil {
		return nil, err
	}
	if m == nil {
		m = map[string]interface{}{}
	}
	for _, k := range keys {
		subv, ok := m[k]
		if !ok {
			subv = map[string]interface{}{}
			m[k] = subv
		}
		subm, ok := subv.(map[string]interface{})
		if !ok {
			return nil, ErrNotFound
		}
		m = subm
	}
	return m, err
}

func (y *Yaml) setMap(m map[string]interface{}) error {
	b, err := yaml.Marshal(m)
	if err != nil {
		return err
	}
	return y.set(b)
}

func setInMap(m map[string]interface{}, keys []string, key string, v interface{}) error {
	if len(keys) == 0 {
		m[key] = v
		if v == nil {
			delete(m, key)
		}
		return nil
	}
	subv, ok := m[keys[0]]
	if !ok {
		subv = map[string]interface{}{}
	}
	subm, ok := subv.(map[string]interface{})
	if !ok {
		return errors.New("clobber")
	}
	if err := setInMap(subm, keys[1:], key, v); err != nil {
		return err
	}
	m[keys[0]] = subm
	return nil
}

func (y *Yaml) get() ([]byte, error) {
	b, err := ioutil.ReadFile(y.path)
	if os.IsNotExist(err) {
		return []byte{}, nil
	}
	return b, err
}

func (y *Yaml) set(b []byte) error {
	return ioutil.WriteFile(y.path, b, os.ModePerm)
}

func keysDFS(m map[string]interface{}) ([][]string, error) {
	keys, _, err := _keysDFS(m)
	return keys, err
}

func _keysDFS(m map[string]interface{}) ([][]string, bool, error) {
	keys := make([][]string, 0)
	hasNonMaps := false
	for k, v := range m {
		if subm, ok := v.(map[string]interface{}); ok {
			subkeys, hasElements, err := _keysDFS(subm)
			if err != nil {
				return nil, false, err
			}
			if hasElements {
				subkeys = append(subkeys, []string{})
			}
			for i := range subkeys {
				subkeys[i] = append([]string{k}, subkeys[i]...)
				keys = append(keys, subkeys[i])
			}
		} else {
			hasNonMaps = true
		}
	}
	return keys, hasNonMaps, nil
}

func mbadToM(mBad map[interface{}]interface{}) (map[string]interface{}, error) {
	m := map[string]interface{}{}
	for k, v := range mBad {
		s, ok := k.(string)
		if !ok {
			s = fmt.Sprint(k)
		}
		m[s] = v
		if m2, ok := v.(map[interface{}]interface{}); ok {
			v2, err := mbadToM(m2)
			if err != nil {
				return nil, err
			}
			m[s] = v2
		}
	}
	return m, nil
}