Skip to content

Commit

Permalink
Batch get entity endpoints (#10880)
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevin1chun
kevin1chun authored Jul 10, 2024
1 parent cd932c3 commit 623b6f9
Show file tree
Hide file tree
Showing 6 changed files with 146 additions and 71 deletions.
4 changes: 2 additions & 2 deletions .../openapi-servlet/models/src/main/java/io/datahubproject/openapi/client/OpenApiClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@

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 io.datahubproject.openapi.models.BatchGetUrnRequest;
import io.datahubproject.openapi.models.BatchGetUrnResponse;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
Expand Down
2 changes: 1 addition & 1 deletion ...openapi/v2/models/BatchGetUrnRequest.java → ...ct/openapi/models/BatchGetUrnRequest.java
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package io.datahubproject.openapi.v2.models;
package io.datahubproject.openapi.models;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
Expand Down
6 changes: 3 additions & 3 deletions ...penapi/v2/models/BatchGetUrnResponse.java → ...t/openapi/models/BatchGetUrnResponse.java
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package io.datahubproject.openapi.v2.models;
package io.datahubproject.openapi.models;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
Expand All @@ -13,8 +13,8 @@
@Builder
@JsonInclude(JsonInclude.Include.NON_NULL)
@JsonDeserialize(builder = BatchGetUrnResponse.BatchGetUrnResponseBuilder.class)
public class BatchGetUrnResponse implements Serializable {
public class BatchGetUrnResponse<T extends GenericEntity> implements Serializable {
@JsonProperty("entities")
@Schema(description = "List of entity responses")
List<GenericEntityV2> entities;
List<T> entities;
}
49 changes: 46 additions & 3 deletions ...servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableSet;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.data.ByteString;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.entity.EnvelopedAspect;
Expand Down Expand Up @@ -48,6 +49,8 @@
import io.datahubproject.metadata.context.RequestContext;
import io.datahubproject.openapi.exception.InvalidUrnException;
import io.datahubproject.openapi.exception.UnauthorizedException;
import io.datahubproject.openapi.models.BatchGetUrnRequest;
import io.datahubproject.openapi.models.BatchGetUrnResponse;
import io.datahubproject.openapi.models.GenericEntity;
import io.datahubproject.openapi.models.GenericEntityScrollResult;
import io.swagger.v3.oas.annotations.Operation;
Expand All @@ -56,9 +59,7 @@
import java.lang.reflect.InvocationTargetException;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.annotation.Nonnull;
Expand Down Expand Up @@ -426,6 +427,48 @@ public ResponseEntity<List<E>> createEntity(
}
}

@Tag(name = "Generic Entities")
@PostMapping(value = "/batch/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE)
@Operation(summary = "Get a batch of entities")
public ResponseEntity<BatchGetUrnResponse<E>> getEntityBatch(
HttpServletRequest httpServletRequest,
@PathVariable("entityName") String entityName,
@RequestBody BatchGetUrnRequest request)
throws URISyntaxException {

List<Urn> urns = request.getUrns().stream().map(UrnUtils::getUrn).collect(Collectors.toList());

Authentication authentication = AuthenticationContext.getAuthentication();
if (!AuthUtil.isAPIAuthorizedEntityUrns(authentication, authorizationChain, READ, urns)) {
throw new UnauthorizedException(
authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities.");
}
OperationContext opContext =
OperationContext.asSession(
systemOperationContext,
RequestContext.builder()
.buildOpenapi(
authentication.getActor().toUrnStr(),
httpServletRequest,
"getEntityBatch",
entityName),
authorizationChain,
authentication,
true);

return ResponseEntity.of(
Optional.of(
BatchGetUrnResponse.<E>builder()
.entities(
new ArrayList<>(
buildEntityList(
opContext,
urns,
new HashSet<>(request.getAspectNames()),
request.isWithSystemMetadata())))
.build()));
}

@Tag(name = "Generic Aspects")
@DeleteMapping(value = "/{entityName}/{entityUrn:urn:li:.+}/{aspectName}")
@Operation(summary = "Delete an entity aspect.")
Expand Down
62 changes: 0 additions & 62 deletions ...enapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,9 @@
package io.datahubproject.openapi.v2.controller;

import static com.linkedin.metadata.authorization.ApiOperation.READ;

import com.datahub.authentication.Actor;
import com.datahub.authentication.Authentication;
import com.datahub.authentication.AuthenticationContext;
import com.datahub.authorization.AuthUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.data.ByteString;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.entity.EnvelopedAspect;
Expand All @@ -28,37 +22,23 @@
import com.linkedin.mxe.SystemMetadata;
import com.linkedin.util.Pair;
import io.datahubproject.metadata.context.OperationContext;
import io.datahubproject.metadata.context.RequestContext;
import io.datahubproject.openapi.controller.GenericEntitiesController;
import io.datahubproject.openapi.exception.InvalidUrnException;
import io.datahubproject.openapi.exception.UnauthorizedException;
import io.datahubproject.openapi.v2.models.BatchGetUrnRequest;
import io.datahubproject.openapi.v2.models.BatchGetUrnResponse;
import io.datahubproject.openapi.v2.models.GenericEntityScrollResultV2;
import io.datahubproject.openapi.v2.models.GenericEntityV2;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.servlet.http.HttpServletRequest;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

Expand All @@ -84,48 +64,6 @@ public GenericEntityScrollResultV2<GenericEntityV2> buildScrollResult(
.build();
}

@Tag(name = "Generic Entities")
@PostMapping(value = "/batch/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE)
@Operation(summary = "Get a batch of entities")
public ResponseEntity<BatchGetUrnResponse> getEntityBatch(
HttpServletRequest httpServletRequest,
@PathVariable("entityName") String entityName,
@RequestBody BatchGetUrnRequest request)
throws URISyntaxException {

List<Urn> urns = request.getUrns().stream().map(UrnUtils::getUrn).collect(Collectors.toList());

Authentication authentication = AuthenticationContext.getAuthentication();
if (!AuthUtil.isAPIAuthorizedEntityUrns(authentication, authorizationChain, READ, urns)) {
throw new UnauthorizedException(
authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities.");
}
OperationContext opContext =
OperationContext.asSession(
systemOperationContext,
RequestContext.builder()
.buildOpenapi(
authentication.getActor().toUrnStr(),
httpServletRequest,
"getEntityBatch",
entityName),
authorizationChain,
authentication,
true);

return ResponseEntity.of(
Optional.of(
BatchGetUrnResponse.builder()
.entities(
new ArrayList<>(
buildEntityList(
opContext,
urns,
new HashSet<>(request.getAspectNames()),
request.isWithSystemMetadata())))
.build()));
}

@Override
protected AspectsBatch toMCPBatch(
@Nonnull OperationContext opContext, String entityArrayList, Actor actor)
Expand Down
94 changes: 94 additions & 0 deletions ...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,8 @@ public static OpenAPI generateOpenApiSpec(EntityRegistry entityRegistry) {
entityName + ENTITY_RESPONSE_SUFFIX, buildEntitySchema(e, aspectNames, true));
components.addSchemas(
"Scroll" + entityName + ENTITY_RESPONSE_SUFFIX, buildEntityScrollSchema(e));
components.addSchemas(
"BatchGet" + entityName + ENTITY_RESPONSE_SUFFIX, buildEntityBatchGetSchema(e));
});
// Parameters
entityRegistry.getEntitySpecs().values().stream()
Expand All @@ -127,6 +129,9 @@ public static OpenAPI generateOpenApiSpec(EntityRegistry entityRegistry) {
paths.addPathItem(
String.format("/v3/entity/%s", e.getName().toLowerCase()),
buildListEntityPath(e));
paths.addPathItem(
String.format("/v3/entity/batch/%s", e.getName().toLowerCase()),
buildBatchGetEntityPath(e));
paths.addPathItem(
String.format("/v3/entity/%s/{urn}", e.getName().toLowerCase()),
buildSingleEntityPath(e));
Expand Down Expand Up @@ -360,6 +365,74 @@ private static PathItem buildListEntityPath(final EntitySpec entity) {
return result;
}

private static PathItem buildBatchGetEntityPath(final EntitySpec entity) {
final String upperFirst = toUpperFirst(entity.getName());
final PathItem result = new PathItem();
// Post Operation
final List<String> aspectNames =
entity.getAspectSpecs().stream()
.map(AspectSpec::getName)
.distinct()
.collect(Collectors.toList());
if (aspectNames.isEmpty()) {
aspectNames.add(entity.getKeyAspectName());
}
final Schema aspectNamesSchema =
new Schema()
.type(TYPE_ARRAY)
.description("List of aspect names to get")
.items(
new Schema()
.type(TYPE_STRING)
._enum(aspectNames)
._default(aspectNames.stream().findFirst().orElse(null)));
final Content requestBatch =
new Content()
.addMediaType(
"application/json",
new MediaType()
.schema(
new Schema()
.properties(
Map.of(
"urns",
new Schema()
.type(TYPE_ARRAY)
.items(
new Schema()
.type(TYPE_STRING)
.description("List of urns to get")),
"aspectNames", aspectNamesSchema,
"withSystemMetadata",
new Schema().type(TYPE_BOOLEAN)._default(true)))));
final ApiResponse apiResponse =
new ApiResponse()
.description("Create a batch of " + entity.getName() + " entities.")
.content(
new Content()
.addMediaType(
"application/json",
new MediaType()
.schema(
new Schema<>()
.$ref(
String.format(
"#/components/schemas/BatchGet%s%s",
upperFirst, ENTITY_RESPONSE_SUFFIX)))));
result.setPost(
new Operation()
.summary("Batch Get " + upperFirst + " entities.")
.tags(List.of(entity.getName() + " Entity"))
.requestBody(
new RequestBody()
.description("Batch Get " + entity.getName() + " entities.")
.required(true)
.content(requestBatch))
.responses(new ApiResponses().addApiResponse("200", apiResponse)));

return result;
}

private static void addExtraParameters(final Components components) {
components.addParameters(
"ScrollId" + MODEL_VERSION,
Expand Down Expand Up @@ -556,6 +629,27 @@ private static Schema buildEntityScrollSchema(final EntitySpec entity) {
toUpperFirst(entity.getName()), ENTITY_RESPONSE_SUFFIX))));
}

private static Schema buildEntityBatchGetSchema(final EntitySpec entity) {
return new Schema<>()
.type(TYPE_OBJECT)
.description("Batch get " + toUpperFirst(entity.getName()) + " objects.")
.required(List.of("entities"))
.addProperty(
"entities",
new Schema()
.properties(
Map.of(
"urn",
new Schema().type(TYPE_STRING).description("Entity key urn"),
"aspects",
new Schema<>()
.description(toUpperFirst(entity.getName()) + " object.")
.$ref(
String.format(
"#/components/schemas/%s%s",
toUpperFirst(entity.getName()), ENTITY_RESPONSE_SUFFIX)))));
}

private static Schema buildAspectRef(final String aspect, final boolean withSystemMetadata) {
// A non-required $ref property must be wrapped in a { allOf: [ $ref ] }
// object to allow the
Expand Down

0 comments on commit 623b6f9

Please sign in to comment.