Updating Schema Registry to cache responses. Adding tests

sdk/schemaregistry/azure-data-schemaregistry/CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## 1.0.0-beta.5 (Unreleased)
 
+### Features added
+- Caching schemas based on schemaId.
+
+### Breaking changes
+- 4xx responses return their respective `HttpResponseException` rather than `IllegalStateException`
+- Removed `Response<T>` overloads for getSchema and getSchemaId because response could be cached.
+- `SchemaRegistryClientBuilder.maxCacheSize` is package-private.
 
 ## 1.0.0-beta.4 (2020-09-21)
 - Minor code cleanup and refactoring

sdk/schemaregistry/azure-data-schemaregistry/pom.xml

@@ -87,6 +87,12 @@
       <version>1.10.1</version> <!-- {x-version-update;org.apache.avro:avro;external_dependency} -->
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-core-test</artifactId>
+      <version>1.7.0</version> <!-- {x-version-update;com.azure:azure-core-test;dependency} -->
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-identity</artifactId>

sdk/schemaregistry/azure-data-schemaregistry/src/main/java/com/azure/data/schemaregistry/SchemaRegistryAsyncClient.java

@@ -9,9 +9,13 @@
 import com.azure.core.http.rest.Response;
 import com.azure.core.http.rest.SimpleResponse;
 import com.azure.core.util.Context;
+import com.azure.core.util.FluxUtil;
 import com.azure.core.util.logging.ClientLogger;
+import com.azure.data.schemaregistry.implementation.AzureSchemaRegistry;
+import com.azure.data.schemaregistry.implementation.models.SchemaId;
 import com.azure.data.schemaregistry.models.SchemaProperties;
 import com.azure.data.schemaregistry.models.SerializationType;
+import reactor.core.publisher.Mono;
 
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
@@ -21,23 +25,16 @@
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.function.Function;
 
