blob: bc8415147feb668368e47bd4c2373303cded7d42 [file] [log] [blame]
// Copyright 2015 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package datastore
import (
"bytes"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"strings"
"google.golang.org/protobuf/proto"
pb "go.chromium.org/luci/gae/service/datastore/internal/protos/datastore"
)
// KeyTok is a single token from a multi-part Key.
type KeyTok struct {
Kind string
IntID int64
StringID string
}
// IsIncomplete returns true iff this token doesn't define either a StringID or
// an IntID.
func (k KeyTok) IsIncomplete() bool {
return k.StringID == "" && k.IntID == 0
}
// Special returns true iff this token begins and ends with "__"
func (k KeyTok) Special() bool {
return len(k.Kind) >= 2 && k.Kind[:2] == "__" && k.Kind[len(k.Kind)-2:] == "__"
}
// ID returns the 'active' id as a Property (either the StringID or the IntID).
func (k KeyTok) ID() Property {
if k.StringID != "" {
return MkProperty(k.StringID)
}
return MkProperty(k.IntID)
}
// Less returns true iff k would sort before other.
func (k KeyTok) Less(other KeyTok) bool {
if k.Kind < other.Kind {
return true
} else if k.Kind > other.Kind {
return false
}
a, b := k.ID(), other.ID()
return a.Less(&b)
}
// KeyContext is the context in which a key is generated.
type KeyContext struct {
AppID string
Namespace string
}
// MkKeyContext is a helper function to create a new KeyContext.
//
// It is preferable to field-based struct initialization because, as a function,
// it has the ability to enforce an exact number of parameters.
func MkKeyContext(appID, namespace string) KeyContext {
return KeyContext{AppID: appID, Namespace: namespace}
}
// Matches returns true iff the AppID and Namespace parameters are the same for
// the two KeyContext instances.
func (kc KeyContext) Matches(o KeyContext) bool {
return (kc.AppID == o.AppID && kc.Namespace == o.Namespace)
}
// NewKeyToks creates a new Key. It is the Key implementation returned from
// the various PropertyMap serialization routines, as well as the native key
// implementation for the in-memory implementation of gae.
//
// See NewKeyToks for a version of this function which automatically
// provides aid and ns.
func (kc KeyContext) NewKeyToks(toks []KeyTok) *Key {
if len(toks) == 0 {
return nil
}
newToks := make([]KeyTok, len(toks))
copy(newToks, toks)
return &Key{kc, newToks}
}
// NewKey is a wrapper around NewToks which has an interface similar to NewKey
// in the SDK.
//
// See NewKey for a version of this function which automatically provides aid
// and ns.
func (kc KeyContext) NewKey(kind, stringID string, intID int64, parent *Key) *Key {
if parent == nil {
return &Key{kc, []KeyTok{{kind, intID, stringID}}}
}
toks := parent.toks
newToks := make([]KeyTok, len(toks), len(toks)+1)
copy(newToks, toks)
newToks = append(newToks, KeyTok{kind, intID, stringID})
return &Key{kc, newToks}
}
// MakeKey is a convenience function for manufacturing a *Key. It should only
// be used when elems... is known statically (e.g. in the code) to be correct.
//
// elems is pairs of (string, string|int|int32|int64) pairs, which correspond to
// Kind/id pairs. Example:
//
// MkKeyContext("aid", "namespace").MakeKey("Parent", 1, "Child", "id")
//
// Would create the key:
//
// aid:namespace:/Parent,1/Child,id
//
// If elems is not parsable (e.g. wrong length, wrong types, etc.) this method
// will panic.
//
// See MakeKey for a version of this function which automatically
// provides aid and ns.
func (kc KeyContext) MakeKey(elems ...any) *Key {
if len(elems) == 0 {
return nil
}
if len(elems)%2 != 0 {
panic(fmt.Errorf("datastore.MakeKey: odd number of tokens: %v", elems))
}
toks := make([]KeyTok, len(elems)/2)
for i := 0; len(elems) > 0; i, elems = i+1, elems[2:] {
knd, ok := elems[0].(string)
if !ok {
panic(fmt.Errorf("datastore.MakeKey: bad kind: %v", elems[i]))
}
t := &toks[i]
t.Kind = knd
switch x := elems[1].(type) {
case string:
t.StringID = x
case int:
t.IntID = int64(x)
case int32:
t.IntID = int64(x)
case int64:
t.IntID = int64(x)
case uint16:
t.IntID = int64(x)
case uint32:
t.IntID = int64(x)
default:
panic(fmt.Errorf("datastore.MakeKey: bad id: %v", x))
}
}
return kc.NewKeyToks(toks)
}
// NewKeyFromMeta constructs a key (potentially partial) based on meta fields.
//
// Looks at `$key`, `$kind`, `$id`, `$parent`. Returns an error if necessary
// fields are missing.
func (kc KeyContext) NewKeyFromMeta(mgs MetaGetterSetter) (*Key, error) {
if key, _ := GetMetaDefault(mgs, "key", nil).(*Key); key != nil {
return key, nil
}
kind := GetMetaDefault(mgs, "kind", "").(string)
if kind == "" {
return nil, errors.New("unable to extract $kind")
}
// get id - allow both to be default for default keys
sid := GetMetaDefault(mgs, "id", "").(string)
iid := GetMetaDefault(mgs, "id", 0).(int64)
par, _ := GetMetaDefault(mgs, "parent", nil).(*Key)
return kc.NewKey(kind, sid, iid, par), nil
}
// Key is the type used for all datastore operations.
type Key struct {
kc KeyContext
toks []KeyTok
}
var _ interface {
json.Marshaler
json.Unmarshaler
} = (*Key)(nil)
// NewKeyEncoded decodes and returns a *Key
func NewKeyEncoded(encoded string) (ret *Key, err error) {
ret = &Key{}
// Re-add padding
if m := len(encoded) % 4; m != 0 {
encoded += strings.Repeat("=", 4-m)
}
b, err := base64.URLEncoding.DecodeString(encoded)
if err != nil {
return
}
r := &pb.Reference{}
if err = proto.Unmarshal(b, r); err != nil {
return
}
ret.kc = MkKeyContext(r.GetApp(), r.GetNameSpace())
ret.toks = make([]KeyTok, len(r.Path.Element))
for i, e := range r.Path.Element {
ret.toks[i] = KeyTok{
Kind: e.GetType(),
IntID: e.GetId(),
StringID: e.GetName(),
}
}
return
}
// LastTok returns the last KeyTok in this Key. Non-nil Keys are always guaranteed
// to have at least one token.
func (k *Key) LastTok() KeyTok {
return k.toks[len(k.toks)-1]
}
// AppID returns the application ID that this Key is for.
func (k *Key) AppID() string { return k.kc.AppID }
// Namespace returns the namespace that this Key is for.
func (k *Key) Namespace() string { return k.kc.Namespace }
// KeyContext returns the KeyContext that this Key is using.
func (k *Key) KeyContext() *KeyContext {
kc := k.kc
return &kc
}
// Kind returns the Kind of the child KeyTok
func (k *Key) Kind() string { return k.toks[len(k.toks)-1].Kind }
// StringID returns the StringID of the child KeyTok
func (k *Key) StringID() string { return k.toks[len(k.toks)-1].StringID }
// IntID returns the IntID of the child KeyTok
func (k *Key) IntID() int64 { return k.toks[len(k.toks)-1].IntID }
// String returns a human-readable representation of the key in the form of
//
// AID:NS:/Kind,id/Kind,id/...
func (k *Key) String() string {
b := bytes.NewBuffer(make([]byte, 0, 512))
fmt.Fprintf(b, "%s:%s:", k.kc.AppID, k.kc.Namespace)
for _, t := range k.toks {
if t.StringID != "" {
fmt.Fprintf(b, "/%s,%q", t.Kind, t.StringID)
} else {
fmt.Fprintf(b, "/%s,%d", t.Kind, t.IntID)
}
}
return b.String()
}
// IsIncomplete returns true iff the last token of this Key doesn't define
// either a StringID or an IntID.
func (k *Key) IsIncomplete() bool {
return k.LastTok().IsIncomplete()
}
// Valid determines if a key is valid, according to a couple of rules:
// - k is not nil
// - every token of k:
// - (if !allowSpecial) token's kind doesn't start with '__'
// - token's kind and appid are non-blank
// - token is not incomplete
// - all tokens have the same namespace and appid
func (k *Key) Valid(allowSpecial bool, kc KeyContext) bool {
if !kc.Matches(k.kc) {
return false
}
for _, t := range k.toks {
if t.IsIncomplete() {
return false
}
if !allowSpecial && t.Special() {
return false
}
if t.Kind == "" {
return false
}
if t.StringID != "" && t.IntID != 0 {
return false
}
}
return true
}
// PartialValid returns true iff this key is suitable for use in a Put
// operation. This is the same as Valid(k, false, ...), but also allowing k to
// be IsIncomplete().
func (k *Key) PartialValid(kc KeyContext) bool {
if k.IsIncomplete() {
if !kc.Matches(k.kc) {
return false
}
k = kc.NewKey(k.Kind(), "", 1, k.Parent())
}
return k.Valid(false, kc)
}
// Parent returns the parent Key of this *Key, or nil. The parent
// will always have the concrete type of *Key.
func (k *Key) Parent() *Key {
if len(k.toks) <= 1 {
return nil
}
return k.kc.NewKeyToks(k.toks[:len(k.toks)-1])
}
// MarshalJSON allows this key to be automatically marshaled by encoding/json.
func (k *Key) MarshalJSON() ([]byte, error) {
return []byte(`"` + k.Encode() + `"`), nil
}
// Encode encodes the provided key as a base64-encoded protobuf.
//
// This encoding is compatible with the SDK-provided encoding and is agnostic
// to the underlying implementation of the Key.
//
// It's encoded with the urlsafe base64 table without padding.
func (k *Key) Encode() string {
e := make([]*pb.Path_Element, len(k.toks))
for i, t := range k.toks {
t := t
e[i] = &pb.Path_Element{
Type: &t.Kind,
}
if t.StringID != "" {
e[i].Name = &t.StringID
} else {
e[i].Id = &t.IntID
}
}
var namespace *string
if ns := k.kc.Namespace; ns != "" {
namespace = &ns
}
r, err := proto.Marshal(&pb.Reference{
App: &k.kc.AppID,
NameSpace: namespace,
Path: &pb.Path{
Element: e,
},
})
if err != nil {
panic(err)
}
// trim padding
return strings.TrimRight(base64.URLEncoding.EncodeToString(r), "=")
}
// UnmarshalJSON allows this key to be automatically unmarshaled by encoding/json.
func (k *Key) UnmarshalJSON(buf []byte) error {
if len(buf) < 2 || buf[0] != '"' || buf[len(buf)-1] != '"' {
return errors.New("datastore: bad JSON key")
}
nk, err := NewKeyEncoded(string(buf[1 : len(buf)-1]))
if err != nil {
return err
}
*k = *nk
return nil
}
// GobEncode allows the Key to be encoded in a Gob struct.
func (k *Key) GobEncode() ([]byte, error) {
return []byte(k.Encode()), nil
}
// GobDecode allows the Key to be decoded in a Gob struct.
func (k *Key) GobDecode(buf []byte) error {
nk, err := NewKeyEncoded(string(buf))
if err != nil {
return err
}
*k = *nk
return nil
}
// Root returns the entity root for the given key.
func (k *Key) Root() *Key {
if len(k.toks) > 1 {
ret := *k
ret.toks = ret.toks[:1]
return &ret
}
return k
}
// Less returns true iff k would sort before other.
func (k *Key) Less(other *Key) bool {
if k.kc.AppID < other.kc.AppID {
return true
} else if k.kc.AppID > other.kc.AppID {
return false
}
if k.kc.Namespace < other.kc.Namespace {
return true
} else if k.kc.Namespace > other.kc.Namespace {
return false
}
lim := len(k.toks)
if len(other.toks) < lim {
lim = len(other.toks)
}
for i := 0; i < lim; i++ {
a, b := k.toks[i], other.toks[i]
if a.Less(b) {
return true
} else if b.Less(a) {
return false
}
}
return len(k.toks) < len(other.toks)
}
// HasAncestor returns true iff other is an ancestor of k (or if other == k).
func (k *Key) HasAncestor(other *Key) bool {
if !k.kc.Matches(other.kc) {
return false
}
if len(k.toks) < len(other.toks) {
return false
}
for i, tok := range other.toks {
if tok != k.toks[i] {
return false
}
}
return true
}
// GQL returns a correctly formatted Cloud Datastore GQL key literal.
//
// The flavor of GQL that this emits is defined here:
//
// https://cloud.google.com/datastore/docs/apis/gql/gql_reference
func (k *Key) GQL() string {
ret := &bytes.Buffer{}
fmt.Fprintf(ret, "KEY(DATASET(%s)", gqlQuoteString(k.kc.AppID))
if ns := k.kc.Namespace; ns != "" {
fmt.Fprintf(ret, ", NAMESPACE(%s)", gqlQuoteString(ns))
}
for _, t := range k.toks {
if t.IntID != 0 {
fmt.Fprintf(ret, ", %s, %d", gqlQuoteString(t.Kind), t.IntID)
} else {
fmt.Fprintf(ret, ", %s, %s", gqlQuoteString(t.Kind), gqlQuoteString(t.StringID))
}
}
if _, err := ret.WriteString(")"); err != nil {
panic(err)
}
return ret.String()
}
// Equal returns true iff the two keys represent identical key values.
func (k *Key) Equal(other *Key) bool {
if k == nil || other == nil {
return k == nil && other == nil
}
return k.IncompleteEqual(other) && (k.LastTok() == other.LastTok())
}
// IncompleteEqual asserts that, were the two keys incomplete, they would be
// equal.
//
// This asserts equality for the full lineage of the key, except for its last
// token ID.
func (k *Key) IncompleteEqual(other *Key) (ret bool) {
ret = (k.kc.Matches(other.kc) &&
len(k.toks) == len(other.toks))
if ret {
for i, t := range k.toks {
if i == len(k.toks)-1 {
// Last token: check only Kind.
if ret = (t.Kind == other.toks[i].Kind); !ret {
return
}
} else {
if ret = t == other.toks[i]; !ret {
return
}
}
}
}
return
}
// Incomplete returns an incomplete version of the key. The ID fields of the
// last token will be set to zero/empty.
func (k *Key) Incomplete() *Key {
if k.IsIncomplete() {
return k
}
return k.kc.NewKey(k.Kind(), "", 0, k.Parent())
}
// WithID returns the key generated by setting the ID of its last token to
// the specified value.
//
// To generate this, k is reduced to its Incomplete form, then populated with a
// new ID. The resulting key will have the same token linage as k (i.e., will
// be IncompleteEqual).
func (k *Key) WithID(stringID string, intID int64) *Key {
if k.StringID() == stringID && k.IntID() == intID {
return k
}
return k.kc.NewKey(k.Kind(), stringID, intID, k.Parent())
}
// Split componentizes the key into pieces (AppID, Namespace and tokens)
//
// Each token represents one piece of they key's 'path'.
//
// toks is guaranteed to be empty if and only if k is nil. If k is non-nil then
// it contains at least one token.
func (k *Key) Split() (appID, namespace string, toks []KeyTok) {
appID = k.kc.AppID
namespace = k.kc.Namespace
toks = make([]KeyTok, len(k.toks))
copy(toks, k.toks)
return
}
// EstimateSize estimates the size of a Key.
//
// It uses https://cloud.google.com/appengine/articles/storage_breakdown?csw=1
// as a guide for these values.
func (k *Key) EstimateSize() int64 {
ret := int64(len(k.kc.AppID))
ret += int64(len(k.kc.Namespace))
for _, t := range k.toks {
ret += int64(len(t.Kind))
if t.StringID != "" {
ret += int64(len(t.StringID))
} else {
ret += 8
}
}
return ret
}