pastebin/store/boltdb.go

package store

import (
	"encoding/json"
	"time"

	"go.etcd.io/bbolt"
)

var (
	pastesBucket = []byte("pastes")
	hashesBucket = []byte("hashes")
)

type BoltStore struct {
	db *bbolt.DB
}

func NewBoltStore(path string) (*BoltStore, error) {
	db, err := bbolt.Open(path, 0600, nil)
	if err != nil {
		return nil, err
	}

	err = db.Update(func(tx *bbolt.Tx) error {
		if _, err := tx.CreateBucketIfNotExists(pastesBucket); err != nil {
			return err
		}
		_, err := tx.CreateBucketIfNotExists(hashesBucket)
		return err
	})
	if err != nil {
		return nil, err
	}

	return &BoltStore{db: db}, nil
}

func (s *BoltStore) Get(id string) (*Paste, bool, error) {
	var paste Paste
	exists := false
	err := s.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket(pastesBucket)
		val := b.Get([]byte(id))
		if val == nil {
			return nil
		}
		exists = true
		return json.Unmarshal(val, &paste)
	})
	if err != nil {
		return nil, false, err
	}
	return &paste, exists, nil
}

func (s *BoltStore) GetIDByHash(hash string) (string, bool, error) {
	var id string
	exists := false
	err := s.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket(hashesBucket)
		val := b.Get([]byte(hash))
		if val == nil {
			return nil
		}
		exists = true
		id = string(val)
		return nil
	})
	return id, exists, err
}

func (s *BoltStore) Set(id, hash, content string) error {
	return s.db.Update(func(tx *bbolt.Tx) error {
		pb := tx.Bucket(pastesBucket)
		hb := tx.Bucket(hashesBucket)

		paste := Paste{
			Content:   content,
			CreatedAt: time.Now(),
		}
		encoded, err := json.Marshal(paste)
		if err != nil {
			return err
		}

		if err := pb.Put([]byte(id), encoded); err != nil {
			return err
		}
		return hb.Put([]byte(hash), []byte(id))
	})
}

func (s *BoltStore) Del(id string) error {
	return s.db.Update(func(tx *bbolt.Tx) error {
		return tx.Bucket(pastesBucket).Delete([]byte(id))
	})
}

func (s *BoltStore) Cleanup(maxAge time.Duration) {
	var keysToDelete [][]byte
	s.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket(pastesBucket)
		c := b.Cursor()
		for k, v := c.First(); k != nil; k, v = c.Next() {
			var p Paste
			if err := json.Unmarshal(v, &p); err == nil {
				if time.Since(p.CreatedAt) > maxAge {
					keysToDelete = append(keysToDelete, k)
				}
			}
		}
		return nil
	})

	if len(keysToDelete) > 0 {
		s.db.Update(func(tx *bbolt.Tx) error {
			b := tx.Bucket(pastesBucket)
			for _, k := range keysToDelete {
				b.Delete(k)
			}
			return nil
		})
	}
}

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