fix(kv): update auto migration store to return migratable store

bolt/kv.go

@@ -18,6 +18,9 @@ import (
 // check that *KVStore implement kv.Store interface.
 var _ kv.Store = (*KVStore)(nil)
 
+// ensure *KVStore implements kv.AutoMigrationStore.
+var _ kv.AutoMigrationStore = (*KVStore)(nil)
+
 // KVStore is a kv.Store backed by boltdb.
 type KVStore struct {
 	path string
@@ -34,9 +37,9 @@ func NewKVStore(log *zap.Logger, path string) *KVStore {
 	}
 }
 
-// AutoMigrate returns true as the bolt KVStore is safe to migrate on initialize.
-func (s *KVStore) AutoMigrate() bool {
-	return true
+// AutoMigrate returns itself as it is safe to automatically apply migrations on initialization.
+func (s *KVStore) AutoMigrate() kv.Store {
+	return s
 }
 
 // Open creates boltDB file it doesn't exists and opens it otherwise.

inmem/kv.go

@@ -14,6 +14,9 @@ import (
 // ensure *KVStore implement kv.Store interface
 var _ kv.Store = (*KVStore)(nil)
 
+// ensure *KVStore implements kv.AutoMigrationStore
+var _ kv.AutoMigrationStore = (*KVStore)(nil)
+
 // cursorBatchSize is the size of a batch sent by a forward cursors
 // tree iterator
 const cursorBatchSize = 1000
@@ -61,9 +64,9 @@ func (s *KVStore) Backup(ctx context.Context, w io.Writer) error {
 	panic("not implemented")
 }
 
-// AutoMigrate returns true as inmem KVStore is safe to migrate on initialize.
-func (s *KVStore) AutoMigrate() bool {
-	return true
+// AutoMigrate returns itlsef as *KVStore is safe to migrate automically on initialize.
+func (s *KVStore) AutoMigrate() kv.Store {
+	return s
 }
 
 // Flush removes all data from the buckets.  Used for testing.

kv/service.go

@@ -101,11 +101,12 @@ type ServiceConfig struct {
 
 // AutoMigrationStore is a Store which also describes whether or not
 // migrations can be applied automatically.
-// Given the AutoMigrate method is defined and it returns true then migrations
-// will automatically be applied on Service.Initialize(...).
+// Given the AutoMigrate method is defined and it returns a non-nil kv.Store
+// implementation, then it will automatically invoke migrator.Up(store)
+// on the returned kv.Store during Service.Initialize(...).
 type AutoMigrationStore interface {
 	Store
-	AutoMigrate() bool
+	AutoMigrate() Store
 }
 
 // Initialize creates Buckets needed.
@@ -114,8 +115,12 @@ func (s *Service) Initialize(ctx context.Context) error {
 		return err
 	}
 
-	if store, ok := s.kv.(AutoMigrationStore); ok && store.AutoMigrate() {
-		return s.Migrator.Up(ctx, s.kv)
+	// if store implements auto migrate and the resulting Store from
+	// AutoMigrate() is non-nil, apply migrator.Up() to the resulting store.
+	if store, ok := s.kv.(AutoMigrationStore); ok {
+		if migrateStore := store.AutoMigrate(); migrateStore != nil {
+			return s.Migrator.Up(ctx, migrateStore)
+		}
 	}
 
 	return nil