Add Custom Header - In AddHeaderPolicy and RequestIDPolicy

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/AddHeadersPolicy.java

@@ -10,23 +10,36 @@
 import com.azure.core.http.HttpResponse;
 import reactor.core.publisher.Mono;
 
+import java.util.Objects;
+import java.util.function.Supplier;
+
 /**
  * The pipeline policy that adds a particular set of headers to HTTP requests.
  */
 public class AddHeadersPolicy implements HttpPipelinePolicy {
-    private final HttpHeaders headers;
+    private final Supplier<HttpHeaders> headersSupplier;
+
+    /**
+     * Creates a AddHeadersPolicy provided {@link Supplier}.
+     *
+     * @param headersSupplier The {@link Supplier} to provide headers to add to outgoing requests.
+     */
+    public AddHeadersPolicy(Supplier<HttpHeaders> headersSupplier) {
+        this.headersSupplier = Objects.requireNonNull(headersSupplier, "headersSuppliercan not be null.");
+    }
 
     /**
      * Creates a AddHeadersPolicy.
      *
      * @param headers The headers to add to outgoing requests.
      */
     public AddHeadersPolicy(HttpHeaders headers) {
-        this.headers = headers;
+        this.headersSupplier = () -> headers;
     }
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        HttpHeaders headers = headersSupplier.get();
         for (HttpHeader header : headers) {
             context.getHttpRequest().setHeader(header.getName(), header.getValue());
         }

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/OverrideHeaderPolicy.java

@@ -0,0 +1,33 @@
+package com.azure.core.http.policy;
+
+import com.azure.core.http.HttpHeader;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.util.Context;
+import com.azure.core.http.HttpRequest;
+import static com.azure.core.util.Context.AZURE_REQUEST_HTTP_HEADERS_KEY;
+import reactor.core.publisher.Mono;
+
+import java.util.Objects;
+
+/**
+ * The pipeline policy that override or add  {@link HttpHeader} in {@link HttpRequest} by reading values from
+ * {@link Context} with key  'azure-request-override-http-headers-key'.
+ */
+public class OverrideHeaderPolicy implements HttpPipelinePolicy {
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        if (!Objects.isNull(context.getData(AZURE_REQUEST_HTTP_HEADERS_KEY))) {
+            HttpHeaders customHttpHeaders = (HttpHeaders) context.getData(AZURE_REQUEST_HTTP_HEADERS_KEY).get();
+            // loop through customHttpHeaders and add header in HttpRequest
+            for (HttpHeader httpHeader : customHttpHeaders) {
+                context.getHttpRequest().getHeaders().put(httpHeader.getName(), httpHeader.getValue());
+            }
+        }
+        return next.process();
+
+    }
+}

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicy.java

@@ -6,22 +6,43 @@
 import com.azure.core.http.HttpPipelineCallContext;
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpResponse;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpHeader;
 import reactor.core.publisher.Mono;
 
+import java.util.Objects;
 import java.util.UUID;
 
 /**
- * The pipeline policy that puts a UUID in the request header. Azure uses the request id as
- * the unique identifier for the request.
+ * The pipeline policy that puts a UUID in the request header. The default {@link HttpHeader} name can be overwritten.
+ * Azure uses the request id as the unique identifier for the request.
  */
 public class RequestIdPolicy implements HttpPipelinePolicy {
+
     private static final String REQUEST_ID_HEADER = "x-ms-client-request-id";
+    private final String requestIdHeaderName;
+
+    /**
+     * Creates  {@link RequestIdPolicy} with provided {@code requestIdHeaderName}.
+     * @param requestIdHeaderName to be used to set in {@link HttpRequest}.
+     */
+    public RequestIdPolicy(String requestIdHeaderName) {
+        this.requestIdHeaderName = Objects.requireNonNull(requestIdHeaderName,
+            "messageIdHeaderName can not be null.");
+    }
+
+    /**
+     * Creates default {@link RequestIdPolicy} with default header name 'x-ms-client-request-id'.
+     */
+    public RequestIdPolicy() {
+        requestIdHeaderName = REQUEST_ID_HEADER;
+    }
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
-        String requestId = context.getHttpRequest().getHeaders().getValue(REQUEST_ID_HEADER);
+        String requestId = context.getHttpRequest().getHeaders().getValue(requestIdHeaderName);
         if (requestId == null) {
-            context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString());
+            context.getHttpRequest().getHeaders().put(requestIdHeaderName, UUID.randomUUID().toString());
         }
         return next.process();
     }

sdk/core/azure-core/src/main/java/com/azure/core/util/Context.java

@@ -5,6 +5,7 @@
 
 import com.azure.core.annotation.Immutable;
 import com.azure.core.util.logging.ClientLogger;
+import com.azure.core.http.HttpHeaders;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -23,6 +24,9 @@
 public class Context {
     private final ClientLogger logger = new ClientLogger(Context.class);
 
+    /**Key used to override headers in HttpRequest. The Value for this key should be {@link HttpHeaders}.*/
+    public static String AZURE_REQUEST_HTTP_HEADERS_KEY = "azure-request-override-http-headers-key";
+
     // All fields must be immutable.
     //
     /**