mattermost-community-enterprise/public/pluginapi/kv.go

package pluginapi

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/pkg/errors"

	"github.com/mattermost/mattermost/server/public/model"
	"github.com/mattermost/mattermost/server/public/plugin"
)

// numRetries is the number of times the setAtomicWithRetries will retry before returning an error.
const numRetries = 5

// KVService exposes methods to read and write key-value pairs for the active plugin.
//
// This service cannot be used to read or write key-value pairs for other plugins.
type KVService struct {
	api plugin.API
}

type KVSetOptions struct {
	model.PluginKVSetOptions
	oldValue any
}

// KVSetOption is an option passed to Set() operation.
type KVSetOption func(*KVSetOptions)

// SetAtomic guarantees the write will occur only when the current value of matches the given old
// value. A client is expected to read the old value first, then pass it back to ensure the value
// has not since been modified.
func SetAtomic(oldValue any) KVSetOption {
	return func(o *KVSetOptions) {
		o.Atomic = true
		o.oldValue = oldValue
	}
}

// SetExpiry configures a key value to expire after the given duration relative to now.
func SetExpiry(ttl time.Duration) KVSetOption {
	return func(o *KVSetOptions) {
		o.ExpireInSeconds = int64(ttl / time.Second)
	}
}

// Set stores a key-value pair, unique per plugin.
// Keys prefixed with `mmi_` are reserved for internal use and will fail to be set.
//
// Returns (false, err) if DB error occurred
// Returns (false, nil) if the value was not set
// Returns (true, nil) if the value was set
//
// Minimum server version: 5.18
func (k *KVService) Set(key string, value any, options ...KVSetOption) (bool, error) {
	if strings.HasPrefix(key, internalKeyPrefix) {
		return false, errors.Errorf("'%s' prefix is not allowed for keys", internalKeyPrefix)
	}

	opts := KVSetOptions{}
	for _, o := range options {
		o(&opts)
	}

	var valueBytes []byte
	if value != nil {
		// Assume JSON encoding, unless explicitly given a byte slice.
		var isValueInBytes bool
		valueBytes, isValueInBytes = value.([]byte)
		if !isValueInBytes {
			var err error
			valueBytes, err = json.Marshal(value)
			if err != nil {
				return false, errors.Wrapf(err, "failed to marshal value %v", value)
			}
		}
	}

	downstreamOpts := model.PluginKVSetOptions{
		Atomic:          opts.Atomic,
		ExpireInSeconds: opts.ExpireInSeconds,
	}

	if opts.oldValue != nil {
		oldValueBytes, isOldValueInBytes := opts.oldValue.([]byte)
		if isOldValueInBytes {
			downstreamOpts.OldValue = oldValueBytes
		} else {
			data, err := json.Marshal(opts.oldValue)
			if err != nil {
				return false, errors.Wrapf(err, "failed to marshal value %v", opts.oldValue)
			}

			downstreamOpts.OldValue = data
		}
	}

	written, appErr := k.api.KVSetWithOptions(key, valueBytes, downstreamOpts)
	return written, normalizeAppErr(appErr)
}

// SetAtomicWithRetries will set a key-value pair atomically using compare and set semantics:
// it will read key's value (to get oldValue), perform valueFunc (to get newValue),
// and compare and set (comparing oldValue and setting newValue).
//
// Parameters:
//
//	`key`              is the key to get and set.
//	`valueFunc`        is a user-provided function that will take the old value as a []byte and
//	                   return the new value or an error. If valueFunc needs to operate on
//	                   oldValue, it will need to use the oldValue as a []byte, or convert
//	                   oldValue into the expected type (e.g., by parsing it, or marshaling it
//	                   into the expected struct). It should then return the newValue as the type
//	                   expected to be stored.
//
// Returns:
//
//	Returns err if the key could not be retrieved (DB error), valueFunc returned an error,
//	if the key could not be set (DB error), or if the key could not be set (after retries).
//	Returns nil if the value was set.
//
// Minimum server version: 5.18
func (k *KVService) SetAtomicWithRetries(key string, valueFunc func(oldValue []byte) (newValue any, err error)) error {
	for range numRetries {
		var oldVal []byte
		if err := k.Get(key, &oldVal); err != nil {
			return errors.Wrapf(err, "failed to get value for key %s", key)
		}

		newVal, err := valueFunc(oldVal)
		if err != nil {
			return errors.Wrap(err, "valueFunc failed")
		}

		if saved, err := k.Set(key, newVal, SetAtomic(oldVal)); err != nil {
			return errors.Wrapf(err, "DB failed to set value for key %s", key)
		} else if saved {
			return nil
		}

		// small delay to allow cooperative scheduling to do its thing
		time.Sleep(10 * time.Millisecond)
	}
	return fmt.Errorf("failed to set value after %d retries", numRetries)
}

