Add support for MaxIntegratedCacheStaleness option (#23406)

* Add support for MaxIntegratedCacheStaleness option

* Add Changelog entry

* Fix lints

sdk/data/azcosmos/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 * Added support for OpenTelemetry trace spans. See [PR 23268](https://github.com/Azure/azure-sdk-for-go/pull/23268)
+* Added support for MaxIntegratedCacheStaleness option See [PR 23406](https://github.com/Azure/azure-sdk-for-go/pull/23406)
 
 ### Breaking Changes
 

sdk/data/azcosmos/cosmos_dedicated_gateway_request_options.go

@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import "time"
+
+// DedicatedGatewayRequestOptions includes options for operations in the dedicated gateway.
+type DedicatedGatewayRequestOptions struct {
+	// Gets or sets the staleness value associated with the request in the Azure Cosmos DB service.
+	// For requests where the ConsistencyLevel is ConsistencyLevel.Eventual or ConsistencyLevel.Session,
+	// responses from the integrated cache are guaranteed to be no staler than value indicated by this MaxIntegratedCacheStaleness.
+	// Cache Staleness is supported in milliseconds granularity. Anything smaller than milliseconds will be ignored.
+	MaxIntegratedCacheStaleness *time.Duration
+}

sdk/data/azcosmos/cosmos_http_constants.go

@@ -78,6 +78,7 @@ const (
 	cosmosHeaderIsPartitionKeyDeletePending        string = "x-ms-cosmos-is-partition-key-delete-pending"
 	cosmosHeaderQueryExecutionInfo                 string = "x-ms-cosmos-query-execution-info"
 	headerXmsItemCount                             string = "x-ms-item-count"
+	headerDedicatedGatewayMaxAge                   string = "x-ms-dedicatedgateway-max-age"
 )
 
 const (

sdk/data/azcosmos/cosmos_item_request_options.go

@@ -4,6 +4,7 @@
 package azcosmos
 
 import (
+	"strconv"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
@@ -34,6 +35,8 @@ type ItemOptions struct {
 	// IfMatchEtag is used to ensure optimistic concurrency control.
 	// https://docs.microsoft.com/azure/cosmos-db/sql/database-transactions-optimistic-concurrency#optimistic-concurrency-control
 	IfMatchEtag *azcore.ETag
+	// Options for operations in the dedicated gateway.
+	DedicatedGatewayRequestOptions *DedicatedGatewayRequestOptions
 }
 
 func (options *ItemOptions) toHeaders() *map[string]string {
@@ -63,5 +66,14 @@ func (options *ItemOptions) toHeaders() *map[string]string {
 		headers[headerIfMatch] = string(*options.IfMatchEtag)
 	}
 
+	if options.DedicatedGatewayRequestOptions != nil {
+		dedicatedGatewayRequestOptions := options.DedicatedGatewayRequestOptions
+
+		if dedicatedGatewayRequestOptions.MaxIntegratedCacheStaleness != nil {
+			milliseconds := dedicatedGatewayRequestOptions.MaxIntegratedCacheStaleness.Milliseconds()
+			headers[headerDedicatedGatewayMaxAge] = strconv.FormatInt(milliseconds, 10)
+		}
+	}
+
 	return &headers
 }

sdk/data/azcosmos/cosmos_item_request_options_test.go

@@ -4,7 +4,9 @@
 package azcosmos
 
 import (
+	"strconv"
 	"testing"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 )
@@ -19,6 +21,10 @@ func TestItemRequestOptionsToHeaders(t *testing.T) {
 	options.IndexingDirective = IndexingDirectiveInclude.ToPtr()
 	etagValue := azcore.ETag("someEtag")
 	options.IfMatchEtag = &etagValue
+	maxIntegratedCacheStalenessDuration := time.Duration(5 * time.Minute)
+	options.DedicatedGatewayRequestOptions = &DedicatedGatewayRequestOptions{
+		MaxIntegratedCacheStaleness: &maxIntegratedCacheStalenessDuration,
+	}
 	header := options.toHeaders()
 	if header == nil {
 		t.Fatal("toHeaders should return non-nil")
@@ -43,4 +49,7 @@ func TestItemRequestOptionsToHeaders(t *testing.T) {
 	if headers[headerIfMatch] != string(*options.IfMatchEtag) {
 		t.Errorf("IfMatchEtag should be someEtag but got %v", headers[headerIfMatch])
 	}
+	if headers[headerDedicatedGatewayMaxAge] != strconv.FormatInt(300000, 10) {
+		t.Errorf("headerDedicatedGatewayMaxAge should be 300000 but got %v", headers[headerDedicatedGatewayMaxAge])
+	}
 }

sdk/data/azcosmos/cosmos_query_request_options.go

@@ -3,7 +3,10 @@
 
 package azcosmos
 
-import "fmt"
+import (
+	"fmt"
+	"strconv"
+)
 
 // QueryOptions includes options for query operations on items.
 type QueryOptions struct {
@@ -34,6 +37,8 @@ type QueryOptions struct {
 	// QueryParameters allows execution of parametrized queries.
 	// See https://docs.microsoft.com/azure/cosmos-db/sql/sql-query-parameterized-queries
 	QueryParameters []QueryParameter
+	// Options for operations in the dedicated gateway.
+	DedicatedGatewayRequestOptions *DedicatedGatewayRequestOptions
 }
 
 func (options *QueryOptions) toHeaders() *map[string]string {
@@ -67,6 +72,15 @@ func (options *QueryOptions) toHeaders() *map[string]string {
 		headers[cosmosHeaderContinuationToken] = *options.ContinuationToken
 	}
 
+	if options.DedicatedGatewayRequestOptions != nil {
+		dedicatedGatewayRequestOptions := options.DedicatedGatewayRequestOptions
+
+		if dedicatedGatewayRequestOptions.MaxIntegratedCacheStaleness != nil {
+			milliseconds := dedicatedGatewayRequestOptions.MaxIntegratedCacheStaleness.Milliseconds()
+			headers[headerDedicatedGatewayMaxAge] = strconv.FormatInt(milliseconds, 10)
+		}
+	}
+
 	headers[cosmosHeaderPopulateQueryMetrics] = "true"
 
 	return &headers

sdk/data/azcosmos/cosmos_query_request_options_test.go

@@ -4,7 +4,9 @@
 package azcosmos
 
 import (
+	"strconv"
 	"testing"
+	"time"
 )
 
 func TestQueryRequestOptionsToHeaders(t *testing.T) {
@@ -18,6 +20,10 @@ func TestQueryRequestOptionsToHeaders(t *testing.T) {
 	options.PopulateIndexMetrics = true
 	continuation := "continuationToken"
 	options.ContinuationToken = &continuation
+	maxIntegratedCacheStalenessDuration := time.Duration(5 * time.Minute)
+	options.DedicatedGatewayRequestOptions = &DedicatedGatewayRequestOptions{
+		MaxIntegratedCacheStaleness: &maxIntegratedCacheStalenessDuration,
+	}
 	header := options.toHeaders()
 	if header == nil {
 		t.Fatal("toHeaders should return non-nil")
@@ -48,4 +54,7 @@ func TestQueryRequestOptionsToHeaders(t *testing.T) {
 	if headers[cosmosHeaderPopulateQueryMetrics] != "true" {
 		t.Errorf("PopulateQueryMetrics should be true but got %v", headers[cosmosHeaderPopulateQueryMetrics])
 	}
+	if headers[headerDedicatedGatewayMaxAge] != strconv.FormatInt(300000, 10) {
+		t.Errorf("headerDedicatedGatewayMaxAge should be 300000 but got %v", headers[headerDedicatedGatewayMaxAge])
+	}
 }