go-mockid/kvdb/kvdb.go

package kvdb

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

type KVDB struct {
	Prefix string
	Ext    string
}

func (kv *KVDB) Load(
	keyif interface{},
	typ ...interface{},
) (value interface{}, ok bool, err error) {
	key, _ := keyif.(string)
	if "" == key || strings.Contains(key, "..") || strings.ContainsAny(key, "$#!:| \n") {
		return nil, false, nil
	}

	userFile := filepath.Join(kv.Prefix, key+"."+kv.Ext)
	fmt.Println("Debug user file:", userFile)
	b, err := ioutil.ReadFile(userFile)
	if nil != err {
		if os.IsNotExist(err) {
			return nil, false, nil
		}
		fmt.Println("kvdb debug read:", err)
		return nil, false, errors.New("database read failed")
	}

	ok = true
	value = b
	if 1 == len(typ) {
		err := json.Unmarshal(b, typ[0])
		if nil != err {
			return nil, false, err
		}
		value = typ[0]
	} else if len(b) > 0 && '"' == b[0] {
		var str string
		err := json.Unmarshal(b, &str)
		if nil == err {
			value = str
		}
	}

	return value, ok, nil
}

func (kv *KVDB) Store(keyif interface{}, value interface{}) (err error) {
	key, _ := keyif.(string)
	if "" == key || strings.Contains(key, "..") || strings.ContainsAny(key, "$#! \n") {
		return errors.New("invalid key name")
	}

	keypath := filepath.Join(kv.Prefix, key+"."+kv.Ext)
	f, err := os.Open(keypath)
	if nil == err {
		s, err := f.Stat()
		if nil != err {
			// if we can open, we should be able to stat
			return errors.New("database connection failure")
		}
		ts := strconv.FormatInt(s.ModTime().Unix(), 10)
		bakpath := filepath.Join(kv.Prefix, key+"."+ts+"."+kv.Ext)
		if err := os.Rename(keypath, bakpath); nil != err {
			// keep the old record as a backup
			return errors.New("database write failure")
		}
	}

	var b []byte
	switch v := value.(type) {
	case []byte:
		b = v
	case string:
		b, _ = json.Marshal(v)
	default:
		fmt.Println("kvdb: not []byte or string:", v)
		jsonb, err := json.Marshal(v)
		if nil != err {
			return err
		}
		b = jsonb
	}

	if err := ioutil.WriteFile(
		keypath,
		b,
		os.FileMode(0600),
	); nil != err {
		fmt.Println("write failure:", err)
		return errors.New("database write failed")
	}

	return nil
}

func (kv *KVDB) Delete(keyif interface{}) (err error) {
	key, _ := keyif.(string)
	if "" == key || strings.Contains(key, "..") || strings.ContainsAny(key, "$#! \n") {
		return errors.New("invalid key name")
	}

	keypath := filepath.Join(kv.Prefix, key+"."+kv.Ext)
	f, err := os.Open(keypath)
	if nil == err {
		s, err := f.Stat()
		if nil != err {
			return errors.New("database connection failure")
		}
		ts := strconv.FormatInt(s.ModTime().Unix(), 64)
		if err := os.Rename(keypath, filepath.Join(kv.Prefix, key+"."+ts+"."+kv.Ext)); nil != err {
			return errors.New("database connection failure")
		}
	}

	return nil
}

func (kv *KVDB) Vacuum() (err error) {
	return nil
}