speedtest-go/database/bolt/bolt.go

83 lines
1.7 KiB
Go
Raw Normal View History

package bolt
import (
"encoding/json"
"errors"
"time"
"github.com/librespeed/speedtest/database/schema"
log "github.com/sirupsen/logrus"
"go.etcd.io/bbolt"
)
const (
bucketName = `speedtest`
)
type Bolt struct {
db *bbolt.DB
}
func Open(databaseFile string) *Bolt {
db, err := bbolt.Open(databaseFile, 0666, nil)
if err != nil {
log.Fatalf("Cannot open BoltDB database file: %s", err)
}
return &Bolt{db: db}
}
func (p *Bolt) Insert(data *schema.TelemetryData) error {
return p.db.Update(func(tx *bbolt.Tx) error {
data.Timestamp = time.Now()
b, _ := json.Marshal(data)
bucket, err := tx.CreateBucketIfNotExists([]byte(bucketName))
if err != nil {
return err
}
return bucket.Put([]byte(data.UUID), b)
})
}
func (p *Bolt) FetchByUUID(uuid string) (*schema.TelemetryData, error) {
var record schema.TelemetryData
err := p.db.View(func(tx *bbolt.Tx) error {
bucket := tx.Bucket([]byte(bucketName))
if bucket == nil {
return errors.New("data bucket doesn't exist yet")
}
b := bucket.Get([]byte(uuid))
return json.Unmarshal(b, &record)
})
return &record, err
}
func (p *Bolt) FetchLast100() ([]schema.TelemetryData, error) {
var records []schema.TelemetryData
err := p.db.View(func(tx *bbolt.Tx) error {
var record schema.TelemetryData
bucket := tx.Bucket([]byte(bucketName))
if bucket == nil {
return errors.New("data bucket doesn't exist yet")
}
cursor := bucket.Cursor()
_, b := cursor.Last()
for len(records) < 100 {
if err := json.Unmarshal(b, &record); err != nil {
return err
}
records = append(records, record)
_, b = cursor.Prev()
if b == nil {
break
}
}
return nil
})
return records, err
}