datastoredb.go

package datastoredb

import (
	"context"

	"cloud.google.com/go/datastore"
	"go.opentelemetry.io/otel/metric"
	"google.golang.org/api/option"
)

// DatastoreDB implements the slackscot StringStorer interface. It maps
// the given name (usually a plugin name) to the datastore entity Kind
// to isolate data between different plugins
type DatastoreDB struct {
	datastorer
	kind string
}

// EntryValue represents an entity/entry value mapped to a datastore key
type EntryValue struct {
	Value string `datastore:",noindex"`
}

const (
	// Try operations that could fail at most twice. The first time is assummed to potentially fail because
	// of authentication errors when credentials have expired. The second time, a failure is probably
	// something to report back
	maxAttemptCount = 2
)

// New returns a new instance of DatastoreDB for the given name (which maps to the datastore entity "Kind" and can
// be thought of as the namespace). This function also requires a gcloudProjectID as well as at least one option to provide gcloud client credentials.
// Note that in order to support a deployment where credentials can get updated, the gcloudClientOpts should use
// something like option.WithCredentialsFile with the credentials file being updated on disk so that when reconnecting
// on a failure, the updated credentials are visible through the same gcloud client options
func New(name string, gcloudProjectID string, gcloudClientOpts ...option.ClientOption) (dsdb *DatastoreDB, err error) {
	ds := new(gcdatastore)
	ds.gcloudProjectID = gcloudProjectID
	ds.gcloudClientOpts = gcloudClientOpts

	return newWithDatastorer(name, ds)
}

// NewWithTelemetry returns a new instance of DatastoreDB similar to New but with added opentelemetry
// instrumentation. Note that even if no metrics exporter is configured, it is safe to use telemetry
// and should result with NOOP operations with a small overhead.
//
// For the given name (which maps to the datastore entity "Kind" and can
// be thought of as the namespace). This function also requires a gcloudProjectID as well as at least one option to provide gcloud client credentials.
// Note that in order to support a deployment where credentials can get updated, the gcloudClientOpts should use
// something like option.WithCredentialsFile with the credentials file being updated on disk so that when reconnecting
// on a failure, the updated credentials are visible through the same gcloud client options
func NewWithTelemetry(appName string, meter metric.Meter, kindName string, gcloudProjectID string, gcloudClientOpts ...option.ClientOption) (dsdb *DatastoreDB, err error) {
	ds := new(gcdatastore)
	ds.gcloudProjectID = gcloudProjectID
	ds.gcloudClientOpts = gcloudClientOpts

	return newWithDatastorer(kindName, NewdatastorerWithTelemetry(ds, appName, meter))
}

// newWithDatastorer returns a new instance of DatastoreDB using the provided datastorer
func newWithDatastorer(name string, datastorer datastorer) (dsdb *DatastoreDB, err error) {
	dsdb = new(DatastoreDB)
	dsdb.kind = name
	dsdb.datastorer = datastorer

	err = dsdb.connect()
	if err != nil {
		return nil, err
	}

	err = dsdb.testDB()
	if err = dsdb.testDB(); err != nil {
		dsdb.Close()
		return nil, err
	}

	return dsdb, nil
}

// testDB makes a lightweight call to the datastore to validate connectivity and credentials
func (dsdb *DatastoreDB) testDB() (err error) {
	err = dsdb.Get(context.Background(), datastore.NameKey(dsdb.kind, "testConnectivity", nil), &EntryValue{})

	if err != nil && err != datastore.ErrNoSuchEntity {
		return err
	}

	return nil
}

// GetString returns the value associated to a given key. If the value is not
// found or an error occurred, the zero-value string is returned along with
// the error
func (dsdb *DatastoreDB) GetString(key string) (value string, err error) {
	return dsdb.GetSiloString("", key)
}

// GetSiloString returns the value associated to a given key within the silo provided. If the value is not
// found or an error occurred, the zero-value string is returned along with the error
func (dsdb *DatastoreDB) GetSiloString(silo string, key string) (value string, err error) {
	ctx := context.Background()

	var e EntryValue
	k := newKeyWithNamespace(silo, dsdb.kind, key)

	// Retry once and try a reconnect if the error is recoverable (like unauthenticated error)
	var attempt int
	for attempt, err = 1, dsdb.Get(ctx, k, &e); attempt < maxAttemptCount && err != nil && shouldRetry(err); attempt, err = attempt+1, dsdb.Get(ctx, k, &e) {
		dsdb.connect()
	}

	if err != nil {
		return "", err
	}

	return e.Value, nil
}

// PutString stores the key/value to the database
func (dsdb *DatastoreDB) PutString(key string, value string) (err error) {
	return dsdb.PutSiloString("", key, value)
}

// PutSiloString stores the key/value to the database in the given silo
func (dsdb *DatastoreDB) PutSiloString(silo string, key string, value string) (err error) {
	ctx := context.Background()
	k := newKeyWithNamespace(silo, dsdb.kind, key)

	// Execute first attempt
	_, err = dsdb.Put(ctx, k, &EntryValue{Value: value})

	// Retry once and try a reconnect if the error is recoverable (like unauthenticated error)
	for attempt := 1; attempt < maxAttemptCount && err != nil && shouldRetry(err); attempt = attempt + 1 {
		dsdb.connect()

		_, err = dsdb.Put(ctx, k, &EntryValue{Value: value})
	}

	return err
}

