diff --git a/sdk/data/azcosmos/CHANGELOG.md b/sdk/data/azcosmos/CHANGELOG.md index 505edfb1cf87..241f8171cef4 100644 --- a/sdk/data/azcosmos/CHANGELOG.md +++ b/sdk/data/azcosmos/CHANGELOG.md @@ -4,6 +4,7 @@ ### Features Added * Set all Telemetry spans to have the Kind of SpanKindClient +* Set request_charge and status_code on all trace spans ### Breaking Changes diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go index cef7d1c8d7fc..8e605b3315ad 100644 --- a/sdk/data/azcosmos/cosmos_client.go +++ b/sdk/data/azcosmos/cosmos_client.go @@ -18,6 +18,7 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime" "github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/tracing" "github.com/Azure/azure-sdk-for-go/sdk/internal/log" ) @@ -327,7 +328,7 @@ func (c *Client) sendPostRequest( return nil, err } - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) sendQueryRequest( @@ -356,7 +357,7 @@ func (c *Client) sendQueryRequest( // Override content type for query req.Raw().Header.Set(headerContentType, cosmosHeaderValuesQuery) - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) sendPutRequest( @@ -376,7 +377,7 @@ func (c *Client) sendPutRequest( return nil, err } - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) sendGetRequest( @@ -390,7 +391,7 @@ func (c *Client) sendGetRequest( return nil, err } - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) sendDeleteRequest( @@ -404,7 +405,7 @@ func (c *Client) sendDeleteRequest( return nil, err } - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) sendBatchRequest( @@ -424,7 +425,7 @@ func (c *Client) sendBatchRequest( return nil, err } - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) sendPatchRequest( @@ -444,7 +445,7 @@ func (c *Client) sendPatchRequest( return nil, err } - return c.executeAndEnsureSuccessResponse(req) + return c.executeAndEnsureSuccessResponse(ctx, req) } func (c *Client) createRequest( @@ -507,13 +508,15 @@ func (c *Client) attachContent(content interface{}, req *policy.Request) error { return nil } -func (c *Client) executeAndEnsureSuccessResponse(request *policy.Request) (*http.Response, error) { +func (c *Client) executeAndEnsureSuccessResponse(ctx context.Context, request *policy.Request) (*http.Response, error) { log.Write(azlog.EventResponse, fmt.Sprintf("\n===== Client preferred regions:\n%v\n=====\n", c.gem.preferredLocations)) response, err := c.internal.Pipeline().Do(request) if err != nil { return nil, err } + c.addResponseValuesToSpan(ctx, response) + successResponse := (response.StatusCode >= 200 && response.StatusCode < 300) || response.StatusCode == 304 if successResponse { return response, nil @@ -526,6 +529,14 @@ func (c *Client) accountEndpointUrl() *url.URL { return c.endpointUrl } +func (c *Client) addResponseValuesToSpan(ctx context.Context, resp *http.Response) { + span := c.internal.Tracer().SpanFromContext(ctx) + span.SetAttributes( + tracing.Attribute{Key: "db.cosmosdb.request_charge", Value: newResponse(resp).RequestCharge}, + tracing.Attribute{Key: "db.cosmosdb.status_code", Value: resp.StatusCode}, + ) +} + type pipelineRequestOptions struct { headerOptionsOverride *headerOptionsOverride resourceType resourceType diff --git a/sdk/data/azcosmos/cosmos_client_test.go b/sdk/data/azcosmos/cosmos_client_test.go index f9840b11f8a6..ee0196c9182c 100644 --- a/sdk/data/azcosmos/cosmos_client_test.go +++ b/sdk/data/azcosmos/cosmos_client_test.go @@ -681,6 +681,54 @@ func TestQueryDatabases(t *testing.T) { } } +func TestSpanResponseAttributes(t *testing.T) { + srv, close := mock.NewTLSServer() + defer close() + srv.SetResponse( + mock.WithStatusCode(200), + mock.WithHeader(cosmosHeaderRequestCharge, "13.42"), + ) + + matcher := &spanMatcher{ + ExpectedSpans: []string{"test_span"}, + } + tp := newSpanValidator(t, matcher) + internalClient, _ := azcore.NewClient( + "azcosmostest", "v1.0.0", + azruntime.PipelineOptions{Tracing: azruntime.TracingOptions{Namespace: "Microsoft.DocumentDB"}}, + &policy.ClientOptions{Transport: srv, TracingProvider: tp}, + ) + gem := &globalEndpointManager{preferredLocations: []string{}} + client := &Client{endpoint: srv.URL(), internal: internalClient, gem: gem} + operationContext := pipelineRequestOptions{ + resourceType: resourceTypeDatabase, + resourceAddress: "", + } + + ctx := context.Background() + ctx, endSpan := azruntime.StartSpan(ctx, "test_span", client.internal.Tracer(), &azruntime.StartSpanOptions{}) + _, err := client.sendGetRequest("/", ctx, operationContext, &DeleteDatabaseOptions{}, nil) + endSpan(err) + if err != nil { + t.Fatal(err) + } + + if len(matcher.MatchedSpans) != 1 { + t.Errorf("Unexpected number of spans") + } + + span := matcher.MatchedSpans[0] + status_value := attributeValueForKey(span.attributes, "db.cosmosdb.status_code") + if status_value != 200 { + t.Fatalf("Expected db.cosmosdb.status_code attribute with 200 value, got %v", status_value) + } + + charge_value := attributeValueForKey(span.attributes, "db.cosmosdb.request_charge") + if charge_value != float32(13.42) { + t.Fatalf("Expected db.cosmosdb.request_charge attribute with 13.42 value, got %v", charge_value) + } +} + type pipelineVerifier struct { requests []pipelineVerifierRequest } diff --git a/sdk/data/azcosmos/emulator_cosmos_aad_test.go b/sdk/data/azcosmos/emulator_cosmos_aad_test.go index 8796b67c12f6..674b22608a54 100644 --- a/sdk/data/azcosmos/emulator_cosmos_aad_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_aad_test.go @@ -11,7 +11,7 @@ import ( func TestAAD(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) @@ -29,7 +29,7 @@ func TestAAD(t *testing.T) { t.Fatalf("Failed to create container: %v", err) } - aadClient := emulatorTests.getAadClient(t, newSpanValidator(t, spanMatcher{ + aadClient := emulatorTests.getAadClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_item aContainer", "read_item aContainer", "replace_item aContainer", "upsert_item aContainer", "delete_item aContainer"}, })) diff --git a/sdk/data/azcosmos/emulator_cosmos_batch_test.go b/sdk/data/azcosmos/emulator_cosmos_batch_test.go index c3d86f4a2637..25fb349b5d34 100644 --- a/sdk/data/azcosmos/emulator_cosmos_batch_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_batch_test.go @@ -11,7 +11,7 @@ import ( func TestItemTransactionalBatch(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"execute_batch aContainer"}, })) @@ -171,7 +171,7 @@ func TestItemTransactionalBatch(t *testing.T) { func TestItemTransactionalBatchError(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"execute_batch aContainer"}, })) diff --git a/sdk/data/azcosmos/emulator_cosmos_container_test.go b/sdk/data/azcosmos/emulator_cosmos_container_test.go index cf33a09b9fdb..aa73e05adbb5 100644 --- a/sdk/data/azcosmos/emulator_cosmos_container_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_container_test.go @@ -10,7 +10,7 @@ import ( func TestContainerCRUD(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_container aContainer", "read_container aContainer", "replace_container aContainer", "read_container_throughput aContainer", "replace_container_throughput aContainer", "delete_container aContainer"}, })) @@ -131,7 +131,7 @@ func TestContainerCRUD(t *testing.T) { func TestContainerAutoscaleCRUD(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_container aContainer", "read_container aContainer", "read_container_throughput aContainer", "replace_container_throughput aContainer", "delete_container aContainer"}, })) diff --git a/sdk/data/azcosmos/emulator_cosmos_database_test.go b/sdk/data/azcosmos/emulator_cosmos_database_test.go index 4e90d8482cf6..c84ff7e44b06 100644 --- a/sdk/data/azcosmos/emulator_cosmos_database_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_database_test.go @@ -10,7 +10,7 @@ import ( func TestDatabaseCRUD(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_database baseDbTest", "read_database baseDbTest", "delete_database baseDbTest", "read_database_throughput baseDbTest"}, })) @@ -70,7 +70,7 @@ func TestDatabaseCRUD(t *testing.T) { func TestDatabaseWithOfferCRUD(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_database baseDbTest", "read_database baseDbTest", "delete_database baseDbTest", "read_database_throughput baseDbTest", "replace_database_throughput baseDbTest"}, })) 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 a698038e85da..225eeec9eb85 100644 --- a/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_global_endpoint_manager_test.go @@ -14,7 +14,7 @@ import ( func TestGlobalEndpointManagerEmulator(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) emulatorRegionName := "South Central US" @@ -79,7 +79,7 @@ func TestGlobalEndpointManagerEmulator(t *testing.T) { func TestGlobalEndpointManagerPolicyEmulator(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) emulatorRegionName := "South Central US" diff --git a/sdk/data/azcosmos/emulator_cosmos_item_test.go b/sdk/data/azcosmos/emulator_cosmos_item_test.go index 10fd0f7b7e4d..9b0c59949d0d 100644 --- a/sdk/data/azcosmos/emulator_cosmos_item_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_item_test.go @@ -17,7 +17,7 @@ import ( func TestItemCRUD(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_item aContainer", "read_item aContainer", "replace_item aContainer", "upsert_item aContainer", "delete_item aContainer", "patch_item aContainer"}, })) @@ -193,7 +193,7 @@ func TestItemCRUD(t *testing.T) { func TestItemCRUDforNullPartitionKey(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{"create_item aContainer", "read_item aContainer", "replace_item aContainer", "upsert_item aContainer", "delete_item aContainer", "patch_item aContainer"}, })) @@ -370,7 +370,7 @@ func TestItemCRUDforNullPartitionKey(t *testing.T) { func TestItemConcurrent(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) @@ -423,7 +423,7 @@ func TestItemConcurrent(t *testing.T) { func TestItemIdEncodingRoutingGW(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) @@ -464,7 +464,7 @@ func TestItemIdEncodingRoutingGW(t *testing.T) { func TestItemIdEncodingComputeGW(t *testing.T) { emulatorTests := newEmulatorTestsWithComputeGateway(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) diff --git a/sdk/data/azcosmos/emulator_cosmos_query_test.go b/sdk/data/azcosmos/emulator_cosmos_query_test.go index d6e8781dadda..992b4ec3e670 100644 --- a/sdk/data/azcosmos/emulator_cosmos_query_test.go +++ b/sdk/data/azcosmos/emulator_cosmos_query_test.go @@ -12,7 +12,7 @@ import ( func TestSinglePartitionQueryWithIndexMetrics(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) @@ -87,7 +87,7 @@ func TestSinglePartitionQueryWithIndexMetrics(t *testing.T) { func TestSinglePartitionQuery(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) @@ -175,7 +175,7 @@ func TestSinglePartitionQuery(t *testing.T) { func TestSinglePartitionQueryWithParameters(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) @@ -227,7 +227,7 @@ func TestSinglePartitionQueryWithParameters(t *testing.T) { func TestSinglePartitionQueryWithProjection(t *testing.T) { emulatorTests := newEmulatorTests(t) - client := emulatorTests.getClient(t, newSpanValidator(t, spanMatcher{ + client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{ ExpectedSpans: []string{}, })) diff --git a/sdk/data/azcosmos/tracing_test.go b/sdk/data/azcosmos/tracing_test.go index 69d8b8c3bb18..69b36ff72a26 100644 --- a/sdk/data/azcosmos/tracing_test.go +++ b/sdk/data/azcosmos/tracing_test.go @@ -13,8 +13,10 @@ import ( "github.com/stretchr/testify/require" ) +type spanContextKey struct{} + // newSpanValidator creates a tracing.Provider that verifies a span was created that matches the specified SpanMatcher. -func newSpanValidator(t *testing.T, matcher spanMatcher) tracing.Provider { +func newSpanValidator(t *testing.T, matcher *spanMatcher) tracing.Provider { return tracing.NewProvider(func(name, version string) tracing.Tracer { tt := matchingTracer{ matcher: matcher, @@ -23,7 +25,7 @@ func newSpanValidator(t *testing.T, matcher spanMatcher) tracing.Provider { t.Cleanup(func() { for _, expectedSpan := range matcher.ExpectedSpans { found := false - for _, match := range tt.matches { + for _, match := range matcher.MatchedSpans { if match.name == expectedSpan { found = true require.True(t, match.ended, "span %s wasn't ended", match.name) @@ -40,18 +42,25 @@ func newSpanValidator(t *testing.T, matcher spanMatcher) tracing.Provider { kind = options.Kind } return tt.Start(ctx, spanName, kind) - }, nil) + }, &tracing.TracerOptions{ + SpanFromContext: func(ctx context.Context) tracing.Span { + if span, ok := ctx.Value(spanContextKey{}).(tracing.Span); ok { + return span + } + return tracing.Span{} + }, + }) }, nil) } // SpanMatcher contains the values to match when a span is created. type spanMatcher struct { ExpectedSpans []string + MatchedSpans []*matchingSpan } type matchingTracer struct { - matcher spanMatcher - matches []*matchingSpan + matcher *spanMatcher } func (mt *matchingTracer) Start(ctx context.Context, spanName string, kind tracing.SpanKind) (context.Context, tracing.Span) { @@ -63,18 +72,22 @@ func (mt *matchingTracer) Start(ctx context.Context, spanName string, kind traci newSpan := &matchingSpan{ name: spanName, } - mt.matches = append(mt.matches, newSpan) - return ctx, tracing.NewSpan(tracing.SpanImpl{ - End: newSpan.End, - SetStatus: newSpan.SetStatus, + mt.matcher.MatchedSpans = append(mt.matcher.MatchedSpans, newSpan) + tracingSpan := tracing.NewSpan(tracing.SpanImpl{ + End: newSpan.End, + SetStatus: newSpan.SetStatus, + SetAttributes: newSpan.SetAttributes, }) + ctx = context.WithValue(ctx, spanContextKey{}, tracingSpan) + return ctx, tracingSpan } type matchingSpan struct { - name string - status tracing.SpanStatus - desc string - ended bool + name string + status tracing.SpanStatus + desc string + attributes []tracing.Attribute + ended bool } func (s *matchingSpan) End() { @@ -86,3 +99,19 @@ func (s *matchingSpan) SetStatus(code tracing.SpanStatus, desc string) { s.desc = desc s.ended = true } + +func (s *matchingSpan) SetAttributes(attrs ...tracing.Attribute) { + s.attributes = append(s.attributes, attrs...) +} + +func attributeValueForKey(attributes []tracing.Attribute, key string) any { + i := slices.IndexFunc(attributes, func(attr tracing.Attribute) bool { + return attr.Key == key + }) + + if i < 0 { + return nil + } + + return attributes[i].Value +}