1package database
  2
  3import (
  4	"bytes"
  5	"errors"
  6	"fmt"
  7	"math"
  8	"path"
  9	"strconv"
 10	"time"
 11
 12	"git.icyphox.sh/paprika/config"
 13	"github.com/dgraph-io/badger/v3"
 14)
 15
 16// Use this as a global DB handle.
 17var DB database
 18
 19type database struct {
 20	*badger.DB
 21}
 22
 23// see: https://dgraph.io/docs/badger/get-started/#garbage-collection
 24func gc() {
 25	ticker := time.NewTicker(5 * time.Minute)
 26	defer ticker.Stop()
 27	for range ticker.C {
 28	again:
 29		err := DB.DB.RunValueLogGC(0.7)
 30		if err == nil {
 31			goto again
 32		}
 33	}
 34	panic("Unreachable!")
 35}
 36
 37func Open() (*badger.DB, error) {
 38	db, err := badger.Open(
 39		badger.DefaultOptions(path.Join(config.C.DbPath, "badger")),
 40	)
 41	if err != nil {
 42		return nil, err
 43	}
 44	go gc()
 45	return db, nil
 46}
 47
 48// Wrapper function to simplify setting a key/val
 49// in badger.
 50func (d *database) Set(key, val []byte) error {
 51	err := d.Update(func(txn *badger.Txn) error {
 52		err := txn.Set(key, val)
 53		return err
 54	})
 55
 56	return err
 57}
 58
 59// Wrapper function to simplify setting a key/val with a duration.
 60// see Set
 61func (d *database) SetWithTTL(key, val []byte, duration time.Duration) error {
 62	err := d.Update(func(txn *badger.Txn) error {
 63		err := txn.SetEntry(badger.NewEntry(key, val).WithTTL(duration))
 64		return err
 65	})
 66
 67	return err
 68}
 69
 70// Wrapper function to simplify getting a key from badger.
 71func (d *database) Get(key []byte) ([]byte, error) {
 72	var val []byte
 73	err := d.View(func(txn *badger.Txn) error {
 74		item, err := txn.Get(key)
 75		if err != nil {
 76			return err
 77		}
 78
 79		err = item.Value(func(v []byte) error {
 80			val = v
 81			return nil
 82		})
 83		return err
 84	})
 85	if err != nil {
 86		return nil, err
 87	}
 88	return val, nil
 89}
 90
 91type KVPair struct {
 92	a, b []byte
 93}
 94
 95// Wrapper function to simplify getting a range of keys with a given prefix.
 96func (d *database) GetRange(prefix []byte) ([]KVPair, error) {
 97	var vals []KVPair
 98	err := d.View(func(txn *badger.Txn) error {
 99		iter := txn.NewIterator(badger.DefaultIteratorOptions)
100		defer iter.Close()
101
102		for iter.Seek(prefix); iter.ValidForPrefix(prefix); iter.Next() {
103			item := iter.Item()
104
105			var pair KVPair
106			key := item.Key()
107			pair.a = key
108			err := item.Value(func(val []byte) error {
109				pair.b = val
110				return nil
111			})
112
113			if err != nil {
114				return err
115			}
116
117			vals = append(vals, pair)
118		}
119
120		return nil
121	})
122	if err != nil {
123		return nil, err
124	}
125	return vals, nil
126}
127
128// Key delete wrapper. Returns deleted Value
129func (d *database) Delete(key []byte) ([]byte, error) {
130	var val []byte
131    err := d.DB.Update(func(txn *badger.Txn) error {
132		item, err := txn.Get(key)
133		if err != nil {
134			return err
135		}
136		err = item.Value(func(v []byte) error {
137			val = v
138			return nil
139		})
140		return txn.Delete(key)
141	})
142	return val, err
143}
144
145var NumTooBig = errors.New("Number Too Big")
146var InvalidNumber = errors.New("Invalid Number")
147
148// encode number so it sorts lexicographically, while being semi readable.
149func EncodeNumber(n int) ([]byte, error) {
150	neg := false
151	num := n
152	if n < 0 {
153		neg = true
154		num = -num
155	}
156
157	digits := int(math.Trunc(math.Log10(float64(num))))+1
158	if digits > 93 {
159		return []byte{}, NumTooBig
160	}
161
162	if !neg {
163		lenCode := 33 + digits
164		return []byte(fmt.Sprintf("%c %d", lenCode, n)), nil
165	} else {
166		lenCode := 127 - digits
167		return []byte(fmt.Sprintf("!%c %d", lenCode, n)), nil
168	}
169}
170
171func ToKey(prefix, key string) []byte {
172	return []byte(fmt.Sprintf("%s/%s", prefix, key))
173}
174
175// encode number so it sorts lexicographically, while being semi readable.
176func DecodeNumber(n []byte) (int, error) {
177	if len(n) < 3 {
178		return 0, InvalidNumber
179	}
180
181	// No digit padding
182	if n[0] < 33 || n[0] > 126 {
183		return 0, InvalidNumber
184	}
185
186	num := bytes.SplitN(n, []byte{' '}, 2)
187	if len(num) != 2 {
188		return 0, InvalidNumber
189	}
190
191	number, err := strconv.Atoi(string(num[1]))
192	if err != nil {
193		return number, err
194	}
195
196	if n[0] == '!' {
197		return -number, nil
198	} else {
199		return number, nil
200	}
201}