Add Tracing context in Key Vault keys package

eng/code-quality-reports/src/main/resources/spotbugs/spotbugs-exclude.xml

@@ -445,6 +445,16 @@
     <Bug pattern="NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE"/>
   </Match>
 
+  <!-- Incorrect flagging, if the response is null a HttpResponseException should be thrown
+  https://github.com/azure/azure-sdk-for-java/issues/4556 -->
+  <Match>
+    <Or>
+      <Class name="com.azure.security.keyvault.keys.KeyClient"/>
+      <Class name="com.azure.security.keyvault.keys.SecretClient"/>
+    </Or>
+    <Bug pattern="NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE"/>
+  </Match>
+
   <!-- Super doesn't need to be cloned since it is Object -->
   <Match>
     <Class name="com.azure.core.util.configuration.Configuration"/>

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxUtil.java

@@ -3,6 +3,7 @@
 
 package com.azure.core.implementation.util;
 
+import com.azure.core.http.rest.Response;
 import com.azure.core.util.Context;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
@@ -199,6 +200,16 @@ public static <T> Mono<T> withContext(Function<Context, Mono<T>> serviceCall) {
             .flatMap(serviceCall);
     }
 
+    /**
+     * Converts the incoming content to Mono.
+     *
+     * @param response whose {@link Response#value() value} is to be converted
+     * @return The converted {@link Mono}
+     */
+    public static <T> Mono<T> toMono(Response<T> response) {
+        return Mono.justOrEmpty(response.value());
+    }
+
     /**
      * This method converts the incoming {@code subscriberContext} from {@link reactor.util.context.Context Reactor
      * Context} to {@link Context Azure Context} and calls the given lambda function with this context and returns a

sdk/core/azure-core/src/samples/java/com/azure/core/implementation/util/FluxUtilJavaDocCodeSnippets.java

@@ -3,6 +3,9 @@
 
 package com.azure.core.implementation.util;
 
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.rest.Response;
 import com.azure.core.util.Context;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -54,4 +57,33 @@ private Mono<String> serviceCallReturnsSingle(String prefix, Context context) {
         return Mono.empty();
     }
 
+    /**
+     * Implementation not provided
+     * @param value The value
+     * @return A {@link Mono} containing a {@link Response} containing a {@link Response#value() value}.
+     */
+    private <T> Mono<Response<T>> getMonoRestResponse(T value) {
+        Response<T> response = new Response<T>() {
+            @Override
+            public int statusCode() {
+                return 200;
+            }
+
+            @Override
+            public HttpHeaders headers() {
+                return null;
+            }
+
+            @Override
+            public HttpRequest request() {
+                return null;
+            }
+
+            @Override
+            public T value() {
+                return value;
+            }
+        };
+        return Mono.just(response);
+    }
 }

sdk/core/azure-core/src/test/java/com/azure/core/implementation/util/FluxUtilTests.java

@@ -12,6 +12,7 @@
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.rest.PagedFlux;
 import com.azure.core.http.rest.PagedResponse;
+import com.azure.core.http.rest.Response;
 import com.azure.core.implementation.http.PagedResponseBase;
 import com.azure.core.util.Context;
 import io.netty.buffer.ByteBuf;
@@ -278,6 +279,13 @@ public void testCollectByteBufStream() {
             0, 0, 1, 0}, bytes);
     }
 
+    @Test
+    public void testToMono() {
+        String value = "test";
+        Assert.assertEquals(getMonoRestResponse(value).flatMap(FluxUtil::toMono).block(), value);
+        Assert.assertEquals(getMonoRestResponse("").flatMap(FluxUtil::toMono).block(), "");
+    }
+
     @Test
     public void testCallWithContextGetSingle() {
         String response = getSingle("Hello, ")
@@ -401,4 +409,29 @@ private File createFileIfNotExist(String fileName) throws IOException {
         return file;
     }
 
+    private <T> Mono<Response<T>> getMonoRestResponse(T value) {
+        Response<T> response = new Response<T>() {
+            @Override
+            public int statusCode() {
+                return 200;
+            }
+
+            @Override
+            public HttpHeaders headers() {
+                return null;
+            }
+
+            @Override
+            public HttpRequest request() {
+                return null;
+            }
+
+            @Override
+            public T value() {
+                return value;
+            }
+        };
+        return Mono.just(response);
+    }
+
 }

sdk/keyvault/azure-keyvault-keys/src/main/java/com/azure/security/keyvault/keys/KeyClientBuilder.java

