Cosmos DB - Adds TransactionalBatch support (Azure#17795)

Adding support for TransactionalBatch

Author: ealsur

sdk/data/azcosmos/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+* Added Transactional Batch support
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/data/azcosmos/README.md

@@ -8,7 +8,7 @@ This client library enables client applications to connect to Azure Cosmos via t
 
 ### Prerequisites
 
-* Go versions 1.16 or higher
+* Go versions 1.18 or higher
 * An Azure subscription or free Azure Cosmos DB trial account
 
 Note: If you don't have an Azure subscription, create a free account before you begin.

sdk/data/azcosmos/cosmos_client.go

@@ -238,6 +238,26 @@ func (c *Client) sendDeleteRequest(
 	return c.executeAndEnsureSuccessResponse(req)
 }
 
+func (c *Client) sendBatchRequest(
+	ctx context.Context,
+	path string,
+	batch []batchOperation,
+	operationContext pipelineRequestOptions,
+	requestOptions cosmosRequestOptions,
+	requestEnricher func(*policy.Request)) (*http.Response, error) {
+	req, err := c.createRequest(path, ctx, http.MethodPost, operationContext, requestOptions, requestEnricher)
+	if err != nil {
+		return nil, err
+	}
+
+	err = c.attachContent(batch, req)
+	if err != nil {
+		return nil, err
+	}
+
+	return c.executeAndEnsureSuccessResponse(req)
+}
+
 func (c *Client) createRequest(
 	path string,
 	ctx context.Context,
@@ -288,7 +308,6 @@ func (c *Client) attachContent(content interface{}, req *policy.Request) error {
 	default:
 		// Otherwise, we need to marshal it
 		err = azruntime.MarshalAsJSON(req, content)
-
 	}
 
 	if err != nil {

sdk/data/azcosmos/cosmos_client_test.go

@@ -354,6 +354,50 @@ func TestSendQuery(t *testing.T) {
 	}
 }
 
+func TestSendBatch(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(
+		mock.WithStatusCode(200))
+	verifier := pipelineVerifier{}
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{&verifier}}, &policy.ClientOptions{Transport: srv})
+	client := &Client{endpoint: srv.URL(), pipeline: pl}
+	operationContext := pipelineRequestOptions{
+		resourceType:    resourceTypeDocument,
+		resourceAddress: "",
+	}
+
+	batch := TransactionalBatch{}
+	batch.partitionKey = NewPartitionKeyString("foo")
+
+	body := map[string]string{
+		"foo": "bar",
+	}
+
+	itemMarshall, _ := json.Marshal(body)
+
+	batch.CreateItem(itemMarshall, nil)
+	batch.ReadItem("someId", nil)
+
+	marshalled, err := json.Marshal(batch.operations)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = client.sendBatchRequest(context.Background(), "/", batch.operations, operationContext, &TransactionalBatchOptions{}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if verifier.requests[0].method != http.MethodPost {
+		t.Errorf("Expected %v, but got %v", http.MethodPost, verifier.requests[0].method)
+	}
+
+	if verifier.requests[0].body != string(marshalled) {
+		t.Errorf("Expected %v, but got %v", string(marshalled), verifier.requests[0].body)
+	}
+}
+
 func TestCreateScopeFromEndpoint(t *testing.T) {
 	url := "https://foo.documents.azure.com:443/"
 	scope, err := createScopeFromEndpoint(url)

sdk/data/azcosmos/cosmos_container.go

@@ -5,6 +5,7 @@ package azcosmos
 
 import (
 	"context"
+	"errors"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
@@ -456,6 +457,66 @@ func (c *ContainerClient) NewQueryItemsPager(query string, partitionKey Partitio
 	})
 }
 
+// NewTransactionalBatch creates a batch of operations to be committed as a single unit.
+// See https://docs.microsoft.com/azure/cosmos-db/sql/transactional-batch
+func (c *ContainerClient) NewTransactionalBatch(partitionKey PartitionKey) TransactionalBatch {
+	return TransactionalBatch{partitionKey: partitionKey}
+}
+
+// ExecuteTransactionalBatch executes a transactional batch.
+// Once executed, verify the Success property of the response to determine if the batch was committed
+func (c *ContainerClient) ExecuteTransactionalBatch(ctx context.Context, b TransactionalBatch, o *TransactionalBatchOptions) (TransactionalBatchResponse, error) {
+	if len(b.operations) == 0 {
+		return TransactionalBatchResponse{}, errors.New("no operations in batch")
+	}
+
+	h := headerOptionsOverride{
+		partitionKey: &b.partitionKey,
+	}
+
+	if o == nil {
+		o = &TransactionalBatchOptions{}
+	} else {
+		h.enableContentResponseOnWrite = &o.EnableContentResponseOnWrite
+	}
+
+	// If contentResponseOnWrite is not enabled at the client level the
+	// service will not even send a batch response payload
+	// Instead we should automatically enforce contentResponseOnWrite for all
+	// batch requests whenever at least one of the item operations requires a content response (read operation)
+	enableContentResponseOnWriteForReadOperations := true
+	for _, op := range b.operations {
+		if op.getOperationType() == operationTypeRead {
+			h.enableContentResponseOnWrite = &enableContentResponseOnWriteForReadOperations
+			break
+		}
+	}
+
+	operationContext := pipelineRequestOptions{
+		resourceType:          resourceTypeDocument,
+		resourceAddress:       c.link,
+		isWriteOperation:      true,
+		headerOptionsOverride: &h}
+
+	path, err := generatePathForNameBased(resourceTypeDocument, operationContext.resourceAddress, true)
+	if err != nil {
+		return TransactionalBatchResponse{}, err
+	}
+
+	azResponse, err := c.database.client.sendBatchRequest(
+		ctx,
+		path,
+		b.operations,
+		operationContext,
+		o,
+		nil)
+	if err != nil {
+		return TransactionalBatchResponse{}, err
+	}
+
+	return newTransactionalBatchResponse(azResponse)
+}
+
 func (c *ContainerClient) getRID(ctx context.Context) (string, error) {
 	containerResponse, err := c.Read(ctx, nil)
 	if err != nil {

sdk/data/azcosmos/cosmos_container_test.go

@@ -520,3 +520,72 @@ func TestContainerQueryItems(t *testing.T) {
 		}
 	}
 }
+
+func TestContainerExecuteBatch(t *testing.T) {
+	batchResponseRaw := []map[string]interface{}{
+		{"statusCode": 200, "requestCharge": 10.0, "eTag": "someETag", "resourceBody": "someBody"},
+		{"statusCode": 201, "requestCharge": 11.0, "eTag": "someETag2"},
+	}
+
+	jsonString, err := json.Marshal(batchResponseRaw)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(
+		mock.WithBody(jsonString),
+		mock.WithStatusCode(http.StatusOK),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"))
+
+	verifier := pipelineVerifier{}
+
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{&verifier}}, &policy.ClientOptions{Transport: srv})
+	client := &Client{endpoint: srv.URL(), pipeline: pl}
+
+	database, _ := newDatabase("databaseId", client)
+	container, _ := newContainer("containerId", database)
+
+	pk := NewPartitionKeyString("pk")
+	batch := container.NewTransactionalBatch(pk)
+	_, err = container.ExecuteTransactionalBatch(context.TODO(), batch, nil)
+	if err == nil {
+		t.Fatal("Expected error, but got nil")
+	}
+
+	batch.ReadItem("someId", nil)
+
+	body := map[string]string{
+		"foo": "bar",
+	}
+
+	itemMarshall, _ := json.Marshal(body)
+	batch.CreateItem(itemMarshall, nil)
+
+	_, err = container.ExecuteTransactionalBatch(context.TODO(), batch, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(verifier.requests) != 1 {
+		t.Fatalf("Expected 1 request, got %d", len(verifier.requests))
+	}
+
+	request := verifier.requests[0]
+
+	if request.method != http.MethodPost {
+		t.Errorf("Expected method to be %s, but got %s", http.MethodPost, request.method)
+	}
+
+	if request.url.RequestURI() != "/dbs/databaseId/colls/containerId/docs" {
+		t.Errorf("Expected url to be %s, but got %s", "/dbs/databaseId/colls/containerId/docs", request.url.RequestURI())
+	}
+
+	marshalledOperations, _ := json.Marshal(batch.operations)
+	if request.body != string(marshalledOperations) {
+		t.Errorf("Expected %v, but got %v", string(marshalledOperations), request.body)
+	}
+}

sdk/data/azcosmos/cosmos_headers.go

@@ -30,6 +30,9 @@ const (
 	cosmosHeaderQueryMetrics                       string = "x-ms-documentdb-query-metrics"
 	cosmosHeaderIndexUtilization                   string = "x-ms-cosmos-index-utilization"
 	cosmosHeaderCorrelatedActivityId               string = "x-ms-cosmos-correlated-activityid"
+	cosmosHeaderIsBatchRequest                     string = "x-ms-cosmos-is-batch-request"
+	cosmosHeaderIsBatchAtomic                      string = "x-ms-cosmos-batch-atomic"
+	cosmosHeaderIsBatchOrdered                     string = "x-ms-cosmos-batch-ordered"
 	headerXmsDate                                  string = "x-ms-date"
 	headerAuthorization                            string = "Authorization"
 	headerContentType                              string = "Content-Type"