Add Custom Header - In AddHeaderPolicy and RequestIDPolicy

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

@@ -0,0 +1,39 @@
+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;
+import java.util.Optional;
+
+/**
+ * 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'. The value for this key should be of type
+ * {@link HttpHeaders} for it to be added in {@link HttpRequest}.
+ */
+public class OverrideHeaderPolicy implements HttpPipelinePolicy {
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+
+        Optional<Object> customHttpHeadersObject = context.getData(AZURE_REQUEST_HTTP_HEADERS_KEY);
+        if (!Objects.isNull(customHttpHeadersObject) && customHttpHeadersObject.get() instanceof HttpHeaders) {
+            HttpHeaders customHttpHeaders = (HttpHeaders) customHttpHeadersObject.get();
+            // loop through customHttpHeaders and add header in HttpRequest
+            for (HttpHeader httpHeader : customHttpHeaders) {
+                if (!Objects.isNull(httpHeader.getName()) && !Objects.isNull(httpHeader.getValue())) {
+                    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.
     //
     /**