Add db.cosmosdb.partition_key as a span attr on item ops

sdk/data/azcosmos/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Features Added
 * Set all Telemetry spans to have the Kind of SpanKindClient
 * Set request_charge and status_code on all trace spans
+* Set partition_key on Item based spans. See [PR 23664](https://github.com/Azure/azure-sdk-for-go/pull/23664)
 
 ### Breaking Changes
 

sdk/data/azcosmos/cosmos_container.go

@@ -229,7 +229,7 @@ func (c *ContainerClient) CreateItem(
 	item []byte,
 	o *ItemOptions) (ItemResponse, error) {
 	var err error
-	spanName, err := c.getSpanForItems(operationTypeCreate)
+	spanName, err := c.getSpanForItems(operationTypeCreate, partitionKey)
 	if err != nil {
 		return ItemResponse{}, err
 	}
@@ -282,7 +282,7 @@ func (c *ContainerClient) UpsertItem(
 	item []byte,
 	o *ItemOptions) (ItemResponse, error) {
 	var err error
-	spanName, err := c.getSpanForItems(operationTypeUpsert)
+	spanName, err := c.getSpanForItems(operationTypeUpsert, partitionKey)
 	if err != nil {
 		return ItemResponse{}, err
 	}
@@ -341,7 +341,7 @@ func (c *ContainerClient) ReplaceItem(
 	item []byte,
 	o *ItemOptions) (ItemResponse, error) {
 	var err error
-	spanName, err := c.getSpanForItems(operationTypeReplace)
+	spanName, err := c.getSpanForItems(operationTypeReplace, partitionKey)
 	if err != nil {
 		return ItemResponse{}, err
 	}
@@ -394,7 +394,7 @@ func (c *ContainerClient) ReadItem(
 	itemId string,
 	o *ItemOptions) (ItemResponse, error) {
 	var err error
-	spanName, err := c.getSpanForItems(operationTypeRead)
+	spanName, err := c.getSpanForItems(operationTypeRead, partitionKey)
 	if err != nil {
 		return ItemResponse{}, err
 	}
@@ -443,7 +443,7 @@ func (c *ContainerClient) DeleteItem(
 	itemId string,
 	o *ItemOptions) (ItemResponse, error) {
 	var err error
-	spanName, err := c.getSpanForItems(operationTypeDelete)
+	spanName, err := c.getSpanForItems(operationTypeDelete, partitionKey)
 	if err != nil {
 		return ItemResponse{}, err
 	}
@@ -516,7 +516,7 @@ func (c *ContainerClient) NewQueryItemsPager(query string, partitionKey Partitio
 		Fetcher: func(ctx context.Context, page *QueryItemsResponse) (QueryItemsResponse, error) {
 			var err error
 			// Move the span to the pager once https://github.com/Azure/azure-sdk-for-go/issues/23294 is fixed
-			spanName, err := c.getSpanForItems(operationTypeQuery)
+			spanName, err := c.getSpanForItems(operationTypeQuery, partitionKey)
 			if err != nil {
 				return QueryItemsResponse{}, err
 			}
@@ -561,7 +561,7 @@ func (c *ContainerClient) PatchItem(
 	ops PatchOperations,
 	o *ItemOptions) (ItemResponse, error) {
 	var err error
-	spanName, err := c.getSpanForItems(operationTypePatch)
+	spanName, err := c.getSpanForItems(operationTypePatch, partitionKey)
 	if err != nil {
 		return ItemResponse{}, err
 	}
@@ -684,6 +684,6 @@ func (c *ContainerClient) getSpanForContainer(operationType operationType, resou
 	return getSpanNameForContainers(c.database.client.accountEndpointUrl(), operationType, resourceType, c.database.id, id)
 }
 
-func (c *ContainerClient) getSpanForItems(operationType operationType) (span, error) {
-	return getSpanNameForItems(c.database.client.accountEndpointUrl(), operationType, c.database.id, c.id)
+func (c *ContainerClient) getSpanForItems(operationType operationType, pk PartitionKey) (span, error) {
+	return getSpanNameForItems(c.database.client.accountEndpointUrl(), operationType, c.database.id, c.id, pk)
 }

sdk/data/azcosmos/otel_constants.go

@@ -116,7 +116,7 @@ func getSpanNameForContainers(endpoint *url.URL, operationType operationType, re
 	return span{name: fmt.Sprintf("%s %s", spanName, id), options: getSpanPropertiesForContainer(endpoint, spanName, database, id)}, nil
 }
 
-func getSpanNameForItems(endpoint *url.URL, operationType operationType, database string, id string) (span, error) {
+func getSpanNameForItems(endpoint *url.URL, operationType operationType, database string, id string, pk PartitionKey) (span, error) {
 	var spanName string
 	switch operationType {
 	case operationTypeCreate:
@@ -139,7 +139,13 @@ func getSpanNameForItems(endpoint *url.URL, operationType operationType, databas
 		return span{}, fmt.Errorf("undefined telemetry span for operationType %v and resourceType %v", operationType, resourceTypeDocument)
 	}
 
-	return span{name: fmt.Sprintf("%s %s", spanName, id), options: getSpanPropertiesForContainer(endpoint, spanName, database, id)}, nil
+	options := getSpanPropertiesForContainer(endpoint, spanName, database, id)
+
+	if pkString, err := pk.toJsonString(); err == nil {
+		options.Attributes = append(options.Attributes, tracing.Attribute{Key: "db.cosmosdb.partition_key", Value: pkString})
+	}
+
+	return span{name: fmt.Sprintf("%s %s", spanName, id), options: options}, nil
 }
 
 func getSpanPropertiesForClient(endpoint *url.URL, operationName string) runtime.StartSpanOptions {

sdk/data/azcosmos/otel_constants_test.go

@@ -172,3 +172,72 @@ func TestSpanForContainers(t *testing.T) {
 		t.Fatalf("Expected error, but got none")
 	}
 }
+
+func TestSpanForItems(t *testing.T) {
+	endpoint, _ := url.Parse("https://localhost:8081/")
+
+	pk := NewPartitionKeyString("aPartitionKey")
+
+	aSpan, err := getSpanNameForItems(endpoint, operationTypeCreate, "db", "test", pk)
+	if err != nil {
+		t.Fatalf("Failed to get span name: %v", err)
+	}
+	if aSpan.name != "create_item test" {
+		t.Fatalf("Expected span name to be 'create_item test', but got %s", aSpan.name)
+	}
+
+	if aSpan.options.Kind != tracing.SpanKindClient {
+		t.Fatalf("Expected span kind to be 'SpanKindClient (%v)', got %v", tracing.SpanKindClient, aSpan.options.Kind)
+	}
+
+	if len(aSpan.options.Attributes) == 0 {
+		t.Fatalf("Expected span options to have attributes, but got none")
+	}
+
+	idx := slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "db.system" && a.Value == "cosmosdb" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'db.system' with value 'cosmosdb', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "db.cosmosdb.connection_mode" && a.Value == "gateway" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'db.cosmosdb.connection_mode' with value 'gateway', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "server.address" && a.Value == "localhost" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'server.address' with value 'localhost', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "server.port" && a.Value == "8081" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'server.port' with value '8081', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "db.operation.name" && a.Value == "create_item" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'db.operation.name' with value 'create_item', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "db.namespace" && a.Value == "db" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'db.namespace' with value 'db', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool { return a.Key == "db.collection.name" && a.Value == "test" })
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'db.collection.name' with value 'test', but got none")
+	}
+
+	idx = slices.IndexFunc(aSpan.options.Attributes, func(a tracing.Attribute) bool {
+		return a.Key == "db.cosmosdb.partition_key" && a.Value == "[\"aPartitionKey\"]"
+	})
+	if idx == -1 {
+		t.Fatalf("Expected attribute 'db.cosmosdb.partition_key' with value '[\"aPartitionKey\"]', but got none")
+	}
+
+	aSpan, err = getSpanNameForItems(endpoint, -1, "db", "test", pk)
+	if err == nil {
+		t.Fatalf("Expected error, but got none")
+	}
+}