From d55210663f9cb591b11448c379efa842d58577b7 Mon Sep 17 00:00:00 2001 From: RyanHolstien Date: Thu, 21 Mar 2024 15:15:29 -0500 Subject: [PATCH] feeat(openapi): add batch endpoint to v2 using requestbody (#10100) --- metadata-service/openapi-servlet/build.gradle | 1 + .../openapi-servlet/models/build.gradle | 16 ++++ .../openapi/client/OpenApiClient.java | 88 +++++++++++++++++++ .../openapi/v2/models/BatchGetUrnRequest.java | 30 +++++++ .../v2/models/BatchGetUrnResponse.java | 20 +++++ .../openapi/v2/models/GenericEntity.java | 12 +++ .../v2/models/GenericRelationship.java | 0 .../v2/models/GenericScrollResult.java | 0 .../v2/models/GenericTimeseriesAspect.java | 0 .../openapi/v2/models/PatchOperation.java | 0 .../v2/controller/EntityController.java | 45 +++++++++- settings.gradle | 1 + 12 files changed, 211 insertions(+), 2 deletions(-) create mode 100644 metadata-service/openapi-servlet/models/build.gradle create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/client/OpenApiClient.java create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnRequest.java create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnResponse.java rename metadata-service/openapi-servlet/{ => models}/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java (82%) rename metadata-service/openapi-servlet/{ => models}/src/main/java/io/datahubproject/openapi/v2/models/GenericRelationship.java (100%) rename metadata-service/openapi-servlet/{ => models}/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java (100%) rename metadata-service/openapi-servlet/{ => models}/src/main/java/io/datahubproject/openapi/v2/models/GenericTimeseriesAspect.java (100%) rename metadata-service/openapi-servlet/{ => models}/src/main/java/io/datahubproject/openapi/v2/models/PatchOperation.java (100%) diff --git a/metadata-service/openapi-servlet/build.gradle b/metadata-service/openapi-servlet/build.gradle index 32c40c31df42d..04ed75d8800c2 100644 --- a/metadata-service/openapi-servlet/build.gradle +++ b/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 diff --git a/metadata-service/openapi-servlet/models/build.gradle b/metadata-service/openapi-servlet/models/build.gradle new file mode 100644 index 0000000000000..e4100b2d094e0 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/client/OpenApiClient.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/client/OpenApiClient.java new file mode 100644 index 0000000000000..267d95f1dddbf --- /dev/null +++ b/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; + } +} diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnRequest.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnRequest.java new file mode 100644 index 0000000000000..02afe8d40dd52 --- /dev/null +++ b/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 urns; + + @JsonProperty("aspectNames") + @Schema(required = true, description = "The list of aspect names to get") + List aspectNames; + + @JsonProperty("withSystemMetadata") + @Schema(required = true, description = "Whether or not to retrieve system metadata") + boolean withSystemMetadata; +} diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnResponse.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnResponse.java new file mode 100644 index 0000000000000..628733e4fd4ae --- /dev/null +++ b/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 entities; +} diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java similarity index 82% rename from metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java index f1e965ca05464..cb049c5ba131a 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java +++ b/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 aspects; public static class GenericEntityBuilder { diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericRelationship.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericRelationship.java similarity index 100% rename from metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericRelationship.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericRelationship.java diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java similarity index 100% rename from metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericTimeseriesAspect.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericTimeseriesAspect.java similarity index 100% rename from metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/GenericTimeseriesAspect.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericTimeseriesAspect.java diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/PatchOperation.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/PatchOperation.java similarity index 100% rename from metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/models/PatchOperation.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/PatchOperation.java diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java index 7a11d60a567f9..55bb8ebe625ae 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java @@ -38,6 +38,8 @@ 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; @@ -45,9 +47,11 @@ 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> 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 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 getEntity( @PathVariable("entityName") String entityName, @PathVariable("entityUrn") String entityUrn, diff --git a/settings.gradle b/settings.gradle index 57d7d1dbc36f2..27928ae7446a1 100644 --- a/settings.gradle +++ b/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'