AWS SDK instrumentation - DynamoDB attributes

instrumentation/aws-sdk/aws-sdk-2.2/library/aws-sdk-2.2-library.gradle

@@ -4,6 +4,10 @@ dependencies {
   implementation deps.opentelemetryExtAws
 
   library group: 'software.amazon.awssdk', name: 'aws-core', version: '2.2.0'
+  library group: 'software.amazon.awssdk', name: 'aws-json-protocol', version: '2.2.0'
 
   testImplementation project(':instrumentation:aws-sdk:aws-sdk-2.2:testing')
+
+  testImplementation deps.assertj
+  testImplementation group: 'org.mockito', name: 'mockito-core', version: '3.6.0'
 }

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkRequest.java

@@ -0,0 +1,155 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.DynamoDB;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.Kinesis;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.S3;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.SQS;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.FieldMapping.request;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.FieldMapping.response;
+
+import java.util.List;
+import java.util.Map;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import software.amazon.awssdk.core.SdkRequest;
+
+/**
+ * Temporary solution - maps only DynamoDB attributes. Final solution should be generated from AWS
+ * SDK automatically
+ * (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/2291).
+ */
+enum AwsSdkRequest {
+  // generic requests
+  DynamoDbRequest(DynamoDB, "DynamoDbRequest"),
+  S3Request(S3, "S3Request"),
+  SqsRequest(SQS, "SqsRequest"),
+  KinesisRequest(Kinesis, "KinesisRequest"),
+  // specific requests
+  BatchGetItem(
+      DynamoDB,
+      "BatchGetItemRequest",
+      request("aws.dynamodb.table_names", "RequestItems"),
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity")),
+  BatchWriteItem(
+      DynamoDB,
+      "BatchWriteItemRequest",
+      request("aws.dynamodb.table_names", "RequestItems"),
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity"),
+      response("aws.dynamodb.item_collection_metrics", "ItemCollectionMetrics")),
+  CreateTable(
+      DynamoDB,
+      "CreateTableRequest",
+      request("aws.dynamodb.global_secondary_indexes", "GlobalSecondaryIndexes"),
+      request("aws.dynamodb.local_secondary_indexes", "LocalSecondaryIndexes"),
+      request(
+          "aws.dynamodb.provisioned_throughput.read_capacity_units",
+          "ProvisionedThroughput.ReadCapacityUnits"),
+      request(
+          "aws.dynamodb.provisioned_throughput.write_capacity_units",
+          "ProvisionedThroughput.WriteCapacityUnits")),
+  DeleteItem(
+      DynamoDB,
+      "DeleteItemRequest",
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity"),
+      response("aws.dynamodb.item_collection_metrics", "ItemCollectionMetrics")),
+  GetItem(
+      DynamoDB,
+      "GetItemRequest",
+      request("aws.dynamodb.projection_expression", "ProjectionExpression"),
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity"),
+      request("aws.dynamodb.consistent_read", "ConsistentRead")),
+  ListTables(
+      DynamoDB,
+      "ListTablesRequest",
+      request("aws.dynamodb.exclusive_start_table_name", "ExclusiveStartTableName"),
+      response("aws.dynamodb.table_count", "TableNames"),
+      request("aws.dynamodb.limit", "Limit")),
+  PutItem(
+      DynamoDB,
+      "PutItemRequest",
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity"),
+      response("aws.dynamodb.item_collection_metrics", "ItemCollectionMetrics")),
+  Query(
+      DynamoDB,
+      "QueryRequest",
+      request("aws.dynamodb.attributes_to_get", "AttributesToGet"),
+      request("aws.dynamodb.consistent_read", "ConsistentRead"),
+      request("aws.dynamodb.index_name", "IndexName"),
+      request("aws.dynamodb.limit", "Limit"),
+      request("aws.dynamodb.projection_expression", "ProjectionExpression"),
+      request("aws.dynamodb.scan_index_forward", "ScanIndexForward"),
+      request("aws.dynamodb.select", "Select"),
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity")),
+  Scan(
+      DynamoDB,
+      "ScanRequest",
+      request("aws.dynamodb.attributes_to_get", "AttributesToGet"),
+      request("aws.dynamodb.consistent_read", "ConsistentRead"),
+      request("aws.dynamodb.index_name", "IndexName"),
+      request("aws.dynamodb.limit", "Limit"),
+      request("aws.dynamodb.projection_expression", "ProjectionExpression"),
+      request("aws.dynamodb.segment", "Segment"),
+      request("aws.dynamodb.select", "Select"),
+      request("aws.dynamodb.total_segments", "TotalSegments"),
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity"),
+      response("aws.dynamodb.count", "Count"),
+      response("aws.dynamodb.scanned_count", "ScannedCount")),
+  UpdateItem(
+      DynamoDB,
+      "UpdateItemRequest",
+      response("aws.dynamodb.consumed_capacity", "ConsumedCapacity"),
+      response("aws.dynamodb.item_collection_metrics", "ItemCollectionMetrics")),
+  UpdateTable(
+      DynamoDB,
+      "UpdateTableRequest",
+      request("aws.dynamodb.attribute_definitions", "AttributeDefinitions"),
+      request("aws.dynamodb.global_secondary_index_updates", "GlobalSecondaryIndexUpdates"),
+      request(
+          "aws.dynamodb.provisioned_throughput.read_capacity_units",
+          "ProvisionedThroughput.ReadCapacityUnits"),
+      request(
+          "aws.dynamodb.provisioned_throughput.write_capacity_units",
+          "ProvisionedThroughput.WriteCapacityUnits"));
+
+  private final AwsSdkRequestType type;
+  private final String requestClass;
+  private final Map<FieldMapping.Type, List<FieldMapping>> fields;
+
+  AwsSdkRequest(AwsSdkRequestType type, String requestClass, FieldMapping... fields) {
+    this.type = type;
+    this.requestClass = requestClass;
+    this.fields = FieldMapping.groupByType(fields);
+  }
+
+  @Nullable
+  static AwsSdkRequest ofSdkRequest(SdkRequest request) {
+    // try request type
+    AwsSdkRequest result = ofType(request.getClass().getSimpleName());
+    // try parent - generic
+    if (result == null) {
+      result = ofType(request.getClass().getSuperclass().getSimpleName());
+    }
+    return result;
+  }
+
+  private static AwsSdkRequest ofType(String typeName) {
+    for (AwsSdkRequest type : values()) {
+      if (type.requestClass.equals(typeName)) {
+        return type;
+      }
+    }
+    return null;
+  }
+
+  List<FieldMapping> fields(FieldMapping.Type type) {
+    return fields.get(type);
+  }
+
+  AwsSdkRequestType type() {
+    return type;
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkRequestType.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import static io.opentelemetry.instrumentation.awssdk.v2_2.FieldMapping.request;
+
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
+import java.util.List;
+import java.util.Map;
+
+enum AwsSdkRequestType {
+  S3(request("aws.bucket.name", "Bucket")),
+  SQS(request("aws.queue.url", "QueueUrl"), request("aws.queue.name", "QueueName")),
+  Kinesis(request("aws.stream.name", "StreamName")),
+  DynamoDB(
+      request("aws.table.name", "TableName"),
+      request(SemanticAttributes.DB_NAME.getKey(), "TableName"));
+
+  private final Map<FieldMapping.Type, List<FieldMapping>> fields;
+
+  AwsSdkRequestType(FieldMapping... fieldMappings) {
+    this.fields = FieldMapping.groupByType(fieldMappings);
+  }
+
+  List<FieldMapping> fields(FieldMapping.Type type) {
+    return fields.get(type);
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapper.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import io.opentelemetry.api.trace.Span;
+import java.util.List;
+import java.util.function.Function;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.core.SdkResponse;
+import software.amazon.awssdk.utils.StringUtils;
+
+class FieldMapper {
+
+  private final Serializer serializer;
+  private final MethodHandleFactory methodHandleFactory;
+
+  FieldMapper() {
+    serializer = new Serializer();
+    methodHandleFactory = new MethodHandleFactory();
+  }
+
+  FieldMapper(Serializer serializer, MethodHandleFactory methodHandleFactory) {
+    this.methodHandleFactory = methodHandleFactory;
+    this.serializer = serializer;
+  }
+
+  void mapToAttributes(SdkRequest sdkRequest, AwsSdkRequest request, Span span) {
+    mapToAttributes(
+        field -> sdkRequest.getValueForField(field, Object.class).orElse(null),
+        FieldMapping.Type.REQUEST,
+        request,
+        span);
+  }
+
+  void mapToAttributes(SdkResponse sdkResponse, AwsSdkRequest request, Span span) {
+    mapToAttributes(
+        field -> sdkResponse.getValueForField(field, Object.class).orElse(null),
+        FieldMapping.Type.RESPONSE,
+        request,
+        span);
+  }
+
+  private void mapToAttributes(
+      Function<String, Object> fieldValueProvider,
+      FieldMapping.Type type,
+      AwsSdkRequest request,
+      Span span) {
+    for (FieldMapping fieldMapping : request.fields(type)) {
+      mapToAttributes(fieldValueProvider, fieldMapping, span);
+    }
+    for (FieldMapping fieldMapping : request.type().fields(type)) {
+      mapToAttributes(fieldValueProvider, fieldMapping, span);
+    }
+  }
+
+  private void mapToAttributes(
+      Function<String, Object> fieldValueProvider, FieldMapping fieldMapping, Span span) {
+    // traverse path
+    List<String> path = fieldMapping.getFields();
+    Object target = fieldValueProvider.apply(path.get(0));
+    for (int i = 1; i < path.size() && target != null; i++) {
+      target = next(target, path.get(i));
+    }
+    if (target != null) {
+      String value = serializer.serialize(target);
+      if (!StringUtils.isEmpty(value)) {
+        span.setAttribute(fieldMapping.getAttribute(), value);
+      }
+    }
+  }
+
+  @Nullable
+  private Object next(Object current, String fieldName) {
+    try {
+      return methodHandleFactory.forField(current.getClass(), fieldName).invoke(current);
+    } catch (Throwable t) {
+      // ignore
+    }
+    return null;
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapping.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.EnumMap;
+import java.util.List;
+import java.util.Map;
+
+class FieldMapping {
+
+  enum Type {
+    REQUEST,
+    RESPONSE;
+  }
+
+  private final Type type;
+  private final String attribute;
+  private final List<String> fields;
+
+  static FieldMapping request(String attribute, String fieldPath) {
+    return new FieldMapping(Type.REQUEST, attribute, fieldPath);
+  }
+
+  static FieldMapping response(String attribute, String fieldPath) {
+    return new FieldMapping(Type.RESPONSE, attribute, fieldPath);
+  }
+
+  FieldMapping(Type type, String attribute, String fieldPath) {
+    this.type = type;
+    this.attribute = attribute;
+    this.fields = Collections.unmodifiableList(Arrays.asList(fieldPath.split("\\.")));
+  }
+
+  String getAttribute() {
+    return attribute;
+  }
+
+  List<String> getFields() {
+    return fields;
+  }
+
+  Type getType() {
+    return type;
+  }
+
+  static final Map<Type, List<FieldMapping>> groupByType(FieldMapping[] fieldMappings) {
+
+    EnumMap<Type, List<FieldMapping>> fields = new EnumMap<>(Type.class);
+    for (FieldMapping.Type type : FieldMapping.Type.values()) {
+      fields.put(type, new ArrayList<>());
+    }
+    for (FieldMapping fieldMapping : fieldMappings) {
+      fields.get(fieldMapping.getType()).add(fieldMapping);
+    }
+    return fields;
+  }
+}