// DeleteString deletes the entry for the given key. If the entry is not found
// an error is returned
func (dsdb *DatastoreDB) DeleteString(key string) (err error) {
	return dsdb.DeleteSiloString("", key)
}

// DeleteSiloString deletes the entry for the given key in the given silo. If the entry is not found
// an error is returned
func (dsdb *DatastoreDB) DeleteSiloString(silo string, key string) (err error) {
	ctx := context.Background()
	k := newKeyWithNamespace(silo, dsdb.kind, key)

	// Retry once and try a reconnect if the error is recoverable (like unauthenticated error)
	var attempt int
	for attempt, err = 1, dsdb.Delete(ctx, k); attempt < maxAttemptCount && err != nil && shouldRetry(err); attempt, err = attempt+1, dsdb.Delete(ctx, k) {
		dsdb.connect()
	}

	return err
}

// Scan returns all key/values from the database
func (dsdb *DatastoreDB) Scan() (entries map[string]string, err error) {
	return dsdb.ScanSilo("")
}

// ScanSilo returns all key/values from the database in the given silo
func (dsdb *DatastoreDB) ScanSilo(silo string) (entries map[string]string, err error) {
	entries = make(map[string]string)

	keys, vals, err := dsdb.scan(datastore.NewQuery(dsdb.kind).Namespace(silo))

	if err != nil {
		return nil, err
	}

	for i, key := range keys {
		entries[key.Name] = vals[i].Value
	}

	return entries, nil
}

// GlobalScan returns all key/values for all silos keyed by silo name
func (dsdb *DatastoreDB) GlobalScan() (entries map[string]map[string]string, err error) {
	entries = make(map[string]map[string]string)

	namespaces, err := dsdb.listNamespaces()
	if err != nil {
		return nil, err
	}

	for _, ns := range namespaces {
		keys, vals, err := dsdb.scan(datastore.NewQuery(dsdb.kind).Namespace(ns))
		if err != nil {
			return nil, err
		}

		for i, key := range keys {
			if _, ok := entries[key.Namespace]; !ok {
				entries[key.Namespace] = make(map[string]string)
			}

			entries[key.Namespace][key.Name] = vals[i].Value
		}
	}

	return entries, nil
}

// listNamespaces lists all namespaces available
func (dsdb *DatastoreDB) listNamespaces() (namespaces []string, err error) {
	q := datastore.NewQuery("__namespace__").KeysOnly()
	keys, err := dsdb.scanKeys(q)
	if err != nil {
		return nil, err
	}

	namespaces = make([]string, 0)

	for _, nsKey := range keys {
		namespaces = append(namespaces, nsKey.Name)
	}

	return namespaces, nil
}

// scanKeys runs an internal datastore query and returns the raw keys only
func (dsdb *DatastoreDB) scanKeys(query *datastore.Query) (keys []*datastore.Key, err error) {
	ctx := context.Background()

	// Run first attempt before looping
	keys, err = dsdb.GetAll(ctx, query, nil)

	// Retry once and try a reconnect if the error is recoverable (like unauthenticated error)
	for attempt := 1; attempt < maxAttemptCount && err != nil && shouldRetry(err); attempt = attempt + 1 {
		dsdb.connect()

		keys, err = dsdb.GetAll(ctx, query, nil)
	}

	return keys, err
}

// scan runs an internal datastore query and returns the raw keys and values for post-processing
func (dsdb *DatastoreDB) scan(query *datastore.Query) (keys []*datastore.Key, vals []*EntryValue, err error) {
	ctx := context.Background()

	vals = make([]*EntryValue, 0)

	// Run first attempt before looping
	keys, err = dsdb.GetAll(ctx, query, &vals)

	// Retry once and try a reconnect if the error is recoverable (like unauthenticated error)
	for attempt := 1; attempt < maxAttemptCount && err != nil && shouldRetry(err); attempt = attempt + 1 {
		dsdb.connect()

		keys, err = dsdb.GetAll(ctx, query, &vals)
	}

	return keys, vals, err
}

// shouldRetry returns true if the given error should be retried or false if not.
// In order to determine this, one approach would be to only retry on a
// statusError (https://github.com/grpc/grpc-go/blob/master/status/status.go#L43)
// with code Unauthenticated (https://godoc.org/google.golang.org/grpc/codes) but that's made
// trickier by the statusError not being promoted outside the package (checking for the Error string
// would be reasonable but a bit dirty).
// Alternatively, and what's done here is to be a little conservative and retry on everything except
// ErrNoSuchEntity, ErrInvalidEntityType and ErrInvalidKey which are not things retries would help
// with. This means we could still retry when it's pointless to do so at the expense of added latency.
func shouldRetry(err error) bool {
	return err != datastore.ErrNoSuchEntity && err != datastore.ErrInvalidEntityType && err != datastore.ErrInvalidKey
}

// newKeyWithNamespace returns a new datastore key for the given kind and key name within the
// specified namespace
func newKeyWithNamespace(namespace string, kind string, key string) (k *datastore.Key) {
	k = datastore.NameKey(kind, key, nil)
	k.Namespace = namespace

	return k
}