@@ -4,6 +4,7 @@
 package com.azure.security.keyvault.keys;
 
 import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.implementation.http.policy.spi.HttpPolicyProviders;
 import com.azure.core.implementation.util.ImplUtils;
 import com.azure.core.util.configuration.Configuration;
 import com.azure.core.credentials.TokenCredential;
@@ -89,6 +90,7 @@ public KeyClientBuilder() {
     public KeyClient buildClient() {
         return new KeyClient(buildAsyncClient());
     }
+
     /**
      * Creates a {@link KeyAsyncClient} based on options set in the builder.
      * Every time {@code buildAsyncClient()} is called, a new instance of {@link KeyAsyncClient} is created.
@@ -122,9 +124,11 @@ public KeyAsyncClient buildAsyncClient() {
         // Closest to API goes first, closest to wire goes last.
         final List<HttpPipelinePolicy> policies = new ArrayList<>();
         policies.add(new UserAgentPolicy(AzureKeyVaultConfiguration.SDK_NAME, AzureKeyVaultConfiguration.SDK_VERSION, buildConfiguration));
+        HttpPolicyProviders.addBeforeRetryPolicies(policies);
         policies.add(retryPolicy);
         policies.add(new BearerTokenAuthenticationPolicy(credential, KeyAsyncClient.KEY_VAULT_SCOPE));
         policies.addAll(this.policies);
+        HttpPolicyProviders.addAfterRetryPolicies(policies);
         policies.add(new HttpLoggingPolicy(httpLogDetailLevel));
 
         HttpPipeline pipeline = new HttpPipelineBuilder()

sdk/keyvault/azure-keyvault-keys/src/samples/java/com/azure/security/keyvault/keys/BackupAndRestoreOperations.java

@@ -45,7 +45,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
         // Backups are good to have, if in case keys get accidentally deleted by you.
         // For long term storage, it is ideal to write the backup to a file.
         String backupFilePath = "YOUR_BACKUP_FILE_PATH";
-        byte[] keyBackup = keyClient.backupKey("CloudRsaKey").value();
+        byte[] keyBackup = keyClient.backupKey("CloudRsaKey");
         writeBackupToFile(keyBackup, backupFilePath);
 
         // The Cloud Rsa key is no longer in use, so you delete it.
@@ -62,7 +62,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
 
         // After sometime, the key is required again. We can use the backup value to restore it in the key vault.
         byte[] backupFromFile = Files.readAllBytes(new File(backupFilePath).toPath());
-        Key restoredKey = keyClient.restoreKey(backupFromFile).value();
+        Key restoredKey = keyClient.restoreKey(backupFromFile);
     }
 
     private static void writeBackupToFile(byte[] bytes, String filePath) {

sdk/keyvault/azure-keyvault-keys/src/samples/java/com/azure/security/keyvault/keys/BackupAndRestoreOperationsAsync.java

@@ -41,23 +41,23 @@ public static void main(String[] args) throws IOException, InterruptedException,
                 .expires(OffsetDateTime.now().plusYears(1))
                 .keySize(2048))
                 .subscribe(keyResponse ->
-                        System.out.printf("Key is created with name %s and type %s \n", keyResponse.value().name(), keyResponse.value().keyMaterial().kty()));
+                        System.out.printf("Key is created with name %s and type %s \n", keyResponse.name(), keyResponse.keyMaterial().kty()));
 
         Thread.sleep(2000);
 
         // Backups are good to have, if in case keys get accidentally deleted by you.
         // For long term storage, it is ideal to write the backup to a file.
         String backupFilePath = "YOUR_BACKUP_FILE_PATH";
         keyAsyncClient.backupKey("CloudRsaKey").subscribe(backupResponse -> {
-            byte[] backupBytes = backupResponse.value();
+            byte[] backupBytes = backupResponse;
             writeBackupToFile(backupBytes, backupFilePath);
         });
 
         Thread.sleep(7000);
 
         // The Cloud Rsa key is no longer in use, so you delete it.
         keyAsyncClient.deleteKey("CloudRsaKey").subscribe(deletedKeyResponse ->
-                System.out.printf("Deleted Key's Recovery Id %s \n", deletedKeyResponse.value().recoveryId()));
+                System.out.printf("Deleted Key's Recovery Id %s \n", deletedKeyResponse.recoveryId()));
 
         //To ensure file is deleted on server side.
         Thread.sleep(30000);
@@ -72,7 +72,7 @@ public static void main(String[] args) throws IOException, InterruptedException,
         // After sometime, the key is required again. We can use the backup value to restore it in the key vault.
         byte[] backupFromFile = Files.readAllBytes(new File(backupFilePath).toPath());
         keyAsyncClient.restoreKey(backupFromFile).subscribe(keyResponse ->
-            System.out.printf("Restored Key with name %s \n", keyResponse.value().name()));
+            System.out.printf("Restored Key with name %s \n", keyResponse.name()));
 
         //To ensure key is restored on server side.
         Thread.sleep(15000);

sdk/keyvault/azure-keyvault-keys/src/samples/java/com/azure/security/keyvault/keys/HelloWorld.java

@@ -38,13 +38,14 @@ public static void main(String[] args) throws InterruptedException, IllegalArgum
                 .keySize(2048));
 
         // Let's Get the Cloud Rsa Key from the key vault.
-        Key cloudRsaKey = keyClient.getKey("CloudRsaKey").value();
-        System.out.printf("Key is returned with name %s and type %s \n", cloudRsaKey.name(), cloudRsaKey.keyMaterial().kty());
+        Key cloudRsaKey = keyClient.getKey("CloudRsaKey");
+        System.out.printf("Key is returned with name %s and type %s \n", cloudRsaKey.name(),
+            cloudRsaKey.keyMaterial().kty());
 
         // After one year, the Cloud Rsa Key is still required, we need to update the expiry time of the key.
         // The update method can be used to update the expiry attribute of the key.
         cloudRsaKey.expires(cloudRsaKey.expires().plusYears(1));
-        Key updatedKey = keyClient.updateKey(cloudRsaKey).value();
+        Key updatedKey = keyClient.updateKey(cloudRsaKey);
         System.out.printf("Key's updated expiry time %s \n", updatedKey.expires());
 
         // We need the Cloud Rsa key with bigger key size, so you want to update the key in key vault to ensure it has the required size.

sdk/keyvault/azure-keyvault-keys/src/samples/java/com/azure/security/keyvault/keys/HelloWorldAsync.java

@@ -36,25 +36,25 @@ public static void main(String[] args) throws InterruptedException {
                 .expires(OffsetDateTime.now().plusYears(1))
                 .keySize(2048))
                 .subscribe(keyResponse ->
-                        System.out.printf("Key is created with name %s and type %s \n", keyResponse.value().name(), keyResponse.value().keyMaterial().kty()));
+                        System.out.printf("Key is created with name %s and type %s \n", keyResponse.name(), keyResponse.keyMaterial().kty()));
 
         Thread.sleep(2000);
 
         // Let's Get the Cloud Rsa Key from the key vault.
         keyAsyncClient.getKey("CloudRsaKey").subscribe(keyResponse ->
-                System.out.printf("Key returned with name %s and type %s \n", keyResponse.value().name(), keyResponse.value().keyMaterial().kty()));
+                System.out.printf("Key returned with name %s and type %s \n", keyResponse.name(), keyResponse.keyMaterial().kty()));
 
         Thread.sleep(2000);
 
 
         // After one year, the Cloud Rsa Key is still required, we need to update the expiry time of the key.
         // The update method can be used to update the expiry attribute of the key.
         keyAsyncClient.getKey("CloudRsaKey").subscribe(keyResponse -> {
-            Key key = keyResponse.value();
+            Key key = keyResponse;
             //Update the expiry time of the key.
             key.expires(key.expires().plusYears(1));
             keyAsyncClient.updateKey(key).subscribe(updatedKeyResponse ->
-                System.out.printf("Key's updated expiry time %s \n", updatedKeyResponse.value().expires().toString()));
+                System.out.printf("Key's updated expiry time %s \n", updatedKeyResponse.expires().toString()));
         });
 
         Thread.sleep(2000);
@@ -65,13 +65,13 @@ public static void main(String[] args) throws InterruptedException {
                 .expires(OffsetDateTime.now().plusYears(1))
                 .keySize(4096))
                 .subscribe(keyResponse ->
-                        System.out.printf("Key is created with name %s and type %s \n", keyResponse.value().name(), keyResponse.value().keyMaterial().kty()));
+                        System.out.printf("Key is created with name %s and type %s \n", keyResponse.name(), keyResponse.keyMaterial().kty()));
 
         Thread.sleep(2000);
 
         // The Cloud Rsa Key is no longer needed, need to delete it from the key vault.
         keyAsyncClient.deleteKey("CloudRsaKey").subscribe(deletedKeyResponse ->
-            System.out.printf("Deleted Key's Recovery Id %s \n", deletedKeyResponse.value().recoveryId()));
+            System.out.printf("Deleted Key's Recovery Id %s \n", deletedKeyResponse.recoveryId()));
 
         //To ensure key is deleted on server side.
         Thread.sleep(30000);