ModelStore Put/Get implmentations (#261)

* Simple Get and Put implementations

Get implementation is currently the one that uses the kopia ID of the
model/manifest.

* Basic tests for ModelStore Get/Put

src/internal/kopia/model_store.go

@@ -2,10 +2,22 @@ package kopia
 
 import (
 	"context"
+
+	"github.com/google/uuid"
+	"github.com/kopia/kopia/repo"
+	"github.com/kopia/kopia/repo/manifest"
+	"github.com/pkg/errors"
+
+	"github.com/alcionai/corso/internal/model"
+)
+
+const stableIDKey = "stableID"
+
+var (
+	errNoModelStoreID = errors.New("model has no ModelStoreID")
+	errBadTagKey      = errors.New("tag key overlaps with required key")
 )
 
-// ID of the manifest in kopia. This is not guaranteed to be stable.
-type ID string
 type modelType int
 
 //go:generate stringer -type=modelType
@@ -14,22 +26,119 @@ const (
 	BackupOpModel
 	RestoreOpModel
 	RestorePointModel
-
-	errStoreFlush = "flushing manifest store"
 )
 
-type ModelStore struct{}
+func NewModelStore(kw *KopiaWrapper) *ModelStore {
+	return &ModelStore{wrapper: kw}
+}
+
+// ModelStore must not be accessed after the given KopiaWrapper is closed.
+type ModelStore struct {
+	wrapper *KopiaWrapper
+}
+
+// tagsForModel creates a copy of tags and adds a tag for the model type to it.
+// Returns an error if another tag has the same key as the model type or if a
+// bad model type is given.
+func tagsForModel(t modelType, tags map[string]string) (map[string]string, error) {
+	if t == UnknownModel {
+		return nil, errors.New("bad model type")
+	}
+
+	if _, ok := tags[manifest.TypeLabelKey]; ok {
+		return nil, errors.WithStack(errBadTagKey)
+	}
+
+	res := make(map[string]string, len(tags)+1)
+	res[manifest.TypeLabelKey] = t.String()
+	for k, v := range tags {
+		res[k] = v
+	}
+
+	return res, nil
+}
+
+// tagsForModelWithID creates a copy of tags and adds tags for the model type
+// StableID to it. Returns an error if another tag has the same key as the model
+// type or if a bad model type is given.
+func tagsForModelWithID(
+	t modelType,
+	id model.ID,
+	tags map[string]string,
+) (map[string]string, error) {
+	if len(id) == 0 {
+		return nil, errors.New("missing ID for model")
+	}
+
+	res, err := tagsForModel(t, tags)
+	if err != nil {
+		return nil, err
+	}
+
+	if _, ok := res[stableIDKey]; ok {
+		return nil, errors.WithStack(errBadTagKey)
+	}
+
+	res[stableIDKey] = string(id)
+
+	return res, nil
+}
+
+// putInner contains logic for adding a model to the store. However, it does not
+// issue a flush operation.
+func putInner(
+	ctx context.Context,
+	w repo.RepositoryWriter,
+	t modelType,
+	tags map[string]string,
+	m model.Model,
+	create bool,
+) error {
+	// ModelStoreID does not need to be persisted in the model itself.
+	m.SetModelStoreID("")
+	if create {
+		m.SetStableID(model.ID(uuid.NewString()))
+	}
+
+	tmpTags, err := tagsForModelWithID(t, m.GetStableID(), tags)
+	if err != nil {
+		// Will be wrapped at a higher layer.
+		return err
+	}
+
+	id, err := w.PutManifest(ctx, tmpTags, m)
+	if err != nil {
+		// Will be wrapped at a higher layer.
+		return err
+	}
+
+	m.SetModelStoreID(id)
+	return nil
+}
 
 // Put adds a model of the given type to the persistent model store. Any tags
-// given to this function can later be used to help lookup the model. The
+// given to this function can later be used to help lookup the model.
-// returned ID can be used for subsequent Get, Update, or Delete calls.
 func (ms *ModelStore) Put(
 	ctx context.Context,
 	t modelType,
 	tags map[string]string,
-	m any,
+	m model.Model,
-) (ID, error) {
+) error {
-	return "", nil
+	err := repo.WriteSession(
+		ctx,
+		ms.wrapper.rep,
+		repo.WriteSessionOptions{Purpose: "ModelStorePut"},
+		func(innerCtx context.Context, w repo.RepositoryWriter) error {
+			err := putInner(innerCtx, w, t, tags, m, true)
+			if err != nil {
+				return err
+			}
+
+			return nil
+		},
+	)
+
+	return errors.Wrap(err, "putting model")
 }
 
 // GetIDsForType returns all IDs for models that match the given type and have
@@ -39,12 +148,35 @@ func (ms *ModelStore) GetIDsForType(
 	ctx context.Context,
 	t modelType,
 	tags map[string]string,
-) ([]ID, error) {
+) ([]model.ID, error) {
 	return nil, nil
 }
 
 // Get deserializes the model with the given ID into data.
-func (ms *ModelStore) Get(ctx context.Context, id ID, data any) error {
+func (ms *ModelStore) Get(ctx context.Context, id model.ID, data any) error {
+	return nil
+}
+
+// GetWithModelStoreID deserializes the model with the given ModelStoreID into
+// data. Returns github.com/kopia/kopia/repo/manifest.ErrNotFound if no model
+// was found.
+func (ms *ModelStore) GetWithModelStoreID(
+	ctx context.Context,
+	id manifest.ID,
+	data model.Model,
+) error {
+	if len(id) == 0 {
+		return errors.WithStack(errNoModelStoreID)
+	}
+
+	_, err := ms.wrapper.rep.GetManifest(ctx, id, data)
+	// TODO(ashmrtnz): Should probably return some recognizable, non-kopia error
+	// if not found. That way kopia doesn't need to be imported to higher layers.
+	if err != nil {
+		return errors.Wrap(err, "getting model data")
+	}
+
+	data.SetModelStoreID(id)
 	return nil
 }
 
@@ -54,14 +186,14 @@ func (ms *ModelStore) Get(ctx context.Context, id ID, data any) error {
 func (ms *ModelStore) Update(
 	ctx context.Context,
 	t modelType,
-	oldID ID,
+	oldID model.ID,
 	tags map[string]string,
 	m any,
-) (ID, error) {
+) (model.ID, error) {
 	return "", nil
 }
 
 // Delete deletes the model with the given ID from the model store.
-func (ms *ModelStore) Delete(ctx context.Context, id ID) error {
+func (ms *ModelStore) Delete(ctx context.Context, id model.ID) error {
 	return nil
 }

src/internal/kopia/model_store_test.go

@@ -0,0 +1,175 @@
+package kopia
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/kopia/kopia/repo/manifest"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+
+	"github.com/alcionai/corso/internal/model"
+	ctesting "github.com/alcionai/corso/internal/testing"
+)
+
+type fooModel struct {
+	model.BaseModel
+	Bar string
+}
+
+func getModelStore(t *testing.T, ctx context.Context) *ModelStore {
+	kw, err := openKopiaRepo(t, ctx)
+	require.NoError(t, err)
+
+	return NewModelStore(kw)
+}
+
+// ---------------
+// integration tests that use kopia
+// ---------------
+type ModelStoreIntegrationSuite struct {
+	suite.Suite
+}
+
+func TestModelStoreIntegrationSuite(t *testing.T) {
+	if err := ctesting.RunOnAny(
+		ctesting.CorsoCITests,
+		ctesting.CorsoModelStoreTests,
+	); err != nil {
+		t.Skip()
+	}
+
+	suite.Run(t, new(ModelStoreIntegrationSuite))
+}
+
+func (suite *ModelStoreIntegrationSuite) SetupSuite() {
+	_, err := ctesting.GetRequiredEnvVars(ctesting.AWSStorageCredEnvs...)
+	require.NoError(suite.T(), err)
+}
+
+func (suite *ModelStoreIntegrationSuite) TestBadTagsErrors() {
+	table := []struct {
+		name string
+		tags map[string]string
+	}{
+		{
+			name: "StableIDTag",
+			tags: map[string]string{
+				stableIDKey: "foo",
+			},
+		},
+		{
+			name: "manifestTypeTag",
+			tags: map[string]string{
+				manifest.TypeLabelKey: "foo",
+			},
+		},
+	}
+
+	ctx := context.Background()
+	t := suite.T()
+
+	m := getModelStore(t, ctx)
+	defer func() {
+		assert.NoError(t, m.wrapper.Close(ctx))
+	}()
+
+	foo := &fooModel{Bar: uuid.NewString()}
+	for _, test := range table {
+		suite.T().Run(test.name, func(t *testing.T) {
+			assert.Error(t, m.Put(ctx, BackupOpModel, test.tags, foo))
+		})
+	}
+}
+
+func (suite *ModelStoreIntegrationSuite) TestNoIDsErrors() {
+	ctx := context.Background()
+	t := suite.T()
+
+	m := getModelStore(t, ctx)
+	defer func() {
+		assert.NoError(t, m.wrapper.Close(ctx))
+	}()
+
+	noStableID := &fooModel{Bar: uuid.NewString()}
+	noStableID.SetStableID("")
+	noStableID.SetModelStoreID(manifest.ID(uuid.NewString()))
+
+	noModelStoreID := &fooModel{Bar: uuid.NewString()}
+	noModelStoreID.SetStableID(model.ID(uuid.NewString()))
+	noModelStoreID.SetModelStoreID("")
+
+	assert.Error(t, m.GetWithModelStoreID(ctx, "", nil))
+}
+
+func (suite *ModelStoreIntegrationSuite) TestPutGet() {
+	table := []struct {
+		t      modelType
+		check  require.ErrorAssertionFunc
+		hasErr bool
+	}{
+		{
+			t:      UnknownModel,
+			check:  require.Error,
+			hasErr: true,
+		},
+		{
+			t:      BackupOpModel,
+			check:  require.NoError,
+			hasErr: false,
+		},
+		{
+			t:      RestoreOpModel,
+			check:  require.NoError,
+			hasErr: false,
+		},
+		{
+			t:      RestorePointModel,
+			check:  require.NoError,
+			hasErr: false,
+		},
+	}
+
+	ctx := context.Background()
+	t := suite.T()
+
+	m := getModelStore(t, ctx)
+	defer func() {
+		assert.NoError(t, m.wrapper.Close(ctx))
+	}()
+
+	for _, test := range table {
+		suite.T().Run(test.t.String(), func(t *testing.T) {
+			foo := &fooModel{Bar: uuid.NewString()}
+
+			err := m.Put(ctx, test.t, nil, foo)
+			test.check(t, err)
+
+			if test.hasErr {
+				return
+			}
+
+			require.NotEmpty(t, foo.GetModelStoreID())
+			require.NotEmpty(t, foo.GetStableID())
+
+			returned := &fooModel{}
+			err = m.GetWithModelStoreID(ctx, foo.GetModelStoreID(), returned)
+			require.NoError(t, err)
+			assert.Equal(t, foo, returned)
+		})
+	}
+}
+
+func (suite *ModelStoreIntegrationSuite) TestGet_NotFoundErrors() {
+	ctx := context.Background()
+	t := suite.T()
+
+	m := getModelStore(t, ctx)
+	defer func() {
+		assert.NoError(t, m.wrapper.Close(ctx))
+	}()
+
+	assert.ErrorIs(t, m.GetWithModelStoreID(ctx, "baz", nil), manifest.ErrNotFound)
+}

src/internal/model/model.go

@@ -0,0 +1,45 @@
+package model
+
+import (
+	"github.com/kopia/kopia/repo/manifest"
+)
+
+type ID string
+
+type Model interface {
+	GetStableID() ID
+	SetStableID(id ID)
+	GetModelStoreID() manifest.ID
+	SetModelStoreID(id manifest.ID)
+}
+
+// BaseModel defines required fields for models stored in ModelStore. Structs
+// that wish to be stored should embed this struct.
+type BaseModel struct {
+	// StableID is an identifier that other objects can use to refer to this
+	// object in the ModelStore.
+	// Once generated (during Put), it is guaranteed not to change. This field
+	// should be treated as read-only by users.
+	StableID ID `json:"stableID,omitempty"`
+	// ModelStoreID is an internal ID for the model in the store. If present it
+	// can be used for efficient lookups, but should not be used by other models
+	// to refer to this one. This field may change if the model is updated. This
+	// field should be treated as read-only by users.
+	ModelStoreID manifest.ID `json:"modelStoreID,omitempty"`
+}
+
+func (bm *BaseModel) GetStableID() ID {
+	return bm.StableID
+}
+
+func (bm *BaseModel) SetStableID(id ID) {
+	bm.StableID = id
+}
+
+func (bm *BaseModel) GetModelStoreID() manifest.ID {
+	return bm.ModelStoreID
+}
+
+func (bm *BaseModel) SetModelStoreID(id manifest.ID) {
+	bm.ModelStoreID = id
+}

src/internal/testing/integration_runners.go

@@ -13,6 +13,7 @@ const (
 	CorsoCLIConfigTests      = "CORSO_CLI_CONFIG_TESTS"
 	CorsoGraphConnectorTests = "CORSO_GRAPH_CONNECTOR_TESTS"
 	CorsoKopiaWrapperTests   = "CORSO_KOPIA_WRAPPER_TESTS"
+	CorsoModelStoreTests     = "CORSO_MODEL_STORE_TESTS"
 	CorsoRepositoryTests     = "CORSO_REPOSITORY_TESTS"
 )