Fixing S3 bucket name extraction for AWS SDK V2 (#393)

* Fixing S3 bucket name extraction for AWS SDK V2
aws · Dec 12, 2023 · 3b0868d · 3b0868d
1 parent d1c7d7b
commit 3b0868d
Show file tree

Hide file tree

Showing 4 changed files with 222 additions and 76 deletions.
diff --git a/aws-xray-recorder-sdk-aws-sdk-v2/build.gradle.kts b/aws-xray-recorder-sdk-aws-sdk-v2/build.gradle.kts
@@ -15,6 +15,7 @@ dependencies {
     testImplementation("software.amazon.awssdk:lambda:2.15.20")
     testImplementation("software.amazon.awssdk:sqs:2.15.20")
     testImplementation("software.amazon.awssdk:sns:2.15.20")
+    testImplementation("software.amazon.awssdk:s3:2.15.20")
 }
 
 tasks.jar {

diff --git a/...rder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java b/...rder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java
@@ -149,7 +149,8 @@ private HashMap<String, Object> extractRequestParameters(
                 SdkRequest request = context.request();
                 Optional<Object> parameterValue = request.getValueForField(parameterName, Object.class);
                 if (parameterValue.isPresent()) {
-                    parameters.put(SNAKE_CASE_NAMING_STRATEGY.translate(parameterName), parameterValue.get());
+                    parameters.put(normalizeParameterName(
+                            SNAKE_CASE_NAMING_STRATEGY.translate(parameterName)), parameterValue.get());
                 }
             });
         }
@@ -161,14 +162,16 @@ private HashMap<String, Object> extractRequestParameters(
                     Optional<Map> parameterValue = request.getValueForField(key, Map.class);
                     if (parameterValue.isPresent()) {
                         String renameTo = descriptor.getRenameTo() != null ? descriptor.getRenameTo() : key;
-                        parameters.put(SNAKE_CASE_NAMING_STRATEGY.translate(renameTo), parameterValue.get().keySet());
+                        parameters.put(normalizeParameterName(
+                                SNAKE_CASE_NAMING_STRATEGY.translate(renameTo)), parameterValue.get().keySet());
                     }
                 } else if (descriptor.isList() && descriptor.shouldGetCount()) {
                     SdkRequest request = context.request();
                     Optional<List> parameterValue = request.getValueForField(key, List.class);
                     if (parameterValue.isPresent()) {
                         String renameTo = descriptor.getRenameTo() != null ? descriptor.getRenameTo() : key;
-                        parameters.put(SNAKE_CASE_NAMING_STRATEGY.translate(renameTo), parameterValue.get().size());
+                        parameters.put(normalizeParameterName(
+                                SNAKE_CASE_NAMING_STRATEGY.translate(renameTo)), parameterValue.get().size());
                     }
                 }
             });
@@ -177,6 +180,20 @@ private HashMap<String, Object> extractRequestParameters(
         return parameters;
     }
 
+    /**
+     * @param parameterName
+     *      The name of the parameter to normalize.
+     * @return
+     *      The field key for a parameter that the X-Ray backend expects.
+     */
+    private String normalizeParameterName(String parameterName) {
+        // AWS SDK V2 changed the field name from BucketName to Bucket.
+        if (parameterName.equals("bucket")) {
+            return "bucket_name";
+        }
+        return parameterName;
+    }
+
     private HashMap<String, Object> extractResponseParameters(
         Context.AfterExecution context, ExecutionAttributes executionAttributes) {
         HashMap<String, Object> parameters = new HashMap<>();
@@ -191,7 +208,8 @@ private HashMap<String, Object> extractResponseParameters(
                 SdkResponse response = context.response();
                 Optional<Object> parameterValue = response.getValueForField(parameterName, Object.class);
                 if (parameterValue.isPresent()) {
-                    parameters.put(SNAKE_CASE_NAMING_STRATEGY.translate(parameterName), parameterValue.get());
+                    parameters.put(normalizeParameterName(
+                            SNAKE_CASE_NAMING_STRATEGY.translate(parameterName)), parameterValue.get());
                 }
             });
         }
@@ -203,14 +221,16 @@ private HashMap<String, Object> extractResponseParameters(
                     Optional<Map> parameterValue = response.getValueForField(key, Map.class);
                     if (parameterValue.isPresent()) {
                         String renameTo = descriptor.getRenameTo() != null ? descriptor.getRenameTo() : key;
-                        parameters.put(SNAKE_CASE_NAMING_STRATEGY.translate(renameTo), parameterValue.get().keySet());
+                        parameters.put(normalizeParameterName(
+                                SNAKE_CASE_NAMING_STRATEGY.translate(renameTo)), parameterValue.get().keySet());
                     }
                 } else if (descriptor.isList() && descriptor.shouldGetCount()) {
                     SdkResponse response = context.response();
                     Optional<List> parameterValue = response.getValueForField(key, List.class);
                     if (parameterValue.isPresent()) {
                         String renameTo = descriptor.getRenameTo() != null ? descriptor.getRenameTo() : key;
-                        parameters.put(SNAKE_CASE_NAMING_STRATEGY.translate(renameTo), parameterValue.get().size());
+                        parameters.put(normalizeParameterName(
+                                SNAKE_CASE_NAMING_STRATEGY.translate(renameTo)), parameterValue.get().size());
                     }
                 }
             });