Add some response values as attributes to the azcosmos traces

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
 

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

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
 }

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"},
 	}))
 

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"},
 	}))
 

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"},
 	}))
 

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"},
 	}))
 

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"

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{},
 	}))
 

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{},
 	}))