Adding Request ID for trace-ability

Adding Request ID for trace-ability

sdk/search/azure-search/pom.xml

@@ -52,12 +52,12 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core</artifactId>
-      <version>1.1.0</version>
+      <version>1.2.0-beta.1</version>
     </dependency>
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core-http-netty</artifactId>
-      <version>1.1.0</version>
+      <version>1.2.0-beta.1</version>
     </dependency>
     <dependency>
       <groupId>org.apache.commons</groupId>

sdk/search/azure-search/src/main/java/com/azure/search/SearchIndexClientBuilder.java

@@ -5,14 +5,21 @@
 
 import com.azure.core.annotation.ServiceClientBuilder;
 import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.policy.AddDatePolicy;
+import com.azure.core.http.policy.AddHeadersPolicy;
 import com.azure.core.http.policy.HttpLogDetailLevel;
 import com.azure.core.http.policy.HttpLogOptions;
 import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.http.policy.HttpPolicyProviders;
+import com.azure.core.http.policy.RequestIdPolicy;
+import com.azure.core.http.policy.RetryPolicy;
 import com.azure.core.util.Configuration;
 import com.azure.core.util.logging.ClientLogger;
-
 import org.apache.commons.lang3.StringUtils;
 
+import java.util.List;
+
 /**
  * This class provides a fluent builder API to help aid the configuration and instantiation of
  * {@link SearchIndexClient SearchIndexClients} and {@link SearchIndexAsyncClient SearchIndexAsyncClients}.
@@ -31,14 +38,23 @@
 @ServiceClientBuilder(serviceClients = {SearchIndexClient.class, SearchIndexAsyncClient.class})
 public class SearchIndexClientBuilder extends SearchClientBuilder {
 
-    private final ClientLogger logger = new ClientLogger(SearchIndexClientBuilder.class);
+    // This header tells the server to return the request id in the HTTP response. Useful for correlation with what
+    // request was sent.
+    private static final String ECHO_REQUEST_ID_HEADER = "x-ms-return-client-request-id";
+
     private String indexName;
+    private final HttpHeaders headers;
+    private RetryPolicy retryPolicy;
+    private final ClientLogger logger = new ClientLogger(SearchIndexClientBuilder.class);
 
     /**
      * Default Constructor
      */
     public SearchIndexClientBuilder() {
         init();
+        headers = new HttpHeaders()
+            .put(ECHO_REQUEST_ID_HEADER, "true");
+
     }
 
     /**
@@ -52,6 +68,15 @@ public SearchIndexClientBuilder apiVersion(SearchServiceVersion apiVersion) {
         return this;
     }
 
+    /**
+     *  Returns the list of policies configured for this builder.
+     *
+     *  @return List of HttpPipelinePolicy
+     *  */
+    public List<HttpPipelinePolicy> getPolicies() {
+        return this.policies;
+    }
+
     /**
      * Sets the Azure Cognitive Search service endpoint
      *
@@ -102,6 +127,19 @@ public SearchIndexClientBuilder credential(SearchApiKeyCredential searchApiKeyCr
         return this;
     }
 
+    /**
+     * Sets the configuration store that is used during construction of the service client.
+     * The default configuration store is a clone of the {@link Configuration#getGlobalConfiguration() global
+     * configuration store}, use {@link Configuration#NONE} to bypass using configuration settings during construction.
+     *
+     * @param configuration The configuration store.
+     * @return The updated SearchIndexClientBuilder object.
+     */
+    public SearchIndexClientBuilder configuration(Configuration configuration) {
+        this.configuration = configuration;
+        return this;
+    }
+
     /**
      * Http Pipeline policy
      *
@@ -114,16 +152,16 @@ public SearchIndexClientBuilder addPolicy(HttpPipelinePolicy policy) {
     }
 
     /**
-     * Sets the configuration store that is used during construction of the service client.
+     * Sets the {@link RetryPolicy} that is used when each request is sent.
+     * <p>
+     * The default retry policy will be used if not provided {@link SearchIndexClientBuilder#buildAsyncClient()}
+     * to build {@link SearchServiceAsyncClient} or {@link SearchServiceClient}.
      *
-     * The default configuration store is a clone of the {@link Configuration#getGlobalConfiguration() global
-     * configuration store}, use {@link Configuration#NONE} to bypass using configuration settings during construction.
-     *
-     * @param configuration The configuration store used to
+     * @param retryPolicy RetryPolicy applied to each request.
      * @return The updated SearchIndexClientBuilder object.
      */
