DGS-23340 Ensure output schema is cached if it is returned

Author: rayokota

client/src/main/java/io/confluent/kafka/schemaregistry/client/CachedSchemaRegistryClient.java

@@ -543,12 +543,16 @@ private RegisterSchemaResponse registerWithResponse(
             ? registerAndGetId(subject, schema, version, id, normalize, propagateSchemaTags)
             : registerAndGetId(subject, schema, normalize, propagateSchemaTags);
         schemaResponseMap.put(schema, retrievedResponse);
-        String context = toQualifiedContext(subject);
-        final Cache<Integer, ParsedSchema> idSchemaMap = idToSchemaCache.get(
-            context, () -> CacheBuilder.newBuilder()
-                .maximumSize(cacheCapacity)
-                .build());
-        idSchemaMap.put(retrievedResponse.getId(), schema);
+        if (retrievedResponse.getSchema() != null) {
+          String context = toQualifiedContext(subject);
+          final Cache<Integer, ParsedSchema> idSchemaMap = idToSchemaCache.get(
+              context, () -> CacheBuilder.newBuilder()
+                  .maximumSize(cacheCapacity)
+                  .build());
+          ParsedSchema retrievedSchema =
+              parseSchemaOrElseThrow(new Schema(null, retrievedResponse));
+          idSchemaMap.put(retrievedResponse.getId(), retrievedSchema);
+        }
         return retrievedResponse;
       }
     } catch (ExecutionException e) {
@@ -826,14 +830,20 @@ public int getId(String subject, ParsedSchema schema, boolean normalize)
           return cachedId;
         }
 
-        final int retrievedId = getIdFromRegistry(subject, schema, normalize);
+        final RegisterSchemaResponse retrievedResponse =
+            getIdWithResponseFromRegistry(subject, schema, normalize, false);
+        final int retrievedId = retrievedResponse.getId();
         schemaIdMap.put(schema, retrievedId);
-        String context = toQualifiedContext(subject);
-        final Cache<Integer, ParsedSchema> idSchemaMap = idToSchemaCache.get(
-            context, () -> CacheBuilder.newBuilder()
-                .maximumSize(cacheCapacity)
-                .build());
-        idSchemaMap.put(retrievedId, schema);
+        if (retrievedResponse.getSchema() != null) {
+          String context = toQualifiedContext(subject);
+          final Cache<Integer, ParsedSchema> idSchemaMap = idToSchemaCache.get(
+              context, () -> CacheBuilder.newBuilder()
+                  .maximumSize(cacheCapacity)
+                  .build());
+          ParsedSchema retrievedSchema =
+              parseSchemaOrElseThrow(new Schema(null, retrievedResponse));
+          idSchemaMap.put(retrievedId, retrievedSchema);
+        }
         return retrievedId;
       }
     } catch (ExecutionException e) {
@@ -908,12 +918,16 @@ public RegisterSchemaResponse getIdWithResponse(
         final RegisterSchemaResponse retrievedResponse =
             getIdWithResponseFromRegistry(subject, schema, normalize, false);
         schemaResponseMap.put(schema, retrievedResponse);
-        String context = toQualifiedContext(subject);
-        final Cache<Integer, ParsedSchema> idSchemaMap = idToSchemaCache.get(
-            context, () -> CacheBuilder.newBuilder()
-                .maximumSize(cacheCapacity)
-                .build());
-        idSchemaMap.put(retrievedResponse.getId(), schema);
+        if (retrievedResponse.getSchema() != null) {
+          String context = toQualifiedContext(subject);
+          final Cache<Integer, ParsedSchema> idSchemaMap = idToSchemaCache.get(
+              context, () -> CacheBuilder.newBuilder()
+                  .maximumSize(cacheCapacity)
+                  .build());
+          ParsedSchema retrievedSchema =
+              parseSchemaOrElseThrow(new Schema(null, retrievedResponse));
+          idSchemaMap.put(retrievedResponse.getId(), retrievedSchema);
+        }
         return retrievedResponse;
       }
     } catch (ExecutionException e) {

client/src/test/java/io/confluent/kafka/schemaregistry/client/CachedSchemaRegistryClientTest.java

@@ -69,6 +69,8 @@ public class CachedSchemaRegistryClientTest {
   private static final int CACHE_CAPACITY = 5;
   private static final String SCHEMA_STR_0 = avroSchemaString(0);
   private static final AvroSchema AVRO_SCHEMA_0 = avroSchema(0);
+  private static final AvroSchema AVRO_SCHEMA_0_WITH_METADATA =
+      AVRO_SCHEMA_0.copy(new Metadata(null, ImmutableMap.of("confluent:version", "1"), null), null);
   private static final AvroSchema SCHEMA_WITH_DECIMAL = new AvroSchema(
           "{\n"
               + "    \"type\": \"record\",\n"
@@ -258,6 +260,21 @@ public void testRegisterFollowedByLookupWillSkipCache() throws Exception {
     verify(restService);
   }
 
+  @Test
+  public void testRegisterFollowedByIdLookupWillReturnRegisteredSchema() throws Exception {
+    expect(restService.registerSchema(anyObject(RegisterSchemaRequest.class),
+        eq(SUBJECT_0), anyBoolean()))
+        .andReturn(new RegisterSchemaResponse(new Schema(SUBJECT_0, 1, ID_25, AVRO_SCHEMA_0_WITH_METADATA)))
+        .once();
+
+    replay(restService);
+
+    assertEquals(ID_25, client.register(SUBJECT_0, AVRO_SCHEMA_0, VERSION_1, ID_25));
+    assertEquals(AVRO_SCHEMA_0_WITH_METADATA, client.getSchemaById(ID_25));
+
+    verify(restService);
+  }
+
   @Test
   public void testRegisterOverCapacity() throws Exception {
     expect(restService.registerSchema(anyObject(RegisterSchemaRequest.class),