Implement the Multi() interface method adding transaction support in … (#372)
* Implement the Multi() interface method adding transaction support in the cosmos db component. The implementation uses a stored procedure to do this. The stored procedure is registered in Init(). * linter * code review comments * code review * remove comment * casing and de-dup * use 'in' in query Co-authored-by: Yaron Schneider <yaronsc@microsoft.com>
This commit is contained in:
Leon Mai
GitHub
parent
4007ba3622
commit
a086a42f2b
|
|
@ -7,12 +7,13 @@ package cosmosdb
|
|||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/dapr/dapr/pkg/logger"
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
"strings"
|
||||
|
||||
"github.com/dapr/components-contrib/state"
|
||||
"github.com/dapr/dapr/pkg/logger"
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
|
||||
"github.com/a8m/documentdb"
|
||||
)
|
||||
|
|
@ -22,6 +23,7 @@ type StateStore struct {
|
|||
client *documentdb.DocumentDB
|
||||
collection *documentdb.Collection
|
||||
db *documentdb.Database
|
||||
sp *documentdb.Sproc
|
||||
|
||||
logger logger.Logger
|
||||
}
|
||||
|
|
@ -36,10 +38,22 @@ type credentials struct {
|
|||
// CosmosItem is a wrapper around a CosmosDB document
|
||||
type CosmosItem struct {
|
||||
documentdb.Document
|
||||
ID string `json:"id"`
|
||||
Value interface{} `json:"value"`
|
||||
ID string `json:"id"`
|
||||
Value interface{} `json:"value"`
|
||||
PartitionKey string `json:"partitionKey"`
|
||||
}
|
||||
|
||||
type storedProcedureDefinition struct {
|
||||
ID string `json:"id"`
|
||||
Body string `json:"body"`
|
||||
}
|
||||
|
||||
const (
|
||||
storedProcedureName = "__dapr__"
|
||||
metadataPartitionKey = "partitionKey"
|
||||
unknownPartitionKey = "__UNKNOWN__"
|
||||
)
|
||||
|
||||
// NewCosmosDBStateStore returns a new CosmosDB state store
|
||||
func NewCosmosDBStateStore(logger logger.Logger) *StateStore {
|
||||
return &StateStore{logger: logger}
|
||||
|
|
@ -47,6 +61,8 @@ func NewCosmosDBStateStore(logger logger.Logger) *StateStore {
|
|||
|
||||
// Init does metadata and connection parsing
|
||||
func (c *StateStore) Init(metadata state.Metadata) error {
|
||||
c.logger.Debugf("CosmosDB init start")
|
||||
|
||||
connInfo := metadata.Properties
|
||||
b, err := json.Marshal(connInfo)
|
||||
if err != nil {
|
||||
|
|
@ -87,12 +103,38 @@ func (c *StateStore) Init(metadata state.Metadata) error {
|
|||
if err != nil {
|
||||
return err
|
||||
} else if len(colls) == 0 {
|
||||
return fmt.Errorf("collection %s for CosmosDB state store not found", creds.Collection)
|
||||
return fmt.Errorf("collection %s for CosmosDB state store not found. This must be created before Dapr uses it", creds.Collection)
|
||||
}
|
||||
|
||||
c.collection = &colls[0]
|
||||
c.client = client
|
||||
|
||||
sps, err := c.client.ReadStoredProcedures(c.collection.Self)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// get a link to the sp
|
||||
for _, proc := range sps {
|
||||
if proc.Id == storedProcedureName {
|
||||
c.sp = &proc
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if c.sp == nil {
|
||||
// register the stored procedure
|
||||
createspBody := storedProcedureDefinition{ID: storedProcedureName, Body: spDefinition}
|
||||
c.sp, err = c.client.CreateStoredProcedure(c.collection.Self, createspBody)
|
||||
if err != nil {
|
||||
// if it already exists that is success
|
||||
if !strings.HasPrefix(err.Error(), "Conflict") {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
c.logger.Debug("cosmos Init done")
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -100,8 +142,10 @@ func (c *StateStore) Init(metadata state.Metadata) error {
|
|||
func (c *StateStore) Get(req *state.GetRequest) (*state.GetResponse, error) {
|
||||
key := req.Key
|
||||
|
||||
partitionKey := populatePartitionMetadata(req.Key, req.Metadata)
|
||||
|
||||
items := []CosmosItem{}
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(req.Key)}
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(partitionKey)}
|
||||
if req.Options.Consistency == state.Strong {
|
||||
options = append(options, documentdb.ConsistencyLevel(documentdb.Strong))
|
||||
}
|
||||
|
|
@ -139,7 +183,8 @@ func (c *StateStore) Set(req *state.SetRequest) error {
|
|||
return err
|
||||
}
|
||||
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(req.Key)}
|
||||
partitionKey := populatePartitionMetadata(req.Key, req.Metadata)
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(partitionKey)}
|
||||
|
||||
if req.ETag != "" {
|
||||
options = append(options, documentdb.IfMatch((req.ETag)))
|
||||
|
|
@ -151,7 +196,7 @@ func (c *StateStore) Set(req *state.SetRequest) error {
|
|||
options = append(options, documentdb.ConsistencyLevel(documentdb.Eventual))
|
||||
}
|
||||
|
||||
_, err = c.client.UpsertDocument(c.collection.Self, CosmosItem{ID: req.Key, Value: req.Value}, options...)
|
||||
_, err = c.client.UpsertDocument(c.collection.Self, CosmosItem{ID: req.Key, Value: req.Value, PartitionKey: partitionKey}, options...)
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
|
|
@ -179,9 +224,21 @@ func (c *StateStore) Delete(req *state.DeleteRequest) error {
|
|||
return err
|
||||
}
|
||||
|
||||
selfLink := fmt.Sprintf("dbs/%s/colls/%s/docs/%s", c.db.Id, c.collection.Id, req.Key)
|
||||
partitionKey := populatePartitionMetadata(req.Key, req.Metadata)
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(partitionKey)}
|
||||
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(req.Key)}
|
||||
items := []CosmosItem{}
|
||||
_, err = c.client.QueryDocuments(
|
||||
c.collection.Self,
|
||||
documentdb.NewQuery("SELECT * FROM ROOT r WHERE r.id=@id", documentdb.P{Name: "@id", Value: req.Key}),
|
||||
&items,
|
||||
options...,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
} else if len(items) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if req.ETag != "" {
|
||||
options = append(options, documentdb.IfMatch((req.ETag)))
|
||||
|
|
@ -193,7 +250,10 @@ func (c *StateStore) Delete(req *state.DeleteRequest) error {
|
|||
options = append(options, documentdb.ConsistencyLevel(documentdb.Eventual))
|
||||
}
|
||||
|
||||
_, err = c.client.DeleteDocument(selfLink, options...)
|
||||
_, err = c.client.DeleteDocument(items[0].Self, options...)
|
||||
if err != nil {
|
||||
c.logger.Debugf("Error from cosmos.DeleteDocument e=%e, e.Error=%s", err, err.Error())
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
|
|
@ -208,3 +268,67 @@ func (c *StateStore) BulkDelete(req []state.DeleteRequest) error {
|
|||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Multi performs a transactional operation. succeeds only if all operations succeed, and fails if one or more operations fail
|
||||
func (c *StateStore) Multi(operations []state.TransactionalRequest) error {
|
||||
upserts := []CosmosItem{}
|
||||
deletes := []CosmosItem{}
|
||||
|
||||
partitionKey := unknownPartitionKey
|
||||
previousPartitionKey := unknownPartitionKey
|
||||
|
||||
for _, o := range operations {
|
||||
t := o.Request.(state.KeyInt)
|
||||
key := t.GetKey()
|
||||
metadata := t.GetMetadata()
|
||||
|
||||
partitionKey = populatePartitionMetadata(key, metadata)
|
||||
if previousPartitionKey != unknownPartitionKey &&
|
||||
partitionKey != previousPartitionKey {
|
||||
return errors.New("all objects used in Multi() must have the same partition key")
|
||||
}
|
||||
previousPartitionKey = partitionKey
|
||||
|
||||
if o.Operation == state.Upsert {
|
||||
req := o.Request.(state.SetRequest)
|
||||
|
||||
upsertOperation := CosmosItem{
|
||||
ID: req.Key,
|
||||
Value: req.Value,
|
||||
PartitionKey: partitionKey}
|
||||
|
||||
upserts = append(upserts, upsertOperation)
|
||||
} else if o.Operation == state.Delete {
|
||||
req := o.Request.(state.DeleteRequest)
|
||||
|
||||
deleteOperation := CosmosItem{
|
||||
ID: req.Key,
|
||||
Value: "", // Value does not need to be specified
|
||||
PartitionKey: partitionKey}
|
||||
deletes = append(deletes, deleteOperation)
|
||||
}
|
||||
}
|
||||
|
||||
c.logger.Debugf("#upserts=%d,#deletes=%d, partitionkey=%s", len(upserts), len(deletes), partitionKey)
|
||||
|
||||
options := []documentdb.CallOption{documentdb.PartitionKey(partitionKey)}
|
||||
var retString string
|
||||
|
||||
// The stored procedure throws if it failed, which sets err to non-nil. It doesn't return anything else.
|
||||
err := c.client.ExecuteStoredProcedure(c.sp.Self, [...]interface{}{upserts, deletes}, &retString, options...)
|
||||
if err != nil {
|
||||
c.logger.Debugf("error=%e", err)
|
||||
return err
|
||||
}
|
||||
|
||||
return 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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,117 @@
|
|||
// ------------------------------------------------------------
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// Licensed under the MIT License.
|
||||
// ------------------------------------------------------------
|
||||
|
||||
package cosmosdb
|
||||
|
||||
const spDefinition string = `// upserts - an array of objects to upsert
|
||||
// deletes - an array of objects to delete
|
||||
|
||||
function dapr_multi(upserts, deletes) {
|
||||
var context = getContext();
|
||||
var container = context.getCollection();
|
||||
var response = context.getResponse();
|
||||
|
||||
if (typeof upserts === "string") {
|
||||
throw new Error("first arg is a string, expected array of objects");
|
||||
}
|
||||
|
||||
if (typeof deletes === "string") {
|
||||
throw new Error("second arg is a string, expected array of objects");
|
||||
}
|
||||
|
||||
// create the query string used to look up deletes
|
||||
var query = "select * from n where n.id in ";
|
||||
if (deletes.length > 0) {
|
||||
query += ("('" + deletes[0].id + "'");
|
||||
|
||||
for (let j = 1; j < deletes.length; j++) {
|
||||
query += ", '" + deletes[j].id + "'"
|
||||
}
|
||||
}
|
||||
|
||||
query += ')'
|
||||
console.log("query" + query)
|
||||
var upsertCount = 0;
|
||||
var deleteCount = 0;
|
||||
|
||||
var collectionLink = container.getSelfLink();
|
||||
|
||||
// do the upserts first
|
||||
if (upserts.length != 0) {
|
||||
tryCreate(upserts[upsertCount], callback);
|
||||
} else {
|
||||
tryQueryAndDelete();
|
||||
}
|
||||
|
||||
function tryCreate(doc, callback) {
|
||||
var isAccepted = container.upsertDocument(collectionLink, doc, callback);
|
||||
|
||||
// fail if we hit execution bounds
|
||||
if (!isAccepted) {
|
||||
throw new Error("upsertDocument() not accepted, please retry");
|
||||
}
|
||||
}
|
||||
|
||||
function callback(err, doc, options) {
|
||||
if (err) throw err;
|
||||
|
||||
upsertCount++;
|
||||
|
||||
if (upsertCount >= upserts.length) {
|
||||
|
||||
// upserts are done, start the deletes, if any
|
||||
if (deletes.length > 0) {
|
||||
tryQueryAndDelete()
|
||||
}
|
||||
} else {
|
||||
tryCreate(upserts[upsertCount], callback);
|
||||
}
|
||||
}
|
||||
|
||||
function tryQueryAndDelete() {
|
||||
var requestOptions = {};
|
||||
var isAccepted = container.queryDocuments(collectionLink, query, requestOptions, function (err, retrievedDocs, responseOptions) {
|
||||
if (err) {
|
||||
throw err;
|
||||
}
|
||||
|
||||
if (retrievedDocs == null) {
|
||||
response.setBody(JSON.stringify("success"));
|
||||
} else if (retrievedDocs.length > 0) {
|
||||
tryDelete(retrievedDocs);
|
||||
} else {
|
||||
// done with all deletes
|
||||
response.setBody(JSON.stringify("success"));
|
||||
}
|
||||
});
|
||||
|
||||
// fail if we hit execution bounds
|
||||
if (!isAccepted) {
|
||||
throw new Error("queryDocuments() not accepted, please retry");
|
||||
}
|
||||
}
|
||||
|
||||
function tryDelete(documents) {
|
||||
if (documents.length > 0) {
|
||||
// Delete the first document in the array.
|
||||
var isAccepted = container.deleteDocument(documents[0]._self, {}, function (err, responseOptions) {
|
||||
if (err) throw err;
|
||||
|
||||
deleteCount++;
|
||||
documents.shift();
|
||||
// Delete the next document in the array.
|
||||
tryDelete(documents);
|
||||
});
|
||||
|
||||
// fail if we hit execution bounds
|
||||
if (!isAccepted) {
|
||||
throw new Error("deleteDocument() not accepted, please retry");
|
||||
}
|
||||
} else {
|
||||
// If the document array is empty, query for more documents.
|
||||
tryQueryAndDelete();
|
||||
}
|
||||
}
|
||||
}`
|
||||
|
|
@ -27,6 +27,16 @@ type DeleteRequest struct {
|
|||
Options DeleteStateOption `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// Key gets the Key on a DeleteRequest
|
||||
func (r DeleteRequest) GetKey() string {
|
||||
return r.Key
|
||||
}
|
||||
|
||||
// Metadata gets the Metadata on a DeleteRequest
|
||||
func (r DeleteRequest) GetMetadata() map[string]string {
|
||||
return r.Metadata
|
||||
}
|
||||
|
||||
// DeleteStateOption controls how a state store reacts to a delete request
|
||||
type DeleteStateOption struct {
|
||||
Concurrency string `json:"concurrency,omitempty"` //"concurrency"
|
||||
|
|
@ -43,6 +53,16 @@ type SetRequest struct {
|
|||
Options SetStateOption `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// Key gets the Key on a SetRequest
|
||||
func (r SetRequest) GetKey() string {
|
||||
return r.Key
|
||||
}
|
||||
|
||||
// Metadata gets the Key on a SetRequest
|
||||
func (r SetRequest) GetMetadata() map[string]string {
|
||||
return r.Metadata
|
||||
}
|
||||
|
||||
//RetryPolicy describes how retries should be handled
|
||||
type RetryPolicy struct {
|
||||
Interval time.Duration `json:"interval"`
|
||||
|
|
@ -72,3 +92,9 @@ type TransactionalRequest struct {
|
|||
Operation OperationType `json:"operation"`
|
||||
Request interface{} `json:"request"`
|
||||
}
|
||||
|
||||
// KeyInt is an interface that allows gets of the Key and Metadata inside requests
|
||||
type KeyInt interface {
|
||||
GetKey() string
|
||||
GetMetadata() map[string]string
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue