Align DynamoDB batch telemetry with database batch semantics

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/internal/DynamoDbAttributesExtractor.java

@@ -49,9 +49,13 @@ public void onStart(AttributesBuilder attributes, Context parentContext, Request
 
     String operation = getOperationName(request.getOriginalRequest());
     Long batchSize = extractBatchSize(operation, request.getOriginalRequest());
+    WriteOperationType writeOpType =
+        "BatchWriteItem".equals(operation)
+            ? extractWriteOperationType(request.getOriginalRequest())
+            : WriteOperationType.NONE;
     if (emitStableDatabaseSemconv()) {
-      attributes.put(DB_OPERATION_NAME, getStableOperationName(operation, batchSize));
-      if (isBatch(batchSize)) {
+      attributes.put(DB_OPERATION_NAME, getStableOperationName(operation, batchSize, writeOpType));
+      if (shouldEmitBatchSize(batchSize)) {
         attributes.put(DB_OPERATION_BATCH_SIZE, batchSize);
       }
     }
@@ -96,30 +100,32 @@ private static String getSingleCollectionName(Map<?, ?> requestItems) {
 
   @Nullable
   private static String getStableOperationName(
-      @Nullable String operation, @Nullable Long batchSize) {
-    if ("BatchGetItem".equals(operation)) {
-      return getStableBatchOperationName(batchSize, "GetItem", operation);
-    }
+      @Nullable String operation, @Nullable Long batchSize, WriteOperationType writeOpType) {
     if ("BatchWriteItem".equals(operation)) {
-      return getStableBatchOperationName(batchSize, "WriteItem", operation);
+      return getStableWriteOperationName(batchSize, writeOpType);
     }
     return operation;
   }
 
-  private static String getStableBatchOperationName(
-      @Nullable Long batchSize, String itemOperation, String batchOperation) {
-    if (batchSize == null || batchSize == 0) {
-      return batchOperation;
+  private static String getStableWriteOperationName(
+      @Nullable Long batchSize, WriteOperationType writeOpType) {
+    if (batchSize == null || batchSize == 0 || writeOpType == WriteOperationType.NONE) {
+      return "BatchWriteItem";
     }
+    String itemOp = writeOpType == WriteOperationType.PUT ? "PutItem" : "DeleteItem";
     if (batchSize == 1) {
-      return itemOperation;
+      return itemOp;
+    }
+    // mixed operations collapse to bare BATCH (consistent with SQL/Cassandra)
+    if (writeOpType == WriteOperationType.MIXED) {
+      return "BATCH";
     }
-    return "BATCH " + itemOperation;
+    return "BATCH " + itemOp;
   }
 
   @Nullable
   private static Long extractBatchSize(@Nullable String operation, Object request) {
-    if (!"BatchGetItem".equals(operation) && !"BatchWriteItem".equals(operation)) {
+    if (!"BatchWriteItem".equals(operation)) {
       return null;
     }
 
@@ -128,36 +134,66 @@ private static Long extractBatchSize(@Nullable String operation, Object request)
       return null;
     }
 
-    long batchSize =
-        "BatchGetItem".equals(operation)
-            ? countBatchGetItems(requestItems)
-            : countBatchWriteItems(requestItems);
-    return batchSize == 0 ? null : batchSize;
+    long batchSize = countBatchWriteItems(requestItems);
+    // return the size for every batch request, including an empty batch with size 0
+    return batchSize;
   }
 
-  private static long countBatchGetItems(Map<?, ?> requestItems) {
+  private static long countBatchWriteItems(Map<?, ?> requestItems) {
     long count = 0;
-    for (Object keysAndAttributes : requestItems.values()) {
-      List<?> keys = RequestAccess.getKeys(keysAndAttributes);
-      if (keys != null) {
-        count += keys.size();
+    for (Object writeRequests : requestItems.values()) {
+      if (writeRequests instanceof Collection) {
+        count += ((Collection<?>) writeRequests).size();
       }
     }
     return count;
   }
 
-  private static long countBatchWriteItems(Map<?, ?> requestItems) {
-    long count = 0;
+  /**
+   * Extracts the write operation type from a BatchWriteItem request. Returns PUT if all requests
+   * are PutRequests, DELETE if all are DeleteRequests, MIXED if both types are present, or NONE if
+   * the request is empty or cannot be inspected.
+   */
+  private static WriteOperationType extractWriteOperationType(Object request) {
+    Map<?, ?> requestItems = RequestAccess.getRequestItems(request);
+    if (requestItems == null) {
+      return WriteOperationType.NONE;
+    }
+
+    WriteOperationType result = WriteOperationType.NONE;
     for (Object writeRequests : requestItems.values()) {
       if (writeRequests instanceof Collection) {
-        count += ((Collection<?>) writeRequests).size();
+        for (Object writeRequest : (Collection<?>) writeRequests) {
+          WriteOperationType opType = classifyWriteRequest(writeRequest);
+          if (opType == WriteOperationType.NONE) {
+            continue;
+          }
+          if (result == WriteOperationType.NONE) {
+            result = opType;
+          } else if (result != opType) {
+            return WriteOperationType.MIXED;
+          }
+        }
       }
     }
-    return count;
+    return result;
+  }
+
+  private static WriteOperationType classifyWriteRequest(Object writeRequest) {
+    // WriteRequest has getPutRequest() and getDeleteRequest() methods; exactly one returns non-null
+    if (RequestAccess.hasPutRequest(writeRequest)) {
+      return WriteOperationType.PUT;
+    }
+    if (RequestAccess.hasDeleteRequest(writeRequest)) {
+      return WriteOperationType.DELETE;
+    }
+    return WriteOperationType.NONE;
   }
 
-  private static boolean isBatch(@Nullable Long batchSize) {
-    return batchSize != null && batchSize > 1;
+  // db.operation.batch.size is captured for every batch request (including an empty batch with
+  // size 0); it is only omitted for a single-item batch, which is reported as a non-batch operation
+  private static boolean shouldEmitBatchSize(@Nullable Long batchSize) {
+    return batchSize != null && batchSize != 1;
   }
 
   @Nullable
@@ -177,4 +213,11 @@ public void onEnd(
       Request<?> request,
       @Nullable Response<?> response,
       @Nullable Throwable error) {}
+
+  private enum WriteOperationType {
+    NONE,
+    PUT,
+    DELETE,
+    MIXED
+  }
 }

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/internal/RequestAccess.java

@@ -8,6 +8,7 @@
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -111,6 +112,24 @@ static List<?> getKeys(Object request) {
     return invokeOrNull(access.getKeys, request, List.class);
   }
 
+  /**
+   * Returns true if the given WriteRequest contains a PutRequest, false otherwise. Uses reflection
+   * to call getPutRequest() on the WriteRequest object.
+   */
+  static boolean hasPutRequest(Object writeRequest) {
+    WriteRequestAccess access = WriteRequestAccess.ACCESSORS.get(writeRequest.getClass());
+    return access.invokeGetPutRequest(writeRequest) != null;
+  }
+
+  /**
+   * Returns true if the given WriteRequest contains a DeleteRequest, false otherwise. Uses
+   * reflection to call getDeleteRequest() on the WriteRequest object.
+   */
+  static boolean hasDeleteRequest(Object writeRequest) {
+    WriteRequestAccess access = WriteRequestAccess.ACCESSORS.get(writeRequest.getClass());
+    return access.invokeGetDeleteRequest(writeRequest) != null;
+  }
+
   @Nullable
   static String getSnsTopicArn(Object request) {
     RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
@@ -220,4 +239,53 @@ private static MethodHandle findGetLambdaArnMethod() {
       }
     }
   }
+
+  private static class WriteRequestAccess {
+    private static final ClassValue<WriteRequestAccess> ACCESSORS =
+        new ClassValue<WriteRequestAccess>() {
+          @Override
+          protected WriteRequestAccess computeValue(Class<?> type) {
+            return new WriteRequestAccess(type);
+          }
+        };
+
+    @Nullable private final Method getPutRequest;
+    @Nullable private final Method getDeleteRequest;
+
+    private WriteRequestAccess(Class<?> clz) {
+      getPutRequest = findMethodOrNull(clz, "getPutRequest");
+      getDeleteRequest = findMethodOrNull(clz, "getDeleteRequest");
+    }
+
+    @Nullable
+    Object invokeGetPutRequest(Object obj) {
+      return invokeMethod(getPutRequest, obj);
+    }
+
+    @Nullable
+    Object invokeGetDeleteRequest(Object obj) {
+      return invokeMethod(getDeleteRequest, obj);
+    }
+
+    @Nullable
+    private static Object invokeMethod(@Nullable Method method, Object obj) {
+      if (method == null) {
+        return null;
+      }
+      try {
+        return method.invoke(obj);
+      } catch (Throwable ignored) {
+        return null;
+      }
+    }
+
+    @Nullable
+    private static Method findMethodOrNull(Class<?> clz, String methodName) {
+      try {
+        return clz.getMethod(methodName);
+      } catch (Throwable ignored) {
+        return null;
+      }
+    }
+  }
 }