datastore/keycompat.go

// Copyright 2019 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package datastore

import (
	"context"
	"sync"

	"google.golang.org/appengine/datastore/internal/cloudkey"
	"google.golang.org/appengine/internal"
)

var keyConversion struct {
	mu    sync.RWMutex
	appID string // read using getKeyConversionAppID
}

// EnableKeyConversion enables encoded key compatibility with the Cloud
// Datastore client library (cloud.google.com/go/datastore). Encoded keys
// generated by the Cloud Datastore client library will be decoded into App
// Engine datastore keys.
//
// The context provided must be an App Engine context if running in App Engine
// first generation runtime. This can be called in the /_ah/start handler. It is
// safe to call multiple times, and is cheap to call, so can also be inserted as
// middleware.
//
// Enabling key compatibility does not affect the encoding format used by
// Key.Encode, it only expands the type of keys that are able to be decoded with
// DecodeKey.
func EnableKeyConversion(ctx context.Context) {
	// Only attempt to set appID if it's unset.
	// If already set, ignore.
	if getKeyConversionAppID() != "" {
		return
	}

	keyConversion.mu.Lock()
	// Check again to avoid race where another goroutine set appID between the call
	// to getKeyConversionAppID above and taking the write lock.
	if keyConversion.appID == "" {
		keyConversion.appID = internal.FullyQualifiedAppID(ctx)
	}
	keyConversion.mu.Unlock()
}

func getKeyConversionAppID() string {
	keyConversion.mu.RLock()
	appID := keyConversion.appID
	keyConversion.mu.RUnlock()
	return appID
}

// decodeCloudKey attempts to decode the given encoded key generated by the
// Cloud Datastore client library (cloud.google.com/go/datastore), returning nil
// if the key couldn't be decoded.
func decodeCloudKey(encoded string) *Key {
	appID := getKeyConversionAppID()
	if appID == "" {
		return nil
	}

	k, err := cloudkey.DecodeKey(encoded)
	if err != nil {
		return nil
	}
	return convertCloudKey(k, appID)
}

// convertCloudKey converts a Cloud Datastore key and converts it to an App
// Engine Datastore key. Cloud Datastore keys don't include the project/app ID,
// so we must add it back in.
func convertCloudKey(key *cloudkey.Key, appID string) *Key {
	if key == nil {
		return nil
	}
	k := &Key{
		intID:     key.ID,
		kind:      key.Kind,
		namespace: key.Namespace,
		parent:    convertCloudKey(key.Parent, appID),
		stringID:  key.Name,
		appID:     appID,
	}
	return k
}