add a structured type to batchGet in OpenAPI V3 spec (#10956)

datahub-project · Jul 22, 2024 · ebe687b · ebe687b
1 parent a090518
commit ebe687b
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 18 deletions.
diff --git a/...servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java b/...servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java
@@ -13,6 +13,7 @@
 import com.datahub.authorization.AuthorizerChain;
 import com.datahub.util.RecordUtils;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableSet;
 import com.linkedin.common.urn.Urn;
@@ -508,10 +509,10 @@ public ResponseEntity<E> createAspect(
       @PathVariable("aspectName") String aspectName,
       @RequestParam(value = "systemMetadata", required = false, defaultValue = "false")
           Boolean withSystemMetadata,
-      @RequestParam(value = "createIfNotExists", required = false, defaultValue = "false")
+      @RequestParam(value = "createIfNotExists", required = false, defaultValue = "true")
           Boolean createIfNotExists,
       @RequestBody @Nonnull String jsonAspect)
-      throws URISyntaxException {
+      throws URISyntaxException, JsonProcessingException {
 
     Urn urn = validatedUrn(entityUrn);
     EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName);
@@ -649,8 +650,8 @@ protected Boolean exists(
    * fixes)
    *
    * @param requestedAspectNames requested aspects
-   * @return updated map
    * @param <T> map values
+   * @return updated map
    */
   protected <T> LinkedHashMap<Urn, Map<String, T>> resolveAspectNames(
       LinkedHashMap<Urn, Map<String, T>> requestedAspectNames, @Nonnull T defaultValue) {
@@ -732,15 +733,17 @@ protected ChangeMCP toUpsertItem(
       Boolean createIfNotExists,
       String jsonAspect,
       Actor actor)
-      throws URISyntaxException {
+      throws JsonProcessingException {
+    JsonNode jsonNode = objectMapper.readTree(jsonAspect);
+    String aspectJson = jsonNode.get("value").toString();
     return ChangeItemImpl.builder()
         .urn(entityUrn)
         .aspectName(aspectSpec.getName())
         .changeType(Boolean.TRUE.equals(createIfNotExists) ? ChangeType.CREATE : ChangeType.UPSERT)
         .auditStamp(AuditStampUtils.createAuditStamp(actor.toUrnStr()))
         .recordTemplate(
             GenericRecordUtils.deserializeAspect(
-                ByteString.copyString(jsonAspect, StandardCharsets.UTF_8),
+                ByteString.copyString(aspectJson, StandardCharsets.UTF_8),
                 GenericRecordUtils.JSON,
                 aspectSpec))
         .build(aspectRetriever);

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
@@ -82,6 +82,20 @@ public static OpenAPI generateOpenApiSpec(EntityRegistry entityRegistry) {
         "SystemMetadata", new Schema().type(TYPE_OBJECT).additionalProperties(true));
     components.addSchemas("SortOrder", new Schema()._enum(List.of("ASCENDING", "DESCENDING")));
     components.addSchemas("AspectPatch", buildAspectPatchSchema());
+    components.addSchemas(
+        "BatchGetRequestBody",
+        new Schema<>()
+            .type(TYPE_OBJECT)
+            .description("Request body for batch get aspects.")
+            .properties(
+                Map.of(
+                    "headers",
+                    new Schema<>()
+                        .type(TYPE_OBJECT)
+                        .additionalProperties(new Schema<>().type(TYPE_STRING))
+                        .description("System headers for the operation.")
+                        .nullable(true)))
+            .nullable(true));
     entityRegistry
         .getAspectSpecs()
         .values()
@@ -645,28 +659,19 @@ private static Schema buildEntityScrollSchema(final EntitySpec entity) {
   private static Schema buildEntityBatchGetRequestSchema(
       final EntitySpec entity, Set<String> aspectNames) {
 
-    final Schema stringTypeSchema = new Schema<>();
-    stringTypeSchema.setType(TYPE_STRING);
-    final Map<String, Schema> headers =
-        Map.of(
-            "headers",
-            new Schema<>()
-                .type(TYPE_OBJECT)
-                .additionalProperties(stringTypeSchema)
-                .description("System headers for the operation.")
-                .nullable(true));
-
     final Map<String, Schema> properties =
         entity.getAspectSpecMap().entrySet().stream()
             .filter(a -> aspectNames.contains(a.getValue().getName()))
             .collect(
                 Collectors.toMap(
-                    Map.Entry::getKey, a -> new Schema().type(TYPE_OBJECT).properties(headers)));
+                    Map.Entry::getKey,
+                    a -> new Schema().$ref("#/components/schemas/BatchGetRequestBody")));
     properties.put(
         PROPERTY_URN,
         new Schema<>().type(TYPE_STRING).description("Unique id for " + entity.getName()));
 
-    properties.put(entity.getKeyAspectName(), new Schema().type(TYPE_OBJECT).properties(headers));
+    properties.put(
+        entity.getKeyAspectName(), new Schema().$ref("#/components/schemas/BatchGetRequestBody"));
 
     return new Schema<>()
         .type(TYPE_OBJECT)

diff --git a/...enapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java b/...enapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java
@@ -37,6 +37,7 @@
 import io.datahubproject.openapi.v3.models.GenericAspectV3;
 import io.datahubproject.openapi.v3.models.GenericEntityScrollResultV3;
 import io.datahubproject.openapi.v3.models.GenericEntityV3;
+import io.swagger.v3.oas.annotations.Hidden;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import jakarta.servlet.http.HttpServletRequest;
@@ -67,6 +68,7 @@
 @RequiredArgsConstructor
 @RequestMapping("/v3/entity")
 @Slf4j
+@Hidden
 public class EntityController
     extends GenericEntitiesController<
         GenericAspectV3, GenericEntityV3, GenericEntityScrollResultV3> {

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
@@ -88,5 +88,18 @@ public void testOpenApiSpecBuilder() throws Exception {
     Schema fabricType = openAPI.getComponents().getSchemas().get("FabricType");
     assertEquals("string", fabricType.getType());
     assertFalse(fabricType.getEnum().isEmpty());
+
+    Map<String, Schema> batchProperties =
+        openAPI
+            .getComponents()
+            .getSchemas()
+            .get("BatchGetContainerEntityRequest_v3")
+            .getProperties();
+    batchProperties.entrySet().stream()
+        .filter(entry -> !entry.getKey().equals("urn"))
+        .forEach(
+            entry ->
+                assertEquals(
+                    "#/components/schemas/BatchGetRequestBody", entry.getValue().get$ref()));
   }
 }