-    public SearchIndexClientBuilder configuration(Configuration configuration) {
-        this.configuration = configuration;
+    public SearchIndexClientBuilder retryPolicy(RetryPolicy retryPolicy) {
+        this.retryPolicy = retryPolicy;
         return this;
     }
 
@@ -151,6 +189,16 @@ public SearchIndexClient buildClient() {
      * @return a {@link SearchIndexAsyncClient} created from the configurations in this builder.
      */
     public SearchIndexAsyncClient buildAsyncClient() {
+
+        policies.add(new AddHeadersPolicy(headers));
+        // We need to add RequestId and override the default header, with the one
+        // That the service expects, in order to capture the request ids.
+        policies.add(new RequestIdPolicy("client-request-id"));
+        policies.add(new AddDatePolicy());
+        HttpPolicyProviders.addBeforeRetryPolicies(policies);
+        policies.add(retryPolicy == null ? new RetryPolicy() : retryPolicy);
+        HttpPolicyProviders.addAfterRetryPolicies(policies);
+
         return new SearchIndexAsyncClient(endpoint, indexName, apiVersion, prepareForBuildClient());
     }
 }

sdk/search/azure-search/src/test/java/com/azure/search/SearchIndexClientBuilderTests.java

@@ -3,6 +3,12 @@
 
 package com.azure.search;
 
+import com.azure.core.http.policy.AddDatePolicy;
+import com.azure.core.http.policy.AddHeadersPolicy;
+import com.azure.core.http.policy.HttpLoggingPolicy;
+import com.azure.core.http.policy.RequestIdPolicy;
+import com.azure.core.http.policy.RetryPolicy;
+import com.azure.core.http.policy.UserAgentPolicy;
 import org.junit.Assert;
 import org.junit.jupiter.api.Test;
 
@@ -294,6 +300,46 @@ public void verifyEmptyVersionThrowsIllegalArgumentExceptionAsync() {
         );
     }
 
+    @Test
+    public void whenCreateUsingClientBuilderThenDefaultPoliciesExists() {
+        SearchIndexClientBuilder searchIndexClientBuilder = new SearchIndexClientBuilder();
+        searchIndexClientBuilder
+            .endpoint(searchEndpoint)
+            .indexName(indexName)
+            .credential(searchApiKeyCredential)
+            .apiVersion(apiVersion)
+            .buildAsyncClient();
+
+        int policyCount = searchIndexClientBuilder.getPolicies().size();
+
+        Assert.assertEquals(7, policyCount);
+
+        Assert.assertEquals(1,
+            searchIndexClientBuilder.getPolicies().stream()
+                .filter(p -> p.getClass() == RetryPolicy.class).count()
+        );
+        Assert.assertEquals(1,
+            searchIndexClientBuilder.getPolicies().stream()
+                .filter(p -> p.getClass() == RequestIdPolicy.class).count()
+        );
+        Assert.assertEquals(1,
+            searchIndexClientBuilder.getPolicies().stream()
+                .filter(p -> p.getClass() == UserAgentPolicy.class).count()
+        );
+        Assert.assertEquals(1,
+            searchIndexClientBuilder.getPolicies().stream()
+                .filter(p -> p.getClass() == AddHeadersPolicy.class).count()
+        );
+        Assert.assertEquals(1,
+            searchIndexClientBuilder.getPolicies().stream()
+                .filter(p -> p.getClass() == AddDatePolicy.class).count()
+        );
+        Assert.assertEquals(1,
+            searchIndexClientBuilder.getPolicies().stream()
+                .filter(p -> p.getClass() == HttpLoggingPolicy.class).count()
+        );
+    }
+
     private void expectThrowsWithMessage(String expectedMessage, Runnable runnable) {
         try {
             runnable.run();