Add Custom Header - In AddHeaderPolicy and RequestIDPolicy

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

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+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 HttpHeaders} 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 AddHeadersFromContextPolicy implements HttpPipelinePolicy {
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+
+        Optional<Object> customHttpHeadersObject = context.getData(AZURE_REQUEST_HTTP_HEADERS_KEY);
+        if (customHttpHeadersObject.isPresent()
+            && customHttpHeadersObject.get() instanceof HttpHeaders) {
+
+            HttpHeaders customHttpHeaders = (HttpHeaders) customHttpHeadersObject.get();
+            // loop through customHttpHeaders and add headers 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,47 @@
 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.
+ *
+ * <p>The default {@link HttpHeader} name can be overwritten as shown below
+ * <p><strong>Code sample</strong></p>
+ * {@codesnippet com.azure.core.http.policy.RequestIdPolicy.constructor.overrideRequestIdHeaderName}
  */
 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 final String AZURE_REQUEST_HTTP_HEADERS_KEY = "azure-request-override-http-headers-key";
+
     // All fields must be immutable.
     //
     /**

sdk/core/azure-core/src/samples/java/com/azure/core/http/policy/RequestIdPolicyJavaDocCodeSnippet.java

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.http.policy;
+
+/**
+ * Code snippets for {@link RequestIdPolicy}
+ */
+public class RequestIdPolicyJavaDocCodeSnippet {
+
+    /**
+     * Code snippets for using {@link RequestIdPolicy#RequestIdPolicy(String)} }
+     */
+    public void overrideRequestIdHeaderName() {
+
+        // BEGIN: com.azure.core.http.policy.RequestIdPolicy.constructor.overrideRequestIdHeaderName
+        RequestIdPolicy requestIdPolicy =  new RequestIdPolicy("x-ms-my-custom-request-id");
+        // END: com.azure.core.http.policy.RequestIdPolicy.constructor.overrideRequestIdHeaderName
+    }
+}