reveil/model/database.go

package reveil

import (
	"encoding/json"
	"fmt"
	"log"
	"math/rand"

	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/errors"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/util"
)

type LevelDBStorage struct {
	db *leveldb.DB
}

// NewMySQLStorage establishes the connection to the database
func NewLevelDBStorage(path string) (s *LevelDBStorage, err error) {
	var db *leveldb.DB

	db, err = leveldb.OpenFile(path, nil)
	if err != nil {
		if _, ok := err.(*errors.ErrCorrupted); ok {
			log.Printf("LevelDB was corrupted; attempting recovery (%s)", err.Error())
			_, err = leveldb.RecoverFile(path, nil)
			if err != nil {
				return
			}
			log.Println("LevelDB recovery succeeded!")
		} else {
			return
		}
	}

	s = &LevelDBStorage{db}
	return
}

func (s *LevelDBStorage) Close() error {
	return s.db.Close()
}

func decodeData(data []byte, v interface{}) error {
	return json.Unmarshal(data, v)
}

func (s *LevelDBStorage) get(key string, v interface{}) error {
	data, err := s.db.Get([]byte(key), nil)
	if err != nil {
		return err
	}

	return decodeData(data, v)
}

func (s *LevelDBStorage) put(key string, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}

	return s.db.Put([]byte(key), data, nil)
}

func (s *LevelDBStorage) findBytesKey(prefix string, len int) (key string, id Identifier, err error) {
	id = make([]byte, len)
	found := true
	for found {
		if _, err = rand.Read(id); err != nil {
			return
		}
		key = fmt.Sprintf("%s%s", prefix, id.ToString())

		found, err = s.db.Has([]byte(key), nil)
		if err != nil {
			return
		}
	}
	return
}

func (s *LevelDBStorage) delete(key string) error {
	return s.db.Delete([]byte(key), nil)
}

func (s *LevelDBStorage) search(prefix string) iterator.Iterator {
	return s.db.NewIterator(util.BytesPrefix([]byte(prefix)), nil)
}