mirror of
https://github.com/micromdm/micromdm/
synced 2026-08-11 20:45:32 +08:00
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:
committed by
Victor Vrantchan
parent
47cd6b7679
commit
9669878d40
28
device/db.go
28
device/db.go
@@ -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 ¬Found{"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 ¬Found{"Device", fmt.Sprintf("serial %s", serial)}
|
||||
}
|
||||
|
||||
74
device/db_test.go
Normal file
74
device/db_test.go
Normal 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
|
||||
}
|
||||
Reference in New Issue
Block a user