add InitStore

diff --git a/account/accounts.go b/account/accounts.go
index a6e92a6..8e60f47 100644 (file)
--- a/account/accounts.go
+++ b/account/accounts.go
@@ -115,17 +115,15 @@ func CreateAccount(xpubs []chainkd.XPub, quorum int, alias string, acctIndex uin
 }
 
 func (m *Manager) saveAccount(account *Account) error {
-       if err := m.store.InitBatch(); err != nil {
-               return err
-       }
+       newStore := m.store.InitStore()
 
        // update account index
-       m.store.SetAccountIndex(account)
-       if err := m.store.SetAccount(account); err != nil {
+       newStore.SetAccountIndex(account)
+       if err := newStore.SetAccount(account); err != nil {
                return err
        }
 
-       if err := m.store.CommitBatch(); err != nil {
+       if err := newStore.CommitBatch(); err != nil {
                return err
        }
 
@@ -214,19 +212,17 @@ func (m *Manager) UpdateAccountAlias(accountID string, newAlias string) error {
 
        account.Alias = normalizedAlias
 
-       if err := m.store.InitBatch(); err != nil {
-               return err
-       }
+       newStore := m.store.InitStore()
 
-       if err := m.store.DeleteAccount(&oldAccount); err != nil {
+       if err := newStore.DeleteAccount(&oldAccount); err != nil {
                return err
        }
 
-       if err := m.store.SetAccount(account); err != nil {
+       if err := newStore.SetAccount(account); err != nil {
                return err
        }
 
-       if err := m.store.CommitBatch(); err != nil {
+       if err := newStore.CommitBatch(); err != nil {
                return err
        }
 
@@ -621,24 +617,22 @@ func (m *Manager) saveControlProgram(prog *CtrlProgram, updateIndex bool) error
                return err
        }
 
-       if err := m.store.InitBatch(); err != nil {
-               return err
-       }
+       newStore := m.store.InitStore()
 
-       if err := m.store.SetControlProgram(bc.NewHash(hash), prog); err != nil {
+       if err := newStore.SetControlProgram(bc.NewHash(hash), prog); err != nil {
                return nil
        }
 
        if updateIndex {
                switch acct.DeriveRule {
                case signers.BIP0032:
-                       m.store.SetContractIndex(acct.ID, prog.KeyIndex)
+                       newStore.SetContractIndex(acct.ID, prog.KeyIndex)
                case signers.BIP0044:
-                       m.store.SetBip44ContractIndex(acct.ID, prog.Change, prog.KeyIndex)
+                       newStore.SetBip44ContractIndex(acct.ID, prog.Change, prog.KeyIndex)
                }
        }
 
-       return m.store.CommitBatch()
+       return newStore.CommitBatch()
 }
 
 // SaveControlPrograms save account control programs

diff --git a/account/image.go b/account/image.go
index 6e456c8..ec2b3e5 100644 (file)
--- a/account/image.go
+++ b/account/image.go
@@ -44,12 +44,10 @@ func (m *Manager) Restore(image *Image) error {
        m.accountMu.Lock()
        defer m.accountMu.Unlock()
 
-       if err := m.store.InitBatch(); err != nil {
-               return err
-       }
+       newStore := m.store.InitStore()
 
        for _, slice := range image.Slice {
-               if _, err := m.store.GetAccountByID(slice.Account.ID); err != nil {
+               if _, err := newStore.GetAccountByID(slice.Account.ID); err != nil {
                        log.WithFields(log.Fields{
                                "module": logModule,
                                "alias":  slice.Account.Alias,
@@ -58,14 +56,14 @@ func (m *Manager) Restore(image *Image) error {
                        continue
                }
 
-               if _, err := m.store.GetAccountByAlias(slice.Account.Alias); err != nil {
+               if _, err := newStore.GetAccountByAlias(slice.Account.Alias); err != nil {
                        return err
                }
 
-               if err := m.store.SetAccount(slice.Account); err != nil {
+               if err := newStore.SetAccount(slice.Account); err != nil {
                        return err
                }
        }
 
-       return m.store.CommitBatch()
+       return newStore.CommitBatch()
 }

diff --git a/account/store.go b/account/store.go
index 6ec4ecf..928938a 100644 (file)
--- a/account/store.go
+++ b/account/store.go
@@ -7,7 +7,7 @@ import (
 
 // AccountStore interface contains account storage functions.
 type AccountStore interface {
-       InitBatch() error
+       InitStore() AccountStore
        CommitBatch() error
        DeleteAccount(*Account) error
        DeleteStandardUTXO(bc.Hash)

diff --git a/account/utxo_keeper_test.go b/account/utxo_keeper_test.go
index 9e0d1ce..78cb3b2 100644 (file)
--- a/account/utxo_keeper_test.go
+++ b/account/utxo_keeper_test.go
@@ -1380,7 +1380,7 @@ func ContractUTXOKey(id bc.Hash) []byte {
        return append(SUTXOPrefix, id.Bytes()...)
 }
 
-func (store *mockAccountStore) InitBatch() error                                { return nil }
+func (store *mockAccountStore) InitStore() AccountStore                         { return nil }
 func (store *mockAccountStore) CommitBatch() error                              { return nil }
 func (store *mockAccountStore) DeleteAccount(*Account) error                    { return nil }
 func (store *mockAccountStore) GetAccountByAlias(string) (*Account, error)      { return nil, nil }

diff --git a/database/account_store.go b/database/account_store.go
index 4ec66e3..4e4c851 100644 (file)
--- a/database/account_store.go
+++ b/database/account_store.go
@@ -82,13 +82,11 @@ func NewAccountStore(db dbm.DB) *AccountStore {
        }
 }
 
-// InitBatch initial batch
-func (store *AccountStore) InitBatch() error {
-       if store.batch != nil {
-               return errors.New("AccountStore initail fail, store batch is not nil.")
-       }
-       store.batch = store.db.NewBatch()
-       return nil
+// InitStore initial batch
+func (store *AccountStore) InitStore() acc.AccountStore {
+       newStore := NewAccountStore(store.db)
+       newStore.batch = newStore.db.NewBatch()
+       return newStore
 }
 
 // CommitBatch commit batch

diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go
index a140b54..43cb78c 100644 (file)
--- a/wallet/wallet_test.go
+++ b/wallet/wallet_test.go
@@ -696,13 +696,11 @@ func NewMockAccountStore(db dbm.DB) *MockAccountStore {
        }
 }
 
-// InitBatch initial batch
-func (store *MockAccountStore) InitBatch() error {
-       if store.batch != nil {
-               return errors.New("MockAccountStore initail fail, store batch is not nil.")
-       }
-       store.batch = store.db.NewBatch()
-       return nil
+// InitStore initial batch
+func (store *MockAccountStore) InitStore() acc.AccountStore {
+       newStore := NewMockAccountStore(store.db)
+       newStore.batch = newStore.db.NewBatch()
+       return newStore
 }
 
 // CommitBatch commit batch