2022-11-30 18:15:47 +00:00
|
|
|
package database
|
2022-02-08 14:13:45 +00:00
|
|
|
|
|
|
|
import (
|
|
|
|
"fmt"
|
|
|
|
"io"
|
|
|
|
"os"
|
|
|
|
"path/filepath"
|
|
|
|
|
2023-11-10 20:36:15 +00:00
|
|
|
"git.sr.ht/~ashkeel/strimertul/utils"
|
2023-02-23 09:50:52 +00:00
|
|
|
|
2023-11-11 14:13:25 +00:00
|
|
|
pebble_driver "git.sr.ht/~ashkeel/kilovolt-driver-pebble"
|
2022-02-08 14:13:45 +00:00
|
|
|
"github.com/cockroachdb/pebble"
|
2023-11-03 15:55:57 +00:00
|
|
|
kv "github.com/strimertul/kilovolt/v11"
|
2022-12-04 13:45:34 +00:00
|
|
|
"go.uber.org/zap"
|
2022-02-08 14:13:45 +00:00
|
|
|
)
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
type PebbleDatabase struct {
|
2022-11-30 18:15:47 +00:00
|
|
|
db *pebble.DB
|
|
|
|
hub *kv.Hub
|
|
|
|
logger *zap.Logger
|
2022-11-18 16:37:30 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
// NewPebble creates a new database driver instance with an underlying Pebble database
|
2022-11-30 18:15:47 +00:00
|
|
|
func NewPebble(directory string, logger *zap.Logger) (*PebbleDatabase, error) {
|
2022-11-16 11:23:54 +00:00
|
|
|
db, err := pebble.Open(directory, &pebble.Options{})
|
|
|
|
if err != nil {
|
2022-11-18 16:37:30 +00:00
|
|
|
return nil, fmt.Errorf("could not open DB: %w", err)
|
2022-11-16 11:23:54 +00:00
|
|
|
}
|
2022-02-08 14:13:45 +00:00
|
|
|
|
|
|
|
// Create file for autodetect
|
2022-11-18 16:37:30 +00:00
|
|
|
err = os.WriteFile(filepath.Join(directory, "stul-driver"), []byte("pebble"), 0o644)
|
2022-11-16 11:23:54 +00:00
|
|
|
if err != nil {
|
2022-11-18 16:37:30 +00:00
|
|
|
return nil, fmt.Errorf("could not write driver file: %w", err)
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
p := &PebbleDatabase{
|
2022-11-30 18:15:47 +00:00
|
|
|
db: db,
|
|
|
|
hub: nil,
|
|
|
|
logger: logger,
|
2022-11-18 16:37:30 +00:00
|
|
|
}
|
2022-02-08 14:13:45 +00:00
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
return p, nil
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
func (p *PebbleDatabase) Hub() *kv.Hub {
|
|
|
|
if p.hub == nil {
|
2022-11-30 18:15:47 +00:00
|
|
|
p.hub, _ = kv.NewHub(pebble_driver.NewPebbleBackend(p.db, true), kv.HubOptions{}, p.logger)
|
2022-11-16 11:23:54 +00:00
|
|
|
}
|
2022-11-18 16:37:30 +00:00
|
|
|
return p.hub
|
|
|
|
}
|
2022-11-16 11:23:54 +00:00
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
func (p *PebbleDatabase) Close() error {
|
2023-04-14 18:04:21 +00:00
|
|
|
if p.hub != nil {
|
|
|
|
p.hub.Close()
|
|
|
|
}
|
2022-11-18 16:37:30 +00:00
|
|
|
err := p.db.Close()
|
2022-11-16 11:23:54 +00:00
|
|
|
if err != nil {
|
2023-02-23 09:50:52 +00:00
|
|
|
return fmt.Errorf("could not close database: %w", err)
|
2022-11-16 11:23:54 +00:00
|
|
|
}
|
|
|
|
return nil
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
func (p *PebbleDatabase) Import(entries map[string]string) error {
|
|
|
|
batch := p.db.NewBatch()
|
|
|
|
for key, value := range entries {
|
2022-12-02 22:52:45 +00:00
|
|
|
err := batch.Set([]byte(key), []byte(value), &pebble.WriteOptions{})
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
2022-11-18 16:37:30 +00:00
|
|
|
return batch.Commit(&pebble.WriteOptions{})
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
func (p *PebbleDatabase) Export(file io.Writer) error {
|
|
|
|
return p.Backup(file)
|
|
|
|
}
|
2022-11-16 11:23:54 +00:00
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
func (p *PebbleDatabase) Restore(file io.Reader) error {
|
2022-02-08 14:13:45 +00:00
|
|
|
in := make(map[string]string)
|
2022-11-23 15:34:49 +00:00
|
|
|
err := json.NewDecoder(file).Decode(&in)
|
2022-02-08 14:13:45 +00:00
|
|
|
if err != nil {
|
2023-04-23 00:26:46 +00:00
|
|
|
return fmt.Errorf("could not decode backup: %w", err)
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
b := p.db.NewBatch()
|
2022-02-08 14:13:45 +00:00
|
|
|
for k, v := range in {
|
|
|
|
err = b.Set([]byte(k), []byte(v), nil)
|
|
|
|
if err != nil {
|
2023-04-23 00:26:46 +00:00
|
|
|
return fmt.Errorf("could not set key %s: %w", k, err)
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
return b.Commit(&pebble.WriteOptions{Sync: true})
|
2022-11-16 11:23:54 +00:00
|
|
|
}
|
|
|
|
|
2022-11-18 16:37:30 +00:00
|
|
|
func (p *PebbleDatabase) Backup(file io.Writer) error {
|
2023-02-23 09:50:52 +00:00
|
|
|
snapshot := p.db.NewSnapshot()
|
|
|
|
defer utils.Close(snapshot, p.logger)
|
|
|
|
|
2023-11-05 11:45:53 +00:00
|
|
|
iter, err := snapshot.NewIter(&pebble.IterOptions{})
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2023-02-23 09:50:52 +00:00
|
|
|
defer utils.Close(iter, p.logger)
|
|
|
|
|
2022-11-16 11:23:54 +00:00
|
|
|
out := make(map[string]string)
|
|
|
|
for iter.First(); iter.Valid(); iter.Next() {
|
2023-02-15 12:46:56 +00:00
|
|
|
val, err := iter.ValueAndErr()
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
out[string(iter.Key())] = string(val)
|
2022-11-16 11:23:54 +00:00
|
|
|
}
|
2022-11-23 15:34:49 +00:00
|
|
|
return json.NewEncoder(file).Encode(out)
|
2022-02-08 14:13:45 +00:00
|
|
|
}
|