fix(openapi): openapi 3.1.0 nullable refs to use oneOf(ref, null) (#14079)

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

@@ -1187,13 +1187,18 @@ public class OpenAPIV3Generator {
         NAME_SYSTEM_METADATA,
         newSchema()
             .types(TYPE_OBJECT_NULLABLE)
-            .$ref(PATH_DEFINITIONS + "SystemMetadata")
+            .oneOf(
+                List.of(
+                    newSchema().$ref(PATH_DEFINITIONS + "SystemMetadata"),
+                    newSchema().type(TYPE_NULL)))
             .description("System metadata for the aspect."));
     result.addProperty(
         NAME_AUDIT_STAMP,
         newSchema()
             .types(TYPE_OBJECT_NULLABLE)
-            .$ref(PATH_DEFINITIONS + "AuditStamp")
+            .oneOf(
+                List.of(
+                    newSchema().$ref(PATH_DEFINITIONS + "AuditStamp"), newSchema().type(TYPE_NULL)))
             .description("Audit stamp for the aspect."));
     return result;
   }
@@ -1210,7 +1215,10 @@ public class OpenAPIV3Generator {
         NAME_SYSTEM_METADATA,
         newSchema()
             .types(TYPE_OBJECT_NULLABLE)
-            .$ref(PATH_DEFINITIONS + "SystemMetadata")
+            .oneOf(
+                List.of(
+                    newSchema().$ref(PATH_DEFINITIONS + "SystemMetadata"),
+                    newSchema().type(TYPE_NULL)))
             .description("System metadata for the aspect."));
 
     Schema stringTypeSchema = newSchema();
