Add ability to set execution attributes as part of client and request override configurations

Author: Nicholas DeJaco

core/aws-core/src/main/java/software/amazon/awssdk/awscore/client/handler/AwsClientHandlerUtils.java

@@ -100,6 +100,10 @@ static <InputT extends SdkRequest, OutputT extends SdkResponse> ExecutionContext
             .putAttribute(SdkInternalExecutionAttribute.DISABLE_HOST_PREFIX_INJECTION,
                           clientConfig.option(SdkAdvancedClientOption.DISABLE_HOST_PREFIX_INJECTION));
 
+        executionAttributes.putAllAttributes(clientConfig.option(SdkClientOption.EXECUTION_ATTRIBUTES));
+        originalRequest.overrideConfiguration().ifPresent(requestOverrideConfiguration ->
+                executionAttributes.putAllAttributes(requestOverrideConfiguration.executionAttributes()));
+
         ExecutionInterceptorChain executionInterceptorChain =
                 new ExecutionInterceptorChain(clientConfig.option(SdkClientOption.EXECUTION_INTERCEPTORS));
         return ExecutionContext.builder()

core/sdk-core/src/main/java/software/amazon/awssdk/core/RequestOverrideConfiguration.java

@@ -27,6 +27,7 @@
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.Immutable;
 import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.utils.CollectionUtils;
