components-contrib/state/azure/cosmosdb/cosmosdb.go

/*
Copyright 2021 The Dapr 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 cosmosdb

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
	"github.com/google/uuid"
	jsoniter "github.com/json-iterator/go"

	"github.com/dapr/components-contrib/common/authentication/azure"
	"github.com/dapr/components-contrib/contenttype"
	contribmeta "github.com/dapr/components-contrib/metadata"
	"github.com/dapr/components-contrib/state"
	"github.com/dapr/components-contrib/state/query"
	stateutils "github.com/dapr/components-contrib/state/utils"
	"github.com/dapr/kit/logger"
	kitmd "github.com/dapr/kit/metadata"
	"github.com/dapr/kit/ptr"
)

// StateStore is a CosmosDB state store.
type StateStore struct {
	state.BulkStore

	client      *azcosmos.ContainerClient
	metadata    metadata
	contentType string
	logger      logger.Logger
}

type metadata struct {
	URL         string `json:"url"`
	MasterKey   string `json:"masterKey"`
	Database    string `json:"database"`
	Collection  string `json:"collection"`
	ContentType string `json:"contentType"`
}

type cosmosOperationType string

// CosmosOperation is a wrapper around a CosmosDB operation.
type CosmosOperation struct {
	Item CosmosItem          `json:"item"`
	Type cosmosOperationType `json:"type"`
}

// CosmosItem is a wrapper around a CosmosDB document.
type CosmosItem struct {
	ID           string      `json:"id"`
	Value        interface{} `json:"value"`
	IsBinary     bool        `json:"isBinary"`
	PartitionKey string      `json:"partitionKey"`
	TTL          *int        `json:"ttl,omitempty"`
	Etag         string      `json:"_etag"`
	TS           int64       `json:"_ts"`
}

const (
	metadataPartitionKey = "partitionKey"
	defaultTimeout       = 20 * time.Second
	statusNotFound       = "NotFound"
)

// Policy that makes all queries cross-partition
type crossPartitionQueryPolicy struct{}

func (p crossPartitionQueryPolicy) Do(req *policy.Request) (*http.Response, error) {
	raw := req.Raw()
	// Check if we're performing a query
	// In that case, remove the partitionkey header and enable cross-partition queries
	if strings.ToLower(raw.Header.Get("x-ms-documentdb-query")) == "true" {
		// Only when the partitionKey is fake (true), it will be removed amd enabled the cross partition
		if strings.ToLower(raw.Header.Get("x-ms-documentdb-partitionkey")) == "[true]" {
			raw.Header.Add("x-ms-documentdb-query-enablecrosspartition", "true")
			raw.Header.Del("x-ms-documentdb-partitionkey")
		}
	}
	return req.Next()
}

// NewCosmosDBStateStore returns a new CosmosDB state store.
func NewCosmosDBStateStore(logger logger.Logger) state.Store {
	s := &StateStore{
		logger: logger,
	}
	s.BulkStore = state.NewDefaultBulkStore(s)
	return s
}

func (c *StateStore) GetComponentMetadata() (metadataInfo contribmeta.MetadataMap) {
	metadataStruct := metadata{}
	contribmeta.GetMetadataInfoFromStructType(reflect.TypeOf(metadataStruct), &metadataInfo, contribmeta.StateStoreType)
	return
}

// Init does metadata and connection parsing.
func (c *StateStore) Init(ctx context.Context, meta state.Metadata) error {
	c.logger.Debugf("CosmosDB init start")

	m := metadata{
		ContentType: "application/json",
	}
	errDecode := kitmd.DecodeMetadata(meta.Properties, &m)
	if errDecode != nil {
		return errDecode
	}

	if m.URL == "" {
		return errors.New("url is required")
	}
	if m.Database == "" {
		return errors.New("database is required")
	}
	if m.Collection == "" {
		return errors.New("collection is required")
	}
	if m.ContentType == "" {
		return errors.New("contentType is required")
	}

	// Internal query policy was created due to lack of cross partition query capability in the current Go sdk
	opts := azcosmos.ClientOptions{
		ClientOptions: policy.ClientOptions{
			PerCallPolicies: []policy.Policy{
				crossPartitionQueryPolicy{},
			},
			Telemetry: policy.TelemetryOptions{
				ApplicationID: "dapr-" + logger.DaprVersion,
			},
		},
	}

	// Create the client; first, try authenticating with a master key, if present
	var client *azcosmos.Client
	if m.MasterKey != "" {
		var cred azcosmos.KeyCredential
		cred, err := azcosmos.NewKeyCredential(m.MasterKey)
		if err != nil {
			return err
		}
		client, err = azcosmos.NewClientWithKey(m.URL, cred, &opts)
		if err != nil {
			return err
		}
	} else {
		// Fallback to using Azure AD
		var env azure.EnvironmentSettings
		env, err := azure.NewEnvironmentSettings(meta.Properties)
		if err != nil {
			return err
		}
		token, tokenErr := env.GetTokenCredential()
		if tokenErr != nil {
			return tokenErr
		}
		client, err = azcosmos.NewClient(m.URL, token, &opts)
		if err != nil {
			return err
		}
	}
	// Create a container client
	dbClient, err := client.NewDatabase(m.Database)
	if err != nil {
		return err
	}
	// Container is synonymous with collection.
	dbContainer, err := dbClient.NewContainer(m.Collection)
	if err != nil {
		return err
	}
	c.client = dbContainer

	c.metadata = m
	c.contentType = m.ContentType

	readCtx, cancel := context.WithTimeout(ctx, defaultTimeout)
	defer cancel()
	_, err = c.client.Read(readCtx, nil)
	return err
}

// Features returns the features available in this state store.
func (c *StateStore) Features() []state.Feature {
	return []state.Feature{
		state.FeatureETag,
		state.FeatureTransactional,
		state.FeatureQueryAPI,
		state.FeatureTTL,
		state.FeaturePartitionKey,
	}
}

// Get retrieves a CosmosDB item.
func (c *StateStore) Get(ctx context.Context, req *state.GetRequest) (*state.GetResponse, error) {
	partitionKey := populatePartitionMetadata(req.Key, req.Metadata)

	options := azcosmos.ItemOptions{}
	if req.Options.Consistency == state.Strong {
		options.ConsistencyLevel = azcosmos.ConsistencyLevelSession.ToPtr()
	} else if req.Options.Consistency == state.Eventual {
		options.ConsistencyLevel = azcosmos.ConsistencyLevelEventual.ToPtr()
	}

	readCtx, cancel := context.WithTimeout(ctx, defaultTimeout)
	defer cancel()
	readItem, err := c.client.ReadItem(readCtx, azcosmos.NewPartitionKeyString(partitionKey), req.Key, &options)
	if err != nil {
		var responseErr *azcore.ResponseError
		if errors.As(err, &responseErr) && responseErr.ErrorCode == "NotFound" {
			return &state.GetResponse{}, nil
		}
		return nil, err
	}

	item, err := NewCosmosItemFromResponse(readItem.Value, c.logger)
	if err != nil {
		return nil, err
	}

	var metadata map[string]string
	if item.TTL != nil {
		metadata = map[string]string{
			state.GetRespMetaKeyTTLExpireTime: time.Unix(item.TS+int64(*item.TTL), 0).UTC().Format(time.RFC3339),
		}
	}

	// We are sure this is a []byte if not nil
	b, _ := item.Value.([]byte)
	return &state.GetResponse{
		Data:     b,
		ETag:     ptr.Of(item.Etag),
		Metadata: metadata,
	}, nil
}

// getMulti retrieves multiple items with a cross-partition query, retrieving multiple records with a single query.
func (c *StateStore) getMulti(ctx context.Context, req []state.GetRequest) ([]state.BulkGetResponse, error) {
	// The partition key doesn't matter since it will be removed for a cross-partition query
	pk := azcosmos.NewPartitionKeyBool(true)

	// Build the query
	var consistency azcosmos.ConsistencyLevel
	keys := make([]string, len(req))
	for i, r := range req {
		keys[i] = r.Key

		// Use the specified consistency if empty
		// If there's a strong consistency, it overrides any eventual
		if r.Options.Consistency == state.Strong {
			consistency = azcosmos.ConsistencyLevelStrong
		} else if r.Options.Consistency == state.Eventual && consistency == "" {
			consistency = azcosmos.ConsistencyLevelEventual
		}
	}

	// Execute the query
	// Source for "ARRAY_CONTAINS": https://stackoverflow.com/a/44639998/192024
	queryOpts := &azcosmos.QueryOptions{
		QueryParameters: []azcosmos.QueryParameter{
			{Name: "@keys", Value: keys},
		},
	}
	if consistency != "" {
		queryOpts.ConsistencyLevel = &consistency
	}
	pager := c.client.NewQueryItemsPager(
		"SELECT * FROM r WHERE ARRAY_CONTAINS(@keys, r.id)",
		pk, queryOpts,
	)
	result := make([]state.BulkGetResponse, len(req))
	n := 0
	for pager.More() {
		queryResponse, err := pager.NextPage(ctx)
		if err != nil {
			return nil, err
		}

		for _, value := range queryResponse.Items {
			item, err := NewCosmosItemFromResponse(value, c.logger)
			if err != nil {
				// If the error was while unserializing JSON, we don't have an ID, so exit right away
				// This should never happen, hopefully
				if item.ID == "" {
					return nil, err
				}
				result[n] = state.BulkGetResponse{
					Key:   item.ID,
					Error: err.Error(),
				}
				n++
				continue
			}

			var metadata map[string]string
			if item.TTL != nil {
				metadata = map[string]string{
					state.GetRespMetaKeyTTLExpireTime: time.Unix(item.TS+int64(*item.TTL), 0).UTC().Format(time.RFC3339),
				}
			}

			// We are sure this is a []byte if not nil
			b, _ := item.Value.([]byte)

			result[n] = state.BulkGetResponse{
				Key:      item.ID,
				Data:     b,
				ETag:     &item.Etag,
				Metadata: metadata,
			}
			n++
		}
	}

	return result[:n], nil
}

// BulkGet performs a Get operation in bulk.
func (c *StateStore) BulkGet(ctx context.Context, req []state.GetRequest, _ state.BulkGetOpts) ([]state.BulkGetResponse, error) {
	if len(req) == 0 {
		return []state.BulkGetResponse{}, nil
	}

	// If we have a single request, execute it as a regular Get request which is more efficient
	if len(req) == 1 {
		item, err := c.Get(ctx, &req[0])
		res := state.BulkGetResponse{
			Key: req[0].Key,
		}
		if err != nil {
			res.Error = err.Error()
		} else if item != nil {
			res.Data = json.RawMessage(item.Data)
			res.ETag = item.ETag
			res.Metadata = item.Metadata
		}
		return []state.BulkGetResponse{res}, nil
	}

	// Execute a single query to retrieve all values, as a cross-partition query
	items, err := c.getMulti(ctx, req)
	if err != nil {
		return nil, err
	}

	// Save all results
	// For consistency with when using Get, we need to include keys that are missing and set an empty value
	if len(items) > len(req) {
		return nil, errors.New("received more results than expected")
	}
	result := make([]state.BulkGetResponse, len(req))
	foundKeys := make(map[string]struct{}, len(items))
	n := 0
	for _, res := range items {
		result[n] = res
		foundKeys[res.Key] = struct{}{}
		n++
	}
	for _, r := range req {
		if _, ok := foundKeys[r.Key]; !ok {
			// Prevents a panic if out of bounds
			if n >= len(result) {
				return nil, errors.New("received more results than expected")
			}
			result[n] = state.BulkGetResponse{
				Key: r.Key,
			}
			n++
		}
	}
	// Prevents a panic if out of bounds
	if n > len(result) {
		n = len(result)
	}
	return result[:n], nil
}

// Set saves a CosmosDB item.
func (c *StateStore) Set(ctx context.Context, req *state.SetRequest) error {
	err := state.CheckRequestOptions(req.Options)
	if err != nil {
		return err
	}

	partitionKey := populatePartitionMetadata(req.Key, req.Metadata)
	options := azcosmos.ItemOptions{}

	if req.HasETag() {
		etag := azcore.ETag(*req.ETag)
		options.IfMatchEtag = &etag
	} else if req.Options.Concurrency == state.FirstWrite {
		var u uuid.UUID
		u, err = uuid.NewRandom()
		if err != nil {
			return err
		}
		options.IfMatchEtag = ptr.Of(azcore.ETag(u.String()))
	}
	// Consistency levels can only be relaxed so the session level is used here
	if req.Options.Consistency == state.Strong {
		options.ConsistencyLevel = azcosmos.ConsistencyLevelSession.ToPtr()
	} else if req.Options.Consistency == state.Eventual {
		options.ConsistencyLevel = azcosmos.ConsistencyLevelEventual.ToPtr()
	}

	doc, err := createUpsertItem(c.contentType, *req, partitionKey)
	if err != nil {
		return err
	}

	marsh, err := json.Marshal(doc)
	if err != nil {
		return err
	}

	upsertCtx, cancel := context.WithTimeout(ctx, defaultTimeout)
	defer cancel()
	pk := azcosmos.NewPartitionKeyString(partitionKey)
	_, err = c.client.UpsertItem(upsertCtx, pk, marsh, &options)
	if err != nil {
		resErr := &azcore.ResponseError{}
		if errors.As(err, &resErr) && resErr.StatusCode == http.StatusPreconditionFailed {
			return state.NewETagError(state.ETagMismatch, err)
		}
		return err
	}
	return nil
}

// Delete performs a delete operation.
func (c *StateStore) Delete(ctx context.Context, req *state.DeleteRequest) error {
	err := state.CheckRequestOptions(req.Options)
	if err != nil {
		return err
	}
	partitionKey := populatePartitionMetadata(req.Key, req.Metadata)
	options := azcosmos.ItemOptions{}

	if req.HasETag() {
		etag := azcore.ETag(*req.ETag)
		options.IfMatchEtag = &etag
	} else if req.Options.Concurrency == state.FirstWrite {
		var u uuid.UUID
		u, err = uuid.NewRandom()
		if err != nil {
			return err
		}
		options.IfMatchEtag = ptr.Of(azcore.ETag(u.String()))
	}

	if req.Options.Consistency == state.Strong {
		options.ConsistencyLevel = azcosmos.ConsistencyLevelSession.ToPtr()
	} else if req.Options.Consistency == state.Eventual {
		options.ConsistencyLevel = azcosmos.ConsistencyLevelEventual.ToPtr()
	}

	deleteCtx, cancel := context.WithTimeout(ctx, defaultTimeout)
	defer cancel()
	pk := azcosmos.NewPartitionKeyString(partitionKey)
	_, err = c.client.DeleteItem(deleteCtx, pk, req.Key, &options)
	if err != nil && !isNotFoundError(err) {
		resErr := &azcore.ResponseError{}
		if errors.As(err, &resErr) && resErr.StatusCode == http.StatusPreconditionFailed {
			return state.NewETagError(state.ETagMismatch, err)
		}
		return err
	}

	return nil
}

// MultiMaxSize returns the maximum number of operations allowed in a transaction.
// For Azure Cosmos DB, that's 100.
func (c StateStore) MultiMaxSize() int {
	return 100
}

// Multi performs a transactional operation. Succeeds only if all operations succeed, and fails if one or more operations fail.
// Note that all operations must be in the same partition.
func (c *StateStore) Multi(ctx context.Context, request *state.TransactionalStateRequest) (err error) {
	if len(request.Operations) == 0 {
		c.logger.Debugf("No operations provided")
		return nil
	}

	partitionKey := request.Metadata[metadataPartitionKey]
	batch := c.client.NewTransactionalBatch(azcosmos.NewPartitionKeyString(partitionKey))

	numOperations := 0
	// Loop through the list of operations. Create and add the operation to the batch
	for _, o := range request.Operations {
		options := &azcosmos.TransactionalBatchItemOptions{}

		switch req := o.(type) {
		case state.SetRequest:
			var doc CosmosItem
			doc, err = createUpsertItem(c.contentType, req, partitionKey)
			if err != nil {
				return err
			}
			doc.PartitionKey = partitionKey

			if req.HasETag() {
				etag := azcore.ETag(*req.ETag)
				options.IfMatchETag = &etag
			} else if req.Options.Concurrency == state.FirstWrite {
				var u uuid.UUID
				u, err = uuid.NewRandom()
				if err != nil {
					return err
				}
				options.IfMatchETag = ptr.Of(azcore.ETag(u.String()))
			}

			var marsh []byte
			marsh, err = json.Marshal(doc)
			if err != nil {
				return err
			}
			batch.UpsertItem(marsh, options)
			numOperations++
		case state.DeleteRequest:
			if req.HasETag() {
				etag := azcore.ETag(*req.ETag)
				options.IfMatchETag = &etag
			} else if req.Options.Concurrency == state.FirstWrite {
				var u uuid.UUID
				u, err = uuid.NewRandom()
				if err != nil {
					return err
				}
				options.IfMatchETag = ptr.Of(azcore.ETag(u.String()))
			}

			batch.DeleteItem(req.Key, options)
			numOperations++
		}
	}

	c.logger.Debugf("#operations=%d,partitionkey=%s", numOperations, partitionKey)

	execCtx, cancel := context.WithTimeout(ctx, defaultTimeout)
	defer cancel()
	batchResponse, err := c.client.ExecuteTransactionalBatch(execCtx, batch, nil)
	if err != nil {
		return err
	}

	if !batchResponse.Success {
		// Transaction failed, look for the offending operation
		for index, operation := range batchResponse.OperationResults {
			if operation.StatusCode != http.StatusFailedDependency {
				c.logger.Errorf("Transaction failed due to operation %v which failed with status code %d", index, operation.StatusCode)
				return fmt.Errorf("transaction failed due to operation %v which failed with status code %d", index, operation.StatusCode)
			}
		}
		return errors.New("transaction failed")
	}

	// Transaction succeeded
	// We can inspect the individual operation results
	if c.logger.IsOutputLevelEnabled(logger.DebugLevel) {
		for index, operation := range batchResponse.OperationResults {
			c.logger.Debugf("Operation %v completed with status code %d", index, operation.StatusCode)
		}
	}

	return nil
}

func (c *StateStore) Query(ctx context.Context, req *state.QueryRequest) (*state.QueryResponse, error) {
	q := &Query{}

	qbuilder := query.NewQueryBuilder(q)
	if err := qbuilder.BuildQuery(&req.Query); err != nil {
		return &state.QueryResponse{}, err
	}

	// If present partitionKey, the value will be used in the query disabling the cross partition
	q.partitionKey = ""
	if val, found := req.Metadata[metadataPartitionKey]; found {
		q.partitionKey = val
	}

	data, token, err := q.execute(ctx, c.client)
	if err != nil {
		return nil, err
	}

	return &state.QueryResponse{
		Results: data,
		Token:   token,
	}, nil
}

func (c *StateStore) Ping(ctx context.Context) error {
	pingCtx, cancel := context.WithTimeout(ctx, defaultTimeout)
	defer cancel()
	_, err := c.client.Read(pingCtx, nil)
	if err != nil {
		return err
	}
	return nil
}

func createUpsertItem(contentType string, req state.SetRequest, partitionKey string) (CosmosItem, error) {
	byteArray, isBinary := req.Value.([]byte)
	if len(byteArray) == 0 {
		isBinary = false
	}

	ttl, err := stateutils.ParseTTL(req.Metadata)
	if err != nil {
		return CosmosItem{}, fmt.Errorf("error parsing TTL from metadata: %s", err)
	}

	if isBinary {
		if contenttype.IsJSONContentType(contentType) {
			var value map[string]interface{}
			err := json.Unmarshal(byteArray, &value)
			// if byte array is not a valid JSON, so keep it as-is to be Base64 encoded in CosmosDB.
			// otherwise, we save it as JSON
			if err == nil {
				item := CosmosItem{
					ID:           req.Key,
					Value:        value,
					PartitionKey: partitionKey,
					IsBinary:     false,
					TTL:          ttl,
				}
				return item, nil
			}
		} else if contenttype.IsStringContentType(contentType) {
			item := CosmosItem{
				ID:           req.Key,
				Value:        string(byteArray),
				PartitionKey: partitionKey,
				IsBinary:     false,
				TTL:          ttl,
			}
			return item, nil
		}
	}

	item := CosmosItem{
		ID:           req.Key,
		Value:        req.Value,
		PartitionKey: partitionKey,
		IsBinary:     isBinary,
		TTL:          ttl,
	}
	return item, nil
}

// This is a helper to return the partition key to use.  If if metadata["partitionkey"] is present,
// use that, otherwise use what's in "key".
func populatePartitionMetadata(key string, requestMetadata map[string]string) string {
	if val, found := requestMetadata[metadataPartitionKey]; found {
		return val
	}

	return key
}

func isNotFoundError(err error) bool {
	if err == nil {
		return false
	}

	if requestError, ok := err.(*azcore.ResponseError); ok {
		if requestError.ErrorCode == statusNotFound {
			return true
		}
	}

	return false
}

func NewCosmosItemFromResponse(value []byte, logger logger.Logger) (item CosmosItem, err error) {
	err = jsoniter.ConfigFastest.Unmarshal(value, &item)
	if err != nil {
		return item, err
	}

	if item.IsBinary {
		if item.Value == nil {
			return item, nil
		}

		strValue, ok := item.Value.(string)
		if !ok || strValue == "" {
			return item, nil
		}

		item.Value, err = base64.StdEncoding.DecodeString(strValue)
		if err != nil {
			logger.Warnf("CosmosDB state store Get request could not decode binary string: %v. Returning raw string instead.", err)
			item.Value = []byte(strValue)
		}

		return item, nil
	}

	item.Value, err = jsoniter.ConfigFastest.Marshal(&item.Value)
	if err != nil {
		return item, err
	}

	return item, nil
}