Add support for tracking business metrics from resolved endpoints (#6263)

* Ignore unknown properties in endpoint tests

* Add support for tracking business metrics from resolved endpoints

* Remove unused import

* Fix syntax error + add functional test

* Improve comments + fix test

* Cleanups from PR + expand tests

Author: alextwoods

.changes/next-release/feature-AWSSDKforJavaV2-6e40ae3.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java V2",
+    "contributor": "",
+    "description": "Add support for tracking business metrics from resolved endpoints."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointResolverInterceptorSpec.java

@@ -160,6 +160,8 @@ public TypeSpec poetSpec() {
         }
 
         endpointParamsKnowledgeIndex.addAccountIdMethodsIfPresent(b);
+
+        b.addMethod(setMetricValuesMethod());
         return b.build();
     }
 
@@ -255,6 +257,7 @@ private MethodSpec modifyRequestMethod(String endpointAuthSchemeStrategyFieldNam
         }
 
         b.addStatement("executionAttributes.putAttribute(SdkInternalExecutionAttribute.RESOLVED_ENDPOINT, endpoint)");
+        b.addStatement("setMetricValues(endpoint, executionAttributes)");
         b.addStatement("return result");
         b.endControlFlow();
         b.beginControlFlow("catch ($T e)", CompletionException.class);
@@ -905,4 +908,19 @@ private MethodSpec constructorMethodSpec(String endpointAuthSchemeFieldName) {
         b.addStatement("this.$N = $N.endpointAuthSchemeStrategy()", endpointAuthSchemeFieldName, factoryLocalVarName);
         return b.build();
     }
