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(); + } +}