Wrap non-required $ref properties in an object to mark as nullable (d…

…atahub-project#10514)

Co-authored-by: tjin <tjin@netflix.com>

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java

@@ -20,8 +20,8 @@
 import io.swagger.v3.oas.models.responses.ApiResponse;
 import io.swagger.v3.oas.models.responses.ApiResponses;
 import java.math.BigDecimal;
-import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -441,16 +441,23 @@ private static void addAspectSchemas(final Components components, final AspectSp
                   final String newDefinition =
                       definition.replaceAll("definitions", "components/schemas");
                   Schema s = Json.mapper().readValue(newDefinition, Schema.class);
-                  // Set nullable attribute
-                  Optional.ofNullable(s.getProperties())
-                      .orElse(new HashMap())
-                      .forEach(
-                          (name, schema) ->
-                              ((Schema) schema)
-                                  .setNullable(
-                                      !Optional.ofNullable(s.getRequired())
-                                          .orElse(new ArrayList())
-                                          .contains(name)));
+                  Set<String> requiredNames = Optional.ofNullable(s.getRequired())
+                      .map(names -> Set.copyOf(names))
+                      .orElse(new HashSet());
+                  Map<String, Schema> properties = Optional.ofNullable(s.getProperties()).orElse(new HashMap<>());
+                  properties.forEach((name, schema) -> {
+                    String $ref = schema.get$ref();
+                    boolean isNameRequired = requiredNames.contains(name);
+                    if ($ref != null && !isNameRequired) {
+                      // A non-required $ref property must be wrapped in a { allOf: [ $ref ] } object to allow the
+                      // property to be marked as nullable
+                      schema.setType("object");
+                      schema.set$ref(null);
+                      schema.setAllOf(List.of(new Schema().$ref($ref)));
+                    }
+                    schema.setNullable(!isNameRequired);
+                  });
+
                   components.addSchemas(n, s);
                 } catch (Exception e) {
                   throw new RuntimeException(e);

metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/OpenAPIV3GeneratorTest.java

@@ -1,5 +1,8 @@
 package io.datahubproject.openapi.v3;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 
 import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor;
@@ -9,36 +12,65 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import io.swagger.v3.oas.models.media.Schema;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
 public class OpenAPIV3GeneratorTest {
   public static final String BASE_DIRECTORY =
-      "../../entity-registry/custom-test-model/build/plugins/models";
+          "../../entity-registry/custom-test-model/build/plugins/models";
 
   @BeforeTest
   public void disableAssert() {
     PathSpecBasedSchemaAnnotationVisitor.class
-        .getClassLoader()
-        .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false);
+            .getClassLoader()
+            .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false);
   }
 
   @Test
   public void testOpenApiSpecBuilder() throws Exception {
     ConfigEntityRegistry er =
-        new ConfigEntityRegistry(
-            OpenAPIV3GeneratorTest.class
-                .getClassLoader()
-                .getResourceAsStream("entity-registry.yml"));
+            new ConfigEntityRegistry(
+                    OpenAPIV3GeneratorTest.class
+                            .getClassLoader()
+                            .getResourceAsStream("entity-registry.yml"));
 
     OpenAPI openAPI = OpenAPIV3Generator.generateOpenApiSpec(er);
     String openapiYaml = Yaml.pretty(openAPI);
     Files.write(
-        Path.of(getClass().getResource("/").getPath(), "open-api.yaml"),
-        openapiYaml.getBytes(StandardCharsets.UTF_8));
+            Path.of(getClass().getResource("/").getPath(), "open-api.yaml"),
+            openapiYaml.getBytes(StandardCharsets.UTF_8));
 
     assertTrue(openAPI.getComponents().getSchemas().size() > 900);
     assertTrue(openAPI.getComponents().getParameters().size() > 50);
     assertTrue(openAPI.getPaths().size() > 500);
+
+    Schema datasetPropertiesSchema = openAPI.getComponents().getSchemas().get("DatasetProperties");
+    List<String> requiredNames = datasetPropertiesSchema.getRequired();
+    Map<String, Schema> properties = datasetPropertiesSchema.getProperties();
+
+    // Assert required properties are non-nullable
+    Schema customProperties = properties.get("customProperties");
+    assertTrue(requiredNames.contains("customProperties"));
+    assertFalse(customProperties.getNullable());
+
+    // Assert non-required properties are nullable
+    Schema name = properties.get("name");
+    assertFalse(requiredNames.contains("name"));
+    assertTrue(name.getNullable());
+
+    // Assert non-required $ref properties are replaced by nullable { allOf: [ $ref ] } objects
+    Schema created = properties.get("created");
+    assertFalse(requiredNames.contains("created"));
+    assertEquals("object", created.getType());
+    assertNull(created.get$ref());
+    assertEquals(List.of(new Schema().$ref("#/components/schemas/TimeStamp")), created.getAllOf());
+    assertTrue(created.getNullable());
   }
 }