feat(openapi-v3): support async and createIfNotExists params on aspect (#11609)

metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java

@@ -1357,6 +1357,7 @@ public class EntityServiceImpl implements EntityService<ChangeItemImpl> {
               return IngestResult.builder()
                   .urn(item.getUrn())
                   .request(item)
+                  .result(result)
                   .publishedMCL(result.getMclFuture() != null)
                   .sqlCommitted(true)
                   .isUpdate(result.getOldValue() != null)

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java

@@ -170,6 +170,9 @@ public abstract class GenericEntitiesController<
       @Nonnull UpdateAspectResult updateAspectResult,
       boolean withSystemMetadata);
 
+  protected abstract E buildGenericEntity(
+      @Nonnull String aspectName, @Nonnull IngestResult ingestResult, boolean withSystemMetadata);
+
   protected abstract AspectsBatch toMCPBatch(
       @Nonnull OperationContext opContext, String entityArrayList, Actor actor)
       throws JsonProcessingException, InvalidUrnException;
@@ -560,8 +563,11 @@ public abstract class GenericEntitiesController<
       @PathVariable("entityName") String entityName,
       @PathVariable("entityUrn") String entityUrn,
       @PathVariable("aspectName") String aspectName,
+      @RequestParam(value = "async", required = false, defaultValue = "false") Boolean async,
       @RequestParam(value = "systemMetadata", required = false, defaultValue = "false")
           Boolean withSystemMetadata,
+      @RequestParam(value = "createIfEntityNotExists", required = false, defaultValue = "false")
+          Boolean createIfEntityNotExists,
       @RequestParam(value = "createIfNotExists", required = false, defaultValue = "true")
           Boolean createIfNotExists,
       @RequestBody @Nonnull String jsonAspect)
@@ -591,24 +597,38 @@ public abstract class GenericEntitiesController<
             opContext.getRetrieverContext().get().getAspectRetriever(),
             urn,
             aspectSpec,
+            createIfEntityNotExists,
             createIfNotExists,
             jsonAspect,
             authentication.getActor());
 
-    List<UpdateAspectResult> results =
+    Set<IngestResult> results =
-        entityService.ingestAspects(
+        entityService.ingestProposal(
             opContext,
             AspectsBatchImpl.builder()
                 .retrieverContext(opContext.getRetrieverContext().get())
                 .items(List.of(upsert))
                 .build(),
-            true,
+            async);
-            true);
 
-    return ResponseEntity.of(
+    if (!async) {
-        results.stream()
+      return ResponseEntity.of(
-            .findFirst()
+          results.stream()
-            .map(result -> buildGenericEntity(aspectName, result, withSystemMetadata)));
+              .filter(item -> aspectName.equals(item.getRequest().getAspectName()))
+              .findFirst()
+              .map(
+                  result ->
+                      buildGenericEntity(aspectName, result.getResult(), withSystemMetadata)));
+    } else {
+      return results.stream()
+          .filter(item -> aspectName.equals(item.getRequest().getAspectName()))
+          .map(
+              result ->
+                  ResponseEntity.accepted()
+                      .body(buildGenericEntity(aspectName, result, withSystemMetadata)))
+          .findFirst()
+          .orElse(ResponseEntity.accepted().build());
+    }
   }
 
   @Tag(name = "Generic Aspects")
@@ -789,6 +809,7 @@ public abstract class GenericEntitiesController<
       @Nonnull AspectRetriever aspectRetriever,
       Urn entityUrn,
       AspectSpec aspectSpec,
+      Boolean createIfEntityNotExists,
       Boolean createIfNotExists,
       String jsonAspect,
       Actor actor)

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java

@@ -232,6 +232,20 @@ public class EntityController
                     withSystemMetadata ? updateAspectResult.getNewSystemMetadata() : null)));
   }
 