// Get gets the value for the given key into the given interface.
//
// An error is returned only if the value cannot be fetched. A non-existent key will return no
// error, with nothing written to the given interface.
//
// Minimum server version: 5.2
func (k *KVService) Get(key string, o any) error {
	data, appErr := k.api.KVGet(key)
	if appErr != nil {
		return normalizeAppErr(appErr)
	}

	if len(data) == 0 {
		return nil
	}

	if bytesOut, ok := o.(*[]byte); ok {
		*bytesOut = data
		return nil
	}

	if err := json.Unmarshal(data, o); err != nil {
		return errors.Wrapf(err, "failed to unmarshal value for key %s", key)
	}

	return nil
}

// Delete deletes the given key-value pair.
//
// An error is returned only if the value failed to be deleted. A non-existent key will return
// no error.
//
// Minimum server version: 5.18
func (k *KVService) Delete(key string) error {
	_, err := k.Set(key, nil)
	return err
}

// DeleteAll removes all key-value pairs.
//
// Minimum server version: 5.6
func (k *KVService) DeleteAll() error {
	return normalizeAppErr(k.api.KVDeleteAll())
}

// ListKeysOption used to configure a ListKeys() operation.
type ListKeysOption func(*listKeysOptions)

// listKeysOptions holds configurations of a ListKeys() operation.
type listKeysOptions struct {
	checkers []func(key string) (keep bool, err error)
}

func (o *listKeysOptions) checkAll(key string) (keep bool, err error) {
	for _, check := range o.checkers {
		keep, err := check(key)
		if err != nil {
			return false, err
		}
		if !keep {
			return false, nil
		}
	}

	// key made it through all checkers
	return true, nil
}

// WithPrefix only return keys that start with the given string.
func WithPrefix(prefix string) ListKeysOption {
	return WithChecker(func(key string) (keep bool, err error) {
		return strings.HasPrefix(key, prefix), nil
	})
}

// WithChecker allows for a custom filter function to determine which keys to return.
// Returning true will keep the key and false will filter it out. Returning an error
// will halt KVListWithOptions immediately and pass the error up (with no other results).
func WithChecker(f func(key string) (keep bool, err error)) ListKeysOption {
	return func(args *listKeysOptions) {
		args.checkers = append(args.checkers, f)
	}
}

// ListKeys lists all keys that match the given options. If no options are provided then all keys are returned.
//
// Minimum server version: 5.6
func (k *KVService) ListKeys(page, count int, options ...ListKeysOption) ([]string, error) {
	// convert functional options into args struct
	args := &listKeysOptions{
		checkers: nil,
	}
	for _, opt := range options {
		opt(args)
	}

	// get our keys a batch at a time, filter out the ones we don't want based on our args
	// any errors will hault the whole process and return the error raw

	keys, appErr := k.api.KVList(page, count)
	if appErr != nil {
		return nil, normalizeAppErr(appErr)
	}

	if len(args.checkers) == 0 {
		// no checkers, just return the keys
		return keys, nil
	}

	ret := make([]string, 0)
	// we have a filter, so check each key, all checkers must say key
	// for us to keep a key
	for _, key := range keys {
		keep, err := args.checkAll(key)
		if err != nil {
			return nil, err
		}

		if !keep {
			continue
		}

		// didn't get filtered out, add to our return
		ret = append(ret, key)
	}

	return ret, nil
}