Add Custom Header - In AddHeaderPolicy and RequestIDPolicy (#6602)

New AddHeadersFromContextPolicy to add multiple headers and values for those headers.
Allowing override header name in RequestIdPolicy.

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

@@ -0,0 +1,61 @@
+// 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 reactor.core.publisher.Mono;
+
+import java.util.Objects;
+
+/**
+ * The pipeline policy that override or add  {@link HttpHeaders} in {@link HttpRequest} by reading values from
+ * {@link Context} with key 'azure-http-headers-key'. The value for this key should be of type {@link HttpHeaders} for
+ * it to be added in {@link HttpRequest}.
+ *
+ * <p><strong>Code Sample: Add multiple HttpHeader in Context and call client</strong></p>
+ * <pre>
+ * // Create ConfigurationClient for example
+ * ConfigurationClient configurationClient = new ConfigurationClientBuilder()
+ *       .connectionString("endpoint={endpoint_value};id={id_value};secret={secret_value}")
+ *       .buildClient();
+ * // Add your headers
+ * HttpHeaders headers = new HttpHeaders();
+ * headers.put("my-header1", "my-header1-value");
+ * headers.put("my-header2", "my-header2-value");
+ * headers.put("my-header3", "my-header3-value");
+ * // Call API by passing headers in Context.
+ * configurationClient.addConfigurationSettingWithResponse(
+ *       new ConfigurationSetting().setKey("key").setValue("value"),
+ *       new Context(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY, headers));
+ * // Above three HttpHeader will be added in outgoing HttpRequest.
+ * </pre>
+ */
+public class AddHeadersFromContextPolicy implements HttpPipelinePolicy {
+
+    /**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-http-headers-key";
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        context.getData(AZURE_REQUEST_HTTP_HEADERS_KEY).ifPresent(headers -> {
+            if (headers instanceof HttpHeaders) {
+                HttpHeaders customHttpHeaders = (HttpHeaders) headers;
+                // 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,
+            "requestIdHeaderName 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/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
+        new RequestIdPolicy("x-ms-my-custom-request-id");
+        // END: com.azure.core.http.policy.RequestIdPolicy.constructor.overrideRequestIdHeaderName
+    }
+}

sdk/core/azure-core/src/test/java/com/azure/core/http/policy/AddHeadersFromContextPolicyTest.java

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.http.policy;
+
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.clients.NoOpHttpClient;
+import com.azure.core.util.Context;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+
+public class AddHeadersFromContextPolicyTest {
+
+    private final HttpResponse mockResponse = new HttpResponse(null) {
+        @Override
+        public int getStatusCode() {
+            return 500;
+        }
+
+        @Override
+        public String getHeaderValue(String name) {
+            return null;
+        }
+
+        @Override
+        public HttpHeaders getHeaders() {
+            return new HttpHeaders();
+        }
+
+        @Override
+        public Mono<byte[]> getBodyAsByteArray() {
+            return Mono.empty();
+        }
+
+        @Override
+        public Flux<ByteBuffer> getBody() {
+            return Flux.empty();
+        }
+
+        @Override
+        public Mono<String> getBodyAsString() {
+            return Mono.empty();
+        }
+
+        @Override
+        public Mono<String> getBodyAsString(Charset charset) {
+            return Mono.empty();
+        }
+    };
+
+    @Test
+    public void clientProvidedMultipleHeaderTest() throws Exception {
+        // Create custom Headers
+        String customRequestId = "request-id-value";
+        final HttpHeaders headers = new HttpHeaders();
+        headers.put("x-ms-client-request-id", customRequestId);
+        headers.put("my-header1", "my-header1-value");
+        headers.put("my-header2", "my-header2-value");
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient() {
+                @Override
+                public Mono<HttpResponse> send(HttpRequest request) {
+                    Assertions.assertEquals(request.getHeaders().getValue("x-ms-client-request-id"), customRequestId);
+                    Assertions.assertEquals(request.getHeaders().getValue("my-header1"), "my-header1-value");
+                    Assertions.assertEquals(request.getHeaders().getValue("my-header2"), "my-header2-value");
+                    return Mono.just(mockResponse);
+                }
+            })
+            .policies(new RequestIdPolicy())
+            .policies(new AddHeadersFromContextPolicy())
+            .build();
+
+        pipeline.send(new HttpRequest(HttpMethod.GET, new URL("http://localhost/")), new Context(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY, headers)).block();
+    }
+}