@@ -46,6 +47,7 @@ public abstract class RequestOverrideConfiguration {
     private final Duration apiCallAttemptTimeout;
     private final Signer signer;
     private final List<MetricPublisher> metricPublishers;
+    private final Map<ExecutionAttribute<?>, Object> executionAttributes;
 
     protected RequestOverrideConfiguration(Builder<?> builder) {
         this.headers = CollectionUtils.deepUnmodifiableMap(builder.headers(), () -> new TreeMap<>(String.CASE_INSENSITIVE_ORDER));
@@ -55,6 +57,7 @@ protected RequestOverrideConfiguration(Builder<?> builder) {
         this.apiCallAttemptTimeout = Validate.isPositiveOrNull(builder.apiCallAttemptTimeout(), "apiCallAttemptTimeout");
         this.signer = builder.signer();
         this.metricPublishers = Collections.unmodifiableList(new ArrayList<>(builder.metricPublishers()));
+        this.executionAttributes = Collections.unmodifiableMap(new HashMap<>(builder.executionAttributes()));
     }
 
     /**
@@ -138,6 +141,14 @@ public List<MetricPublisher> metricPublishers() {
         return metricPublishers;
     }
 
+    /**
+     * Returns the additional execution attributes to be added to this request.
+     * This collection of attributes is added in addition to the attributes set on the client.
+     * An attribute value added on the client within the collection of attributes is superseded by an
+     * attribute value added on the request.
+     */
+    public Map<ExecutionAttribute<?>, Object> executionAttributes() { return executionAttributes; }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -371,6 +382,22 @@ default B putRawQueryParameter(String name, String value) {
 
         List<MetricPublisher> metricPublishers();
 
+        /**
+         * Sets the additional execution attributes collection for this request.
+         * @param executionAttributes Execution attributes map for this request
+         * @return This object for method chaining.
+         */
+        B executionAttributes(Map<ExecutionAttribute<?>, Object> executionAttributes);
+
+        /**
+         * Add an execution attribute to the existing collection of execution attributes.
+         * @param attribute The execution attribute object
+         * @param value The value of the execution attribute.
+         */
+        B addExecutionAttribute(ExecutionAttribute attribute, Object value);
+
+        Map<ExecutionAttribute<?>, Object> executionAttributes();
+
         /**
          * Create a new {@code SdkRequestOverrideConfiguration} with the properties set on this builder.
          *
@@ -387,6 +414,7 @@ protected abstract static class BuilderImpl<B extends Builder> implements Builde
         private Duration apiCallAttemptTimeout;
         private Signer signer;
         private List<MetricPublisher> metricPublishers = new ArrayList<>();
+        private Map<ExecutionAttribute<?>, Object> executionAttributes = new HashMap<>();
 
         protected BuilderImpl() {
         }
@@ -526,5 +554,23 @@ public void setMetricPublishers(List<MetricPublisher> metricPublishers) {
         public List<MetricPublisher> metricPublishers() {
             return metricPublishers;
         }
+
+        @Override
+        public B executionAttributes(Map<ExecutionAttribute<?>, Object> executionAttributes) {
+            Validate.paramNotNull(executionAttributes, "executionAttributes");
+            this.executionAttributes = executionAttributes;
+            return (B) this;
+        }
+
+        @Override
+        public B addExecutionAttribute(ExecutionAttribute executionAttribute, Object value) {
+            this.executionAttributes.put(executionAttribute, value);
+            return (B) this;
+        }
+
+        @Override
+        public Map<ExecutionAttribute<?>, Object> executionAttributes() { return executionAttributes; }
+
+        public void setExecutionAttributes(Map<ExecutionAttribute<?>, Object> executionAttributes) { executionAttributes(executionAttributes); }
     }
 }

core/sdk-core/src/main/java/software/amazon/awssdk/core/client/builder/SdkDefaultClientBuilder.java

@@ -22,19 +22,7 @@
 import static software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.SIGNER;
 import static software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.USER_AGENT_PREFIX;
 import static software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.USER_AGENT_SUFFIX;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.ADDITIONAL_HTTP_HEADERS;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.API_CALL_ATTEMPT_TIMEOUT;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.API_CALL_TIMEOUT;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.ASYNC_HTTP_CLIENT;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.CRC32_FROM_COMPRESSED_DATA_ENABLED;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.ENDPOINT_OVERRIDDEN;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.EXECUTION_INTERCEPTORS;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.METRIC_PUBLISHERS;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.PROFILE_FILE;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.PROFILE_NAME;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.RETRY_POLICY;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.SCHEDULED_EXECUTOR_SERVICE;
-import static software.amazon.awssdk.core.client.config.SdkClientOption.SIGNER_OVERRIDDEN;
+import static software.amazon.awssdk.core.client.config.SdkClientOption.*;
 import static software.amazon.awssdk.core.internal.SdkInternalTestAdvancedClientOption.ENDPOINT_OVERRIDDEN_OVERRIDE;
 import static software.amazon.awssdk.utils.CollectionUtils.mergeLists;
 import static software.amazon.awssdk.utils.Validate.paramNotNull;
@@ -370,6 +358,7 @@ public final B overrideConfiguration(ClientOverrideConfiguration overrideConfig)
         clientConfiguration.option(PROFILE_FILE, overrideConfig.defaultProfileFile().orElse(null));
         clientConfiguration.option(PROFILE_NAME, overrideConfig.defaultProfileName().orElse(null));
         clientConfiguration.option(METRIC_PUBLISHERS, overrideConfig.metricPublishers());
+        clientConfiguration.option(EXECUTION_ATTRIBUTES, overrideConfig.executionAttributes());
         overrideConfig.advancedOption(ENDPOINT_OVERRIDDEN_OVERRIDE).ifPresent(value -> {
             clientConfiguration.option(ENDPOINT_OVERRIDDEN, value);
         });

core/sdk-core/src/main/java/software/amazon/awssdk/core/client/config/ClientOverrideConfiguration.java

@@ -25,6 +25,7 @@
 import java.util.TreeMap;
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.core.retry.RetryPolicy;
@@ -57,6 +58,7 @@ public final class ClientOverrideConfiguration
     private final ProfileFile defaultProfileFile;
     private final String defaultProfileName;
     private final List<MetricPublisher> metricPublishers;
+    private final Map<ExecutionAttribute<?>, Object> executionAttributes;
 
     /**
      * Initialize this configuration. Private to require use of {@link #builder()}.
@@ -71,6 +73,7 @@ private ClientOverrideConfiguration(Builder builder) {
         this.defaultProfileFile = builder.defaultProfileFile();
         this.defaultProfileName = builder.defaultProfileName();
         this.metricPublishers = Collections.unmodifiableList(new ArrayList<>(builder.metricPublishers()));
+        this.executionAttributes = Collections.unmodifiableMap(new HashMap<>(builder.executionAttributes()));
     }
 
     @Override
@@ -196,6 +199,13 @@ public List<MetricPublisher> metricPublishers() {
         return metricPublishers;
     }
 
+    /**
+     *  Returns the additional execution attributes to be added for this client.
+     *
+     * @Return Map of execution attributes.
+     */
+    public Map<ExecutionAttribute<?>, Object> executionAttributes() { return executionAttributes; }
+
     @Override
     public String toString() {
         return ToString.builder("ClientOverrideConfiguration")
@@ -439,6 +449,22 @@ default Builder retryPolicy(RetryMode retryMode) {
         Builder addMetricPublisher(MetricPublisher metricPublisher);
 
         List<MetricPublisher> metricPublishers();
+
+        /**
+         * Sets the additional execution attributes collection for this client.
+         * @param executionAttributes Execution attributes map for this client.
+         * @return This object for method chaining.
+         */
+        Builder executionAttributes(Map<ExecutionAttribute<?>, Object> executionAttributes);
+
+        /**
+         * Add an execution attribute to the existing collection of execution attributes.
+         * @param attribute The execution attribute object
+         * @param value The value of the execution attribute.
+         */
+        Builder addExecutionAttribute(ExecutionAttribute attribute, Object value);
+
+        Map<ExecutionAttribute<?>, Object> executionAttributes();
     }
 
     /**
@@ -454,6 +480,7 @@ private static final class DefaultClientOverrideConfigurationBuilder implements
         private ProfileFile defaultProfileFile;
         private String defaultProfileName;
         private List<MetricPublisher> metricPublishers = new ArrayList<>();
+        private Map<ExecutionAttribute<?>, Object> executionAttributes = new HashMap<>();
 
         @Override
         public Builder headers(Map<String, List<String>> headers) {
@@ -618,6 +645,22 @@ public List<MetricPublisher> metricPublishers() {
             return Collections.unmodifiableList(metricPublishers);
         }
 
+        @Override
+        public Builder executionAttributes(Map<ExecutionAttribute<?>, Object> executionAttributes) {
+            Validate.paramNotNull(executionAttributes, "executionAttributes");
+            this.executionAttributes = executionAttributes;
+            return this;
+        }
+
+        @Override
+        public Builder addExecutionAttribute(ExecutionAttribute executionAttribute, Object value) {
+            this.executionAttributes.put(executionAttribute, value);
+            return this;
+        }
+
+        @Override
+        public Map<ExecutionAttribute<?>, Object> executionAttributes() { return executionAttributes; }
+
         @Override
         public ClientOverrideConfiguration build() {
             return new ClientOverrideConfiguration(this);

core/sdk-core/src/main/java/software/amazon/awssdk/core/client/config/SdkClientOption.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.core.ClientType;
 import software.amazon.awssdk.core.ServiceConfiguration;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.retry.RetryPolicy;
 import software.amazon.awssdk.http.SdkHttpClient;
@@ -138,6 +139,12 @@ public final class SdkClientOption<T> extends ClientOption<T> {
      */
     public static final SdkClientOption<Boolean> SIGNER_OVERRIDDEN = new SdkClientOption<>(Boolean.class);
 
+    /**
+     * Option to specify additional execution attributes to each client call.
+     */
+    public static final SdkClientOption<Map<ExecutionAttribute<?>, Object>> EXECUTION_ATTRIBUTES =
+            new SdkClientOption<>(new UnsafeValueType(Map.class));
+
     private SdkClientOption(Class<T> valueClass) {
         super(valueClass);
     }

core/sdk-core/src/main/java/software/amazon/awssdk/core/interceptor/ExecutionAttributes.java

@@ -48,6 +48,14 @@ public <U> ExecutionAttributes putAttribute(ExecutionAttribute<U> attribute, U v
         return this;
     }
 
+    /**
+     * Add all execution attributes from the provided collection to the internal collection of attributes.
+     */
+    public <U> ExecutionAttributes putAllAttributes(Map<ExecutionAttribute<?>, Object> executionAttributes) {
+        this.attributes.putAll(executionAttributes);
+        return this;
+    }
+
     /**
      * Set the provided attribute in this collection of attributes if it does not already exist in the collection.
      */

core/sdk-core/src/test/java/software/amazon/awssdk/core/RequestOverrideConfigurationTest.java

@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import org.junit.Test;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.utils.ImmutableMap;
 
@@ -165,4 +166,71 @@ public void addMetricPublisher_listPreviouslyAdded_appendedToList() {
 
         assertThat(overrideConfig.metricPublishers()).containsExactly(publishers.get(0), publishers.get(1), thirdAdded);
     }
+
+    @Test
+    public void executionAttributes_createsCopy() {
+        Map<ExecutionAttribute<?>, Object> executionAttributes = new HashMap<>();
+
+        ExecutionAttribute testAttribute = new ExecutionAttribute("TestAttribute");
+        String expectedValue = "Value1";
+        executionAttributes.put(testAttribute, expectedValue);
+
+        SdkRequestOverrideConfiguration overrideConfig = SdkRequestOverrideConfiguration.builder()
+                .executionAttributes(executionAttributes)
+                .build();
+
+        executionAttributes.remove(testAttribute);
+        assertThat(overrideConfig.executionAttributes().get(testAttribute)).isEqualTo(expectedValue);
+    }
+
+    @Test
+    public void executionAttributes_maintainsAllAdded() {
+        Map<ExecutionAttribute, Object> executionAttributeObjectMap = new HashMap<>();
+        for (int i = 0; i < 5; i++) {
+            executionAttributeObjectMap.put(new ExecutionAttribute<>("Attribute" + i), mock(Object.class));
+        }
+
+        SdkRequestOverrideConfiguration.Builder builder = SdkRequestOverrideConfiguration.builder();
+
+        for (Map.Entry<ExecutionAttribute, Object> attributeObjectEntry : executionAttributeObjectMap.entrySet()) {
+            builder.addExecutionAttribute(attributeObjectEntry.getKey(), attributeObjectEntry.getValue());
+        }
+
+        SdkRequestOverrideConfiguration overrideConfig = builder.build();
+        assertThat(overrideConfig.executionAttributes()).isEqualTo(executionAttributeObjectMap);
+    }
+
+    @Test
+    public void executionAttributes_overwritesPreviouslyAdded() {
+        Map<ExecutionAttribute<?>, Object> executionAttributes = new HashMap<>();
+        for (int i = 0; i < 5; i++) {
+            executionAttributes.put(new ExecutionAttribute<>("Attribute" + i), mock(Object.class));
+        }
+
+        SdkRequestOverrideConfiguration.Builder builder = SdkRequestOverrideConfiguration.builder();
+
+        builder.addExecutionAttribute(new ExecutionAttribute("AddedAttribute"), mock(Object.class));
+        builder.executionAttributes(executionAttributes);
+        SdkRequestOverrideConfiguration overrideConfig = builder.build();
+        assertThat(overrideConfig.executionAttributes()).isEqualTo(executionAttributes);
+    }
+
+    @Test
+    public void executionAttributes_listPreviouslyAdded_appendedToList() {
+        Map<ExecutionAttribute<?>, Object> executionAttributes = new HashMap<>();
+        for (int i = 0; i < 5; i++) {
+            executionAttributes.put(new ExecutionAttribute<>("Attribute" + i), mock(Object.class));
+        }
+
+        SdkRequestOverrideConfiguration.Builder builder = SdkRequestOverrideConfiguration.builder();
+
+        builder.executionAttributes(executionAttributes);
+        ExecutionAttribute addedAttribute = new ExecutionAttribute("AddedAttribute");
+        Object addedValue = mock(Object.class);
+
+        builder.addExecutionAttribute(addedAttribute, addedValue);
+
+        SdkRequestOverrideConfiguration overrideConfig = builder.build();
+        assertThat(overrideConfig.executionAttributes().get(addedAttribute)).isEqualTo(addedValue);
+    }
 }