go.mercari.io/datastore@v1.8.2/boom/tx_batch.go (about)

     1  package boom
     2  
     3  import (
     4  	"sync"
     5  
     6  	"go.mercari.io/datastore"
     7  )
     8  
     9  // TransactionBatch can queue operations on Datastore and process them in batch.
    10  // Batch does nothing until you call Exec().
    11  // This helps to reduce the number of RPCs.
    12  type TransactionBatch struct {
    13  	m  sync.Mutex
    14  	bm *Boom
    15  	tx *Transaction
    16  	b  *datastore.TransactionBatch
    17  
    18  	earlyErrors []error
    19  }
    20  
    21  // Boom object that is the source of the TransactionBatch object is returned.
    22  func (b *TransactionBatch) Boom() *Boom {
    23  	return b.bm
    24  }
    25  
    26  // Transaction object that is the source of the TransactionBatch object is returned.
    27  func (b *TransactionBatch) Transaction() *Transaction {
    28  	return b.tx
    29  }
    30  
    31  // Put Entity operation into the queue.
    32  // This operation doesn't Put to Datastore immediately.
    33  // If a h is provided, it passes the processing result to the handler, and treats the return value as the value of the result of Putting.
    34  func (b *TransactionBatch) Put(src interface{}, h datastore.TxBatchPutHandler) {
    35  	keys, err := b.bm.extractKeys([]interface{}{src})
    36  	if err != nil {
    37  		if h != nil {
    38  			err = h(nil, err)
    39  		}
    40  		if err != nil {
    41  			b.m.Lock()
    42  			b.earlyErrors = append(b.earlyErrors, err)
    43  			b.m.Unlock()
    44  		}
    45  		return
    46  	}
    47  
    48  	b.b.Put(keys[0], src, func(pKey datastore.PendingKey, err error) error {
    49  		b.tx.m.Lock()
    50  		defer b.tx.m.Unlock()
    51  		if err != nil {
    52  			if h != nil {
    53  				err = h(pKey, err)
    54  			}
    55  			if err != nil {
    56  				b.m.Lock()
    57  				b.earlyErrors = append(b.earlyErrors, err)
    58  				b.m.Unlock()
    59  			}
    60  			return err
    61  		}
    62  
    63  		if keys[0].Incomplete() {
    64  			b.tx.pendingKeysLater = append(b.tx.pendingKeysLater, &setKeyLater{
    65  				pendingKey: pKey,
    66  				src:        src,
    67  			})
    68  		}
    69  
    70  		if h != nil {
    71  			return h(pKey, nil)
    72  		}
    73  
    74  		return nil
    75  	})
    76  }
    77  
    78  // Get Entity operation into the queue.
    79  func (b *TransactionBatch) Get(dst interface{}, h datastore.BatchErrHandler) {
    80  	keys, err := b.bm.extractKeys([]interface{}{dst})
    81  	if err != nil {
    82  		if h != nil {
    83  			err = h(err)
    84  		}
    85  		if err != nil {
    86  			b.m.Lock()
    87  			b.earlyErrors = append(b.earlyErrors, err)
    88  			b.m.Unlock()
    89  		}
    90  		return
    91  	}
    92  
    93  	b.b.Get(keys[0], dst, h)
    94  }
    95  
    96  // Delete Entity operation into the queue.
    97  func (b *TransactionBatch) Delete(dst interface{}, h datastore.BatchErrHandler) {
    98  	keys, err := b.bm.extractKeys([]interface{}{dst})
    99  	if err != nil {
   100  		if h != nil {
   101  			err = h(err)
   102  		}
   103  		if err != nil {
   104  			b.m.Lock()
   105  			b.earlyErrors = append(b.earlyErrors, err)
   106  			b.m.Unlock()
   107  		}
   108  		return
   109  	}
   110  
   111  	b.b.Delete(keys[0], h)
   112  }
   113  
   114  // Exec will perform all the processing that was queued.
   115  // This process is done recursively until the queue is empty.
   116  // The return value may be MultiError, but the order of contents is not guaranteed.
   117  func (b *TransactionBatch) Exec() error {
   118  	b.m.Lock()
   119  	defer b.m.Unlock()
   120  
   121  	err := b.b.Exec()
   122  
   123  	if merr, ok := err.(datastore.MultiError); ok {
   124  		merr = append(merr, b.earlyErrors...)
   125  		if len(merr) == 0 {
   126  			return nil
   127  		}
   128  		return merr
   129  	} else if err != nil {
   130  		return err
   131  	}
   132  
   133  	b.earlyErrors = nil
   134  
   135  	return nil
   136  }