go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/service/datastore/interface.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  	"container/heap"
    19  	"context"
    20  	"fmt"
    21  	"reflect"
    22  	"sort"
    23  
    24  	"golang.org/x/sync/errgroup"
    25  
    26  	"go.chromium.org/luci/common/data/stringset"
    27  	"go.chromium.org/luci/common/errors"
    28  
    29  	multicursor "go.chromium.org/luci/gae/service/datastore/internal/protos/multicursor"
    30  )
    31  
    32  type resolvedRunCallback func(reflect.Value, CursorCB) error
    33  
    34  func parseRunCallback(cbIface any) (rcb resolvedRunCallback, isKey bool, mat *multiArgType, hasCursorCB bool) {
    35  	badSig := func() {
    36  		panic(fmt.Errorf(
    37  			"cb does not match the required callback signature: `%T` != `func(TYPE, [CursorCB]) [error]`",
    38  			cbIface))
    39  	}
    40  
    41  	if cbIface == nil {
    42  		badSig()
    43  	}
    44  
    45  	// TODO(riannucci): Profile and determine if any of this is causing a real
    46  	// slowdown. Could potentially cache reflection stuff by cbTyp?
    47  	cbVal := reflect.ValueOf(cbIface)
    48  	cbTyp := cbVal.Type()
    49  
    50  	if cbTyp.Kind() != reflect.Func {
    51  		badSig()
    52  	}
    53  
    54  	numIn := cbTyp.NumIn()
    55  	if numIn != 1 && numIn != 2 {
    56  		badSig()
    57  	}
    58  
    59  	firstArg := cbTyp.In(0)
    60  	if firstArg == typeOfKey {
    61  		isKey = true
    62  	} else {
    63  		mat = mustParseArg(firstArg, false)
    64  		if mat.newElem == nil {
    65  			badSig()
    66  		}
    67  	}
    68  
    69  	hasCursorCB = numIn == 2
    70  	if hasCursorCB && cbTyp.In(1) != typeOfCursorCB {
    71  		badSig()
    72  	}
    73  
    74  	if cbTyp.NumOut() > 1 {
    75  		badSig()
    76  	} else if cbTyp.NumOut() == 1 && cbTyp.Out(0) != typeOfError {
    77  		badSig()
    78  	}
    79  	hasErr := cbTyp.NumOut() == 1
    80  
    81  	// Resolve to generic function.
    82  	switch {
    83  	case hasErr && hasCursorCB:
    84  		// func(reflect.Value, CursorCB) error
    85  		rcb = func(v reflect.Value, cb CursorCB) error {
    86  			err := cbVal.Call([]reflect.Value{v, reflect.ValueOf(cb)})[0].Interface()
    87  			if err != nil {
    88  				return err.(error)
    89  			}
    90  			return nil
    91  		}
    92  
    93  	case hasErr && !hasCursorCB:
    94  		// func(reflect.Value) error
    95  		rcb = func(v reflect.Value, _ CursorCB) error {
    96  			err := cbVal.Call([]reflect.Value{v})[0].Interface()
    97  			if err != nil {
    98  				return err.(error)
    99  			}
   100  			return nil
   101  		}
   102  
   103  	case !hasErr && hasCursorCB:
   104  		// func(reflect.Value, CursorCB)
   105  		rcb = func(v reflect.Value, cb CursorCB) error {
   106  			cbVal.Call([]reflect.Value{v, reflect.ValueOf(cb)})
   107  			return nil
   108  		}
   109  
   110  	case !hasErr && !hasCursorCB:
   111  		// func(reflect.Value)
   112  		rcb = func(v reflect.Value, _ CursorCB) error {
   113  			cbVal.Call([]reflect.Value{v})
   114  			return nil
   115  		}
   116  
   117  	default:
   118  		badSig()
   119  	}
   120  
   121  	return
   122  }
   123  
   124  // AllocateIDs allows you to allocate IDs from the datastore without putting
   125  // any data.
   126  //
   127  // A partial valid key will be constructed from each entity's kind and parent,
   128  // if present. An allocation will then be performed against the datastore for
   129  // each key, and the partial key will be populated with a unique integer ID.
   130  // The resulting keys will be applied to their objects using PopulateKey. If
   131  // successful, any existing ID will be destroyed.
   132  //
   133  // If the object is supplied that cannot accept an integer key, this method
   134  // will panic.
   135  //
   136  // ent must be one of:
   137  //   - *S where S is a struct
   138  //   - *P where *P is a concrete type implementing PropertyLoadSaver
   139  //   - []S or []*S where S is a struct
   140  //   - []P or []*P where *P is a concrete type implementing PropertyLoadSaver
   141  //   - []I, where I is some interface type. Each element of the slice must have
   142  //     either *S or *P as its underlying type.
   143  //   - []*Key, to populate a slice of partial-valid keys.
   144  //
   145  // nil values (or interface-typed nils) are not allowed, neither as standalone
   146  // arguments nor inside slices. Passing them will cause a panic.
   147  //
   148  // If an error is encountered, the returned error value will depend on the
   149  // input arguments. If one argument is supplied, the result will be the
   150  // encountered error type. If multiple arguments are supplied, the result will
   151  // be a MultiError whose error index corresponds to the argument in which the
   152  // error was encountered.
   153  //
   154  // If an ent argument is a slice, its error type will be a MultiError. Note
   155  // that in the scenario where multiple slices are provided, this will return a
   156  // MultiError containing a nested MultiError for each slice argument.
   157  func AllocateIDs(c context.Context, ent ...any) error {
   158  	if len(ent) == 0 {
   159  		return nil
   160  	}
   161  
   162  	mma, err := makeMetaMultiArg(ent, mmaWriteKeys)
   163  	if err != nil {
   164  		panic(err)
   165  	}
   166  
   167  	keys, _, et := mma.getKeysPMs(GetKeyContext(c), false)
   168  	if len(keys) == 0 {
   169  		return nil
   170  	}
   171  
   172  	var dat DroppedArgTracker
   173  	dat.MarkNilKeys(keys)
   174  	keys, dal := dat.DropKeys(keys)
   175  
   176  	// Convert each key to be partial valid, assigning an integer ID of 0.
   177  	// Confirm that each object can be populated with such a key.
   178  	for compressedIdx, key := range keys {
   179  		keys[compressedIdx] = key.Incomplete()
   180  	}
   181  
   182  	err = Raw(c).AllocateIDs(keys, func(compressedIdx int, key *Key, err error) {
   183  		idx := dal.OriginalIndex(compressedIdx)
   184  
   185  		index := mma.index(idx)
   186  
   187  		if err != nil {
   188  			et.trackError(index, err)
   189  			return
   190  		}
   191  
   192  		mat, v := mma.get(index)
   193  		if !mat.setKey(v, key) {
   194  			et.trackError(index, MakeErrInvalidKey("failed to export key [%s]", key).Err())
   195  			return
   196  		}
   197  	})
   198  	if err == nil {
   199  		err = et.error()
   200  	}
   201  	return maybeSingleError(err, ent)
   202  }
   203  
   204  // KeyForObj extracts a key from src.
   205  //
   206  // It is the same as KeyForObjErr, except that if KeyForObjErr would have
   207  // returned an error, this method panics. It's safe to use if you know that
   208  // src statically meets the metadata constraints described by KeyForObjErr.
   209  func KeyForObj(c context.Context, src any) *Key {
   210  	ret, err := KeyForObjErr(c, src)
   211  	if err != nil {
   212  		panic(err)
   213  	}
   214  	return ret
   215  }
   216  
   217  // KeyForObjErr extracts a key from src.
   218  //
   219  // src must be one of:
   220  //   - *S, where S is a struct
   221  //   - a PropertyLoadSaver
   222  //
   223  // It is expected that the struct exposes the following metadata (as retrieved
   224  // by MetaGetter.GetMeta):
   225  //   - "key" (type: Key) - The full datastore key to use. Must not be nil.
   226  //     OR
   227  //   - "id" (type: int64 or string) - The id of the Key to create.
   228  //   - "kind" (optional, type: string) - The kind of the Key to create. If
   229  //     blank or not present, KeyForObjErr will extract the name of the src
   230  //     object's type.
   231  //   - "parent" (optional, type: Key) - The parent key to use.
   232  //
   233  // By default, the metadata will be extracted from the struct and its tagged
   234  // properties. However, if the struct implements MetaGetterSetter it is
   235  // wholly responsible for exporting the required fields. A struct that
   236  // implements GetMeta to make some minor tweaks can evoke the defualt behavior
   237  // by using GetPLS(s).GetMeta.
   238  //
   239  // If a required metadata item is missing or of the wrong type, then this will
   240  // return an error.
   241  func KeyForObjErr(c context.Context, src any) (*Key, error) {
   242  	return GetKeyContext(c).NewKeyFromMeta(getMGS(src))
   243  }
   244  
   245  // MakeKey is a convenience method for manufacturing a *Key. It should only be
   246  // used when elems... is known statically (e.g. in the code) to be correct.
   247  //
   248  // elems is pairs of (string, string|int|int32|int64) pairs, which correspond
   249  // to Kind/id pairs. Example:
   250  //
   251  //	dstore.MakeKey("Parent", 1, "Child", "id")
   252  //
   253  // Would create the key:
   254  //
   255  //	<current appID>:<current Namespace>:/Parent,1/Child,id
   256  //
   257  // If elems is not parsable (e.g. wrong length, wrong types, etc.) this method
   258  // will panic.
   259  func MakeKey(c context.Context, elems ...any) *Key {
   260  	kc := GetKeyContext(c)
   261  	return kc.MakeKey(elems...)
   262  }
   263  
   264  // NewKey constructs a new key in the current appID/Namespace, using the
   265  // specified parameters.
   266  func NewKey(c context.Context, kind, stringID string, intID int64, parent *Key) *Key {
   267  	kc := GetKeyContext(c)
   268  	return kc.NewKey(kind, stringID, intID, parent)
   269  }
   270  
   271  // NewIncompleteKeys allocates count incomplete keys sharing the same kind and
   272  // parent. It is useful as input to AllocateIDs.
   273  func NewIncompleteKeys(c context.Context, count int, kind string, parent *Key) (keys []*Key) {
   274  	kc := GetKeyContext(c)
   275  	if count > 0 {
   276  		keys = make([]*Key, count)
   277  		for i := range keys {
   278  			keys[i] = kc.NewKey(kind, "", 0, parent)
   279  		}
   280  	}
   281  	return
   282  }
   283  
   284  // NewKeyToks constructs a new key in the current appID/Namespace, using the
   285  // specified key tokens.
   286  func NewKeyToks(c context.Context, toks []KeyTok) *Key {
   287  	kc := GetKeyContext(c)
   288  	return kc.NewKeyToks(toks)
   289  }
   290  
   291  // PopulateKey loads key into obj.
   292  //
   293  // obj is any object that Interface.Get is able to accept.
   294  //
   295  // Upon successful application, this method will return true. If the key could
   296  // not be applied to the object, this method will return false. It will panic if
   297  // obj is an invalid datastore model.
   298  func PopulateKey(obj any, key *Key) bool {
   299  	return populateKeyMGS(getMGS(obj), key)
   300  }
   301  
   302  func populateKeyMGS(mgs MetaGetterSetter, key *Key) bool {
   303  	setViaKey := mgs.SetMeta("key", key)
   304  
   305  	lst := key.LastTok()
   306  	mgs.SetMeta("kind", lst.Kind)
   307  	mgs.SetMeta("parent", key.Parent())
   308  
   309  	setViaID := false
   310  	if lst.StringID != "" {
   311  		setViaID = mgs.SetMeta("id", lst.StringID)
   312  	} else {
   313  		setViaID = mgs.SetMeta("id", lst.IntID)
   314  	}
   315  
   316  	return setViaKey || setViaID
   317  }
   318  
   319  // RunInTransaction runs f inside of a transaction. See the appengine SDK's
   320  // documentation for full details on the behavior of transactions in the
   321  // datastore.
   322  //
   323  // Note that the behavior of transactions may change depending on what filters
   324  // have been installed. It's possible that we'll end up implementing things
   325  // like nested/buffered transactions as filters.
   326  func RunInTransaction(c context.Context, f func(c context.Context) error, opts *TransactionOptions) error {
   327  	return Raw(c).RunInTransaction(f, opts)
   328  }
   329  
   330  // Run executes the given query, and calls `cb` for each successfully
   331  // retrieved item.
   332  //
   333  // By default, datastore applies a short (~5s) timeout to queries. This can be
   334  // increased, usually to around several minutes, by explicitly setting a
   335  // deadline on the supplied Context.
   336  //
   337  // cb is a callback function whose signature is
   338  //
   339  //	func(obj TYPE[, getCursor CursorCB]) [error]
   340  //
   341  // Where TYPE is one of:
   342  //   - S or *S, where S is a struct
   343  //   - P or *P, where *P is a concrete type implementing PropertyLoadSaver
   344  //   - *Key (implies a keys-only query)
   345  //
   346  // If the error is omitted from the signature, this will run until the query
   347  // returns all its results, or has an error/times out.
   348  //
   349  // If error is in the signature, the query will continue as long as the
   350  // callback returns nil. If it returns `Stop`, the query will stop and Run
   351  // will return nil. Otherwise, the query will stop and Run will return the
   352  // user's error.
   353  //
   354  // Run may also stop on the first datastore error encountered, which can occur
   355  // due to flakiness, timeout, etc. If it encounters such an error, it will
   356  // be returned.
   357  func Run(c context.Context, q *Query, cb any) error {
   358  	rcb, isKey, mat, _ := parseRunCallback(cb)
   359  
   360  	if isKey {
   361  		q = q.KeysOnly(true)
   362  	}
   363  	fq, err := q.Finalize()
   364  	if err != nil {
   365  		return err
   366  	}
   367  
   368  	raw := Raw(c)
   369  
   370  	if isKey {
   371  		err = raw.Run(fq, func(k *Key, _ PropertyMap, gc CursorCB) error {
   372  			return rcb(reflect.ValueOf(k), gc)
   373  		})
   374  	} else {
   375  		err = raw.Run(fq, func(k *Key, pm PropertyMap, gc CursorCB) error {
   376  			itm := mat.newElem()
   377  			if err := mat.setPM(itm, pm); err != nil {
   378  				return err
   379  			}
   380  			mat.setKey(itm, k)
   381  			return rcb(itm, gc)
   382  		})
   383  	}
   384  	return filterStop(err)
   385  }
   386  
   387  // RunMulti executes the logical OR of multiple queries, calling `cb` for each
   388  // unique entity (by *Key) that it finds. Results will be returned in the order
   389  // of the provided queries; All queries must have matching Orders.
   390  //
   391  // cb is a callback function (please refer to the `Run` function comments for
   392  // formats and restrictions for `cb` in this file).
   393  //
   394  // The cursor that is returned by the callback cannot be used on a single query
   395  // by doing `query.Start(cursor)` (In some cases it may not even complain when
   396  // you try to do this. But the results are undefined). Apply the cursor to the
   397  // same list of queries using ApplyCursors.
   398  //
   399  // Note: projection queries are not supported, as they are non-trivial in
   400  // complexity and haven't been needed yet.
   401  //
   402  // Note: The cb is called for every unique entity (by *Key) that is retrieved
   403  // on the current run. It is possible to get the same entity twice over two
   404  // calls to RunMulti with different cursors.
   405  //
   406  // DANGER: Cursors are buggy when using Cloud Datastore production backend.
   407  // Paginated queries skip entities sitting on page boundaries. This doesn't
   408  // happen when using `impl/memory` and thus hard to spot in unit tests. See
   409  // queryIterator doc for more details.
   410  func RunMulti(c context.Context, queries []*Query, cb any) error {
   411  	rcb, isKey, mat, hasCursorCB := parseRunCallback(cb)
   412  
   413  	// A helper that passes an entity to the user callback.
   414  	var dispatchEntity func(key *Key, pm PropertyMap, ccb CursorCB) error
   415  	if isKey {
   416  		dispatchEntity = func(key *Key, _ PropertyMap, ccb CursorCB) error {
   417  			return rcb(reflect.ValueOf(key), ccb)
   418  		}
   419  	} else {
   420  		dispatchEntity = func(key *Key, pm PropertyMap, ccb CursorCB) error {
   421  			itm := mat.newElem()
   422  			if err := mat.setPM(itm, pm); err != nil {
   423  				return err
   424  			}
   425  			mat.setKey(itm, key)
   426  			return rcb(itm, ccb)
   427  		}
   428  	}
   429  
   430  	// Finalize queries and do some basic validation. At very least queries must
   431  	// use the same kind and ordering, otherwise putting their results in a single
   432  	// sorted heap makes no sense.
   433  	finalized := make([]*FinalizedQuery, len(queries))
   434  	overallKind := ""
   435  	overallOrder := ""
   436  	for i, q := range queries {
   437  		if isKey {
   438  			q = q.KeysOnly(true)
   439  		}
   440  		fq, err := q.Finalize()
   441  		if err != nil {
   442  			return err
   443  		}
   444  		finalized[i] = fq
   445  		// Build a string identifying ordering of this query, e.g.
   446  		// "-field1,field2,__key__".
   447  		order := ""
   448  		for j, col := range fq.orders {
   449  			if j != 0 {
   450  				order += ","
   451  			}
   452  			order += col.String()
   453  		}
   454  		switch {
   455  		case i == 0:
   456  			overallKind = fq.kind
   457  			overallOrder = order
   458  		case fq.kind != overallKind:
   459  			return fmt.Errorf("all RunMulti queries should query the same kind, but got %q and %q", fq.kind, overallKind)
   460  		case order != overallOrder:
   461  			return fmt.Errorf("all RunMulti queries should use the same order, but got %q and %q", order, overallOrder)
   462  		}
   463  	}
   464  
   465  	// No queries to run => no results to return. This is an edge case.
   466  	if len(finalized) == 0 {
   467  		return nil
   468  	}
   469  
   470  	// If we have only one query, just run it directly without any extra
   471  	// synchronization overhead. Just make sure to use the correct cursor format.
   472  	// This is worth optimizing since running only one query is a very very
   473  	// common case.
   474  	if len(finalized) == 1 {
   475  		var err error
   476  		if hasCursorCB {
   477  			err = Raw(c).Run(finalized[0], func(key *Key, pm PropertyMap, cursorCB CursorCB) error {
   478  				return dispatchEntity(key, pm, func() (Cursor, error) {
   479  					cur, err := cursorCB()
   480  					if err != nil {
   481  						return nil, err
   482  					}
   483  					cursorStr := ""
   484  					if cur != nil {
   485  						cursorStr = cur.String()
   486  					}
   487  					return multiCursor{
   488  						curs: &multicursor.Cursors{
   489  							Version:     multiCursorVersion,
   490  							MagicNumber: multiCursorMagic,
   491  							Cursors:     []string{cursorStr},
   492  						},
   493  					}, nil
   494  				})
   495  			})
   496  		} else {
   497  			err = Raw(c).Run(finalized[0], func(key *Key, pm PropertyMap, _ CursorCB) error {
   498  				return dispatchEntity(key, pm, nil)
   499  			})
   500  		}
   501  		return filterStop(err)
   502  	}
   503  
   504  	// All iterators (active and exhausted) in some arbitrary order.
   505  	iterators := make([]*queryIterator, 0, len(finalized))
   506  
   507  	c, cancel := context.WithCancel(c)
   508  	eg, ectx := errgroup.WithContext(c)
   509  
   510  	// Make sure all spawned goroutines have fully stopped before returning.
   511  	defer func() {
   512  		// Signal all iterators to stop ASAP.
   513  		cancel()
   514  		// Wait for all of them to stop. Calling Next makes sure internal goroutines
   515  		// are not getting stuck trying to write to a channel that nothing is
   516  		// reading from (this blocks forever).
   517  		for _, iter := range iterators {
   518  			for done := false; !done; done, _ = iter.Next() {
   519  			}
   520  		}
   521  		// All goroutines should be stopping now. Wait until they are fully stopped.
   522  		_ = eg.Wait()
   523  	}()
   524  
   525  	// Launch all queries in parallel. Do it before ordering them as a heap, since
   526  	// to build a heap we need to have the first result from each query. We want
   527  	// all such first results to be fetched *in parallel*.
   528  	for _, fq := range finalized {
   529  		iterators = append(iterators, startQueryIterator(ectx, eg, fq))
   530  	}
   531  
   532  	// Wait for first items from all iterators. Gather all non-exhausted iterators
   533  	// to make a sorted heap out of them.
   534  	iHeap := make(iteratorHeap, 0, len(iterators))
   535  	for _, iter := range iterators {
   536  		switch done, err := iter.Next(); {
   537  		case err != nil:
   538  			return err // the defer will clean up everything
   539  		case !done:
   540  			iHeap = append(iHeap, iter)
   541  		}
   542  	}
   543  	heap.Init(&iHeap)
   544  
   545  	// ccb is the cursor callback for RunMulti. This grabs all the cursors for the
   546  	// queries involved and returns a single cursor. It is only executed if the
   547  	// user callback invokes it.
   548  	var ccb CursorCB
   549  	if hasCursorCB {
   550  		ccb = func() (Cursor, error) {
   551  			// Sort the list of queries. It is OK to update `iterators` in-place here.
   552  			// It is only used in the defer, the order doesn't matter there.
   553  			sort.Slice(iterators, func(i, j int) bool {
   554  				queryI := iterators[i].Query()
   555  				queryJ := iterators[j].Query()
   556  				return queryI.Less(queryJ)
   557  			})
   558  			// Create the cursor. It points to all items currently sitting in heap.
   559  			// We'll need to refetch them all again to repopulate the heap when
   560  			// resuming the query.
   561  			var curs multicursor.Cursors
   562  			curs.MagicNumber = multiCursorMagic
   563  			curs.Version = multiCursorVersion
   564  			for _, iter := range iterators {
   565  				switch cur, err := iter.CurrentCursor(); {
   566  				case err != nil:
   567  					return nil, err
   568  				case cur != nil:
   569  					curs.Cursors = append(curs.Cursors, cur.String())
   570  				default:
   571  					curs.Cursors = append(curs.Cursors, "")
   572  				}
   573  			}
   574  			return multiCursor{curs: &curs}, nil
   575  		}
   576  	}
   577  
   578  	// If queries are ordered only by key, all duplicates will be returned from
   579  	// the heap one after another and we can use a simple check to skip them. This
   580  	// is important for CountMulti(...) that can be visiting tens of thousands
   581  	// of entities: storing them all in a hash map for deduplication is a waste of
   582  	// memory.
   583  	//
   584  	// Use a hash map for any other ordering. There may be weird results if this
   585  	// is running non-transactionally and two different subqueries see two
   586  	// different versions of the same entity (with different values of fields
   587  	// affecting the order). Such entity will appear twice in the output, with
   588  	// some other entities in between these appearances. A simple check will not
   589  	// detect such deduplication.
   590  	var seenKey func(keyStr string) bool
   591  	if overallOrder == "__key__" || overallOrder == "-__key__" {
   592  		lastSeen := ""
   593  		seenKey = func(keyStr string) bool {
   594  			if lastSeen == keyStr {
   595  				return true
   596  			}
   597  			lastSeen = keyStr
   598  			return false
   599  		}
   600  	} else {
   601  		seenKeys := stringset.New(128)
   602  		seenKey = func(keyStr string) bool {
   603  			return !seenKeys.Add(keyStr)
   604  		}
   605  	}
   606  
   607  	// Merge query results.
   608  	for iHeap.Len() > 0 {
   609  		pm, key, keyStr, err := iHeap.nextData()
   610  		if err != nil {
   611  			return err
   612  		}
   613  		if !seenKey(keyStr) {
   614  			if err := dispatchEntity(key, pm, ccb); err != nil {
   615  				return filterStop(err)
   616  			}
   617  		}
   618  	}
   619  	return nil
   620  }
   621  
   622  // Count executes the given query and returns the number of entries which
   623  // match it.
   624  //
   625  // If the query is marked as eventually consistent via EventualConsistency(true)
   626  // will use a fast server-side aggregation, with the downside that such queries
   627  // may return slightly stale results and can't be used inside transactions.
   628  //
   629  // If the query is strongly consistent, will essentially do a full keys-only
   630  // query and count the number of matches locally.
   631  func Count(c context.Context, q *Query) (int64, error) {
   632  	fq, err := q.Finalize()
   633  	if err != nil {
   634  		return 0, err
   635  	}
   636  	v, err := Raw(c).Count(fq)
   637  	return v, filterStop(err)
   638  }
   639  
   640  // CountMulti runs multiple queries in parallel and counts the total number of
   641  // unique entities produced by them.
   642  //
   643  // Unlike Count, this method doesn't support server-side aggregation. It always
   644  // does full keys-only queries. If you have only one query and don't care about
   645  // strong consistency, use `Count(c, q.EventualConsistency(true))`: it will use
   646  // the server-side aggregation which is orders of magnitude faster than the
   647  // local counting.
   648  func CountMulti(c context.Context, queries []*Query) (int64, error) {
   649  	var count int64
   650  	err := RunMulti(c, queries,
   651  		func(_ *Key) error {
   652  			// RunMulti already does deduplication, we just need to count unique hits.
   653  			count++
   654  			return nil
   655  		},
   656  	)
   657  	if err != nil {
   658  		return 0, err
   659  	}
   660  	return count, nil
   661  }
   662  
   663  // DecodeCursor converts a string returned by a Cursor into a Cursor instance.
   664  // It will return an error if the supplied string is not valid, or could not
   665  // be decoded by the implementation.
   666  func DecodeCursor(c context.Context, s string) (Cursor, error) {
   667  	return Raw(c).DecodeCursor(s)
   668  }
   669  
   670  // GetAll retrieves all of the Query results into dst.
   671  //
   672  // By default, datastore applies a short (~5s) timeout to queries. This can be
   673  // increased, usually to around several minutes, by explicitly setting a
   674  // deadline on the supplied Context.
   675  //
   676  // dst must be one of:
   677  //   - *[]S or *[]*S, where S is a struct
   678  //   - *[]P or *[]*P, where *P is a concrete type implementing
   679  //     PropertyLoadSaver
   680  //   - *[]*Key implies a keys-only query.
   681  func GetAll(c context.Context, q *Query, dst any) error {
   682  	return getAllRaw(Raw(c), q, dst)
   683  }
   684  
   685  func getAllRaw(raw RawInterface, q *Query, dst any) error {
   686  	v := reflect.ValueOf(dst)
   687  	if v.Kind() != reflect.Ptr {
   688  		panic(fmt.Errorf("invalid GetAll dst: must have a ptr-to-slice: %T", dst))
   689  	}
   690  	if !v.IsValid() || v.IsNil() {
   691  		panic(errors.New("invalid GetAll dst: <nil>"))
   692  	}
   693  
   694  	if keys, ok := dst.(*[]*Key); ok {
   695  		fq, err := q.KeysOnly(true).Finalize()
   696  		if err != nil {
   697  			return err
   698  		}
   699  
   700  		return raw.Run(fq, func(k *Key, _ PropertyMap, _ CursorCB) error {
   701  			*keys = append(*keys, k)
   702  			return nil
   703  		})
   704  	}
   705  	fq, err := q.Finalize()
   706  	if err != nil {
   707  		return err
   708  	}
   709  
   710  	slice := v.Elem()
   711  	mat := mustParseMultiArg(slice.Type())
   712  	if mat.newElem == nil {
   713  		panic(fmt.Errorf("invalid GetAll dst (non-concrete element type): %T", dst))
   714  	}
   715  
   716  	errs := map[int]error{}
   717  	i := 0
   718  	err = filterStop(raw.Run(fq, func(k *Key, pm PropertyMap, _ CursorCB) error {
   719  		slice.Set(reflect.Append(slice, mat.newElem()))
   720  		itm := slice.Index(i)
   721  		mat.setKey(itm, k)
   722  		err := mat.setPM(itm, pm)
   723  		if err != nil {
   724  			errs[i] = err
   725  		}
   726  		i++
   727  		return nil
   728  	}))
   729  	if err == nil {
   730  		if len(errs) > 0 {
   731  			me := make(errors.MultiError, slice.Len())
   732  			for i, e := range errs {
   733  				me[i] = e
   734  			}
   735  			err = me
   736  		}
   737  	}
   738  	return err
   739  }
   740  
   741  // Exists tests if the supplied objects are present in the datastore.
   742  //
   743  // ent must be one of:
   744  //   - *S, where S is a struct
   745  //   - *P, where *P is a concrete type implementing PropertyLoadSaver
   746  //   - []S or []*S, where S is a struct
   747  //   - []P or []*P, where *P is a concrete type implementing PropertyLoadSaver
   748  //   - []I, where I is some interface type. Each element of the slice must have
   749  //     either *S or *P as its underlying type.
   750  //   - *Key, to check a specific key from the datastore.
   751  //   - []*Key, to check a slice of keys from the datastore.
   752  //
   753  // nil values (or interface-typed nils) are not allowed, neither as standalone
   754  // arguments nor inside slices. Passing them will cause a panic.
   755  //
   756  // If an error is encountered, the returned error value will depend on the
   757  // input arguments. If one argument is supplied, the result will be the
   758  // encountered error type. If multiple arguments are supplied, the result will
   759  // be a MultiError whose error index corresponds to the argument in which the
   760  // error was encountered.
   761  //
   762  // If an ent argument is a slice, its error type will be a MultiError. Note
   763  // that in the scenario, where multiple slices are provided, this will return a
   764  // MultiError containing a nested MultiError for each slice argument.
   765  func Exists(c context.Context, ent ...any) (*ExistsResult, error) {
   766  	if len(ent) == 0 {
   767  		return nil, nil
   768  	}
   769  
   770  	mma, err := makeMetaMultiArg(ent, mmaKeysOnly)
   771  	if err != nil {
   772  		panic(err)
   773  	}
   774  
   775  	keys, _, et := mma.getKeysPMs(GetKeyContext(c), false)
   776  	if len(keys) == 0 {
   777  		return nil, nil
   778  	}
   779  
   780  	var dat DroppedArgTracker
   781  	dat.MarkNilKeys(keys)
   782  	keys, dal := dat.DropKeys(keys)
   783  
   784  	bt := newBoolTracker(mma, et)
   785  	err = Raw(c).GetMulti(keys, nil, func(compressedIdx int, _ PropertyMap, err error) {
   786  		idx := dal.OriginalIndex(compressedIdx)
   787  		bt.trackExistsResult(mma.index(idx), err)
   788  	})
   789  
   790  	if err == nil {
   791  		err = bt.error()
   792  	}
   793  	return bt.result(), maybeSingleError(err, ent)
   794  }
   795  
   796  // Get retrieves objects from the datastore.
   797  //
   798  // Each element in dst must be one of:
   799  //   - *S, where S is a struct
   800  //   - *P, where *P is a concrete type implementing PropertyLoadSaver
   801  //   - []S or []*S, where S is a struct
   802  //   - []P or []*P, where *P is a concrete type implementing PropertyLoadSaver
   803  //   - []I, where I is some interface type. Each element of the slice must have
   804  //     either *S or *P as its underlying type.
   805  //
   806  // nil values (or interface-typed nils) are not allowed, neither as standalone
   807  // arguments nor inside slices. Passing them will cause a panic.
   808  //
   809  // If an error is encountered, the returned error value will depend on the
   810  // input arguments. If one argument is supplied, the result will be the
   811  // encountered error type. If multiple arguments are supplied, the result will
   812  // be a MultiError whose error index corresponds to the argument in which the
   813  // error was encountered.
   814  //
   815  // If a dst argument is a slice, its error type will be a MultiError. Note
   816  // that in the scenario where multiple slices are provided, this will return a
   817  // MultiError containing a nested MultiError for each slice argument.
   818  //
   819  // If there was an issue retrieving the entity, the input `dst` objects will
   820  // not be affected. This means that you can populate an object for dst with some
   821  // values, do a Get, and on an ErrNoSuchEntity, do a Put (inside a transaction,
   822  // of course :)).
   823  func Get(c context.Context, dst ...any) error {
   824  	if len(dst) == 0 {
   825  		return nil
   826  	}
   827  
   828  	mma, err := makeMetaMultiArg(dst, mmaReadWrite)
   829  	if err != nil {
   830  		panic(err)
   831  	}
   832  
   833  	keys, pms, et := mma.getKeysPMs(GetKeyContext(c), true)
   834  	if len(keys) == 0 {
   835  		return nil
   836  	}
   837  
   838  	var dat DroppedArgTracker
   839  	dat.MarkNilKeysVals(keys, pms)
   840  	keys, pms, dal := dat.DropKeysAndVals(keys, pms)
   841  
   842  	meta := NewMultiMetaGetter(pms)
   843  	err = Raw(c).GetMulti(keys, meta, func(compressedIdx int, pm PropertyMap, err error) {
   844  		idx := dal.OriginalIndex(compressedIdx)
   845  		index := mma.index(idx)
   846  		if err != nil {
   847  			et.trackError(index, err)
   848  			return
   849  		}
   850  
   851  		mat, v := mma.get(index)
   852  		if err := mat.setPM(v, pm); err != nil {
   853  			et.trackError(index, err)
   854  			return
   855  		}
   856  	})
   857  
   858  	if err == nil {
   859  		err = et.error()
   860  	}
   861  	return maybeSingleError(err, dst)
   862  }
   863  
   864  // Put writes objects into the datastore.
   865  //
   866  // src must be one of:
   867  //   - *S, where S is a struct
   868  //   - *P, where *P is a concrete type implementing PropertyLoadSaver
   869  //   - []S or []*S, where S is a struct
   870  //   - []P or []*P, where *P is a concrete type implementing PropertyLoadSaver
   871  //   - []I, where I is some interface type. Each element of the slice must have
   872  //     either *S or *P as its underlying type.
   873  //
   874  // nil values (or interface-typed nils) are not allowed, neither as standalone
   875  // arguments nor inside slices. Passing them will cause a panic.
   876  //
   877  // A *Key will be extracted from src via KeyForObj. If
   878  // extractedKey.IsIncomplete() is true, and the object is put to the datastore
   879  // successfully, then Put will write the resolved (datastore-generated) *Key
   880  // back to src.
   881  //
   882  // NOTE: The datastore only autogenerates *Keys with integer IDs. Only models
   883  // which use a raw `$key` or integer-typed `$id` field are elegible for this.
   884  // A model with a string-typed `$id` field will not accept an integer id'd *Key
   885  // and will cause the Put to fail.
   886  //
   887  // If an error is encountered, the returned error value will depend on the
   888  // input arguments. If one argument is supplied, the result will be the
   889  // encountered error type. If multiple arguments are supplied, the result will
   890  // be a MultiError whose error index corresponds to the argument in which the
   891  // error was encountered.
   892  //
   893  // If a src argument is a slice, its error type will be a MultiError. Note
   894  // that in the scenario where multiple slices are provided, this will return a
   895  // MultiError containing a nested MultiError for each slice argument.
   896  func Put(c context.Context, src ...any) error {
   897  	return putRaw(Raw(c), GetKeyContext(c), src)
   898  }
   899  
   900  func putRaw(raw RawInterface, kctx KeyContext, src []any) error {
   901  	if len(src) == 0 {
   902  		return nil
   903  	}
   904  
   905  	mma, err := makeMetaMultiArg(src, mmaReadWrite)
   906  	if err != nil {
   907  		panic(err)
   908  	}
   909  
   910  	keys, vals, et := mma.getKeysPMs(kctx, false)
   911  	if len(keys) == 0 {
   912  		return nil
   913  	}
   914  
   915  	var dat DroppedArgTracker
   916  	dat.MarkNilKeysVals(keys, vals)
   917  	keys, vals, dal := dat.DropKeysAndVals(keys, vals)
   918  
   919  	err = raw.PutMulti(keys, vals, func(compressedIdx int, key *Key, err error) {
   920  		idx := dal.OriginalIndex(compressedIdx)
   921  		index := mma.index(idx)
   922  
   923  		if err != nil {
   924  			et.trackError(index, err)
   925  			return
   926  		}
   927  
   928  		if !key.Equal(keys[compressedIdx]) {
   929  			mat, v := mma.get(index)
   930  			mat.setKey(v, key)
   931  		}
   932  	})
   933  	if err == nil {
   934  		err = et.error()
   935  	}
   936  	return maybeSingleError(err, src)
   937  }
   938  
   939  // Delete removes the supplied entities from the datastore.
   940  //
   941  // ent must be one of:
   942  //   - *S, where S is a struct
   943  //   - *P, where *P is a concrete type implementing PropertyLoadSaver
   944  //   - []S or []*S, where S is a struct
   945  //   - []P or []*P, where *P is a concrete type implementing PropertyLoadSaver
   946  //   - []I, where I is some interface type. Each element of the slice must have
   947  //     either *S or *P as its underlying type.
   948  //   - *Key, to remove a specific key from the datastore.
   949  //   - []*Key, to remove a slice of keys from the datastore.
   950  //
   951  // nil values (or interface-typed nils) are not allowed, neither as standalone
   952  // arguments nor inside slices. Passing them will cause a panic.
   953  //
   954  // If an error is encountered, the returned error value will depend on the
   955  // input arguments. If one argument is supplied, the result will be the
   956  // encountered error type. If multiple arguments are supplied, the result will
   957  // be a MultiError whose error index corresponds to the argument in which the
   958  // error was encountered.
   959  //
   960  // If an ent argument is a slice, its error type will be a MultiError. Note
   961  // that in the scenario where multiple slices are provided, this will return a
   962  // MultiError containing a nested MultiError for each slice argument.
   963  func Delete(c context.Context, ent ...any) error {
   964  	if len(ent) == 0 {
   965  		return nil
   966  	}
   967  
   968  	mma, err := makeMetaMultiArg(ent, mmaKeysOnly)
   969  	if err != nil {
   970  		panic(err)
   971  	}
   972  
   973  	keys, _, et := mma.getKeysPMs(GetKeyContext(c), false)
   974  	if len(keys) == 0 {
   975  		return nil
   976  	}
   977  
   978  	var dat DroppedArgTracker
   979  	dat.MarkNilKeys(keys)
   980  	keys, dal := dat.DropKeys(keys)
   981  
   982  	err = Raw(c).DeleteMulti(keys, func(compressedIdx int, err error) {
   983  		idx := dal.OriginalIndex(compressedIdx)
   984  
   985  		if err != nil {
   986  			index := mma.index(idx)
   987  			et.trackError(index, err)
   988  		}
   989  	})
   990  
   991  	if err == nil {
   992  		err = et.error()
   993  	}
   994  	return maybeSingleError(err, ent)
   995  }
   996  
   997  // GetTestable returns the Testable interface for the implementation, or nil if
   998  // there is none.
   999  func GetTestable(c context.Context) Testable {
  1000  	return Raw(c).GetTestable()
  1001  }
  1002  
  1003  // maybeSingleError normalizes the error experience between single- and
  1004  // multi-element API calls.
  1005  //
  1006  // Single-element API calls will return a single error for that element, while
  1007  // multi-element API calls will return a MultiError, one for each element. This
  1008  // accepts the slice of elements that is being operated on and determines what
  1009  // sort of error to return.
  1010  func maybeSingleError(err error, elems []any) error {
  1011  	if err == nil {
  1012  		return nil
  1013  	}
  1014  	if len(elems) == 1 {
  1015  		return errors.SingleError(err)
  1016  	}
  1017  	return err
  1018  }
  1019  
  1020  func filterStop(err error) error {
  1021  	if err == Stop {
  1022  		err = nil
  1023  	}
  1024  	return err
  1025  }
  1026  
  1027  // a min heap for a slice of queryIterator.
  1028  //
  1029  // All iterators are in "not done" state.
  1030  type iteratorHeap []*queryIterator
  1031  
  1032  var _ heap.Interface = &iteratorHeap{}
  1033  
  1034  func (h iteratorHeap) Len() int { return len(h) }
  1035  
  1036  func (h iteratorHeap) Less(i, j int) bool { return h[i].CurrentItemOrder() < h[j].CurrentItemOrder() }
  1037  
  1038  func (h iteratorHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
  1039  
  1040  func (h *iteratorHeap) Push(x any) {
  1041  	*h = append(*h, x.(*queryIterator))
  1042  }
  1043  
  1044  func (h *iteratorHeap) Pop() any {
  1045  	old := *h
  1046  	n := len(old)
  1047  	item := old[n-1]
  1048  	*h = old[0 : n-1]
  1049  	return item
  1050  }
  1051  
  1052  // nextData returns data of the peak queryIterator, advances the queryIterator
  1053  // and either removes it from the heap (if it has no results left) or adjusts
  1054  // its position in the heap.
  1055  //
  1056  // Must be called only with a non-empty heap.
  1057  func (h *iteratorHeap) nextData() (pm PropertyMap, key *Key, keyStr string, err error) {
  1058  	if len(*h) == 0 {
  1059  		panic("the heap is empty")
  1060  	}
  1061  
  1062  	qi := (*h)[0]
  1063  	key, pm = qi.CurrentItem()
  1064  	keyStr = qi.CurrentItemKey()
  1065  
  1066  	var done bool
  1067  	done, err = qi.Next()
  1068  	if !done {
  1069  		heap.Fix(h, 0)
  1070  	} else {
  1071  		heap.Remove(h, 0)
  1072  	}
  1073  
  1074  	return
  1075  }