feeat(openapi): add batch endpoint to v2 using requestbody (#10100)

metadata-service/openapi-servlet/build.gradle

@@ -9,6 +9,7 @@ dependencies {
   implementation project(':metadata-service:auth-impl')
   implementation project(':metadata-service:factories')
   implementation project(':metadata-service:schema-registry-api')
+  implementation project (':metadata-service:openapi-servlet:models')
 
   implementation externalDependency.reflections
   implementation externalDependency.springBoot

metadata-service/openapi-servlet/models/build.gradle

@@ -0,0 +1,16 @@
+plugins {
+  id 'java'
+}
+
+dependencies {
+  implementation project(':entity-registry')
+  implementation project(':metadata-operation-context')
+  implementation project(':metadata-auth:auth-api')
+
+  implementation externalDependency.jacksonDataBind
+  implementation externalDependency.httpClient
+
+  compileOnly externalDependency.lombok
+
+  annotationProcessor externalDependency.lombok
+}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/client/OpenApiClient.java

@@ -0,0 +1,88 @@
+package io.datahubproject.openapi.client;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.datahubproject.metadata.context.OperationContext;
+import io.datahubproject.openapi.v2.models.BatchGetUrnRequest;
+import io.datahubproject.openapi.v2.models.BatchGetUrnResponse;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.hc.client5.http.classic.methods.HttpPost;
+import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
+import org.apache.hc.client5.http.impl.classic.HttpClientBuilder;
+import org.apache.hc.core5.http.ClassicHttpResponse;
+import org.apache.hc.core5.http.ContentType;
+import org.apache.hc.core5.http.HttpHeaders;
+import org.apache.hc.core5.http.io.entity.StringEntity;
+
+/** TODO: This should be autogenerated from our own OpenAPI */
+@Slf4j
+public class OpenApiClient {
+
+  private final CloseableHttpClient httpClient;
+  private final String gmsHost;
+  private final int gmsPort;
+  private final boolean useSsl;
+  @Getter private final OperationContext systemOperationContext;
+
+  private static final String OPENAPI_PATH = "/openapi/v2/entity/batch/";
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+  public OpenApiClient(
+      String gmsHost, int gmsPort, boolean useSsl, OperationContext systemOperationContext) {
+    this.gmsHost = gmsHost;
+    this.gmsPort = gmsPort;
+    this.useSsl = useSsl;
+    httpClient = HttpClientBuilder.create().build();
+    this.systemOperationContext = systemOperationContext;
+  }
+
+  public BatchGetUrnResponse getBatchUrnsSystemAuth(String entityName, BatchGetUrnRequest request) {
+    return getBatchUrns(
+        entityName,
+        request,
+        systemOperationContext.getSystemAuthentication().get().getCredentials());
+  }
+
+  public BatchGetUrnResponse getBatchUrns(
+      String entityName, BatchGetUrnRequest request, String authCredentials) {
+    String url =
+        (useSsl ? "https://" : "http://") + gmsHost + ":" + gmsPort + OPENAPI_PATH + entityName;
+    HttpPost httpPost = new HttpPost(url);
+    httpPost.setHeader(HttpHeaders.AUTHORIZATION, authCredentials);
+    try {
+      httpPost.setEntity(
+          new StringEntity(
+              OBJECT_MAPPER.writeValueAsString(request), ContentType.APPLICATION_JSON));
+      httpPost.setHeader("Content-type", "application/json");
+      return httpClient.execute(httpPost, OpenApiClient::mapResponse);
+    } catch (IOException e) {
+      log.error("Unable to execute Batch Get request for urn: " + request.getUrns(), e);
+      throw new RuntimeException(e);
+    }
+  }
+
+  private static BatchGetUrnResponse mapResponse(ClassicHttpResponse response) {
+    BatchGetUrnResponse serializedResponse;
+    try {
+      ByteArrayOutputStream result = new ByteArrayOutputStream();
+      InputStream contentStream = response.getEntity().getContent();
+      byte[] buffer = new byte[1024];
+      int length = contentStream.read(buffer);
+      while (length > 0) {
+        result.write(buffer, 0, length);
+        length = contentStream.read(buffer);
+      }
+      serializedResponse =
+          OBJECT_MAPPER.readValue(
+              result.toString(StandardCharsets.UTF_8), BatchGetUrnResponse.class);
+    } catch (IOException e) {
+      log.error("Wasn't able to convert response into expected type.", e);
+      throw new RuntimeException(e);
+    }
+    return serializedResponse;
+  }
+}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnRequest.java

@@ -0,0 +1,30 @@
+package io.datahubproject.openapi.v2.models;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import io.swagger.v3.oas.annotations.media.Schema;
+import java.io.Serializable;
+import java.util.List;
+import lombok.Builder;
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+
+@Value
+@EqualsAndHashCode
+@Builder
+@JsonDeserialize(builder = BatchGetUrnRequest.BatchGetUrnRequestBuilder.class)
+@JsonInclude(JsonInclude.Include.NON_NULL)
+public class BatchGetUrnRequest implements Serializable {
+  @JsonProperty("urns")
+  @Schema(required = true, description = "The list of urns to get.")
+  List<String> urns;
+
+  @JsonProperty("aspectNames")
+  @Schema(required = true, description = "The list of aspect names to get")
+  List<String> aspectNames;
+
+  @JsonProperty("withSystemMetadata")
+  @Schema(required = true, description = "Whether or not to retrieve system metadata")
+  boolean withSystemMetadata;
+}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnResponse.java

@@ -0,0 +1,20 @@
+package io.datahubproject.openapi.v2.models;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import io.swagger.v3.oas.annotations.media.Schema;
+import java.io.Serializable;
+import java.util.List;
+import lombok.Builder;
+import lombok.Value;
+
+@Value
+@Builder
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@JsonDeserialize(builder = BatchGetUrnResponse.BatchGetUrnResponseBuilder.class)
+public class BatchGetUrnResponse implements Serializable {
+  @JsonProperty("entities")
+  @Schema(description = "List of entity responses")
+  List<GenericEntity> entities;
+}

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java → metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java

@@ -2,22 +2,34 @@
 
 import com.datahub.util.RecordUtils;
 import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.util.Pair;
+import io.swagger.v3.oas.annotations.media.Schema;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Map;
 import java.util.stream.Collectors;
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
+import lombok.NoArgsConstructor;
 
 @Data
 @Builder
 @JsonInclude(JsonInclude.Include.NON_NULL)
+@NoArgsConstructor(force = true, access = AccessLevel.PRIVATE)
+@AllArgsConstructor
 public class GenericEntity {
+  @JsonProperty("urn")
+  @Schema(description = "Urn of the entity")
   private String urn;
+
+  @JsonProperty("aspects")
+  @Schema(description = "Map of aspect name to aspect")
   private Map<String, Object> aspects;
 
   public static class GenericEntityBuilder {

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

@@ -38,16 +38,20 @@
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.util.Pair;
 import io.datahubproject.metadata.context.OperationContext;
+import io.datahubproject.openapi.v2.models.BatchGetUrnRequest;
+import io.datahubproject.openapi.v2.models.BatchGetUrnResponse;
 import io.datahubproject.openapi.v2.models.GenericEntity;
 import io.datahubproject.openapi.v2.models.GenericScrollResult;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import java.lang.reflect.InvocationTargetException;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -140,8 +144,45 @@ public ResponseEntity<GenericScrollResult<GenericEntity>> getEntities(
   }
 
   @Tag(name = "Generic Entities")
-  @GetMapping(value = "/{entityName}/{entityUrn}", produces = MediaType.APPLICATION_JSON_VALUE)
-  @Operation(summary = "Get an entity")
+  @PostMapping(value = "/batch/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE)
+  @Operation(summary = "Get a batch of entities")
+  public ResponseEntity<BatchGetUrnResponse> getEntityBatch(
+      @PathVariable("entityName") String entityName, @RequestBody BatchGetUrnRequest request)
+      throws URISyntaxException {
+
+    if (restApiAuthorizationEnabled) {
+      Authentication authentication = AuthenticationContext.getAuthentication();
+      EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName);
+      request
+          .getUrns()
+          .forEach(
+              entityUrn ->
+                  checkAuthorized(
+                      authorizationChain,
+                      authentication.getActor(),
+                      entitySpec,
+                      entityUrn,
+                      ImmutableList.of(PoliciesConfig.GET_ENTITY_PRIVILEGE.getType())));
+    }
+
+    return ResponseEntity.of(
+        Optional.of(
+            BatchGetUrnResponse.builder()
+                .entities(
+                    new ArrayList<>(
+                        toRecordTemplates(
+                            request.getUrns().stream()
+                                .map(UrnUtils::getUrn)
+                                .collect(Collectors.toList()),
+                            new HashSet<>(request.getAspectNames()),
+                            request.isWithSystemMetadata())))
+                .build()));
+  }
+
+  @Tag(name = "Generic Entities")
+  @GetMapping(
+      value = "/{entityName}/{entityUrn:urn:li:.+}",
+      produces = MediaType.APPLICATION_JSON_VALUE)
   public ResponseEntity<GenericEntity> getEntity(
       @PathVariable("entityName") String entityName,
       @PathVariable("entityUrn") String entityUrn,

settings.gradle

@@ -68,3 +68,4 @@ include 'metadata-service:services'
 include 'metadata-service:configuration'
 include ':metadata-jobs:common'
 include ':metadata-operation-context'
+include ':metadata-service:openapi-servlet:models'