diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/AddHeadersFromContextPolicy.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/AddHeadersFromContextPolicy.java
new file mode 100644
index 000000000000..26caf06f4bb6
--- /dev/null
+++ b/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}.
+ *
+ *
Code Sample: Add multiple HttpHeader in Context and call client
+ *
+ * // 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.
+ *
+ */
+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 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();
+
+ }
+}
diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicy.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicy.java
index 389386519cf9..e9634030fe6d 100644
--- a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicy.java
+++ b/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.
+ *
+ * The default {@link HttpHeader} name can be overwritten as shown below
+ *
Code sample
+ * {@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 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();
}
diff --git a/sdk/core/azure-core/src/samples/java/com/azure/core/http/policy/RequestIdPolicyJavaDocCodeSnippet.java b/sdk/core/azure-core/src/samples/java/com/azure/core/http/policy/RequestIdPolicyJavaDocCodeSnippet.java
new file mode 100644
index 000000000000..64eaa194ad07
--- /dev/null
+++ b/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
+ }
+}
diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/AddHeadersFromContextPolicyTest.java b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/AddHeadersFromContextPolicyTest.java
new file mode 100644
index 000000000000..856b49957376
--- /dev/null
+++ b/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 getBodyAsByteArray() {
+ return Mono.empty();
+ }
+
+ @Override
+ public Flux getBody() {
+ return Flux.empty();
+ }
+
+ @Override
+ public Mono getBodyAsString() {
+ return Mono.empty();
+ }
+
+ @Override
+ public Mono 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 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();
+ }
+}