Use type: string for enum schemas (datahub-project#10663)

acryldata · Jun 7, 2024 · 812bcbb · 812bcbb
1 parent 1b6763f
commit 812bcbb
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 21 deletions.
diff --git a/...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java b/...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java
@@ -437,27 +437,31 @@ 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<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(TYPE_OBJECT);
-                          schema.set$ref(null);
-                          schema.setAllOf(List.of(new Schema().$ref($ref)));
-                        }
-                        schema.setNullable(!isNameRequired);
-                      });
-
+                  // Set enums to "string".
+                  if (s.getEnum() != null && !s.getEnum().isEmpty()) {
+                    s.setType("string");
+                  } else {
+                    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(TYPE_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);

diff --git a/...ce/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/OpenAPIV3GeneratorTest.java b/...ce/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/OpenAPIV3GeneratorTest.java
@@ -83,5 +83,10 @@ public void testOpenApiSpecBuilder() throws Exception {
         List.of(new Schema().$ref("#/components/schemas/SystemMetadata")),
         systemMetadata.getAllOf());
     assertTrue(systemMetadata.getNullable());
+
+    // Assert enum property is string.
+    Schema fabricType = openAPI.getComponents().getSchemas().get("FabricType");
+    assertEquals("string", fabricType.getType());
+    assertFalse(fabricType.getEnum().isEmpty());
   }
 }