go.mercari.io/datastore@v1.8.2/boom/transaction.go

package boom

import (
	"reflect"
	"sync"

	"go.mercari.io/datastore"
)

// Transaction represents a set of datastore operations to be committed atomically.
//
// Operations are enqueued by calling the Put and Delete methods on Transaction
// (or their Multi-equivalents).  These operations are only committed when the
// Commit method is invoked. To ensure consistency, reads must be performed by
// using Transaction's Get method or by using the Transaction method when
// building a query.
//
// A Transaction must be committed or rolled back exactly once.
type Transaction struct {
	m                sync.Mutex
	bm               *Boom
	tx               datastore.Transaction
	pendingKeysLater []*setKeyLater
}

// DatastoreTransaction returns datastore.Transaction that contains in boom's Transaction.
// This function should not be used unless you have a special reason.
func DatastoreTransaction(tx *Transaction) datastore.Transaction {
	return tx.tx
}

type setKeyLater struct {
	pendingKey datastore.PendingKey
	src        interface{}
}

// Boom object that is the source of the Batch object is returned.
func (tx *Transaction) Boom() *Boom {
	return tx.bm
}

// Kind retrieves kind name from struct.
func (tx *Transaction) Kind(src interface{}) string {
	return tx.bm.Kind(src)
}

// Key retrieves datastore key from struct without error occurred.
func (tx *Transaction) Key(src interface{}) datastore.Key {
	return tx.bm.Key(src)
}

// KeyError retrieves datastore key from struct with error occurred.
func (tx *Transaction) KeyError(src interface{}) (datastore.Key, error) {
	return tx.bm.KeyError(src)
}

// Get loads the entity stored for key into dst, which must be a struct pointer or implement PropertyLoadSaver.
// key will be extracted from dst.
//
// If there is no such entity for the key, Get returns ErrNoSuchEntity.
// The values of dst's unmatched struct fields are not modified, and matching slice-typed fields are not reset before appending to them.
// In particular, it is recommended to pass a pointer to a zero valued struct on each Get call.
func (tx *Transaction) Get(dst interface{}) error {
	dsts := []interface{}{dst}
	err := tx.GetMulti(dsts)
	if merr, ok := err.(datastore.MultiError); ok {
		return merr[0]
	} else if err != nil {
		return err
	}

	return nil
}

// GetMulti is a batch version of Get.
// key will be extracted from each struct of dst.
//
// dst must be a []S, []*S, []I or []P, for some struct type S, some interface type I, or some non-interface non-pointer type P such that P or *P implements PropertyLoadSaver.
// If an []I, each element must be a valid dst for Get: it must be a struct pointer or implement PropertyLoadSaver.
func (tx *Transaction) GetMulti(dst interface{}) error {
	keys, err := tx.bm.extractKeys(dst)
	if err != nil {
		return err
	}

	return tx.tx.GetMulti(keys, dst)
}

// Put saves the entity src into the datastore.
// key will be extract from src struct.
// src must be a struct pointer or implement PropertyLoadSaver; if a struct pointer then any unexported fields of that struct will be skipped.
// If k is an incomplete key, the returned key will be a unique key generated by the datastore,
// and inject key to src struct.
func (tx *Transaction) Put(src interface{}) (datastore.PendingKey, error) {
	srcs := []interface{}{src}
	keys, err := tx.PutMulti(srcs)
	if merr, ok := err.(datastore.MultiError); ok {
		return nil, merr[0]
	} else if err != nil {
		return nil, err
	}

	return keys[0], nil
}

// PutMulti is a batch version of Put.
//
// src must satisfy the same conditions as the dst argument to GetMulti.
func (tx *Transaction) PutMulti(src interface{}) ([]datastore.PendingKey, error) {
	keys, err := tx.bm.extractKeys(src)
	if err != nil {
		return nil, err
	}

	pKeys, err := tx.tx.PutMulti(keys, src)
	if err != nil {
		return nil, err
	}

	v := reflect.Indirect(reflect.ValueOf(src))
	tx.m.Lock()
	defer tx.m.Unlock()
	for idx, pKey := range pKeys {
		if !keys[idx].Incomplete() {
			continue
		}
		tx.pendingKeysLater = append(tx.pendingKeysLater, &setKeyLater{
			pendingKey: pKey,
			src:        v.Index(idx).Interface(),
		})
	}

	return pKeys, nil
}

// Delete deletes the entity.
// key will be extract from src struct.
func (tx *Transaction) Delete(src interface{}) error {
	srcs := []interface{}{src}
	err := tx.DeleteMulti(srcs)
	if merr, ok := err.(datastore.MultiError); ok {
		return merr[0]
	} else if err != nil {
		return err
	}

	return nil
}

// DeleteMulti is a batch version of Delete.
func (tx *Transaction) DeleteMulti(src interface{}) error {
	keys, err := tx.bm.extractKeys(src)
	if err != nil {
		return err
	}

	return tx.tx.DeleteMulti(keys)
}

// Commit applies the enqueued operations atomically.
func (tx *Transaction) Commit() (datastore.Commit, error) {
	commit, err := tx.tx.Commit()
	if err != nil {
		return nil, err
	}

	tx.m.Lock()
	defer tx.m.Unlock()

	for _, s := range tx.pendingKeysLater {
		key := commit.Key(s.pendingKey)
		err = tx.bm.setStructKey(s.src, key)
		if err != nil {
			return nil, err
		}
	}
	tx.pendingKeysLater = nil

	return commit, nil
}

// Rollback abandons a pending transaction.
func (tx *Transaction) Rollback() error {
	return tx.tx.Rollback()
}

// Batch creates batch mode objects.
func (tx *Transaction) Batch() *TransactionBatch {
	b := tx.tx.Batch()
	return &TransactionBatch{bm: tx.bm, tx: tx, b: b}
}