+  @Override
+  protected GenericEntityV2 buildGenericEntity(
+      @Nonnull String aspectName, @Nonnull IngestResult ingestResult, boolean withSystemMetadata) {
+    return GenericEntityV2.builder()
+        .urn(ingestResult.getUrn().toString())
+        .build(
+            objectMapper,
+            Map.of(
+                aspectName,
+                Pair.of(
+                    ingestResult.getRequest().getRecordTemplate(),
+                    withSystemMetadata ? ingestResult.getRequest().getSystemMetadata() : null)));
+  }
+
   private List<GenericEntityV2> toRecordTemplates(
       @Nonnull OperationContext opContext,
       SearchEntityArray searchEntities,
@@ -278,14 +292,25 @@ public class EntityController
       @Nonnull AspectRetriever aspectRetriever,
       Urn entityUrn,
       AspectSpec aspectSpec,
+      Boolean createIfEntityNotExists,
       Boolean createIfNotExists,
       String jsonAspect,
       Actor actor)
       throws URISyntaxException {
+
+    final ChangeType changeType;
+    if (Boolean.TRUE.equals(createIfEntityNotExists)) {
+      changeType = ChangeType.CREATE_ENTITY;
+    } else if (Boolean.TRUE.equals(createIfNotExists)) {
+      changeType = ChangeType.CREATE;
+    } else {
+      changeType = ChangeType.UPSERT;
+    }
+
     return ChangeItemImpl.builder()
         .urn(entityUrn)
         .aspectName(aspectSpec.getName())
-        .changeType(Boolean.TRUE.equals(createIfNotExists) ? ChangeType.CREATE : ChangeType.UPSERT)
+        .changeType(changeType)
         .auditStamp(AuditStampUtils.createAuditStamp(actor.toUrnStr()))
         .recordTemplate(
             GenericRecordUtils.deserializeAspect(

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

@@ -1100,6 +1100,28 @@ public class OpenAPIV3Generator {
         new Operation()
             .summary(String.format("Create aspect %s on %s ", aspect, upperFirstEntity))
             .tags(tags)
+            .parameters(
+                List.of(
+                    new Parameter()
+                        .in(NAME_QUERY)
+                        .name("async")
+                        .description("Use async ingestion for high throughput.")
+                        .schema(new Schema().type(TYPE_BOOLEAN)._default(false)),
+                    new Parameter()
+                        .in(NAME_QUERY)
+                        .name(NAME_SYSTEM_METADATA)
+                        .description("Include systemMetadata with response.")
+                        .schema(new Schema().type(TYPE_BOOLEAN)._default(false)),
+                    new Parameter()
+                        .in(NAME_QUERY)
+                        .name("createIfEntityNotExists")
+                        .description("Only create the aspect if the Entity doesn't exist.")
+                        .schema(new Schema().type(TYPE_BOOLEAN)._default(false)),
+                    new Parameter()
+                        .in(NAME_QUERY)
+                        .name("createIfNotExists")
+                        .description("Only create the aspect if the Aspect doesn't exist.")
+                        .schema(new Schema().type(TYPE_BOOLEAN)._default(true))))
             .requestBody(requestBody)
             .responses(new ApiResponses().addApiResponse("201", successPostResponse));
     // Patch Operation

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

@@ -328,6 +328,24 @@ public class EntityController
                     .build()));
   }
 
+  @Override
+  protected GenericEntityV3 buildGenericEntity(
+      @Nonnull String aspectName, @Nonnull IngestResult ingestResult, boolean withSystemMetadata) {
+    return GenericEntityV3.builder()
+        .build(
+            objectMapper,
+            ingestResult.getUrn(),
+            Map.of(
+                aspectName,
+                AspectItem.builder()
+                    .aspect(ingestResult.getRequest().getRecordTemplate())
+                    .systemMetadata(
+                        withSystemMetadata ? ingestResult.getRequest().getSystemMetadata() : null)
+                    .auditStamp(
+                        withSystemMetadata ? ingestResult.getRequest().getAuditStamp() : null)
+                    .build()));
+  }
+
   private List<GenericEntityV3> toRecordTemplates(
       @Nonnull OperationContext opContext,
       SearchEntityArray searchEntities,
@@ -472,16 +490,27 @@ public class EntityController
       @Nonnull AspectRetriever aspectRetriever,
       Urn entityUrn,
       AspectSpec aspectSpec,
+      Boolean createIfEntityNotExists,
       Boolean createIfNotExists,
       String jsonAspect,
       Actor actor)
       throws JsonProcessingException {
     JsonNode jsonNode = objectMapper.readTree(jsonAspect);
     String aspectJson = jsonNode.get("value").toString();
+
+    final ChangeType changeType;
+    if (Boolean.TRUE.equals(createIfEntityNotExists)) {
+      changeType = ChangeType.CREATE_ENTITY;
+    } else if (Boolean.TRUE.equals(createIfNotExists)) {
+      changeType = ChangeType.CREATE;
+    } else {
+      changeType = ChangeType.UPSERT;
+    }
+
     return ChangeItemImpl.builder()
         .urn(entityUrn)
         .aspectName(aspectSpec.getName())
-        .changeType(Boolean.TRUE.equals(createIfNotExists) ? ChangeType.CREATE : ChangeType.UPSERT)
+        .changeType(changeType)
         .auditStamp(AuditStampUtils.createAuditStamp(actor.toUrnStr()))
         .recordTemplate(
             GenericRecordUtils.deserializeAspect(

metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java

@@ -2,6 +2,7 @@ package com.linkedin.metadata.entity;
 
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.aspect.batch.BatchItem;
+import javax.annotation.Nullable;
 import lombok.Builder;
 import lombok.Value;
 
@@ -10,6 +11,7 @@ import lombok.Value;
 public class IngestResult {
   Urn urn;
   BatchItem request;
+  @Nullable UpdateAspectResult result;
   boolean publishedMCL;
   boolean processedMCL;
   boolean publishedMCP;