@@ -1359,8 +1367,8 @@ public class OpenAPIV3Generator {
                         .toList()))
         .properties(
             Map.of(
-                "entities", entitiesSchema,
-                "aspects", aspectsSchema));
+                "entities", newSchema().oneOf(List.of(entitiesSchema, newSchema().type(TYPE_NULL))),
+                "aspects", newSchema().oneOf(List.of(aspectsSchema, newSchema().type(TYPE_NULL)))));
   }
 
   private static Schema buildEntitiesPatchRequestSchema(List<EntitySpec> entitySpecs) {
@@ -1434,19 +1442,26 @@ public class OpenAPIV3Generator {
   private static Schema buildEntityBatchGetRequestSchema(
       final EntitySpec entity, Set<String> aspectNames) {
 
-    final Map<String, Schema> properties =
-        entity.getAspectSpecMap().entrySet().stream()
-            .filter(a -> aspectNames.contains(a.getValue().getName()))
-            .collect(
-                Collectors.toMap(
-                    Map.Entry::getKey,
-                    a -> newSchema().$ref("#/components/schemas/BatchGetRequestBody")));
+    Map<String, Schema> properties = new LinkedHashMap<>();
     properties.put(
         PROPERTY_URN,
         newSchema().type(TYPE_STRING).description("Unique id for " + entity.getName()));
 
-    properties.put(
-        entity.getKeyAspectName(), newSchema().$ref("#/components/schemas/BatchGetRequestBody"));
+    entity.getAspectSpecMap().entrySet().stream()
+        .filter(
+            e ->
+                aspectNames.contains(e.getValue().getName())
+                    || e.getKey().equals(entity.getKeyAspectName()))
+        .forEach(
+            e ->
+                properties.put(
+                    e.getKey(),
+                    newSchema()
+                        .types(TYPE_OBJECT_NULLABLE)
+                        .oneOf(
+                            List.of(
+                                newSchema().$ref("#/components/schemas/BatchGetRequestBody"),
+                                newSchema().type(TYPE_NULL)))));
 
     return newSchema()
         .type(TYPE_OBJECT)
@@ -1456,19 +1471,24 @@ public class OpenAPIV3Generator {
   }
 
   private static Schema buildCrossEntityUpsertSchema(List<EntitySpec> entitySpecs) {
+
     Map<String, Schema> props = new LinkedHashMap<>();
+
     entitySpecs.forEach(
-        e ->
-            props.put(
-                e.getName(),
-                newSchema()
-                    .type(TYPE_ARRAY)
-                    .items(
-                        newSchema()
-                            .$ref(
-                                String.format(
-                                    "#/components/schemas/%s%s",
-                                    toUpperFirst(e.getName()), ENTITY_REQUEST_SUFFIX)))));
+        e -> {
+          Schema arraySchema =
+              newSchema()
+                  .type(TYPE_ARRAY)
+                  .items(
+                      newSchema()
+                          .$ref(
+                              String.format(
+                                  "#/components/schemas/%s%s",
+                                  toUpperFirst(e.getName()), ENTITY_REQUEST_SUFFIX)));
+          props.put(
+              e.getName(), newSchema().oneOf(List.of(arraySchema, newSchema().type(TYPE_NULL))));
+        });
+
     return newSchema()
         .type(TYPE_OBJECT)
         .description("Mixed-entity upsert request body.")
@@ -1477,19 +1497,25 @@ public class OpenAPIV3Generator {
   }
 
   private static Schema buildCrossEntityPatchSchema(List<EntitySpec> entitySpecs) {
+
     Map<String, Schema> props = new LinkedHashMap<>();
+
     entitySpecs.forEach(
-        e ->
-            props.put(
-                e.getName(),
-                newSchema()
-                    .type(TYPE_ARRAY)
-                    .items(
-                        newSchema()
-                            .$ref(
-                                String.format(
-                                    "#/components/schemas/%s%s",
-                                    toUpperFirst(e.getName()), ENTITY_REQUEST_PATCH_SUFFIX)))));
+        e -> {
+          Schema arraySchema =
+              newSchema()
+                  .type(TYPE_ARRAY)
+                  .items(
+                      newSchema()
+                          .$ref(
+                              String.format(
+                                  "#/components/schemas/%s%s",
+                                  toUpperFirst(e.getName()), ENTITY_REQUEST_PATCH_SUFFIX)));
+
+          props.put(
+              e.getName(), newSchema().oneOf(List.of(newSchema().type(TYPE_NULL), arraySchema)));
+        });
+
     return newSchema()
         .type(TYPE_OBJECT)
         .description("Mixed-entity patch request body.")
@@ -1499,18 +1525,23 @@ public class OpenAPIV3Generator {
 
   private static Schema buildCrossEntityResponseSchema(List<EntitySpec> entitySpecs) {
     Map<String, Schema> props = new LinkedHashMap<>();
+
     entitySpecs.forEach(
-        e ->
-            props.put(
-                e.getName(),
-                newSchema()
-                    .type(TYPE_ARRAY)
-                    .items(
-                        newSchema()
-                            .$ref(
-                                String.format(
-                                    "#/components/schemas/%s%s",
-                                    toUpperFirst(e.getName()), ENTITY_RESPONSE_SUFFIX)))));
+        e -> {
+          Schema arraySchema =
+              newSchema()
+                  .type(TYPE_ARRAY)
+                  .items(
+                      newSchema()
+                          .$ref(
+                              String.format(
+                                  "#/components/schemas/%s%s",
+                                  toUpperFirst(e.getName()), ENTITY_RESPONSE_SUFFIX)));
+
+          props.put(
+              e.getName(), newSchema().oneOf(List.of(arraySchema, newSchema().type(TYPE_NULL))));
+        });
+
     return newSchema()
         .type(TYPE_OBJECT)
         .description("Mixed-entity upsert / patch response.")
@@ -1519,21 +1550,24 @@ public class OpenAPIV3Generator {
   }
 
   private static Schema buildCrossEntityBatchGetRequestSchema(List<EntitySpec> entitySpecs) {
+
     Map<String, Schema> props = new LinkedHashMap<>();
 
     entitySpecs.forEach(
-        e ->
-            props.put(
-                e.getName(),
-                newSchema()
-                    .type(TYPE_ARRAY)
-                    .items(
-                        newSchema()
-                            .$ref(
-                                String.format(
-                                    "#/components/schemas/%s%s",
-                                    "BatchGet" + toUpperFirst(e.getName()), // BatchGet<Ent>
-                                    ENTITY_REQUEST_SUFFIX)))));
+        e -> {
+          Schema arraySchema =
+              newSchema()
+                  .type(TYPE_ARRAY)
+                  .items(
+                      newSchema()
+                          .$ref(
+                              String.format(
+                                  "#/components/schemas/%s%s",
+                                  "BatchGet" + toUpperFirst(e.getName()), ENTITY_REQUEST_SUFFIX)));
+
+          props.put(
+              e.getName(), newSchema().oneOf(List.of(arraySchema, newSchema().type(TYPE_NULL))));
+        });
 
     return newSchema()
         .type(TYPE_OBJECT)
@@ -1542,29 +1576,6 @@ public class OpenAPIV3Generator {
         .properties(props);
   }
 
-  /** Same structure as buildEntityBatchGetRequestSchema but covers the union of all aspects. */
-  private static Schema buildEntitiesBatchGetRequestSchema(
-      Map<String, AspectSpec> aspectSpecs, Set<String> aspectNames) {
-
-    Map<String, Schema> properties =
-        aspectSpecs.entrySet().stream()
-            .filter(e -> aspectNames.contains(e.getKey()))
-            .collect(
-                Collectors.toMap(
-                    Map.Entry::getKey,
-                    e -> newSchema().$ref("#/components/schemas/BatchGetRequestBody"),
-                    (a, b) -> a, // merge func (won’t actually happen)
-                    LinkedHashMap::new));
-
-    properties.put(PROPERTY_URN, newSchema().type(TYPE_STRING).description("Unique id for entity"));
-
-    return newSchema()
-        .type(TYPE_OBJECT)
-        .description(ENTITIES + " object.")
-        .required(List.of(PROPERTY_URN))
-        .properties(properties);
-  }
-
   private static Schema buildAspectRef(final String aspect, final boolean withSystemMetadata) {
     final Schema result = newSchema();
 
@@ -1968,7 +1979,10 @@ public class OpenAPIV3Generator {
         NAME_SYSTEM_METADATA,
         newSchema()
             .types(TYPE_OBJECT_NULLABLE)
-            .$ref(PATH_DEFINITIONS + "SystemMetadata")
+            .oneOf(
+                List.of(
+                    newSchema().$ref(PATH_DEFINITIONS + "SystemMetadata"),
+                    newSchema().type(TYPE_NULL)))
             .description("System metadata for the aspect."));
     schema.addProperty(
         "headers",

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

@@ -192,18 +192,36 @@ public class OpenAPIV3GeneratorTest {
 
   @Test
   public void testBatchProperties() {
-    Map<String, Schema> batchProperties =
+    Map<String, Schema> batchProps =
         openAPI
             .getComponents()
             .getSchemas()
             .get("BatchGetContainerEntityRequest_v3")
             .getProperties();
-    batchProperties.entrySet().stream()
-        .filter(entry -> !entry.getKey().equals("urn"))
+
+    batchProps.entrySet().stream()
+        .filter(e -> !e.getKey().equals("urn"))
         .forEach(
-            entry ->
-                assertEquals(
-                    "#/components/schemas/BatchGetRequestBody", entry.getValue().get$ref()));
+            e -> {
+              Schema<?> prop = e.getValue();
+
+              // must be wrapped in oneOf
+              assertNull(prop.get$ref());
+              assertNotNull(prop.getOneOf());
+              assertEquals(prop.getOneOf().size(), 2);
+
+              boolean hasRef =
+                  prop.getOneOf().stream()
+                      .anyMatch(
+                          s ->
+                              "#/components/schemas/BatchGetRequestBody"
+                                  .equals(((Schema<?>) s).get$ref()));
+              boolean hasNull =
+                  prop.getOneOf().stream().anyMatch(s -> "null".equals(((Schema<?>) s).getType()));
+
+              assertTrue(hasRef, "oneOf must contain BatchGetRequestBody ref");
+              assertTrue(hasNull, "oneOf must contain null type");
+            });
   }
 
   @Test
@@ -232,10 +250,16 @@ public class OpenAPIV3GeneratorTest {
     assertNull(created.getType());
     assertNull(created.getTypes());
     assertNull(created.get$ref());
-    assertEquals(
-        new HashSet<>(created.getOneOf()),
-        Set.of(new Schema().$ref("#/components/schemas/TimeStamp"), new Schema<>().type("null")));
-    assertNull(created.getNullable());
+    assertEquals(created.getOneOf().size(), 2);
+
+    assertTrue(
+        created.getOneOf().stream()
+            .anyMatch(s -> "#/components/schemas/TimeStamp".equals(((Schema<?>) s).get$ref())));
+    assertTrue(
+        created.getOneOf().stream()
+            .anyMatch(
+                s ->
+                    ((Schema<?>) s).get$ref() == null && "null".equals(((Schema<?>) s).getType())));
 
     // Assert systemMetadata property on response schema is optional per v3.1.0
     Map<String, Schema> datasetPropertiesResponseSchemaProps =
@@ -244,10 +268,23 @@ public class OpenAPIV3GeneratorTest {
             .getSchemas()
             .get("DatasetPropertiesAspectResponse_v3")
             .getProperties();
-    Schema systemMetadata = datasetPropertiesResponseSchemaProps.get("systemMetadata");
-    assertEquals(systemMetadata.getTypes(), Set.of("object", "null"));
-    assertEquals(systemMetadata.get$ref(), "#/components/schemas/SystemMetadata");
-    assertNull(systemMetadata.getNullable());
+    Schema systemMetadataProperty = datasetPropertiesResponseSchemaProps.get("systemMetadata");
+    assertNotNull(systemMetadataProperty);
+
+    assertNull(systemMetadataProperty.get$ref());
+    assertEquals(systemMetadataProperty.getTypes(), Set.of("object", "null"));
+    assertNotNull(systemMetadataProperty.getOneOf());
+    assertEquals(systemMetadataProperty.getOneOf().size(), 2);
+
+    boolean hasSysMetaRef =
+        systemMetadataProperty.getOneOf().stream()
+            .anyMatch(s -> "#/components/schemas/SystemMetadata".equals(((Schema<?>) s).get$ref()));
+    boolean hasNullAlt =
+        systemMetadataProperty.getOneOf().stream()
+            .anyMatch(s -> "null".equals(((Schema<?>) s).getType()));
+
+    assertTrue(hasSysMetaRef, "systemMetadata oneOf must contain SystemMetadata ref");
+    assertTrue(hasNullAlt, "systemMetadata oneOf must contain null type");
   }
 
   @Test
@@ -592,20 +629,22 @@ public class OpenAPIV3GeneratorTest {
 
     // Verify 'systemMetadata' property
     Schema systemMetadataProperty = properties.get("systemMetadata");
-    assertNotNull(
-        systemMetadataProperty, "AspectPatchProperty should have 'systemMetadata' property");
-    assertEquals(
-        systemMetadataProperty.get$ref(),
-        "#/components/schemas/SystemMetadata",
-        "SystemMetadata property should reference SystemMetadata schema");
-    assertEquals(
-        systemMetadataProperty.getTypes(),
-        Set.of("object", "null"),
-        "SystemMetadata property should allow object and null types");
-    assertEquals(
-        systemMetadataProperty.getDescription(),
-        "System metadata for the aspect.",
-        "SystemMetadata property should have correct description");
+    assertNotNull(systemMetadataProperty);
+
+    assertNull(systemMetadataProperty.get$ref());
+    assertEquals(systemMetadataProperty.getTypes(), Set.of("object", "null"));
+    assertNotNull(systemMetadataProperty.getOneOf());
+    assertEquals(systemMetadataProperty.getOneOf().size(), 2);
+
+    boolean hasSysMetaRef =
+        systemMetadataProperty.getOneOf().stream()
+            .anyMatch(s -> "#/components/schemas/SystemMetadata".equals(((Schema<?>) s).get$ref()));
+    boolean hasNullAlt =
+        systemMetadataProperty.getOneOf().stream()
+            .anyMatch(s -> "null".equals(((Schema<?>) s).getType()));
+
+    assertTrue(hasSysMetaRef, "systemMetadata oneOf must contain SystemMetadata ref");
+    assertTrue(hasNullAlt, "systemMetadata oneOf must contain null type");
 
     // Verify 'headers' property
     Schema headersProperty = properties.get("headers");
@@ -730,6 +769,45 @@ public class OpenAPIV3GeneratorTest {
     assertTrue(componentKeys.contains("CrossEntitiesResponse_v3"));
   }
 
+  @Test
+  public void testCrossEntityArraysAreOptional() {
+    String[] schemas = {"CrossEntitiesRequest_v3", "CrossEntitiesPatch_v3"};
+
+    for (String schemaName : schemas) {
+      Schema<?> schema = openAPI.getComponents().getSchemas().get(schemaName);
+      assertNotNull(schema, "Component " + schemaName + " must exist");
+
+      // 1)  No property except urn is required (required list null or empty)
+      assertTrue(
+          schema.getRequired() == null || schema.getRequired().isEmpty(),
+          schemaName + " must not require any entity arrays");
+
+      // 2)  Every property value is oneOf( array , null )
+      schema
+          .getProperties()
+          .forEach(
+              (propName, propSchema) -> {
+                // Property schema should have no direct $ref and be wrapped in oneOf
+                assertNull(
+                    propSchema.get$ref(),
+                    schemaName + "." + propName + " must not have direct $ref");
+                assertNotNull(
+                    propSchema.getOneOf(),
+                    schemaName + "." + propName + " must be defined with oneOf");
+
+                List<Schema<?>> oneOf = propSchema.getOneOf();
+                assertEquals(oneOf.size(), 2, "oneOf must contain exactly two alternatives");
+
+                boolean hasArray = oneOf.stream().anyMatch(s -> "array".equals(s.getType()));
+                boolean hasNull = oneOf.stream().anyMatch(s -> "null".equals(s.getType()));
+
+                assertTrue(
+                    hasArray, schemaName + "." + propName + " oneOf must include array type");
+                assertTrue(hasNull, schemaName + "." + propName + " oneOf must include null type");
+              });
+    }
+  }
+
   private JsonSchema loadOpenAPI31Schema(JsonSchemaFactory schemaFactory) throws Exception {
     URL schemaUrl = new URL("https://spec.openapis.org/oas/3.1/schema/2022-10-07");
     return schemaFactory.getSchema(schemaUrl.openStream());