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

     1  package boom
     2  
     3  import (
     4  	"sync"
     5  
     6  	"go.mercari.io/datastore"
     7  )
     8  
     9  // Batch 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 Batch struct {
    13  	m  sync.Mutex
    14  	bm *Boom
    15  	b  *datastore.Batch
    16  
    17  	earlyErrors []error
    18  }
    19  
    20  // Boom object that is the source of the Batch object is returned.
    21  func (b *Batch) Boom() *Boom {
    22  	return b.bm
    23  }
    24  
    25  // Put Entity operation into the queue.
    26  // This operation doesn't Put to Datastore immediately.
    27  // 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.
    28  func (b *Batch) Put(src interface{}, h datastore.BatchPutHandler) {
    29  	keys, err := b.bm.extractKeys([]interface{}{src})
    30  	if err != nil {
    31  		if h != nil {
    32  			err = h(nil, err)
    33  		}
    34  		if err != nil {
    35  			b.m.Lock()
    36  			b.earlyErrors = append(b.earlyErrors, err)
    37  			b.m.Unlock()
    38  		}
    39  		return
    40  	}
    41  
    42  	b.b.Put(keys[0], src, func(key datastore.Key, err error) error {
    43  		if err != nil {
    44  			if h != nil {
    45  				err = h(key, err)
    46  			}
    47  			return err
    48  		}
    49  
    50  		err = b.bm.setStructKey(src, key)
    51  		if err != nil {
    52  			if h != nil {
    53  				err = h(key, 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 h != nil {
    64  			return h(key, nil)
    65  		}
    66  
    67  		return nil
    68  	})
    69  }
    70  
    71  // Get Entity operation into the queue.
    72  func (b *Batch) Get(dst interface{}, h datastore.BatchErrHandler) {
    73  	keys, err := b.bm.extractKeys([]interface{}{dst})
    74  	if err != nil {
    75  		if h != nil {
    76  			err = h(err)
    77  		}
    78  		if err != nil {
    79  			b.m.Lock()
    80  			b.earlyErrors = append(b.earlyErrors, err)
    81  			b.m.Unlock()
    82  		}
    83  		return
    84  	}
    85  
    86  	b.b.Get(keys[0], dst, h)
    87  }
    88  
    89  // Delete Entity operation into the queue.
    90  func (b *Batch) Delete(dst interface{}, h datastore.BatchErrHandler) {
    91  	keys, err := b.bm.extractKeys([]interface{}{dst})
    92  	if err != nil {
    93  		if h != nil {
    94  			err = h(err)
    95  		}
    96  		if err != nil {
    97  			b.m.Lock()
    98  			b.earlyErrors = append(b.earlyErrors, err)
    99  			b.m.Unlock()
   100  		}
   101  		return
   102  	}
   103  
   104  	b.b.Delete(keys[0], h)
   105  }
   106  
   107  // Exec will perform all the processing that was queued.
   108  // This process is done recursively until the queue is empty.
   109  // The return value may be MultiError, but the order of contents is not guaranteed.
   110  func (b *Batch) Exec() error {
   111  	err := b.b.Exec(b.bm.Context)
   112  
   113  	b.m.Lock()
   114  	defer b.m.Unlock()
   115  
   116  	if merr, ok := err.(datastore.MultiError); ok {
   117  		merr = append(merr, b.earlyErrors...)
   118  		b.earlyErrors = nil
   119  		if len(merr) == 0 {
   120  			return nil
   121  		}
   122  		return merr
   123  	} else if err != nil {
   124  		return err
   125  	} else if len(b.earlyErrors) != 0 {
   126  		errs := b.earlyErrors
   127  		b.earlyErrors = nil
   128  		return datastore.MultiError(errs)
   129  	}
   130  
   131  	return nil
   132  }