kv backend store for inventory and profile subsystems (#61)

subsystem/inventory/storage/diskv/diskv.go

@@ -1,87 +1,27 @@
-// Package diskv implements a diskv-backed inventory subsystem storage backend.
+// Package diskv implements an inventory subsystem backend using diskv.
 package diskv
 
 import (
-	"context"
-	"encoding/json"
-	"fmt"
 	"path/filepath"
 
-	"github.com/micromdm/nanocmd/subsystem/inventory/storage"
+	"github.com/micromdm/nanocmd/subsystem/inventory/storage/kv"
+
+	"github.com/micromdm/nanolib/storage/kv/kvdiskv"
 	"github.com/peterbourgon/diskv/v3"
 )
 
-// Diskv is an on-disk enrollment inventory data store.
+// Diskv is an inventory subsystem backend which uses diskv as the key-value store.
 type Diskv struct {
-	diskv *diskv.Diskv
+	*kv.KV
 }
 
-// New creates a new initialized inventory data store.
+// New creates a new profile store at on disk at path.
 func New(path string) *Diskv {
-	flatTransform := func(s string) []string { return []string{} }
 	return &Diskv{
-		diskv: diskv.New(diskv.Options{
+		KV: kv.New(kvdiskv.New(diskv.New(diskv.Options{
 			BasePath:     filepath.Join(path, "inventory"),
-			Transform:    flatTransform,
+			Transform:    kvdiskv.FlatTransform,
 			CacheSizeMax: 1024 * 1024,
-		}),
-	}
-}
-
-// RetrieveInventory retrieves the inventory data for enrollment IDs.
-func (s *Diskv) RetrieveInventory(ctx context.Context, opt *storage.SearchOptions) (map[string]storage.Values, error) {
-	ret := make(map[string]storage.Values)
-	for _, id := range opt.IDs {
-		if !s.diskv.Has(id) {
-			continue
-		}
-		raw, err := s.diskv.Read(id)
-		if err != nil {
-			return ret, fmt.Errorf("reading values for %s: %w", id, err)
-		}
-		var vals storage.Values
-		if err = json.Unmarshal(raw, &vals); err != nil {
-			return ret, fmt.Errorf("unmarshal values for %s: %w", id, err)
-		}
-		ret[id] = vals
-	}
-	return ret, nil
-}
-
-// StoreInventoryValues stores inventory data about the specified ID.
-func (s *Diskv) StoreInventoryValues(ctx context.Context, id string, values storage.Values) error {
-	var err error
-	var raw []byte
-	var vals storage.Values
-	if s.diskv.Has(id) {
-		// this is likely race-prone as we perform a read-process-write on the same key.
-		if raw, err = s.diskv.Read(id); err != nil {
-			return fmt.Errorf("reading values: %w", err)
-		}
-		if len(raw) > 0 {
-			if err = json.Unmarshal(raw, &vals); err != nil {
-				return fmt.Errorf("unmarshal values: %w", err)
-			}
-			if vals != nil {
-				for k := range values {
-					vals[k] = values[k]
-				}
-			}
-		}
-	}
-	if vals == nil {
-		vals = values
-	}
-	if raw, err = json.Marshal(vals); err != nil {
-		return fmt.Errorf("marshal values: %w", err)
+		}))),
 	}
-	if err = s.diskv.Write(id, raw); err != nil {
-		return fmt.Errorf("write values: %w", err)
-	}
-	return nil
-}
-
-// DeleteInventory deletes all inventory data for an enrollment ID.
-func (s *Diskv) DeleteInventory(ctx context.Context, id string) error {
-	return s.diskv.Erase(id)
 }

subsystem/inventory/storage/diskv/diskv_test.go

@@ -1,14 +1,12 @@
 package diskv
 
 import (
-	"os"
 	"testing"
 
 	"github.com/micromdm/nanocmd/subsystem/inventory/storage"
 	"github.com/micromdm/nanocmd/subsystem/inventory/storage/test"
 )
 
 func TestDiskv(t *testing.T) {
-	test.TestStorage(t, func() storage.Storage { return New("teststor") })
-	os.RemoveAll("teststor")
+	test.TestStorage(t, func() storage.Storage { return New(t.TempDir()) })
 }

subsystem/inventory/storage/inmem/inmem.go

@@ -2,60 +2,17 @@
 package inmem
 
 import (
-	"context"
-	"sync"
+	"github.com/micromdm/nanocmd/subsystem/inventory/storage/kv"
 
-	"github.com/micromdm/nanocmd/subsystem/inventory/storage"
+	"github.com/micromdm/nanolib/storage/kv/kvmap"
 )
 
-// InMem represents the in-memory enrollment inventory data store.
+// InMem is an in-memory inventory subsystem storage system backend.
 type InMem struct {
-	mu  sync.RWMutex
-	inv map[string]storage.Values
+	*kv.KV
 }
 
-// New creates a new initialized inventory data store.
+// New creates a new inventory subsystem storage system backend.
 func New() *InMem {
-	return &InMem{inv: make(map[string]storage.Values)}
-}
-
-// RetrieveInventory retrieves the inventory data for enrollment IDs.
-func (s *InMem) RetrieveInventory(ctx context.Context, opt *storage.SearchOptions) (map[string]storage.Values, error) {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
-	if opt == nil || len(opt.IDs) <= 0 {
-		return nil, nil
-	}
-	ret := make(map[string]storage.Values)
-	for _, id := range opt.IDs {
-		if vals, ok := s.inv[id]; ok {
-			ret[id] = make(storage.Values)
-			for k, v := range vals {
-				ret[id][k] = v
-			}
-		}
-	}
-	return ret, nil
-}
-
-// StoreInventoryValues stores inventory data about the specified ID.
-func (s *InMem) StoreInventoryValues(ctx context.Context, id string, values storage.Values) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if s.inv[id] == nil {
-		s.inv[id] = values
-	} else {
-		for k, v := range values {
-			s.inv[id][k] = v
-		}
-	}
-	return nil
-}
-
-// DeleteInventory deletes all inventory data for an enrollment ID.
-func (s *InMem) DeleteInventory(ctx context.Context, id string) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	delete(s.inv, id)
-	return nil
+	return &InMem{KV: kv.New(kvmap.New())}
 }

