go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/service/datastore/checkfilter.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package datastore
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  
    21  	"go.chromium.org/luci/common/errors"
    22  )
    23  
    24  type checkFilter struct {
    25  	RawInterface
    26  
    27  	kc      KeyContext
    28  	userCtx context.Context
    29  }
    30  
    31  func (tcf *checkFilter) RunInTransaction(f func(c context.Context) error, opts *TransactionOptions) error {
    32  	if f == nil {
    33  		return fmt.Errorf("datastore: RunInTransaction function is nil")
    34  	}
    35  	return tcf.checkCtxDone(tcf.RawInterface.RunInTransaction(f, opts))
    36  }
    37  
    38  func (tcf *checkFilter) Run(fq *FinalizedQuery, cb RawRunCB) error {
    39  	if fq == nil {
    40  		return fmt.Errorf("datastore: Run query is nil")
    41  	}
    42  	if cb == nil {
    43  		return fmt.Errorf("datastore: Run callback is nil")
    44  	}
    45  	return tcf.checkCtxDone(
    46  		tcf.RawInterface.Run(fq, func(key *Key, val PropertyMap, getCursor CursorCB) error {
    47  			if err := tcf.checkCtxDone(nil); err != nil {
    48  				return err
    49  			}
    50  			return cb(key, val, getCursor)
    51  		}))
    52  }
    53  
    54  func (tcf *checkFilter) GetMulti(keys []*Key, meta MultiMetaGetter, cb GetMultiCB) error {
    55  	if cb == nil {
    56  		return fmt.Errorf("datastore: GetMulti callback is nil")
    57  	}
    58  
    59  	var dat DroppedArgTracker
    60  
    61  	for i, k := range keys {
    62  		var err error
    63  		switch {
    64  		case k.IsIncomplete():
    65  			err = MakeErrInvalidKey("key [%s] is incomplete", k).Err()
    66  		case !k.Valid(true, tcf.kc):
    67  			err = MakeErrInvalidKey("key [%s] is not valid in context %s", k, tcf.kc).Err()
    68  		}
    69  		if err != nil {
    70  			cb(i, nil, err)
    71  			dat.MarkForRemoval(i, len(keys))
    72  		}
    73  	}
    74  
    75  	keys, meta, dal := dat.DropKeysAndMeta(keys, meta)
    76  	if len(keys) == 0 {
    77  		return nil
    78  	}
    79  
    80  	return tcf.checkCtxDone(
    81  		tcf.RawInterface.GetMulti(keys, meta, func(idx int, val PropertyMap, err error) {
    82  			cb(dal.OriginalIndex(idx), val, err)
    83  		}))
    84  }
    85  
    86  func (tcf *checkFilter) PutMulti(keys []*Key, vals []PropertyMap, cb NewKeyCB) error {
    87  	if len(keys) != len(vals) {
    88  		return fmt.Errorf("datastore: PutMulti with mismatched keys/vals lengths (%d/%d)", len(keys), len(vals))
    89  	}
    90  	if cb == nil {
    91  		return fmt.Errorf("datastore: PutMulti callback is nil")
    92  	}
    93  
    94  	var dat DroppedArgTracker
    95  	for i, k := range keys {
    96  		if !k.PartialValid(tcf.kc) {
    97  			cb(i, nil, MakeErrInvalidKey("key [%s] is not partially valid in context %s", k, tcf.kc).Err())
    98  			dat.MarkForRemoval(i, len(keys))
    99  			continue
   100  		}
   101  		if vals[i] == nil {
   102  			cb(i, nil, errors.New("datastore: PutMulti got nil vals entry"))
   103  			dat.MarkForRemoval(i, len(keys))
   104  		}
   105  	}
   106  	keys, vals, dal := dat.DropKeysAndVals(keys, vals)
   107  
   108  	if len(keys) == 0 {
   109  		return nil
   110  	}
   111  	return tcf.checkCtxDone(
   112  		tcf.RawInterface.PutMulti(keys, vals, func(idx int, key *Key, err error) {
   113  			cb(dal.OriginalIndex(idx), key, err)
   114  		}))
   115  }
   116  
   117  func (tcf *checkFilter) DeleteMulti(keys []*Key, cb DeleteMultiCB) error {
   118  	if cb == nil {
   119  		return fmt.Errorf("datastore: DeleteMulti callback is nil")
   120  	}
   121  	var dat DroppedArgTracker
   122  	for i, k := range keys {
   123  		var err error
   124  		switch {
   125  		case k.IsIncomplete():
   126  			err = MakeErrInvalidKey("key [%s] is incomplete", k).Err()
   127  		case !k.Valid(false, tcf.kc):
   128  			err = MakeErrInvalidKey("key [%s] is not valid in context %s", k, tcf.kc).Err()
   129  		}
   130  		if err != nil {
   131  			cb(i, err)
   132  			dat.MarkForRemoval(i, len(keys))
   133  		}
   134  	}
   135  	keys, dal := dat.DropKeys(keys)
   136  	if len(keys) == 0 {
   137  		return nil
   138  	}
   139  	return tcf.checkCtxDone(
   140  		tcf.RawInterface.DeleteMulti(keys, func(idx int, err error) {
   141  			cb(dal.OriginalIndex(idx), err)
   142  		}))
   143  }
   144  
   145  // checkCtxDone checks if the user context has done. If done, return ctx.Err().
   146  // Otherwise, return the original error.
   147  func (tcf *checkFilter) checkCtxDone(err error) error {
   148  	select {
   149  	case <-tcf.userCtx.Done():
   150  		return tcf.userCtx.Err()
   151  	default:
   152  		return err
   153  	}
   154  }
   155  
   156  func applyCheckFilter(c context.Context, i RawInterface) RawInterface {
   157  	return &checkFilter{
   158  		RawInterface: i,
   159  		kc:           GetKeyContext(c),
   160  		userCtx:      c,
   161  	}
   162  }