go.mercari.io/datastore@v1.8.2/dsmiddleware/fishbone/fishbone.go (about)

     1  package fishbone
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  
     7  	"go.mercari.io/datastore"
     8  )
     9  
    10  var _ datastore.Middleware = &modifier{}
    11  
    12  // New fishbone middleware creates and returns.
    13  func New() datastore.Middleware {
    14  	return &modifier{}
    15  }
    16  
    17  type contextQuery struct{}
    18  
    19  type modifier struct {
    20  }
    21  
    22  func (m *modifier) AllocateIDs(info *datastore.MiddlewareInfo, keys []datastore.Key) ([]datastore.Key, error) {
    23  	return info.Next.AllocateIDs(info, keys)
    24  }
    25  
    26  func (m *modifier) PutMultiWithoutTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) ([]datastore.Key, error) {
    27  	return info.Next.PutMultiWithoutTx(info, keys, psList)
    28  }
    29  
    30  func (m *modifier) PutMultiWithTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) ([]datastore.PendingKey, error) {
    31  	return info.Next.PutMultiWithTx(info, keys, psList)
    32  }
    33  
    34  func (m *modifier) GetMultiWithoutTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) error {
    35  	return info.Next.GetMultiWithoutTx(info, keys, psList)
    36  }
    37  
    38  func (m *modifier) GetMultiWithTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) error {
    39  	return info.Next.GetMultiWithTx(info, keys, psList)
    40  }
    41  
    42  func (m *modifier) DeleteMultiWithoutTx(info *datastore.MiddlewareInfo, keys []datastore.Key) error {
    43  	return info.Next.DeleteMultiWithoutTx(info, keys)
    44  }
    45  
    46  func (m *modifier) DeleteMultiWithTx(info *datastore.MiddlewareInfo, keys []datastore.Key) error {
    47  	return info.Next.DeleteMultiWithTx(info, keys)
    48  }
    49  
    50  func (m *modifier) PostCommit(info *datastore.MiddlewareInfo, tx datastore.Transaction, commit datastore.Commit) error {
    51  	return info.Next.PostCommit(info, tx, commit)
    52  }
    53  
    54  func (m *modifier) PostRollback(info *datastore.MiddlewareInfo, tx datastore.Transaction) error {
    55  	return info.Next.PostRollback(info, tx)
    56  }
    57  
    58  func (m *modifier) Run(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump) datastore.Iterator {
    59  	if qDump.KeysOnly {
    60  		return info.Next.Run(info, q, qDump)
    61  	}
    62  
    63  	q = q.KeysOnly()
    64  	qDump.KeysOnly = true
    65  	defer func() {
    66  		qDump.KeysOnly = false
    67  	}()
    68  
    69  	mQDumpMap, ok := info.Context.Value(contextQuery{}).(map[*modifier]map[*datastore.QueryDump]bool)
    70  	if !ok {
    71  		mQDumpMap = make(map[*modifier]map[*datastore.QueryDump]bool)
    72  		info.Context = context.WithValue(info.Context, contextQuery{}, mQDumpMap)
    73  	}
    74  	_, ok = mQDumpMap[m]
    75  	if !ok {
    76  		mQDumpMap[m] = make(map[*datastore.QueryDump]bool)
    77  	}
    78  	mQDumpMap[m][qDump] = true
    79  
    80  	return info.Next.Run(info, q, qDump)
    81  }
    82  
    83  func (m *modifier) GetAll(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump, psList *[]datastore.PropertyList) ([]datastore.Key, error) {
    84  	if qDump.KeysOnly {
    85  		return info.Next.GetAll(info, q, qDump, psList)
    86  	}
    87  
    88  	q = q.KeysOnly()
    89  	qDump.KeysOnly = true
    90  	defer func() {
    91  		qDump.KeysOnly = false
    92  	}()
    93  
    94  	keys, err := info.Next.GetAll(info, q, qDump, nil)
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  
    99  	if qDump.Transaction != nil {
   100  		if info.Transaction == nil {
   101  			return nil, errors.New("cacheInfo.Transaction is required")
   102  		}
   103  
   104  		altPsList := make([]datastore.PropertyList, len(keys))
   105  		err := qDump.Transaction.GetMulti(keys, altPsList)
   106  		if err != nil {
   107  			return nil, err
   108  		}
   109  		*psList = altPsList
   110  
   111  		return keys, nil
   112  	}
   113  
   114  	altPsList := make([]datastore.PropertyList, len(keys))
   115  	err = info.Client.GetMulti(info.Context, keys, altPsList)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  	*psList = altPsList
   120  
   121  	return keys, nil
   122  }
   123  
   124  func (m *modifier) Next(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump, iter datastore.Iterator, ps *datastore.PropertyList) (datastore.Key, error) {
   125  	mQDumpMap, ok := info.Context.Value(contextQuery{}).(map[*modifier]map[*datastore.QueryDump]bool)
   126  	if !ok {
   127  		return info.Next.Next(info, q, qDump, iter, ps)
   128  	}
   129  
   130  	qDumpMap, ok := mQDumpMap[m]
   131  	if !ok {
   132  		return info.Next.Next(info, q, qDump, iter, ps)
   133  	}
   134  
   135  	_, ok = qDumpMap[qDump]
   136  	if !ok {
   137  		return info.Next.Next(info, q, qDump, iter, ps)
   138  	}
   139  
   140  	// NOTE: We can't delete(qDumpMap, qDump). q is reusable!
   141  
   142  	key, err := info.Next.Next(info, q, qDump, iter, ps)
   143  
   144  	if err != nil {
   145  		return nil, err
   146  	}
   147  
   148  	if qDump.Transaction != nil {
   149  		if info.Transaction == nil {
   150  			return nil, errors.New("cacheInfo.Transaction is required")
   151  		}
   152  
   153  		altPs := datastore.PropertyList{}
   154  		err := qDump.Transaction.Get(key, &altPs)
   155  		if err != nil {
   156  			return nil, err
   157  		}
   158  		*ps = altPs
   159  
   160  		return key, nil
   161  	}
   162  
   163  	altPs := datastore.PropertyList{}
   164  	err = info.Client.Get(info.Context, key, &altPs)
   165  	if err != nil {
   166  		return nil, err
   167  	}
   168  	*ps = altPs
   169  
   170  	return key, nil
   171  }
   172  
   173  func (m *modifier) Count(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump) (int, error) {
   174  	return info.Next.Count(info, q, qDump)
   175  }