-import com.azure.data.schemaregistry.implementation.AzureSchemaRegistry;
-import com.azure.data.schemaregistry.implementation.models.SchemaId;
-import reactor.core.publisher.Mono;
-
 /**
  * HTTP-based client that interacts with Azure Schema Registry service to store and retrieve schemas on demand.
  * @see SchemaRegistryClientBuilder Follows builder pattern for object instantiation
  */
 @ServiceClient(builder = SchemaRegistryClientBuilder.class, isAsync = true)
 public final class SchemaRegistryAsyncClient {
 
-    private final ClientLogger logger = new ClientLogger(SchemaRegistryAsyncClient.class);
-
     static final Charset SCHEMA_REGISTRY_SERVICE_ENCODING = StandardCharsets.UTF_8;
-    static final int MAX_SCHEMA_MAP_SIZE_DEFAULT = 1000;
-    static final int MAX_SCHEMA_MAP_SIZE_MINIMUM = 10;
 
+    private final ClientLogger logger = new ClientLogger(SchemaRegistryAsyncClient.class);
     private final AzureSchemaRegistry restService;
     private final Integer maxSchemaMapSize;
     private final ConcurrentSkipListMap<String, Function<String, Object>> typeParserMap;
@@ -69,15 +66,6 @@ public final class SchemaRegistryAsyncClient {
     @ServiceMethod(returns = ReturnType.SINGLE)
     public Mono<SchemaProperties> registerSchema(
         String schemaGroup, String schemaName, String schemaString, SerializationType serializationType) {
-
-        if (schemaStringCache.containsKey(getSchemaStringCacheKey(schemaGroup, schemaName, schemaString))) {
-            logger.verbose(
-                "Cache hit schema string. Group: '{}', name: '{}', schema type: '{}', payload: '{}'",
-                schemaGroup, schemaName, serializationType, schemaString);
-            return Mono.fromCallable(
-                () -> schemaStringCache.get(getSchemaStringCacheKey(schemaGroup, schemaName, schemaString)));
-        }
-
         return registerSchemaWithResponse(schemaGroup, schemaName, schemaString, serializationType)
             .map(Response::getValue);
     }
@@ -95,42 +83,29 @@ public Mono<SchemaProperties> registerSchema(
      */
     @ServiceMethod(returns = ReturnType.SINGLE)
     public Mono<Response<SchemaProperties>> registerSchemaWithResponse(String schemaGroup, String schemaName,
-                                   String schemaString, SerializationType serializationType) {
-        return registerSchemaWithResponse(schemaGroup, schemaName, schemaString, serializationType, Context.NONE);
+                                    String schemaString, SerializationType serializationType) {
+        return FluxUtil.withContext(context -> registerSchemaWithResponse(schemaGroup, schemaName, schemaString,
+            serializationType, context));
     }
 
     Mono<Response<SchemaProperties>> registerSchemaWithResponse(String schemaGroup, String schemaName,
                                     String schemaString, SerializationType serializationType, Context context) {
-        logger.verbose(
-            "Registering schema. Group: '{}', name: '{}', serialization type: '{}', payload: '{}'",
+        logger.verbose("Registering schema. Group: '{}', name: '{}', serialization type: '{}', payload: '{}'",
             schemaGroup, schemaName, serializationType, schemaString);
 
-        return this.restService
-            .getSchemas().registerWithResponseAsync(schemaGroup, schemaName,
-                com.azure.data.schemaregistry.implementation.models.SerializationType.AVRO, schemaString)
+        return this.restService.getSchemas().registerWithResponseAsync(schemaGroup, schemaName,
+            com.azure.data.schemaregistry.implementation.models.SerializationType.AVRO, schemaString)
             .handle((response, sink) -> {
-                if (response == null) {
-                    sink.error(logger.logExceptionAsError(
-                        new NullPointerException("Client returned null response")));
-                    return;
-                }
-
-                if (response.getStatusCode() == 400) {
-                    sink.error(logger.logExceptionAsError(
-                        new IllegalStateException("Invalid schema registration attempted")));
-                    return;
-                }
-
                 SchemaId schemaId = response.getValue();
-
                 SchemaProperties registered = new SchemaProperties(schemaId.getId(),
                     serializationType,
                     schemaName,
                     schemaString.getBytes(SCHEMA_REGISTRY_SERVICE_ENCODING));
 
-                resetIfNeeded();
-                schemaStringCache
-                    .putIfAbsent(getSchemaStringCacheKey(schemaGroup, schemaName, schemaString), registered);
+                schemaStringCache.putIfAbsent(getSchemaStringCacheKey(schemaGroup, schemaName, schemaString),
+                    registered);
+                idCache.putIfAbsent(schemaId.getId(), registered);
+
                 logger.verbose("Cached schema string. Group: '{}', name: '{}'", schemaGroup, schemaName);
                 SimpleResponse<SchemaProperties> schemaRegistryObjectSimpleResponse = new SimpleResponse<>(
                     response.getRequest(), response.getStatusCode(),
@@ -161,36 +136,22 @@ public Mono<SchemaProperties> getSchema(String schemaId) {
      * @return The {@link SchemaProperties} associated with the given {@code schemaId} along with the HTTP
      * response.
      */
-    @ServiceMethod(returns = ReturnType.SINGLE)
-    public Mono<Response<SchemaProperties>> getSchemaWithResponse(String schemaId) {
-        return getSchemaWithResponse(schemaId, Context.NONE);
+    Mono<Response<SchemaProperties>> getSchemaWithResponse(String schemaId) {
+        return FluxUtil.withContext(context -> getSchemaWithResponse(schemaId, context));
     }
 
     Mono<Response<SchemaProperties>> getSchemaWithResponse(String schemaId, Context context) {
         Objects.requireNonNull(schemaId, "'schemaId' should not be null");
         return this.restService.getSchemas().getByIdWithResponseAsync(schemaId)
             .handle((response, sink) -> {
-                if (response == null) {
-                    sink.error(logger.logExceptionAsError(
-                        new NullPointerException("Client returned null response")));
-                    return;
-                }
-
-                if (response.getStatusCode() == 404) {
-                    sink.error(logger.logExceptionAsError(
-                        new IllegalStateException(String.format("Schema does not exist, id %s", schemaId))));
-                    return;
-                }
-
                 SerializationType serializationType =
-                    SerializationType.fromString(response.getDeserializedHeaders().getXSchemaType());
+                    SerializationType.fromString(response.getDeserializedHeaders().getSchemaType());
 
                 SchemaProperties schemaObject = new SchemaProperties(schemaId,
                     serializationType,
                     null,
-                    response.getValue().getBytes(SCHEMA_REGISTRY_SERVICE_ENCODING));
+                    response.getValue());
 
-                resetIfNeeded();
                 idCache.putIfAbsent(schemaId, schemaObject);
                 logger.verbose("Cached schema object. Path: '{}'", schemaId);
                 SimpleResponse<SchemaProperties> schemaRegistryObjectSimpleResponse = new SimpleResponse<>(
@@ -201,7 +162,8 @@ Mono<Response<SchemaProperties>> getSchemaWithResponse(String schemaId, Context
     }
 
     /**
-     * Gets the schema identifier associated with the given schema.
+     * Gets the schema identifier associated with the given schema. Gets a cached value if it exists, otherwise makes a
+     * call to the service.
      *
      * @param schemaGroup The schema group.
      * @param schemaName The schema name.
@@ -212,20 +174,23 @@ Mono<Response<SchemaProperties>> getSchemaWithResponse(String schemaId, Context
      */
     @ServiceMethod(returns = ReturnType.SINGLE)
     public Mono<String> getSchemaId(String schemaGroup, String schemaName, String schemaString,
-                                    SerializationType serializationType) {
+        SerializationType serializationType) {
 
         String schemaStringCacheKey = getSchemaStringCacheKey(schemaGroup, schemaName, schemaString);
+
         if (schemaStringCache.containsKey(schemaStringCacheKey)) {
-            logger.verbose("Cache hit schema string. Group: '{}', name: '{}'", schemaGroup, schemaName);
-            return Mono.fromCallable(() -> schemaStringCache.get(schemaStringCacheKey).getSchemaId());
+            return Mono.fromCallable(() -> {
+                logger.verbose("Cache hit schema string. Group: '{}', name: '{}'", schemaGroup, schemaName);
+                return schemaStringCache.get(schemaStringCacheKey).getSchemaId();
+            });
         }
 
         return getSchemaIdWithResponse(schemaGroup, schemaName, schemaString, serializationType)
             .map(response -> response.getValue());
     }
 
     /**
-     * Gets the schema identifier associated with the given schema.
+     * Gets the schema identifier associated with the given schema. Always makes a call to the service.
      *
      * @param schemaGroup The schema group.
      * @param schemaName The schema name.
@@ -236,39 +201,37 @@ public Mono<String> getSchemaId(String schemaGroup, String schemaName, String sc
      */
     @ServiceMethod(returns = ReturnType.SINGLE)
     public Mono<Response<String>> getSchemaIdWithResponse(String schemaGroup, String schemaName, String schemaString,
-                                                          SerializationType serializationType) {
-        return getSchemaIdWithResponse(schemaGroup, schemaName, schemaString, serializationType, Context.NONE);
+        SerializationType serializationType) {
+
+        return FluxUtil.withContext(context ->
+            getSchemaIdWithResponse(schemaGroup, schemaName, schemaString, serializationType, context));
     }
 
+    /**
+     * Gets the schema id associated with the schema name a string representation of the schema.
+     *
+     * @param schemaGroup The schema group.
+     * @param schemaName The schema name.
+     * @param schemaString The string representation of the schema.
+     * @param serializationType The serialization type of this schema.
+     * @param context Context to pass along with this request.
+     * @return A mono that completes with the schema id.
+     */
     Mono<Response<String>> getSchemaIdWithResponse(String schemaGroup, String schemaName, String schemaString,
-                                                   SerializationType serializationType, Context context) {
+        SerializationType serializationType, Context context) {
 
         return this.restService.getSchemas()
             .queryIdByContentWithResponseAsync(schemaGroup, schemaName,
                 com.azure.data.schemaregistry.implementation.models.SerializationType.AVRO, schemaString)
             .handle((response, sink) -> {
-                if (response == null) {
-                    sink.error(logger.logExceptionAsError(
-                        new NullPointerException("Client returned null response")));
-                    return;
-                }
-
-                if (response.getStatusCode() == 404) {
-                    sink.error(
-                        logger.logExceptionAsError(new IllegalStateException("Existing matching schema not found.")));
-                    return;
-                }
-
                 SchemaId schemaId = response.getValue();
+                SchemaProperties properties = new SchemaProperties(schemaId.getId(), serializationType, schemaName,
+                    schemaString.getBytes(SCHEMA_REGISTRY_SERVICE_ENCODING));
 
-                resetIfNeeded();
                 schemaStringCache.putIfAbsent(
-                    getSchemaStringCacheKey(schemaGroup, schemaName, schemaString),
-                    new SchemaProperties(
-                        schemaId.getId(),
-                        serializationType,
-                        schemaName,
-                        schemaString.getBytes(SCHEMA_REGISTRY_SERVICE_ENCODING)));
+                    getSchemaStringCacheKey(schemaGroup, schemaName, schemaString), properties);
+                idCache.putIfAbsent(schemaId.getId(), properties);
+
                 logger.verbose("Cached schema string. Group: '{}', name: '{}'", schemaGroup, schemaName);
 
                 SimpleResponse<String> schemaIdResponse = new SimpleResponse<>(
@@ -287,8 +250,6 @@ void clearCache() {
         typeParserMap.clear();
     }
 
-    // TODO: max age for schema maps? or will schemas always be immutable?
-
     /**
      * Checks if caches should be reinitialized to satisfy initial configuration
      */
@@ -304,7 +265,7 @@ private void resetIfNeeded() {
         }
     }
 
-    private String getSchemaStringCacheKey(String schemaGroup, String schemaName, String schemaString) {
+    private static String getSchemaStringCacheKey(String schemaGroup, String schemaName, String schemaString) {
         return schemaGroup + schemaName + schemaString;
     }
 }

sdk/schemaregistry/azure-data-schemaregistry/src/main/java/com/azure/data/schemaregistry/SchemaRegistryClient.java

@@ -78,7 +78,7 @@ public SchemaProperties getSchema(String schemaId) {
      * response.
      */
     @ServiceMethod(returns = ReturnType.SINGLE)
-    public Response<SchemaProperties> getSchemaWithResponse(String schemaId, Context context) {
+    Response<SchemaProperties> getSchemaWithResponse(String schemaId, Context context) {
         return this.asyncClient.getSchemaWithResponse(schemaId, context).block();
     }
 
@@ -110,7 +110,7 @@ public String getSchemaId(String schemaGroup, String schemaName, String schemaSt
      * @return The unique identifier for this schema.
      */
     @ServiceMethod(returns = ReturnType.SINGLE)
-    public Response<String> getSchemaIdWithResponse(String schemaGroup, String schemaName, String schemaString,
+    Response<String> getSchemaIdWithResponse(String schemaGroup, String schemaName, String schemaString,
         SerializationType serializationType, Context context) {
         return this.asyncClient
             .getSchemaIdWithResponse(schemaGroup, schemaName, schemaString, serializationType, context).block();
@@ -119,5 +119,4 @@ public Response<String> getSchemaIdWithResponse(String schemaGroup, String schem
     void clearCache() {
         this.asyncClient.clearCache();
     }
-
 }

sdk/schemaregistry/azure-data-schemaregistry/src/main/java/com/azure/data/schemaregistry/SchemaRegistryClientBuilder.java

@@ -45,6 +45,9 @@
  */
 @ServiceClientBuilder(serviceClients = SchemaRegistryAsyncClient.class)
 public class SchemaRegistryClientBuilder {
+    static final int MAX_SCHEMA_MAP_SIZE_DEFAULT = 1000;
+    static final int MAX_SCHEMA_MAP_SIZE_MINIMUM = 10;
+
     private final ClientLogger logger = new ClientLogger(SchemaRegistryClientBuilder.class);
 
     private static final String DEFAULT_SCOPE = "https://eventhubs.azure.net/.default";
@@ -125,11 +128,11 @@ public SchemaRegistryClientBuilder endpoint(String endpoint) {
      * @return The updated {@link SchemaRegistryClientBuilder} object.
      * @throws IllegalArgumentException on invalid maxCacheSize value
      */
-    public SchemaRegistryClientBuilder maxCacheSize(int maxCacheSize) {
-        if (maxCacheSize < SchemaRegistryAsyncClient.MAX_SCHEMA_MAP_SIZE_MINIMUM) {
+    SchemaRegistryClientBuilder maxCacheSize(int maxCacheSize) {
+        if (maxCacheSize < MAX_SCHEMA_MAP_SIZE_MINIMUM) {
             throw logger.logExceptionAsError(new IllegalArgumentException(
                 String.format("Schema map size must be greater than %s entries",
-                    SchemaRegistryAsyncClient.MAX_SCHEMA_MAP_SIZE_MINIMUM)));
+                    MAX_SCHEMA_MAP_SIZE_MINIMUM)));
         }
 
         this.maxSchemaMapSize = maxCacheSize;
@@ -319,7 +322,7 @@ public SchemaRegistryAsyncClient buildAsyncClient() {
             .buildClient();
 
         int buildMaxSchemaMapSize = (maxSchemaMapSize == null)
-            ? SchemaRegistryAsyncClient.MAX_SCHEMA_MAP_SIZE_DEFAULT
+            ? MAX_SCHEMA_MAP_SIZE_DEFAULT
             : maxSchemaMapSize;
 
         return new SchemaRegistryAsyncClient(restService, buildMaxSchemaMapSize, typeParserMap);