3
0
mirror of https://github.com/ergochat/ergo.git synced 2024-11-15 00:19:29 +01:00
ergo/irc/bunt/bunt_datastore.go
2023-01-04 05:06:21 -05:00

107 lines
2.7 KiB
Go

// Copyright (c) 2022 Shivaram Lingamneni
// released under the MIT license
package bunt
import (
"fmt"
"strings"
"time"
"github.com/tidwall/buntdb"
"github.com/ergochat/ergo/irc/datastore"
"github.com/ergochat/ergo/irc/logger"
"github.com/ergochat/ergo/irc/utils"
)
// BuntKey yields a string key corresponding to a (table, UUID) pair.
// Ideally this would not be public, but some of the migration code
// needs it.
func BuntKey(table datastore.Table, uuid utils.UUID) string {
return fmt.Sprintf("%x %s", table, uuid.String())
}
// buntdbDatastore implements datastore.Datastore using a buntdb.
type buntdbDatastore struct {
db *buntdb.DB
logger *logger.Manager
}
// NewBuntdbDatastore returns a datastore.Datastore backed by buntdb.
func NewBuntdbDatastore(db *buntdb.DB, logger *logger.Manager) datastore.Datastore {
return &buntdbDatastore{
db: db,
logger: logger,
}
}
func (b *buntdbDatastore) Backoff() time.Duration {
return 0
}
func (b *buntdbDatastore) GetAll(table datastore.Table) (result []datastore.KV, err error) {
tablePrefix := fmt.Sprintf("%x ", table)
err = b.db.View(func(tx *buntdb.Tx) error {
err := tx.AscendGreaterOrEqual("", tablePrefix, func(key, value string) bool {
if !strings.HasPrefix(key, tablePrefix) {
return false
}
uuid, err := utils.DecodeUUID(strings.TrimPrefix(key, tablePrefix))
if err == nil {
result = append(result, datastore.KV{UUID: uuid, Value: []byte(value)})
} else {
b.logger.Error("datastore", "invalid uuid", key)
}
return true
})
return err
})
return
}
func (b *buntdbDatastore) Get(table datastore.Table, uuid utils.UUID) (value []byte, err error) {
buntKey := BuntKey(table, uuid)
var result string
err = b.db.View(func(tx *buntdb.Tx) error {
result, err = tx.Get(buntKey)
return err
})
return []byte(result), err
}
func (b *buntdbDatastore) Set(table datastore.Table, uuid utils.UUID, value []byte, expiration time.Time) (err error) {
buntKey := BuntKey(table, uuid)
var setOptions *buntdb.SetOptions
if !expiration.IsZero() {
ttl := time.Until(expiration)
if ttl > 0 {
setOptions = &buntdb.SetOptions{Expires: true, TTL: ttl}
} else {
return nil // it already expired, i guess?
}
}
strVal := string(value)
err = b.db.Update(func(tx *buntdb.Tx) error {
_, _, err := tx.Set(buntKey, strVal, setOptions)
return err
})
return
}
func (b *buntdbDatastore) Delete(table datastore.Table, key utils.UUID) (err error) {
buntKey := BuntKey(table, key)
err = b.db.Update(func(tx *buntdb.Tx) error {
_, err := tx.Delete(buntKey)
return err
})
// deleting a nonexistent key is not considered an error
switch err {
case buntdb.ErrNotFound:
return nil
default:
return err
}
}