From 70049cb64b6245c537b3927120cece26781eb350 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 9 Jan 2024 13:48:19 -0500 Subject: [PATCH 01/12] added policy and linked gem to client --- sdk/data/azcosmos/cosmos_client.go | 38 +++++++++++++++++-- .../cosmos_global_endpoint_manager_policy.go | 26 +++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) create mode 100644 sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index 41807ab416e0..a940a065a14a 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -24,6 +24,7 @@ import ( type Client struct { endpoint string pipeline azruntime.Pipeline + gem *globalEndpointManager } // Endpoint used to create the client. @@ -36,7 +37,14 @@ func (c *Client) Endpoint() string { // cred - The credential used to authenticate with the cosmos service. // options - Optional Cosmos client options. Pass nil to accept default values. func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) { - return &Client{endpoint: endpoint, pipeline: newPipeline(newSharedKeyCredPolicy(cred), o)}, nil + internalClient := &Client{endpoint: endpoint, pipeline: newInternalPipeline(newSharedKeyCredPolicy(cred), o), gem: &globalEndpointManager{}} + + //need to pass in preferredRegions from options here once those changes are merged + gem, err := newGlobalEndpointManager(internalClient, []string{}, 0) + if err != nil { + return nil, err + } + return &Client{endpoint: endpoint, pipeline: newPipeline(newSharedKeyCredPolicy(cred), gem, o), gem: gem}, nil } // NewClient creates a new instance of Cosmos client with Azure AD access token authentication. It uses the default pipeline configuration. @@ -48,7 +56,16 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) ( if err != nil { return nil, err } - return &Client{endpoint: endpoint, pipeline: newPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o)}, nil + + internalClient := &Client{endpoint: endpoint, pipeline: newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), gem: &globalEndpointManager{}} + + //need to pass in preferredRegions from options here once those changes are merged + gem, err := newGlobalEndpointManager(internalClient, []string{}, 0) + if err != nil { + return nil, err + } + + return &Client{endpoint: endpoint, pipeline: newPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), gem, o), gem: gem}, nil } // NewClientFromConnectionString creates a new instance of Cosmos client from connection string. It uses the default pipeline configuration. @@ -87,7 +104,7 @@ func NewClientFromConnectionString(connectionString string, o *ClientOptions) (* return NewClientWithKey(endpoint, cred, o) } -func newPipeline(authPolicy policy.Policy, options *ClientOptions) azruntime.Pipeline { +func newPipeline(authPolicy policy.Policy, gem *globalEndpointManager, options *ClientOptions) azruntime.Pipeline { if options == nil { options = &ClientOptions{} } @@ -98,7 +115,21 @@ func newPipeline(authPolicy policy.Policy, options *ClientOptions) azruntime.Pip &headerPolicies{ enableContentResponseOnWrite: options.EnableContentResponseOnWrite, }, + &globalEndpointManagerPolicy{gem: gem}, + }, + PerRetry: []policy.Policy{ + authPolicy, }, + }, + &options.ClientOptions) +} + +func newInternalPipeline(authPolicy policy.Policy, options *ClientOptions) azruntime.Pipeline { + if options == nil { + options = &ClientOptions{} + } + return azruntime.NewPipeline("azcosmos", serviceLibVersion, + azruntime.PipelineOptions{ PerRetry: []policy.Policy{ authPolicy, }, @@ -176,6 +207,7 @@ func (c *Client) CreateDatabase( if err != nil { return DatabaseResponse{}, err } + fmt.Printf("- db create succeeded with code %d", azResponse.StatusCode) return newDatabaseResponse(azResponse) } diff --git a/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go b/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go new file mode 100644 index 000000000000..cd5bb4bf5403 --- /dev/null +++ b/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package azcosmos + +import ( + "context" + "fmt" + "net/http" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" +) + +type globalEndpointManagerPolicy struct { + gem *globalEndpointManager +} + +func (p *globalEndpointManagerPolicy) Do(req *policy.Request) (*http.Response, error) { + shouldRefresh := p.gem.ShouldRefresh() + if shouldRefresh { + fmt.Println("should refresh true go routine") + go p.gem.Update(context.Background()) + } + fmt.Println("policy done") + return req.Next() +} From 7b6cd4f24e2b58d25351439ccb15c8620d2e12fe Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 9 Jan 2024 16:01:59 -0500 Subject: [PATCH 02/12] refactor go routine for CI pipeline lint --- sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go b/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go index cd5bb4bf5403..9265c3522663 100644 --- a/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go +++ b/sdk/data/azcosmos/cosmos_global_endpoint_manager_policy.go @@ -5,7 +5,6 @@ package azcosmos import ( "context" - "fmt" "net/http" "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" @@ -18,9 +17,9 @@ type globalEndpointManagerPolicy struct { func (p *globalEndpointManagerPolicy) Do(req *policy.Request) (*http.Response, error) { shouldRefresh := p.gem.ShouldRefresh() if shouldRefresh { - fmt.Println("should refresh true go routine") - go p.gem.Update(context.Background()) + go func() { + _ = p.gem.Update(context.Background()) + }() } - fmt.Println("policy done") return req.Next() } From 8e0eea6c233d15e605fc3440651f357021fa595d Mon Sep 17 00:00:00 2001 From: simorenoh Date: Wed, 10 Jan 2024 09:51:49 -0500 Subject: [PATCH 03/12] remove debugging print statement --- sdk/data/azcosmos/cosmos_client.go | 1 - 1 file changed, 1 deletion(-) diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index a940a065a14a..4ed5680f9cff 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -207,7 +207,6 @@ func (c *Client) CreateDatabase( if err != nil { return DatabaseResponse{}, err } - fmt.Printf("- db create succeeded with code %d", azResponse.StatusCode) return newDatabaseResponse(azResponse) } From f219320dd4ee668ee7008be1feaa33ef7952f0d5 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Thu, 11 Jan 2024 09:11:20 -0500 Subject: [PATCH 04/12] changed to internal pipeline, updated tests --- sdk/data/azcosmos/cosmos_client.go | 9 +--- .../cosmos_global_endpoint_manager.go | 44 ++++++++++++---- .../cosmos_global_endpoint_manager_test.go | 51 ++++++------------- ...tor_cosmos_global_endpoint_manager_test.go | 2 +- 4 files changed, 52 insertions(+), 54 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index 4ed5680f9cff..b897e3bd43c3 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -37,10 +37,8 @@ func (c *Client) Endpoint() string { // cred - The credential used to authenticate with the cosmos service. // options - Optional Cosmos client options. Pass nil to accept default values. func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) { - internalClient := &Client{endpoint: endpoint, pipeline: newInternalPipeline(newSharedKeyCredPolicy(cred), o), gem: &globalEndpointManager{}} - //need to pass in preferredRegions from options here once those changes are merged - gem, err := newGlobalEndpointManager(internalClient, []string{}, 0) + gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), []string{}, 0) if err != nil { return nil, err } @@ -56,11 +54,8 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) ( if err != nil { return nil, err } - - internalClient := &Client{endpoint: endpoint, pipeline: newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), gem: &globalEndpointManager{}} - //need to pass in preferredRegions from options here once those changes are merged - gem, err := newGlobalEndpointManager(internalClient, []string{}, 0) + gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), []string{}, 0) if err != nil { return nil, err } diff --git a/sdk/data/azcosmos/cosmos_global_endpoint_manager.go b/sdk/data/azcosmos/cosmos_global_endpoint_manager.go index 82770d999ec9..5abd950ae65d 100644 --- a/sdk/data/azcosmos/cosmos_global_endpoint_manager.go +++ b/sdk/data/azcosmos/cosmos_global_endpoint_manager.go @@ -17,7 +17,8 @@ import ( const defaultUnavailableLocationRefreshInterval = 5 * time.Minute type globalEndpointManager struct { - client *Client + clientEndpoint string + pipeline azruntime.Pipeline preferredLocations []string locationCache *locationCache refreshTimeInterval time.Duration @@ -25,8 +26,8 @@ type globalEndpointManager struct { lastUpdateTime time.Time } -func newGlobalEndpointManager(client *Client, preferredLocations []string, refreshTimeInterval time.Duration) (*globalEndpointManager, error) { - endpoint, err := url.Parse(client.endpoint) +func newGlobalEndpointManager(clientEndpoint string, pipeline azruntime.Pipeline, preferredLocations []string, refreshTimeInterval time.Duration) (*globalEndpointManager, error) { + endpoint, err := url.Parse(clientEndpoint) if err != nil { return &globalEndpointManager{}, err } @@ -36,7 +37,8 @@ func newGlobalEndpointManager(client *Client, preferredLocations []string, refre } gem := &globalEndpointManager{ - client: client, + clientEndpoint: clientEndpoint, + pipeline: pipeline, preferredLocations: preferredLocations, locationCache: newLocationCache(preferredLocations, *endpoint), refreshTimeInterval: refreshTimeInterval, @@ -115,19 +117,39 @@ func (gem *globalEndpointManager) GetAccountProperties(ctx context.Context) (acc return accountProperties{}, fmt.Errorf("failed to generate path for name-based request: %v", err) } - ctx, cancel := context.WithTimeout(ctx, 60*time.Second) - azResponse, err := gem.client.sendGetRequest(path, ctx, operationContext, nil, nil) - cancel() + finalURL := gem.clientEndpoint + if path != "" { + finalURL = azruntime.JoinPaths(gem.clientEndpoint, path) + } + + ctxt, cancel := context.WithTimeout(ctx, 60*time.Second) + defer cancel() + req, err := azruntime.NewRequest(ctxt, http.MethodGet, finalURL) if err != nil { - return accountProperties{}, fmt.Errorf("failed to retrieve account properties: %v", err) + return accountProperties{}, err } - properties, err := newAccountProperties(azResponse) + req.Raw().Header.Set(headerXmsDate, time.Now().UTC().Format(http.TimeFormat)) + req.Raw().Header.Set(headerXmsVersion, "2020-11-05") + req.Raw().Header.Set(cosmosHeaderSDKSupportedCapabilities, supportedCapabilitiesHeaderValue) + + req.SetOperationValue(operationContext) + + azResponse, err := gem.pipeline.Do(req) if err != nil { - return accountProperties{}, fmt.Errorf("failed to parse account properties: %v", err) + return accountProperties{}, err } - return properties, nil + successResponse := (azResponse.StatusCode >= 200 && azResponse.StatusCode < 300) || azResponse.StatusCode == 304 + if successResponse { + properties, err := newAccountProperties(azResponse) + if err != nil { + return accountProperties{}, fmt.Errorf("failed to parse account properties: %v", err) + } + return properties, nil + } + + return accountProperties{}, newCosmosError(azResponse) } func newAccountProperties(azResponse *http.Response) (accountProperties, error) { diff --git a/sdk/data/azcosmos/cosmos_global_endpoint_manager_test.go b/sdk/data/azcosmos/cosmos_global_endpoint_manager_test.go index 4c0089e55c74..0273aa6c83d4 100644 --- a/sdk/data/azcosmos/cosmos_global_endpoint_manager_test.go +++ b/sdk/data/azcosmos/cosmos_global_endpoint_manager_test.go @@ -34,11 +34,7 @@ func TestGlobalEndpointManagerGetWriteEndpoints(t *testing.T) { pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv}) - client := &Client{endpoint: srv.URL(), pipeline: pl} - - preferredRegions := []string{"West US", "Central US"} - - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute) assert.NoError(t, err) writeEndpoints, err := gem.GetWriteEndpoints() @@ -50,6 +46,7 @@ func TestGlobalEndpointManagerGetWriteEndpoints(t *testing.T) { expectedWriteEndpoints := []url.URL{ *serverEndpoint, } + assert.Equal(t, expectedWriteEndpoints, writeEndpoints) } @@ -60,11 +57,7 @@ func TestGlobalEndpointManagerGetReadEndpoints(t *testing.T) { pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv}) - client := &Client{endpoint: srv.URL(), pipeline: pl} - - preferredRegions := []string{"West US", "Central US"} - - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute) assert.NoError(t, err) readEndpoints, err := gem.GetReadEndpoints() @@ -88,12 +81,10 @@ func TestGlobalEndpointManagerMarkEndpointUnavailableForRead(t *testing.T) { client := &Client{endpoint: srv.URL(), pipeline: pl} - preferredRegions := []string{"West US", "Central US"} - - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + endpoint, err := url.Parse(client.endpoint) assert.NoError(t, err) - endpoint, err := url.Parse(client.endpoint) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute) assert.NoError(t, err) err = gem.MarkEndpointUnavailableForRead(*endpoint) @@ -112,12 +103,10 @@ func TestGlobalEndpointManagerMarkEndpointUnavailableForWrite(t *testing.T) { client := &Client{endpoint: srv.URL(), pipeline: pl} - preferredRegions := []string{"West US", "Central US"} - - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + endpoint, err := url.Parse(client.endpoint) assert.NoError(t, err) - endpoint, err := url.Parse(client.endpoint) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute) assert.NoError(t, err) err = gem.MarkEndpointUnavailableForWrite(*endpoint) @@ -130,7 +119,6 @@ func TestGlobalEndpointManagerMarkEndpointUnavailableForWrite(t *testing.T) { func TestGlobalEndpointManagerGetEndpointLocation(t *testing.T) { srv, close := mock.NewTLSServer() defer close() - srv.SetResponse(mock.WithStatusCode(http.StatusOK)) westRegion := accountRegion{ Name: "West US", @@ -144,19 +132,17 @@ func TestGlobalEndpointManagerGetEndpointLocation(t *testing.T) { } jsonString, err := json.Marshal(properties) - if err != nil { - t.Fatal(err) - } + assert.NoError(t, err) + + srv.SetResponse(mock.WithStatusCode(200)) srv.SetResponse(mock.WithBody(jsonString)) pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv}) - client := &Client{endpoint: srv.URL(), pipeline: pl} - - gem, err := newGlobalEndpointManager(client, []string{}, 5*time.Minute) + serverEndpoint, err := url.Parse(srv.URL()) assert.NoError(t, err) - serverEndpoint, err := url.Parse(srv.URL()) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Minute) assert.NoError(t, err) err = gem.Update(context.Background()) @@ -175,11 +161,7 @@ func TestGlobalEndpointManagerGetAccountProperties(t *testing.T) { pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv}) - client := &Client{endpoint: srv.URL(), pipeline: pl} - - preferredRegions := []string{"West US", "Central US"} - - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute) assert.NoError(t, err) accountProps, err := gem.GetAccountProperties(context.Background()) @@ -212,13 +194,13 @@ func TestGlobalEndpointManagerCanUseMultipleWriteLocations(t *testing.T) { mockLc.useMultipleWriteLocations = true mockGem := globalEndpointManager{ - client: client, + clientEndpoint: client.endpoint, preferredLocations: preferredRegions, locationCache: mockLc, refreshTimeInterval: 5 * time.Minute, } - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Minute) assert.NoError(t, err) // Multiple locations should be false for default GEM @@ -254,9 +236,8 @@ func TestGlobalEndpointManagerConcurrentUpdate(t *testing.T) { srv.SetResponse(mock.WithBody(jsonString)) pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{countPolicy}}, &policy.ClientOptions{Transport: srv}) - client := &Client{endpoint: srv.URL(), pipeline: pl} - gem, err := newGlobalEndpointManager(client, []string{}, 5*time.Second) + gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Second) assert.NoError(t, err) // Call update concurrently and see how many times the policy gets called diff --git a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go index 3409758c4f42..128b7a973ae1 100644 --- a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go @@ -19,7 +19,7 @@ func TestGlobalEndpointManagerEmulator(t *testing.T) { preferredRegions := []string{} emulatorRegion := accountRegion{Name: emulatorRegionName, Endpoint: "https://127.0.0.1:8081/"} - gem, err := newGlobalEndpointManager(client, preferredRegions, 5*time.Minute) + gem, err := newGlobalEndpointManager(client.endpoint, client.pipeline, preferredRegions, 5*time.Minute) assert.NoError(t, err) accountProps, err := gem.GetAccountProperties(context.Background()) From 10757e2bd1b30e8ef0bd046b0fe3ecf8563c9bb9 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Thu, 11 Jan 2024 15:33:03 -0500 Subject: [PATCH 05/12] removed unneeded code --- .../azcosmos/cosmos_global_endpoint_manager.go | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_global_endpoint_manager.go b/sdk/data/azcosmos/cosmos_global_endpoint_manager.go index 5abd950ae65d..2740bac74b46 100644 --- a/sdk/data/azcosmos/cosmos_global_endpoint_manager.go +++ b/sdk/data/azcosmos/cosmos_global_endpoint_manager.go @@ -112,19 +112,9 @@ func (gem *globalEndpointManager) GetAccountProperties(ctx context.Context) (acc resourceAddress: "", } - path, err := generatePathForNameBased(resourceTypeDatabaseAccount, "", false) - if err != nil { - return accountProperties{}, fmt.Errorf("failed to generate path for name-based request: %v", err) - } - - finalURL := gem.clientEndpoint - if path != "" { - finalURL = azruntime.JoinPaths(gem.clientEndpoint, path) - } - ctxt, cancel := context.WithTimeout(ctx, 60*time.Second) defer cancel() - req, err := azruntime.NewRequest(ctxt, http.MethodGet, finalURL) + req, err := azruntime.NewRequest(ctxt, http.MethodGet, gem.clientEndpoint) if err != nil { return accountProperties{}, err } @@ -140,7 +130,7 @@ func (gem *globalEndpointManager) GetAccountProperties(ctx context.Context) (acc return accountProperties{}, err } - successResponse := (azResponse.StatusCode >= 200 && azResponse.StatusCode < 300) || azResponse.StatusCode == 304 + successResponse := (azResponse.StatusCode >= 200 && azResponse.StatusCode < 300) if successResponse { properties, err := newAccountProperties(azResponse) if err != nil { From 2beda11bbcb0fc8530d5c08f7ec1355c992b9a97 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Thu, 11 Jan 2024 16:08:31 -0500 Subject: [PATCH 06/12] saved apiVersion into constant, changed all explicit uses to constant --- sdk/data/azcosmos/cosmos_client.go | 6 +++++- sdk/data/azcosmos/cosmos_client_test.go | 4 ++-- sdk/data/azcosmos/cosmos_global_endpoint_manager.go | 2 +- sdk/data/azcosmos/shared_key_credential_test.go | 6 +++--- 4 files changed, 11 insertions(+), 7 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index b897e3bd43c3..51416fde58a5 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -20,6 +20,10 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming" ) +const ( + apiVersion = "2020-11-05" +) + // Client is used to interact with the Azure Cosmos DB database service. type Client struct { endpoint string @@ -420,7 +424,7 @@ func (c *Client) createRequest( } req.Raw().Header.Set(headerXmsDate, time.Now().UTC().Format(http.TimeFormat)) - req.Raw().Header.Set(headerXmsVersion, "2020-11-05") + req.Raw().Header.Set(headerXmsVersion, apiVersion) req.Raw().Header.Set(cosmosHeaderSDKSupportedCapabilities, supportedCapabilitiesHeaderValue) req.SetOperationValue(operationContext) diff --git a/sdk/data/azcosmos/cosmos_client_test.go b/sdk/data/azcosmos/cosmos_client_test.go index d93f4ca93d7c..b6ceea6de18e 100644 --- a/sdk/data/azcosmos/cosmos_client_test.go +++ b/sdk/data/azcosmos/cosmos_client_test.go @@ -254,8 +254,8 @@ func TestCreateRequest(t *testing.T) { t.Errorf("Expected %v, but got %v", "", req.Raw().Header.Get(headerXmsDate)) } - if req.Raw().Header.Get(headerXmsVersion) != "2020-11-05" { - t.Errorf("Expected %v, but got %v", "2020-11-05", req.Raw().Header.Get(headerXmsVersion)) + if req.Raw().Header.Get(headerXmsVersion) != apiVersion { + t.Errorf("Expected %v, but got %v", apiVersion, req.Raw().Header.Get(headerXmsVersion)) } if req.Raw().Header.Get(cosmosHeaderSDKSupportedCapabilities) != supportedCapabilitiesHeaderValue { diff --git a/sdk/data/azcosmos/cosmos_global_endpoint_manager.go b/sdk/data/azcosmos/cosmos_global_endpoint_manager.go index 2740bac74b46..5383864a4407 100644 --- a/sdk/data/azcosmos/cosmos_global_endpoint_manager.go +++ b/sdk/data/azcosmos/cosmos_global_endpoint_manager.go @@ -120,7 +120,7 @@ func (gem *globalEndpointManager) GetAccountProperties(ctx context.Context) (acc } req.Raw().Header.Set(headerXmsDate, time.Now().UTC().Format(http.TimeFormat)) - req.Raw().Header.Set(headerXmsVersion, "2020-11-05") + req.Raw().Header.Set(headerXmsVersion, apiVersion) req.Raw().Header.Set(cosmosHeaderSDKSupportedCapabilities, supportedCapabilitiesHeaderValue) req.SetOperationValue(operationContext) diff --git a/sdk/data/azcosmos/shared_key_credential_test.go b/sdk/data/azcosmos/shared_key_credential_test.go index 4ac2d3f83570..7e632f563561 100644 --- a/sdk/data/azcosmos/shared_key_credential_test.go +++ b/sdk/data/azcosmos/shared_key_credential_test.go @@ -69,7 +69,7 @@ func Test_buildCanonicalizedAuthHeaderFromRequest(t *testing.T) { } req.Raw().Header.Set(headerXmsDate, xmsDate) - req.Raw().Header.Set(headerXmsVersion, "2020-11-05") + req.Raw().Header.Set(headerXmsVersion, apiVersion) req.SetOperationValue(operationContext) authHeader, _ := cred.buildCanonicalizedAuthHeaderFromRequest(req) @@ -102,7 +102,7 @@ func Test_buildCanonicalizedAuthHeaderFromRequestWithRid(t *testing.T) { } req.Raw().Header.Set(headerXmsDate, xmsDate) - req.Raw().Header.Set(headerXmsVersion, "2020-11-05") + req.Raw().Header.Set(headerXmsVersion, apiVersion) req.SetOperationValue(operationContext) authHeader, _ := cred.buildCanonicalizedAuthHeaderFromRequest(req) @@ -135,7 +135,7 @@ func Test_buildCanonicalizedAuthHeaderFromRequestWithEscapedCharacters(t *testin } req.Raw().Header.Set(headerXmsDate, xmsDate) - req.Raw().Header.Set(headerXmsVersion, "2020-11-05") + req.Raw().Header.Set(headerXmsVersion, apiVersion) req.SetOperationValue(operationContext) authHeader, _ := cred.buildCanonicalizedAuthHeaderFromRequest(req) From 503a684920b7afc1b87297b221c6def6bf6ab2d6 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Fri, 12 Jan 2024 16:56:24 -0500 Subject: [PATCH 07/12] add gem policy emulator test --- ...tor_cosmos_global_endpoint_manager_test.go | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go index 128b7a973ae1..306baed05d20 100644 --- a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go @@ -73,3 +73,63 @@ func TestGlobalEndpointManagerEmulator(t *testing.T) { assert.Equal(t, len(locationInfo.availReadEndpointsByLocation), len(availableEndpointsByLocation)+1) assert.Equal(t, len(locationInfo.availWriteEndpointsByLocation), len(availableEndpointsByLocation)+1) } + +func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { + emulatorTests := newEmulatorTests(t) + client := emulatorTests.getClient(t) + emulatorRegionName := "South Central US" + emulatorRegion := accountRegion{Name: emulatorRegionName, Endpoint: "https://127.0.0.1:8081/"} + + accountProps, err := client.gem.GetAccountProperties(context.Background()) + assert.NoError(t, err) + + // Verify the expected account properties + expectedAccountProps := accountProperties{ + ReadRegions: []accountRegion{emulatorRegion}, + WriteRegions: []accountRegion{emulatorRegion}, + EnableMultipleWriteLocations: false, + } + assert.Equal(t, expectedAccountProps, accountProps) + + emulatorEndpoint, err := url.Parse("https://localhost:8081/") + assert.NoError(t, err) + + // Verify the read endpoints + readEndpoints, err := client.gem.GetReadEndpoints() + assert.NoError(t, err) + + expectedEndpoints := []url.URL{ + *emulatorEndpoint, + } + assert.Equal(t, expectedEndpoints, readEndpoints) + + // Verify the write endpoints + writeEndpoints, err := client.gem.GetWriteEndpoints() + assert.NoError(t, err) + + assert.Equal(t, expectedEndpoints, writeEndpoints) + + // Assert location cache is not populated until update() is called + locationInfo := client.gem.locationCache.locationInfo + availableLocation := []string{} + availableEndpointsByLocation := map[string]url.URL{} + + assert.Equal(t, locationInfo.availReadLocations, availableLocation) + assert.Equal(t, locationInfo.availWriteLocations, availableLocation) + assert.Equal(t, locationInfo.availReadEndpointsByLocation, availableEndpointsByLocation) + assert.Equal(t, locationInfo.availWriteEndpointsByLocation, availableEndpointsByLocation) + + //assert that information gets populated by the gem policy after running an api request + database_properties := DatabaseProperties{ID: "GEMPolicyDB"} + _, err = client.CreateDatabase(context.Background(), database_properties, nil) + assert.NoError(t, err) + + locationInfo = client.gem.locationCache.locationInfo + + assert.Equal(t, len(locationInfo.availReadLocations), len(availableLocation)+1) + assert.Equal(t, len(locationInfo.availWriteLocations), len(availableLocation)+1) + assert.Equal(t, locationInfo.availWriteLocations[0], emulatorRegionName) + assert.Equal(t, locationInfo.availReadLocations[0], emulatorRegionName) + assert.Equal(t, len(locationInfo.availReadEndpointsByLocation), len(availableEndpointsByLocation)+1) + assert.Equal(t, len(locationInfo.availWriteEndpointsByLocation), len(availableEndpointsByLocation)+1) +} From 4291cba655db528103a0cddfe97698727b0dd98a Mon Sep 17 00:00:00 2001 From: simorenoh Date: Fri, 12 Jan 2024 17:13:02 -0500 Subject: [PATCH 08/12] small changes to the test --- .../emulator_cosmos_global_endpoint_manager_test.go | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go index 306baed05d20..e15f7b290dfb 100644 --- a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go @@ -61,7 +61,7 @@ func TestGlobalEndpointManagerEmulator(t *testing.T) { assert.Equal(t, locationInfo.availReadEndpointsByLocation, availableEndpointsByLocation) assert.Equal(t, locationInfo.availWriteEndpointsByLocation, availableEndpointsByLocation) - //update and assert available locations are now populated in location cache + // Run Update() and assert available locations are now populated in location cache err = gem.Update(context.Background()) assert.NoError(t, err) locationInfo = gem.locationCache.locationInfo @@ -109,7 +109,7 @@ func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { assert.Equal(t, expectedEndpoints, writeEndpoints) - // Assert location cache is not populated until update() is called + // Assert location cache is not populated until update() is called within the policy locationInfo := client.gem.locationCache.locationInfo availableLocation := []string{} availableEndpointsByLocation := map[string]url.URL{} @@ -119,8 +119,9 @@ func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { assert.Equal(t, locationInfo.availReadEndpointsByLocation, availableEndpointsByLocation) assert.Equal(t, locationInfo.availWriteEndpointsByLocation, availableEndpointsByLocation) - //assert that information gets populated by the gem policy after running an api request - database_properties := DatabaseProperties{ID: "GEMPolicyDB"} + // Assert that information gets populated by the gem policy after running an api request + database_id := "GEMPolicyTestDB" + database_properties := DatabaseProperties{ID: database_id} _, err = client.CreateDatabase(context.Background(), database_properties, nil) assert.NoError(t, err) @@ -132,4 +133,8 @@ func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { assert.Equal(t, locationInfo.availReadLocations[0], emulatorRegionName) assert.Equal(t, len(locationInfo.availReadEndpointsByLocation), len(availableEndpointsByLocation)+1) assert.Equal(t, len(locationInfo.availWriteEndpointsByLocation), len(availableEndpointsByLocation)+1) + + db, _ := client.NewDatabase(database_id) + _, err = db.Delete(context.TODO(), nil) + assert.NoError(t, err) } From 7cbc68f23df80712be26ec8aedbc8fe66a75e96d Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 16 Jan 2024 10:43:36 -0500 Subject: [PATCH 09/12] clean up test --- ...tor_cosmos_global_endpoint_manager_test.go | 44 +++---------------- 1 file changed, 5 insertions(+), 39 deletions(-) diff --git a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go index e15f7b290dfb..be2a47ca9a04 100644 --- a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go @@ -78,36 +78,6 @@ func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { emulatorTests := newEmulatorTests(t) client := emulatorTests.getClient(t) emulatorRegionName := "South Central US" - emulatorRegion := accountRegion{Name: emulatorRegionName, Endpoint: "https://127.0.0.1:8081/"} - - accountProps, err := client.gem.GetAccountProperties(context.Background()) - assert.NoError(t, err) - - // Verify the expected account properties - expectedAccountProps := accountProperties{ - ReadRegions: []accountRegion{emulatorRegion}, - WriteRegions: []accountRegion{emulatorRegion}, - EnableMultipleWriteLocations: false, - } - assert.Equal(t, expectedAccountProps, accountProps) - - emulatorEndpoint, err := url.Parse("https://localhost:8081/") - assert.NoError(t, err) - - // Verify the read endpoints - readEndpoints, err := client.gem.GetReadEndpoints() - assert.NoError(t, err) - - expectedEndpoints := []url.URL{ - *emulatorEndpoint, - } - assert.Equal(t, expectedEndpoints, readEndpoints) - - // Verify the write endpoints - writeEndpoints, err := client.gem.GetWriteEndpoints() - assert.NoError(t, err) - - assert.Equal(t, expectedEndpoints, writeEndpoints) // Assert location cache is not populated until update() is called within the policy locationInfo := client.gem.locationCache.locationInfo @@ -119,11 +89,11 @@ func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { assert.Equal(t, locationInfo.availReadEndpointsByLocation, availableEndpointsByLocation) assert.Equal(t, locationInfo.availWriteEndpointsByLocation, availableEndpointsByLocation) - // Assert that information gets populated by the gem policy after running an api request - database_id := "GEMPolicyTestDB" - database_properties := DatabaseProperties{ID: database_id} - _, err = client.CreateDatabase(context.Background(), database_properties, nil) - assert.NoError(t, err) + // Assert that information gets populated by the gem policy after running an http request (read item) + db, _ := client.NewDatabase("database_id") + container, _ := db.NewContainer("container_id") + _, err := container.ReadItem(context.TODO(), NewPartitionKeyString("1"), "doc1", nil) + assert.Error(t, err) locationInfo = client.gem.locationCache.locationInfo @@ -133,8 +103,4 @@ func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { assert.Equal(t, locationInfo.availReadLocations[0], emulatorRegionName) assert.Equal(t, len(locationInfo.availReadEndpointsByLocation), len(availableEndpointsByLocation)+1) assert.Equal(t, len(locationInfo.availWriteEndpointsByLocation), len(availableEndpointsByLocation)+1) - - db, _ := client.NewDatabase(database_id) - _, err = db.Delete(context.TODO(), nil) - assert.NoError(t, err) } From af844ef07a3a070475dae0363733881b8585e415 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 16 Jan 2024 10:45:39 -0500 Subject: [PATCH 10/12] pass preferred regions --- sdk/data/azcosmos/cosmos_client.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index 51416fde58a5..16939a221a15 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -42,7 +42,7 @@ func (c *Client) Endpoint() string { // options - Optional Cosmos client options. Pass nil to accept default values. func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) { //need to pass in preferredRegions from options here once those changes are merged - gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), []string{}, 0) + gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), o.PreferredRegions, 0) if err != nil { return nil, err } @@ -59,7 +59,7 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) ( return nil, err } //need to pass in preferredRegions from options here once those changes are merged - gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), []string{}, 0) + gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), o.PreferredRegions, 0) if err != nil { return nil, err } From 29fc6995c738119541a024048fbcbbaabbf17729 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 16 Jan 2024 10:57:44 -0500 Subject: [PATCH 11/12] removed excess comments --- sdk/data/azcosmos/cosmos_client.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index 16939a221a15..b9847d396a64 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -41,7 +41,6 @@ func (c *Client) Endpoint() string { // cred - The credential used to authenticate with the cosmos service. // options - Optional Cosmos client options. Pass nil to accept default values. func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) { - //need to pass in preferredRegions from options here once those changes are merged gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), o.PreferredRegions, 0) if err != nil { return nil, err @@ -58,7 +57,6 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) ( if err != nil { return nil, err } - //need to pass in preferredRegions from options here once those changes are merged gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), o.PreferredRegions, 0) if err != nil { return nil, err From 98647782d090e7e8d93da8203f1b242f0871da32 Mon Sep 17 00:00:00 2001 From: simorenoh Date: Tue, 16 Jan 2024 12:17:21 -0500 Subject: [PATCH 12/12] add nil check to client options for preferred regions --- sdk/data/azcosmos/cosmos_client.go | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index b9847d396a64..ba6ed793e6b0 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -41,7 +41,11 @@ func (c *Client) Endpoint() string { // cred - The credential used to authenticate with the cosmos service. // options - Optional Cosmos client options. Pass nil to accept default values. func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) { - gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), o.PreferredRegions, 0) + preferredRegions := []string{} + if o != nil { + preferredRegions = o.PreferredRegions + } + gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), preferredRegions, 0) if err != nil { return nil, err } @@ -57,7 +61,11 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) ( if err != nil { return nil, err } - gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), o.PreferredRegions, 0) + preferredRegions := []string{} + if o != nil { + preferredRegions = o.PreferredRegions + } + gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), preferredRegions, 0) if err != nil { return nil, err }