fix(openapi-schema): fix openapi schema generator (#12590)

Author: david-leifker

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

@@ -114,7 +114,8 @@ public static OpenAPI generateOpenApiSpec(
 
     // --> Aspect components
     components.addSchemas(
-        ASPECTS + ASPECT_RESPONSE_SUFFIX, buildAspectsRefResponseSchema(entityRegistry));
+        ASPECTS + ASPECT_RESPONSE_SUFFIX,
+        buildAspectsRefResponseSchema(entityRegistry, definitionNames));
     entityRegistry
         .getAspectSpecs()
         .values()
@@ -724,7 +725,8 @@ private static void addAspectSchemas(final Components components, final AspectSp
    * @param entityRegistry entity registry
    * @return schema
    */
-  private static Schema buildAspectsRefResponseSchema(final EntityRegistry entityRegistry) {
+  private static Schema buildAspectsRefResponseSchema(
+      final EntityRegistry entityRegistry, final Set<String> definitionNames) {
     final Schema result =
         new Schema<>()
             .type(TYPE_OBJECT)
@@ -734,7 +736,11 @@ private static Schema buildAspectsRefResponseSchema(final EntityRegistry entityR
     // Create a list of reference schemas for each aspect
     List<Schema> aspectRefs =
         entityRegistry.getAspectSpecs().values().stream()
-            .map(aspect -> new Schema<>().$ref(PATH_DEFINITIONS + toUpperFirst(aspect.getName())))
+            .filter(a -> definitionNames.contains(a.getName()))
+            .map(
+                aspect ->
+                    new Schema<>()
+                        .$ref(PATH_DEFINITIONS + toUpperFirst(aspect.getPegasusSchema().getName())))
             .distinct()
             .collect(Collectors.toList());
 

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

@@ -130,5 +130,16 @@ public void testOpenApiSpecBuilder() throws Exception {
                   Character.isUpperCase(refName.charAt(0)),
                   "schema reference should start with capital letter: " + name);
             });
+
+    // Check for pegasus name example
+    assertTrue(
+        valueProperty.getOneOf().stream()
+            .anyMatch(
+                schema ->
+                    schema.get$ref().equals("#/components/schemas/StructuredPropertyDefinition")));
+    assertTrue(
+        valueProperty.getOneOf().stream()
+            .noneMatch(
+                schema -> schema.get$ref().equals("#/components/schemas/PropertyDefinition")));
   }
 }