+
+    private MethodSpec setMetricValuesMethod() {
+        MethodSpec.Builder b = MethodSpec.methodBuilder("setMetricValues")
+                                         .addModifiers(Modifier.PRIVATE)
+                                         .addParameter(Endpoint.class, "endpoint")
+                                         .addParameter(ExecutionAttributes.class, "executionAttributes")
+                                         .returns(void.class);
+
+        b.beginControlFlow("if (endpoint.attribute($T.METRIC_VALUES) != null)", AwsEndpointAttribute.class);
+        b.addStatement("executionAttributes.getOptionalAttribute($T.BUSINESS_METRICS).ifPresent("
+                       + "metrics -> endpoint.attribute($T.METRIC_VALUES).forEach(v -> metrics.addMetric(v)))",
+                       SdkInternalExecutionAttribute.class, AwsEndpointAttribute.class);
+        b.endControlFlow();
+        return b.build();
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/TestGeneratorUtils.java

@@ -30,6 +30,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.awscore.endpoints.AwsEndpointAttribute;
 import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4AuthScheme;
 import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4aAuthScheme;
@@ -45,6 +47,8 @@
 import software.amazon.awssdk.utils.StringUtils;
 
 public final class TestGeneratorUtils {
+    private static final Logger log = LoggerFactory.getLogger(TestGeneratorUtils.class);
+
     private TestGeneratorUtils() {
     }
 
@@ -111,16 +115,40 @@ private static void addEndpointAttributeBlock(CodeBlock.Builder builder, String
                                                   Map<String, KeyTypePair> knownEndpointAttributes) {
         if ("authSchemes".equals(attrName)) {
             addAuthSchemesBlock(builder, attrValue);
+        } else if ("metricValues".equals(attrName)) {
+            addMetricValuesBlock(builder, attrValue);
         } else if (knownEndpointAttributes.containsKey(attrName)) {
             addAttributeBlock(builder, attrValue, knownEndpointAttributes.get(attrName));
         } else {
-            throw new RuntimeException(
-                String.format("Encountered unknown expected endpoint attribute: %s. Known attributes: %s.",
+            log.warn("Ignoring unknown expected endpoint attribute: {}. Known attributes: {}.",
                               attrName,
-                              knownEndpointAttributes));
+                              knownEndpointAttributes);
         }
     }
 
+    private static void addMetricValuesBlock(CodeBlock.Builder builder, TreeNode attrValue) {
+        CodeBlock keyExpr = CodeBlock.builder()
+                                     .add("$T.METRIC_VALUES", AwsEndpointAttribute.class)
+                                     .build();
+
+        CodeBlock.Builder schemesListExpr = CodeBlock.builder()
+                                                     .add("$T.asList(", Arrays.class);
+
+        JrsArray schemesArray = (JrsArray) attrValue;
+
+        Iterator<JrsValue> elementsIter = schemesArray.elements();
+        while (elementsIter.hasNext()) {
+            schemesListExpr.add("$S", elementsIter.next().asText());
+
+            if (elementsIter.hasNext()) {
+                schemesListExpr.add(",");
+            }
+        }
+        schemesListExpr.add(")");
+
+        builder.add(".putAttribute($L, $L)", keyExpr, schemesListExpr.build());
+    }
+
     private static void addAttributeBlock(CodeBlock.Builder builder, TreeNode attrValue, KeyTypePair keyType) {
         CodeBlock keyExpr = CodeBlock.builder()
                                      .add("$L", keyType.getKey())

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules2/CodeGeneratorVisitor.java

@@ -353,6 +353,8 @@ public Void visitPropertiesExpression(PropertiesExpression e) {
         properties.forEach((k, v) -> {
             if ("authSchemes".equals(k)) {
                 addAuthSchemesBlock(v);
+            } else if ("metricValues".equals(k)) {
+                addMetricValuesBlock(v);
             } else if (knownEndpointAttributes.containsKey(k)) {
                 addAttributeBlock(k, v);
             } else {
@@ -436,6 +438,12 @@ private ClassName authSchemeClass(String name) {
         }
     }
 
+    private void addMetricValuesBlock(RuleExpression v) {
+        builder.add(".putAttribute($T.METRIC_VALUES, ", AwsEndpointAttribute.class);
+        v.accept(this);
+        builder.add(")");
+    }
+
     private void addAttributeBlock(String k, RuleExpression v) {
         KeyTypePair keyType = knownEndpointAttributes.get(k);
         ClassConstant classConstant = parseClassConstant(keyType.getKey());

codegen/src/test/java/software/amazon/awssdk/codegen/poet/ClientTestModels.java

@@ -192,7 +192,27 @@ public static IntermediateModel queryServiceModelsWithUnknownEndpointProperties(
         File endpointRuleSetModel =
             new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-rule-set-unknown-properties.json").getFile());
         File endpointTestsModel =
-            new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-tests.json").getFile());
+            new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-tests-unknown-properties.json").getFile());
+
+        C2jModels models = C2jModels
+            .builder()
+            .serviceModel(getServiceModel(serviceModel))
+            .waitersModel(getWaiters(waitersModel))
+            .customizationConfig(CustomizationConfig.create())
+            .endpointRuleSetModel(getEndpointRuleSet(endpointRuleSetModel))
+            .endpointTestSuiteModel(getEndpointTestSuite(endpointTestsModel))
+            .build();
+
+        return new IntermediateModelBuilder(models).build();
+    }
+
+    public static IntermediateModel queryServiceModelsWithUnknownEndpointMetricValues() {
+        File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/query/service-2.json").getFile());
+        File waitersModel = new File(ClientTestModels.class.getResource("client/c2j/query/waiters-2.json").getFile());
+        File endpointRuleSetModel =
+            new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-rule-set-metric-values.json").getFile());
+        File endpointTestsModel =
+            new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-tests-metric-values.json").getFile());
 
         C2jModels models = C2jModels
             .builder()

codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules/EndpointProviderCompiledRulesClassSpecTest.java

@@ -51,4 +51,11 @@ void endpointProviderClassWithUriCache() {
             new EndpointProviderSpec2(ClientTestModels.queryServiceModelsWithUriCache());
         assertThat(endpointProviderSpec, generatesTo("endpoint-provider-uri-cache-class.java"));
     }
+
+    @Test
+    void endpointProviderClassWithMetricValues() {
+        ClassSpec endpointProviderSpec =
+            new EndpointProviderSpec2(ClientTestModels.queryServiceModelsWithUnknownEndpointMetricValues());
+        assertThat(endpointProviderSpec, generatesTo("endpoint-provider-metric-values-class.java"));
+    }
 }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules/EndpointRulesClientTestSpecTest.java

@@ -34,4 +34,16 @@ public void endpointProviderTestClassWithStringArray() {
         ClassSpec endpointProviderSpec = new EndpointProviderTestSpec(ClientTestModels.stringArrayServiceModels());
         assertThat(endpointProviderSpec, generatesTo("endpoint-rules-stringarray-test-class.java"));
     }
+
+    @Test
+    public void endpointProviderTestClassWithUnknownProperties() {
+        ClassSpec endpointProviderSpec = new EndpointProviderTestSpec(ClientTestModels.queryServiceModelsWithUnknownEndpointProperties());
+        assertThat(endpointProviderSpec, generatesTo("endpoint-rules-unknown-property-test-class.java"));
+    }
+
+    @Test
+    public void endpointProviderTestClassWithMetricValues() {
+        ClassSpec endpointProviderSpec = new EndpointProviderTestSpec(ClientTestModels.queryServiceModelsWithUnknownEndpointMetricValues());
+        assertThat(endpointProviderSpec, generatesTo("endpoint-rules-metric-values-test-class.java"));
+    }
 }