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 5fd67f4d3e4c..615a42ae3f1f 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 @@ -3,25 +3,64 @@ package com.azure.core.http.policy; +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 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 pipeline policy that puts a UUID or user provided id in the request header. Azure uses the request id as * the unique identifier for the request. + * @see RequestIdPolicyOptions */ public class RequestIdPolicy implements HttpPipelinePolicy { private static final String REQUEST_ID_HEADER = "x-ms-client-request-id"; + private final RequestIdPolicyOptions requestIdPolicyOptions; + + /** + * Creates {@link RequestIdPolicy} with default behaviour. + */ + public RequestIdPolicy() { + this.requestIdPolicyOptions = null; + } + + /** + * Creates a {@link RequestIdPolicy} with the provided {@link RequestIdPolicyOptions}. + * User can specify {@link java.util.function.Function} to dynamically generated id for various id headers. + * + * @param requestIdPolicyOptions with given {@link RetryPolicyOptions}. + */ + public RequestIdPolicy(RequestIdPolicyOptions requestIdPolicyOptions) { + this.requestIdPolicyOptions = requestIdPolicyOptions; + } + + @Override public Mono process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) { String requestId = context.getHttpRequest().getHeaders().getValue(REQUEST_ID_HEADER); - if (requestId == null) { - context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString()); + if (!Objects.isNull(requestIdPolicyOptions)) { + HttpHeaders httpHeaders = requestIdPolicyOptions.getIdHeaderSupplier().get(); + if (!Objects.isNull(httpHeaders)) { + httpHeaders.stream().forEach(httpHeader -> { + String requestIdHeaderValue = context.getHttpRequest().getHeaders().getValue(httpHeader.getName()); + if (requestIdHeaderValue == null) { + context.getHttpRequest().getHeaders().put(httpHeader.getName(), httpHeader.getValue()); + } + }); + }else { + if (requestId == null) { + context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString()); + } + } + } else { + if (requestId == null) { + context.getHttpRequest().getHeaders().put(REQUEST_ID_HEADER, UUID.randomUUID().toString()); + } } return next.process(); } diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicyOptions.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicyOptions.java new file mode 100644 index 000000000000..e4860b1d5749 --- /dev/null +++ b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RequestIdPolicyOptions.java @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.core.http.policy; + +import com.azure.core.annotation.Immutable; +import com.azure.core.http.HttpHeaders; + +import java.util.Objects; +import java.util.function.Supplier; + +/** + * Immutable Configuration options for {@link RequestIdPolicy}. + * User can provide a {@link Supplier} to generated id for various id headers to be used in + * {@link com.azure.core.http.HttpRequest}. Examples of such headers are 'x-ms-client-request-id + * or 'x-ms-correlation-request-id' etc. + */ +@Immutable +public class RequestIdPolicyOptions { + + private final Supplier idHeaderSupplier; + + /** + * Creates {@link RequestIdPolicyOptions} with provided {@code idHeaderSupplier }. + * + * @param idHeaderSupplier {@link Supplier} to provide dynamically generated id for various request id headers. + * Examples of such headers are 'x-ms-client-request-id or 'x-ms-correlation-request-id' etc. + */ + public RequestIdPolicyOptions(Supplier idHeaderSupplier) { + this.idHeaderSupplier = Objects.requireNonNull(idHeaderSupplier, + "'idHeaderSupplier' cannot be null."); + } + + /** + * Get the provided {@link Supplier} which dynamically generate id for various request id headers. + * + * @return {@link Supplier} to provide dynamically generated id for various request id headers + */ + public Supplier getIdHeaderSupplier() { + return idHeaderSupplier; + } +}