save device indexes in a separate bucket

This change makes it easier to range over the device bucket.
Also added some necessary tests for this package.
This commit is contained in:
Victor Vrantchan
2017-03-19 23:01:17 +00:00
committed by Victor Vrantchan
parent 47cd6b7679
commit 9669878d40
2 changed files with 97 additions and 5 deletions

View File

@@ -11,7 +11,13 @@ import (
uuid "github.com/satori/go.uuid"
)
const DeviceBucket = "mdm.Devices"
const (
DeviceBucket = "mdm.Devices"
// The deviceIndexBucket index bucket stores serial number and UDID references
// to the device uuid.
deviceIndexBucket = "mdm.DeviceIdx"
)
type DB struct {
*bolt.DB
@@ -19,7 +25,11 @@ type DB struct {
func NewDB(db *bolt.DB, sub pubsub.Subscriber) (*DB, error) {
err := db.Update(func(tx *bolt.Tx) error {
_, err := tx.CreateBucketIfNotExists([]byte(DeviceBucket))
_, err := tx.CreateBucketIfNotExists([]byte(deviceIndexBucket))
if err != nil {
return err
}
_, err = tx.CreateBucketIfNotExists([]byte(DeviceBucket))
return err
})
if err != nil {
@@ -50,18 +60,24 @@ func (db *DB) Save(dev *Device) error {
if err != nil {
return errors.Wrap(err, "marshalling device")
}
// store an array of indices to reference the UUID, which will be the
// key used to store the actual device.
indexes := []string{dev.UDID, dev.SerialNumber}
idxBucket := tx.Bucket([]byte(deviceIndexBucket))
if idxBucket == nil {
return fmt.Errorf("bucket %q not found!", deviceIndexBucket)
}
for _, idx := range indexes {
if idx == "" {
continue
}
key := []byte(idx)
if err := bkt.Put(key, []byte(dev.UUID)); err != nil {
if err := idxBucket.Put(key, []byte(dev.UUID)); err != nil {
return errors.Wrap(err, "put device to boltdb")
}
}
key := []byte(dev.UUID)
if err := bkt.Put(key, devproto); err != nil {
return errors.Wrap(err, "put device to boltdb")
@@ -82,7 +98,8 @@ func (db *DB) DeviceByUDID(udid string) (*Device, error) {
var dev Device
err := db.View(func(tx *bolt.Tx) error {
b := tx.Bucket([]byte(DeviceBucket))
idx := b.Get([]byte(udid))
ib := tx.Bucket([]byte(deviceIndexBucket))
idx := ib.Get([]byte(udid))
if idx == nil {
return &notFound{"Device", fmt.Sprintf("udid %s", udid)}
}
@@ -102,7 +119,8 @@ func (db *DB) DeviceBySerial(serial string) (*Device, error) {
var dev Device
err := db.View(func(tx *bolt.Tx) error {
b := tx.Bucket([]byte(DeviceBucket))
idx := b.Get([]byte(serial))
ib := tx.Bucket([]byte(deviceIndexBucket))
idx := ib.Get([]byte(serial))
if idx == nil {
return &notFound{"Device", fmt.Sprintf("serial %s", serial)}
}

74
device/db_test.go Normal file
View File

@@ -0,0 +1,74 @@
package device
import (
"io/ioutil"
"os"
"testing"
"github.com/boltdb/bolt"
)
func TestSave(t *testing.T) {
db := setupDB(t)
dev := &Device{
UUID: "a-b-c-d",
UDID: "UDID-FOO-BAR-BAZ",
SerialNumber: "foobarbaz",
ProductName: "MacBook",
}
if err := db.Save(dev); err != nil {
t.Fatalf("saving device in datastore: %s", err)
}
byUDID, err := db.DeviceByUDID(dev.UDID)
if err != nil {
t.Fatalf("getting device by UDID: %s", err)
}
bySerial, err := db.DeviceBySerial(dev.SerialNumber)
if err != nil {
t.Fatalf("getting device by UDID: %s", err)
}
// test helper that verifies that the retrieved device is the same
tf := func(haveDev *Device) func(t *testing.T) {
return func(t *testing.T) {
if have, want := haveDev.UDID, dev.UDID; have != want {
t.Errorf("have %s, want %s", have, want)
}
if have, want := haveDev.UUID, dev.UUID; have != want {
t.Errorf("have %s, want %s", have, want)
}
if have, want := haveDev.SerialNumber, dev.SerialNumber; have != want {
t.Errorf("have %s, want %s", have, want)
}
if have, want := haveDev.ProductName, dev.ProductName; have != want {
t.Errorf("have %s, want %s", have, want)
}
}
}
t.Run("byUDID", tf(byUDID))
t.Run("bySerial", tf(bySerial))
}
func setupDB(t *testing.T) *DB {
f, _ := ioutil.TempFile("", "bolt-")
f.Close()
os.Remove(f.Name())
db, err := bolt.Open(f.Name(), 0777, nil)
if err != nil {
t.Fatalf("couldn't open bolt, err %s\n", err)
}
devDB, err := NewDB(db, nil)
if err != nil {
t.Fatalf("couldn't create device DB, err %s\n", err)
}
return devDB
}