subsystem/inventory/storage/kv/kv.go

@@ -0,0 +1,101 @@
+// Package kv implements an inventory subsystem storage backend using a key-value store.
+package kv
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/micromdm/nanocmd/subsystem/inventory/storage"
+
+	"github.com/micromdm/nanolib/storage/kv"
+)
+
+// KV is an inventory subsystem storage backend using a key-value store.
+type KV struct {
+	mu sync.RWMutex
+	b  kv.KeysPrefixTraversingBucket
+}
+
+// New creates a new inventory subsystem backend.
+func New(b kv.KeysPrefixTraversingBucket) *KV {
+	return &KV{b: b}
+}
+
+// RetrieveInventory queries and returns the inventory values by mapped
+// by enrollment ID from the key-value store. Must provide opt and IDs.
+func (s *KV) RetrieveInventory(ctx context.Context, opt *storage.SearchOptions) (map[string]storage.Values, error) {
+	if opt == nil || len(opt.IDs) < 1 {
+		return nil, storage.ErrNoIDs
+	}
+
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	r := make(map[string]storage.Values)
+	for _, id := range opt.IDs {
+		jsonValues, err := s.b.Get(ctx, id)
+		if errors.Is(err, kv.ErrKeyNotFound) {
+			continue
+		} else if err != nil {
+			return r, fmt.Errorf("getting values for %s: %w", id, err)
+		}
+
+		var values storage.Values
+		if err = json.Unmarshal(jsonValues, &values); err != nil {
+			return r, fmt.Errorf("unmarshal values for %s: %w", id, err)
+		}
+		r[id] = values
+	}
+	return r, nil
+}
+
+// StoreInventoryValues stores inventory data about the specified ID.
+func (s *KV) StoreInventoryValues(ctx context.Context, id string, newValues storage.Values) error {
+	if id == "" {
+		return storage.ErrNoIDs
+	}
+	if len(newValues) == 0 {
+		return nil
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	jsonValues, err := s.b.Get(ctx, id)
+	if err != nil && !errors.Is(err, kv.ErrKeyNotFound) {
+		return fmt.Errorf("get values: %w", err)
+	}
+
+	var values storage.Values
+	if len(jsonValues) < 1 {
+		values = newValues
+	} else {
+		// load existing values
+		if err = json.Unmarshal(jsonValues, &values); err != nil {
+			return fmt.Errorf("unmarshal values: %w", err)
+		}
+
+		// merge the new values in
+		for k := range newValues {
+			values[k] = newValues[k]
+		}
+	}
+
+	if jsonValues, err = json.Marshal(&values); err != nil {
+		return fmt.Errorf("marshal values: %w", err)
+	}
+
+	if err = s.b.Set(ctx, id, jsonValues); err != nil {
+		return fmt.Errorf("set values: %w", err)
+	}
+
+	return nil
+}
+
+// DeleteInventory deletes all inventory data for an enrollment ID.
+func (s *KV) DeleteInventory(ctx context.Context, id string) error {
+	return s.b.Delete(ctx, id)
+}

subsystem/inventory/storage/storage.go

@@ -3,6 +3,11 @@ package storage
 
 import (
 	"context"
+	"errors"
+)
+
+var (
+	ErrNoIDs = errors.New("no ids supplied")
 )
 
 // SearchOptions is a basic query for inventory of enrollment IDs.
@@ -15,11 +20,18 @@ type Values map[string]interface{}
 
 type ReadStorage interface {
 	// RetrieveInventory queries and returns the inventory values by mapped by enrollment ID.
+	// If no search opt nor IDs are provided an ErrNoIDs should be returned.
+	// If IDs are have no inventory data then they should be skipped and
+	// omitted from the output with no error.
 	RetrieveInventory(ctx context.Context, opt *SearchOptions) (map[string]Values, error)
 }
 
 type Storage interface {
 	ReadStorage
+
+	// StoreInventoryValues stores inventory data about the specified ID.
 	StoreInventoryValues(ctx context.Context, id string, values Values) error
+
+	// DeleteInventory deletes all inventory data for an enrollment ID.
 	DeleteInventory(ctx context.Context, id string) error
 }

subsystem/inventory/storage/test/test.go

@@ -57,6 +57,6 @@ func TestStorage(t *testing.T, newStorage func() storage.Storage) {
 
 	_, ok = idVals[id]
 	if ok {
-		t.Error("expected id to missing in id values map")
+		t.Error("expected id to be missing in id values map")
 	}
 }