From d8fc016ab3e3d66a84a4dd672b4e9ddd090dd125 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Thu, 6 Jun 2024 15:24:40 -0500 Subject: [PATCH] fix(openapiv3): v3 scroll response fix * Add skipCache flag * Separate v2 and v3 controllers * Return `entities` vs `results` for v3 per spec * Remove `aspects` nesting for v3 per spec * Properly throw error when missing `urn` field v3 * Properly throw error when missing `urn` or `aspects` field v2 --- .../openapi-analytics-servlet/build.gradle | 1 + .../openapi-entity-servlet/build.gradle | 1 + .../openapi-servlet/models/build.gradle | 8 + .../controller/GenericEntitiesController.java | 641 ++++++++++++++++ .../exception/UnauthorizedException.java | 0 .../openapi/models/GenericEntity.java | 7 + .../models/GenericEntityScrollResult.java | 3 + .../{v2 => }/models/GenericScrollResult.java | 2 +- .../v2/models/BatchGetUrnResponse.java | 2 +- .../models/GenericEntityScrollResultV2.java | 15 + ...enericEntity.java => GenericEntityV2.java} | 9 +- .../models/GenericEntityScrollResultV3.java | 15 + .../openapi/v3/models/GenericEntityV3.java | 77 ++ .../v2/controller/EntityController.java | 709 +++--------------- .../v2/controller/RelationshipController.java | 2 +- .../v2/controller/TimeseriesController.java | 2 +- .../v3/controller/EntityController.java | 203 ++++- 17 files changed, 1071 insertions(+), 626 deletions(-) create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java rename metadata-service/openapi-servlet/{ => models}/src/main/java/io/datahubproject/openapi/exception/UnauthorizedException.java (100%) create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntity.java create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntityScrollResult.java rename metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/{v2 => }/models/GenericScrollResult.java (79%) create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityScrollResultV2.java rename metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/{GenericEntity.java => GenericEntityV2.java} (90%) create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityScrollResultV3.java create mode 100644 metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityV3.java diff --git a/metadata-service/openapi-analytics-servlet/build.gradle b/metadata-service/openapi-analytics-servlet/build.gradle index 7c6568fa78f64..3a879cb1b0071 100644 --- a/metadata-service/openapi-analytics-servlet/build.gradle +++ b/metadata-service/openapi-analytics-servlet/build.gradle @@ -9,6 +9,7 @@ dependencies { implementation project(':metadata-service:auth-impl') implementation project(':metadata-service:factories') implementation project(':metadata-service:openapi-servlet') + implementation project(':metadata-service:openapi-servlet:models') implementation project(':metadata-models') implementation externalDependency.springBoot diff --git a/metadata-service/openapi-entity-servlet/build.gradle b/metadata-service/openapi-entity-servlet/build.gradle index 016ac6693f55b..4c2d587a81fd7 100644 --- a/metadata-service/openapi-entity-servlet/build.gradle +++ b/metadata-service/openapi-entity-servlet/build.gradle @@ -9,6 +9,7 @@ dependencies { implementation project(':metadata-service:auth-impl') implementation project(':metadata-service:factories') implementation project(':metadata-service:openapi-servlet') + implementation project(':metadata-service:openapi-servlet:models') implementation project(':metadata-models') implementation externalDependency.servletApi diff --git a/metadata-service/openapi-servlet/models/build.gradle b/metadata-service/openapi-servlet/models/build.gradle index e4100b2d094e0..a0e1a553fe814 100644 --- a/metadata-service/openapi-servlet/models/build.gradle +++ b/metadata-service/openapi-servlet/models/build.gradle @@ -6,6 +6,14 @@ dependencies { implementation project(':entity-registry') implementation project(':metadata-operation-context') implementation project(':metadata-auth:auth-api') + implementation project(':metadata-service:auth-impl') + implementation project(':metadata-io') + + implementation externalDependency.springWeb + implementation(externalDependency.springDocUI) { + exclude group: 'org.springframework.boot' + } + implementation externalDependency.swaggerAnnotations implementation externalDependency.jacksonDataBind implementation externalDependency.httpClient diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java new file mode 100644 index 0000000000000..a68d87434f7aa --- /dev/null +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java @@ -0,0 +1,641 @@ +package io.datahubproject.openapi.controller; + +import static com.linkedin.metadata.authorization.ApiOperation.CREATE; +import static com.linkedin.metadata.authorization.ApiOperation.DELETE; +import static com.linkedin.metadata.authorization.ApiOperation.EXISTS; +import static com.linkedin.metadata.authorization.ApiOperation.READ; +import static com.linkedin.metadata.authorization.ApiOperation.UPDATE; + +import com.datahub.authentication.Actor; +import com.datahub.authentication.Authentication; +import com.datahub.authentication.AuthenticationContext; +import com.datahub.authorization.AuthUtil; +import com.datahub.authorization.AuthorizerChain; +import com.datahub.util.RecordUtils; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +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; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.AspectsBatch; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.patch.GenericJsonPatch; +import com.linkedin.metadata.aspect.patch.template.common.GenericPatchTemplate; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.IngestResult; +import com.linkedin.metadata.entity.UpdateAspectResult; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.query.SearchFlags; +import com.linkedin.metadata.query.filter.SortCriterion; +import com.linkedin.metadata.query.filter.SortOrder; +import com.linkedin.metadata.search.ScrollResult; +import com.linkedin.metadata.search.SearchEntityArray; +import com.linkedin.metadata.search.SearchService; +import com.linkedin.metadata.utils.AuditStampUtils; +import com.linkedin.metadata.utils.GenericRecordUtils; +import com.linkedin.metadata.utils.SearchUtil; +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.exception.UnauthorizedException; +import io.datahubproject.openapi.models.GenericEntity; +import io.datahubproject.openapi.models.GenericEntityScrollResult; +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.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.DeleteMapping; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.PatchMapping; +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.RequestMethod; +import org.springframework.web.bind.annotation.RequestParam; + +public abstract class GenericEntitiesController< + E extends GenericEntity, S extends GenericEntityScrollResult> { + protected static final SearchFlags DEFAULT_SEARCH_FLAGS = + new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true); + + @Autowired protected EntityRegistry entityRegistry; + @Autowired protected SearchService searchService; + @Autowired protected EntityService entityService; + @Autowired protected AuthorizerChain authorizationChain; + @Autowired protected ObjectMapper objectMapper; + + @Qualifier("systemOperationContext") + @Autowired + protected OperationContext systemOperationContext; + + /** + * Returns scroll result entities + * + * @param searchEntities the entities to contain in the result + * @param aspectNames the aspect names present + * @param withSystemMetadata whether to include system metadata in the result + * @param scrollId the pagination token + * @return result containing entities/aspects + * @throws URISyntaxException parsing error + */ + protected abstract S buildScrollResult( + @Nonnull OperationContext opContext, + SearchEntityArray searchEntities, + Set aspectNames, + boolean withSystemMetadata, + @Nullable String scrollId) + throws URISyntaxException; + + protected abstract List buildEntityList( + @Nonnull OperationContext opContext, + List urns, + Set aspectNames, + boolean withSystemMetadata) + throws URISyntaxException; + + protected abstract List buildEntityList( + Set ingestResults, boolean withSystemMetadata); + + protected abstract E buildGenericEntity( + @Nonnull String aspectName, + @Nonnull UpdateAspectResult updateAspectResult, + boolean withSystemMetadata); + + protected abstract AspectsBatch toMCPBatch( + @Nonnull OperationContext opContext, String entityArrayList, Actor actor) + throws JsonProcessingException, URISyntaxException; + + @Tag(name = "Generic Entities", description = "API for interacting with generic entities.") + @GetMapping(value = "/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE) + @Operation(summary = "Scroll entities") + public ResponseEntity getEntities( + @PathVariable("entityName") String entityName, + @RequestParam(value = "aspectNames", defaultValue = "") Set aspectNames, + @RequestParam(value = "count", defaultValue = "10") Integer count, + @RequestParam(value = "query", defaultValue = "*") String query, + @RequestParam(value = "scrollId", required = false) String scrollId, + @RequestParam(value = "sort", required = false, defaultValue = "urn") String sortField, + @RequestParam(value = "sortOrder", required = false, defaultValue = "ASCENDING") + String sortOrder, + @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") + Boolean withSystemMetadata, + @RequestParam(value = "skipCache", required = false, defaultValue = "false") + Boolean skipCache) + throws URISyntaxException { + + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + Authentication authentication = AuthenticationContext.getAuthentication(); + + if (!AuthUtil.isAPIAuthorizedEntityType(authentication, authorizationChain, READ, entityName)) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); + } + + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("getEntities", entityName), + authorizationChain, + authentication, + true); + + // TODO: support additional and multiple sort params + SortCriterion sortCriterion = SearchUtil.sortBy(sortField, SortOrder.valueOf(sortOrder)); + + ScrollResult result = + searchService.scrollAcrossEntities( + opContext + .withSearchFlags(flags -> DEFAULT_SEARCH_FLAGS) + .withSearchFlags(flags -> flags.setSkipCache(skipCache)), + List.of(entitySpec.getName()), + query, + null, + sortCriterion, + scrollId, + null, + count); + + if (!AuthUtil.isAPIAuthorizedResult(authentication, authorizationChain, result)) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); + } + + return ResponseEntity.ok( + buildScrollResult( + opContext, + result.getEntities(), + aspectNames, + withSystemMetadata, + result.getScrollId())); + } + + @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, + @RequestParam(value = "aspectNames", defaultValue = "") Set aspectNames, + @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") + Boolean withSystemMetadata) + throws URISyntaxException { + + Urn urn = UrnUtils.getUrn(entityUrn); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, READ, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("getEntity", entityName), + authorizationChain, + authentication, + true); + + return ResponseEntity.of( + buildEntityList(opContext, List.of(urn), aspectNames, withSystemMetadata).stream() + .findFirst()); + } + + @Tag(name = "Generic Entities") + @RequestMapping( + value = "/{entityName}/{entityUrn}", + method = {RequestMethod.HEAD}) + @Operation(summary = "Entity exists") + public ResponseEntity headEntity( + @PathVariable("entityName") String entityName, @PathVariable("entityUrn") String entityUrn) { + + Urn urn = UrnUtils.getUrn(entityUrn); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, EXISTS, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + EXISTS + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("headEntity", entityName), + authorizationChain, + authentication, + true); + + return exists(opContext, urn, null) + ? ResponseEntity.noContent().build() + : ResponseEntity.notFound().build(); + } + + @Tag(name = "Generic Aspects", description = "API for generic aspects.") + @GetMapping( + value = "/{entityName}/{entityUrn}/{aspectName}", + produces = MediaType.APPLICATION_JSON_VALUE) + @Operation(summary = "Get an entity's generic aspect.") + public ResponseEntity getAspect( + @PathVariable("entityName") String entityName, + @PathVariable("entityUrn") String entityUrn, + @PathVariable("aspectName") String aspectName, + @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") + Boolean withSystemMetadata) + throws URISyntaxException { + + Urn urn = UrnUtils.getUrn(entityUrn); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, READ, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("getAspect", entityName), + authorizationChain, + authentication, + true); + + return ResponseEntity.of( + buildEntityList(opContext, List.of(urn), Set.of(aspectName), withSystemMetadata).stream() + .findFirst() + .flatMap( + e -> + e.getAspects().entrySet().stream() + .filter( + entry -> + entry.getKey().equals(lookupAspectSpec(urn, aspectName).getName())) + .map(Map.Entry::getValue) + .findFirst())); + } + + @Tag(name = "Generic Aspects") + @RequestMapping( + value = "/{entityName}/{entityUrn}/{aspectName}", + method = {RequestMethod.HEAD}) + @Operation(summary = "Whether an entity aspect exists.") + public ResponseEntity headAspect( + @PathVariable("entityName") String entityName, + @PathVariable("entityUrn") String entityUrn, + @PathVariable("aspectName") String aspectName) { + + Urn urn = UrnUtils.getUrn(entityUrn); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, EXISTS, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + EXISTS + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("headAspect", entityName), + authorizationChain, + authentication, + true); + + return exists(opContext, urn, lookupAspectSpec(urn, aspectName).getName()) + ? ResponseEntity.noContent().build() + : ResponseEntity.notFound().build(); + } + + @Tag(name = "Generic Entities") + @DeleteMapping(value = "/{entityName}/{entityUrn}") + @Operation(summary = "Delete an entity") + public void deleteEntity( + @PathVariable("entityName") String entityName, @PathVariable("entityUrn") String entityUrn) { + + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + Urn urn = UrnUtils.getUrn(entityUrn); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, DELETE, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + DELETE + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("deleteEntity", entityName), + authorizationChain, + authentication, + true); + + entityService.deleteAspect(opContext, entityUrn, entitySpec.getKeyAspectName(), Map.of(), true); + } + + @Tag(name = "Generic Entities") + @PostMapping(value = "/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE) + @Operation(summary = "Create a batch of entities.") + public ResponseEntity> createEntity( + @PathVariable("entityName") String entityName, + @RequestParam(value = "async", required = false, defaultValue = "true") Boolean async, + @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") + Boolean withSystemMetadata, + @RequestBody @Nonnull String jsonEntityList) + throws URISyntaxException, JsonProcessingException { + + Authentication authentication = AuthenticationContext.getAuthentication(); + + if (!AuthUtil.isAPIAuthorizedEntityType( + authentication, authorizationChain, CREATE, entityName)) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + CREATE + " entities."); + } + + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("createEntity", entityName), + authorizationChain, + authentication, + true); + + AspectsBatch batch = toMCPBatch(opContext, jsonEntityList, authentication.getActor()); + Set results = entityService.ingestProposal(opContext, batch, async); + + if (!async) { + return ResponseEntity.ok(buildEntityList(results, withSystemMetadata)); + } else { + return ResponseEntity.accepted().body(buildEntityList(results, withSystemMetadata)); + } + } + + @Tag(name = "Generic Aspects") + @DeleteMapping(value = "/{entityName}/{entityUrn}/{aspectName}") + @Operation(summary = "Delete an entity aspect.") + public void deleteAspect( + @PathVariable("entityName") String entityName, + @PathVariable("entityUrn") String entityUrn, + @PathVariable("aspectName") String aspectName) { + + Urn urn = UrnUtils.getUrn(entityUrn); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, DELETE, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + DELETE + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("deleteAspect", entityName), + authorizationChain, + authentication, + true); + + entityService.deleteAspect( + opContext, entityUrn, lookupAspectSpec(urn, aspectName).getName(), Map.of(), true); + } + + @Tag(name = "Generic Aspects") + @PostMapping( + value = "/{entityName}/{entityUrn}/{aspectName}", + produces = MediaType.APPLICATION_JSON_VALUE) + @Operation(summary = "Create an entity aspect.") + public ResponseEntity createAspect( + @PathVariable("entityName") String entityName, + @PathVariable("entityUrn") String entityUrn, + @PathVariable("aspectName") String aspectName, + @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") + Boolean withSystemMetadata, + @RequestParam(value = "createIfNotExists", required = false, defaultValue = "false") + Boolean createIfNotExists, + @RequestBody @Nonnull String jsonAspect) + throws URISyntaxException { + + Urn urn = UrnUtils.getUrn(entityUrn); + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + Authentication authentication = AuthenticationContext.getAuthentication(); + + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, CREATE, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + CREATE + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("createAspect", entityName), + authorizationChain, + authentication, + true); + + AspectSpec aspectSpec = lookupAspectSpec(entitySpec, aspectName); + ChangeMCP upsert = + toUpsertItem( + opContext.getRetrieverContext().get().getAspectRetriever(), + urn, + aspectSpec, + createIfNotExists, + jsonAspect, + authentication.getActor()); + + List results = + entityService.ingestAspects( + opContext, + AspectsBatchImpl.builder() + .retrieverContext(opContext.getRetrieverContext().get()) + .items(List.of(upsert)) + .build(), + true, + true); + + return ResponseEntity.of( + results.stream() + .findFirst() + .map(result -> buildGenericEntity(aspectName, result, withSystemMetadata))); + } + + @Tag(name = "Generic Aspects") + @PatchMapping( + value = "/{entityName}/{entityUrn}/{aspectName}", + consumes = "application/json-patch+json", + produces = MediaType.APPLICATION_JSON_VALUE) + @Operation(summary = "Patch an entity aspect. (Experimental)") + public ResponseEntity patchAspect( + @PathVariable("entityName") String entityName, + @PathVariable("entityUrn") String entityUrn, + @PathVariable("aspectName") String aspectName, + @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") + Boolean withSystemMetadata, + @RequestBody @Nonnull GenericJsonPatch patch) + throws URISyntaxException, + NoSuchMethodException, + InvocationTargetException, + InstantiationException, + IllegalAccessException { + + Urn urn = UrnUtils.getUrn(entityUrn); + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + Authentication authentication = AuthenticationContext.getAuthentication(); + if (!AuthUtil.isAPIAuthorizedEntityUrns( + authentication, authorizationChain, UPDATE, List.of(urn))) { + throw new UnauthorizedException( + authentication.getActor().toUrnStr() + " is unauthorized to " + UPDATE + " entities."); + } + OperationContext opContext = + OperationContext.asSession( + systemOperationContext, + RequestContext.builder().buildOpenapi("patchAspect", entityName), + authorizationChain, + authentication, + true); + + AspectSpec aspectSpec = lookupAspectSpec(entitySpec, aspectName); + RecordTemplate currentValue = entityService.getAspect(opContext, urn, aspectSpec.getName(), 0); + + GenericPatchTemplate genericPatchTemplate = + GenericPatchTemplate.builder() + .genericJsonPatch(patch) + .templateType(aspectSpec.getDataTemplateClass()) + .templateDefault( + aspectSpec.getDataTemplateClass().getDeclaredConstructor().newInstance()) + .build(); + ChangeMCP upsert = + toUpsertItem( + opContext.getRetrieverContext().get().getAspectRetriever(), + UrnUtils.getUrn(entityUrn), + aspectSpec, + currentValue, + genericPatchTemplate, + authentication.getActor()); + + List results = + entityService.ingestAspects( + opContext, + AspectsBatchImpl.builder() + .retrieverContext(opContext.getRetrieverContext().get()) + .items(List.of(upsert)) + .build(), + true, + true); + + return ResponseEntity.of( + results.stream() + .findFirst() + .map(result -> buildGenericEntity(aspectSpec.getName(), result, withSystemMetadata))); + } + + protected Boolean exists(@Nonnull OperationContext opContext, Urn urn, @Nullable String aspect) { + return aspect == null + ? entityService.exists(opContext, urn, true) + : entityService.exists(opContext, urn, aspect, true); + } + + protected Set resolveAspectNames(Set urns, Set requestedAspectNames) { + if (requestedAspectNames.isEmpty()) { + return urns.stream() + .flatMap(u -> entityRegistry.getEntitySpec(u.getEntityType()).getAspectSpecs().stream()) + .collect(Collectors.toSet()); + } else { + // ensure key is always present + return Stream.concat( + urns.stream() + .flatMap( + urn -> + requestedAspectNames.stream() + .map(aspectName -> lookupAspectSpec(urn, aspectName))), + urns.stream() + .map(u -> entityRegistry.getEntitySpec(u.getEntityType()).getKeyAspectSpec())) + .collect(Collectors.toSet()); + } + } + + protected Map> toAspectMap( + Urn urn, List aspects, boolean withSystemMetadata) { + return aspects.stream() + .map( + a -> + Map.entry( + a.getName(), + Pair.of( + toRecordTemplate(lookupAspectSpec(urn, a.getName()), a), + withSystemMetadata ? a.getSystemMetadata() : null))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + protected AspectSpec lookupAspectSpec(Urn urn, String aspectName) { + return lookupAspectSpec(entityRegistry.getEntitySpec(urn.getEntityType()), aspectName); + } + + protected RecordTemplate toRecordTemplate( + AspectSpec aspectSpec, EnvelopedAspect envelopedAspect) { + return RecordUtils.toRecordTemplate( + aspectSpec.getDataTemplateClass(), envelopedAspect.getValue().data()); + } + + protected ChangeMCP toUpsertItem( + @Nonnull AspectRetriever aspectRetriever, + Urn entityUrn, + AspectSpec aspectSpec, + Boolean createIfNotExists, + String jsonAspect, + Actor actor) + throws URISyntaxException { + 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), + GenericRecordUtils.JSON, + aspectSpec)) + .build(aspectRetriever); + } + + protected ChangeMCP toUpsertItem( + @Nonnull AspectRetriever aspectRetriever, + @Nonnull Urn urn, + @Nonnull AspectSpec aspectSpec, + @Nullable RecordTemplate currentValue, + @Nonnull GenericPatchTemplate genericPatchTemplate, + @Nonnull Actor actor) { + return ChangeItemImpl.fromPatch( + urn, + aspectSpec, + currentValue, + genericPatchTemplate, + AuditStampUtils.createAuditStamp(actor.toUrnStr()), + aspectRetriever); + } + + /** + * Case-insensitive fallback + * + * @return + */ + protected static AspectSpec lookupAspectSpec(EntitySpec entitySpec, String aspectName) { + return entitySpec.getAspectSpec(aspectName) != null + ? entitySpec.getAspectSpec(aspectName) + : entitySpec.getAspectSpecs().stream() + .filter(aspec -> aspec.getName().toLowerCase().equals(aspectName)) + .findFirst() + .get(); + } +} diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/exception/UnauthorizedException.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/exception/UnauthorizedException.java similarity index 100% rename from metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/exception/UnauthorizedException.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/exception/UnauthorizedException.java diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntity.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntity.java new file mode 100644 index 0000000000000..f25f8b89f8026 --- /dev/null +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntity.java @@ -0,0 +1,7 @@ +package io.datahubproject.openapi.models; + +import java.util.Map; + +public interface GenericEntity { + Map getAspects(); +} diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntityScrollResult.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntityScrollResult.java new file mode 100644 index 0000000000000..69b97956e0cf2 --- /dev/null +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntityScrollResult.java @@ -0,0 +1,3 @@ +package io.datahubproject.openapi.models; + +public interface GenericEntityScrollResult {} diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericScrollResult.java similarity index 79% rename from metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java rename to metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericScrollResult.java index 2befc83c00363..7864af3bb4cdd 100644 --- a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericScrollResult.java @@ -1,4 +1,4 @@ -package io.datahubproject.openapi.v2.models; +package io.datahubproject.openapi.models; import java.util.List; import lombok.Builder; 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 index 628733e4fd4ae..c1fd809ad3649 100644 --- 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 @@ -16,5 +16,5 @@ public class BatchGetUrnResponse implements Serializable { @JsonProperty("entities") @Schema(description = "List of entity responses") - List entities; + List entities; } diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityScrollResultV2.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityScrollResultV2.java new file mode 100644 index 0000000000000..685f45c60dbdc --- /dev/null +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityScrollResultV2.java @@ -0,0 +1,15 @@ +package io.datahubproject.openapi.v2.models; + +import io.datahubproject.openapi.models.GenericEntity; +import io.datahubproject.openapi.models.GenericEntityScrollResult; +import java.util.List; +import lombok.Builder; +import lombok.Data; + +@Data +@Builder +public class GenericEntityScrollResultV2 + implements GenericEntityScrollResult { + private String scrollId; + private List results; +} diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityV2.java similarity index 90% rename from metadata-service/openapi-servlet/models/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/GenericEntityV2.java index cb049c5ba131a..85d404fb57e0e 100644 --- a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityV2.java @@ -7,6 +7,7 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; +import io.datahubproject.openapi.models.GenericEntity; import io.swagger.v3.oas.annotations.media.Schema; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -23,7 +24,7 @@ @JsonInclude(JsonInclude.Include.NON_NULL) @NoArgsConstructor(force = true, access = AccessLevel.PRIVATE) @AllArgsConstructor -public class GenericEntity { +public class GenericEntityV2 implements GenericEntity { @JsonProperty("urn") @Schema(description = "Urn of the entity") private String urn; @@ -32,9 +33,9 @@ public class GenericEntity { @Schema(description = "Map of aspect name to aspect") private Map aspects; - public static class GenericEntityBuilder { + public static class GenericEntityV2Builder { - public GenericEntity build( + public GenericEntityV2 build( ObjectMapper objectMapper, Map> aspects) { Map jsonObjectMap = aspects.entrySet().stream() @@ -63,7 +64,7 @@ public GenericEntity build( }) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - return new GenericEntity(urn, jsonObjectMap); + return new GenericEntityV2(urn, jsonObjectMap); } } } diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityScrollResultV3.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityScrollResultV3.java new file mode 100644 index 0000000000000..265095f0f2c6e --- /dev/null +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityScrollResultV3.java @@ -0,0 +1,15 @@ +package io.datahubproject.openapi.v3.models; + +import io.datahubproject.openapi.models.GenericEntity; +import io.datahubproject.openapi.models.GenericEntityScrollResult; +import java.util.List; +import lombok.Builder; +import lombok.Data; + +@Data +@Builder +public class GenericEntityScrollResultV3 + implements GenericEntityScrollResult { + private String scrollId; + private List entities; +} diff --git a/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityV3.java b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityV3.java new file mode 100644 index 0000000000000..2e030390dd9cb --- /dev/null +++ b/metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityV3.java @@ -0,0 +1,77 @@ +package io.datahubproject.openapi.v3.models; + +import com.datahub.util.RecordUtils; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; +import io.datahubproject.openapi.models.GenericEntity; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Data; +import lombok.EqualsAndHashCode; + +@EqualsAndHashCode(callSuper = true) +@Data +@Builder +@JsonInclude(JsonInclude.Include.NON_NULL) +@AllArgsConstructor +public class GenericEntityV3 extends LinkedHashMap implements GenericEntity { + + public GenericEntityV3(Map m) { + super(m); + } + + @Override + public Map getAspects() { + return this; + } + + public static class GenericEntityV3Builder { + + public GenericEntityV3 build( + ObjectMapper objectMapper, + @Nonnull Urn urn, + Map> aspects) { + Map jsonObjectMap = + aspects.entrySet().stream() + .map( + e -> { + try { + Map valueMap = + Map.of( + "value", + objectMapper.readTree( + RecordUtils.toJsonString(e.getValue().getFirst()) + .getBytes(StandardCharsets.UTF_8))); + + if (e.getValue().getSecond() != null) { + return Map.entry( + e.getKey(), + Map.of( + "systemMetadata", e.getValue().getSecond(), + "value", valueMap.get("value"))); + } else { + return Map.entry(e.getKey(), Map.of("value", valueMap.get("value"))); + } + } catch (IOException ex) { + throw new RuntimeException(ex); + } + }) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + + GenericEntityV3 genericEntityV3 = new GenericEntityV3(); + genericEntityV3.put("urn", urn.toString()); + genericEntityV3.putAll(jsonObjectMap); + return genericEntityV3; + } + } +} 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 41cf972079c25..23cd89147173a 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 @@ -1,63 +1,42 @@ package io.datahubproject.openapi.v2.controller; -import static com.linkedin.metadata.authorization.ApiOperation.CREATE; -import static com.linkedin.metadata.authorization.ApiOperation.DELETE; -import static com.linkedin.metadata.authorization.ApiOperation.EXISTS; import static com.linkedin.metadata.authorization.ApiOperation.READ; -import static com.linkedin.metadata.authorization.ApiOperation.UPDATE; import com.datahub.authentication.Actor; import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.authorization.AuthUtil; -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.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; -import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.aspect.batch.BatchItem; -import com.linkedin.metadata.aspect.batch.ChangeMCP; -import com.linkedin.metadata.aspect.patch.GenericJsonPatch; -import com.linkedin.metadata.aspect.patch.template.common.GenericPatchTemplate; import com.linkedin.metadata.entity.EntityApiUtils; -import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.IngestResult; import com.linkedin.metadata.entity.UpdateAspectResult; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.query.SearchFlags; -import com.linkedin.metadata.query.filter.SortCriterion; -import com.linkedin.metadata.query.filter.SortOrder; -import com.linkedin.metadata.search.ScrollResult; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchEntityArray; -import com.linkedin.metadata.search.SearchService; import com.linkedin.metadata.utils.AuditStampUtils; import com.linkedin.metadata.utils.GenericRecordUtils; -import com.linkedin.metadata.utils.SearchUtil; 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.UnauthorizedException; 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.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 java.lang.reflect.InvocationTargetException; import java.net.URISyntaxException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; @@ -69,100 +48,38 @@ import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; -import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.springframework.beans.factory.annotation.Autowired; -import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.http.MediaType; import org.springframework.http.ResponseEntity; -import org.springframework.web.bind.annotation.DeleteMapping; -import org.springframework.web.bind.annotation.GetMapping; -import org.springframework.web.bind.annotation.PatchMapping; 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.RequestMethod; -import org.springframework.web.bind.annotation.RequestParam; import org.springframework.web.bind.annotation.RestController; @RestController @RequiredArgsConstructor @RequestMapping("/v2/entity") @Slf4j -public class EntityController { - private static final SearchFlags DEFAULT_SEARCH_FLAGS = - new SearchFlags().setFulltext(false).setSkipAggregates(true).setSkipHighlighting(true); - @Autowired private EntityRegistry entityRegistry; - @Autowired private SearchService searchService; - @Autowired private EntityService entityService; - @Autowired private AuthorizerChain authorizationChain; - @Autowired private ObjectMapper objectMapper; +public class EntityController + extends GenericEntitiesController< + GenericEntityV2, GenericEntityScrollResultV2> { - @Qualifier("systemOperationContext") - @Autowired - private OperationContext systemOperationContext; - - @Tag(name = "Generic Entities", description = "API for interacting with generic entities.") - @GetMapping(value = "/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE) - @Operation(summary = "Scroll entities") - public ResponseEntity> getEntities( - @PathVariable("entityName") String entityName, - @RequestParam(value = "aspectNames", defaultValue = "") Set aspectNames, - @RequestParam(value = "count", defaultValue = "10") Integer count, - @RequestParam(value = "query", defaultValue = "*") String query, - @RequestParam(value = "scrollId", required = false) String scrollId, - @RequestParam(value = "sort", required = false, defaultValue = "urn") String sortField, - @RequestParam(value = "sortOrder", required = false, defaultValue = "ASCENDING") - String sortOrder, - @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") - Boolean withSystemMetadata) + @Override + public GenericEntityScrollResultV2 buildScrollResult( + @Nonnull OperationContext opContext, + SearchEntityArray searchEntities, + Set aspectNames, + boolean withSystemMetadata, + @Nullable String scrollId) throws URISyntaxException { - - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); - Authentication authentication = AuthenticationContext.getAuthentication(); - - if (!AuthUtil.isAPIAuthorizedEntityType(authentication, authorizationChain, READ, entityName)) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); - } - - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("getEntities", entityName), - authorizationChain, - authentication, - true); - - // TODO: support additional and multiple sort params - SortCriterion sortCriterion = SearchUtil.sortBy(sortField, SortOrder.valueOf(sortOrder)); - - ScrollResult result = - searchService.scrollAcrossEntities( - opContext.withSearchFlags(flags -> DEFAULT_SEARCH_FLAGS), - List.of(entitySpec.getName()), - query, - null, - sortCriterion, - scrollId, - null, - count); - - if (!AuthUtil.isAPIAuthorizedResult(authentication, authorizationChain, result)) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); - } - - return ResponseEntity.ok( - GenericScrollResult.builder() - .results( - toRecordTemplates(opContext, result.getEntities(), aspectNames, withSystemMetadata)) - .scrollId(result.getScrollId()) - .build()); + return GenericEntityScrollResultV2.builder() + .results(toRecordTemplates(opContext, searchEntities, aspectNames, withSystemMetadata)) + .scrollId(scrollId) + .build(); } @Tag(name = "Generic Entities") @@ -192,7 +109,7 @@ public ResponseEntity getEntityBatch( BatchGetUrnResponse.builder() .entities( new ArrayList<>( - toRecordTemplates( + buildEntityList( opContext, urns, new HashSet<>(request.getAspectNames()), @@ -200,506 +117,10 @@ public ResponseEntity getEntityBatch( .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, - @RequestParam(value = "aspectNames", defaultValue = "") Set aspectNames, - @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") - Boolean withSystemMetadata) - throws URISyntaxException { - - Urn urn = UrnUtils.getUrn(entityUrn); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, READ, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("getEntity", entityName), - authorizationChain, - authentication, - true); - - return ResponseEntity.of( - toRecordTemplates(opContext, List.of(urn), aspectNames, withSystemMetadata).stream() - .findFirst()); - } - - @Tag(name = "Generic Entities") - @RequestMapping( - value = "/{entityName}/{entityUrn}", - method = {RequestMethod.HEAD}) - @Operation(summary = "Entity exists") - public ResponseEntity headEntity( - @PathVariable("entityName") String entityName, @PathVariable("entityUrn") String entityUrn) { - - Urn urn = UrnUtils.getUrn(entityUrn); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, EXISTS, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + EXISTS + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("headEntity", entityName), - authorizationChain, - authentication, - true); - - return exists(opContext, urn, null) - ? ResponseEntity.noContent().build() - : ResponseEntity.notFound().build(); - } - - @Tag(name = "Generic Aspects", description = "API for generic aspects.") - @GetMapping( - value = "/{entityName}/{entityUrn}/{aspectName}", - produces = MediaType.APPLICATION_JSON_VALUE) - @Operation(summary = "Get an entity's generic aspect.") - public ResponseEntity getAspect( - @PathVariable("entityName") String entityName, - @PathVariable("entityUrn") String entityUrn, - @PathVariable("aspectName") String aspectName, - @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") - Boolean withSystemMetadata) - throws URISyntaxException { - - Urn urn = UrnUtils.getUrn(entityUrn); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, READ, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + READ + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("getAspect", entityName), - authorizationChain, - authentication, - true); - - return ResponseEntity.of( - toRecordTemplates(opContext, List.of(urn), Set.of(aspectName), withSystemMetadata).stream() - .findFirst() - .flatMap( - e -> - e.getAspects().entrySet().stream() - .filter( - entry -> - entry.getKey().equals(lookupAspectSpec(urn, aspectName).getName())) - .map(Map.Entry::getValue) - .findFirst())); - } - - @Tag(name = "Generic Aspects") - @RequestMapping( - value = "/{entityName}/{entityUrn}/{aspectName}", - method = {RequestMethod.HEAD}) - @Operation(summary = "Whether an entity aspect exists.") - public ResponseEntity headAspect( - @PathVariable("entityName") String entityName, - @PathVariable("entityUrn") String entityUrn, - @PathVariable("aspectName") String aspectName) { - - Urn urn = UrnUtils.getUrn(entityUrn); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, EXISTS, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + EXISTS + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("headAspect", entityName), - authorizationChain, - authentication, - true); - - return exists(opContext, urn, lookupAspectSpec(urn, aspectName).getName()) - ? ResponseEntity.noContent().build() - : ResponseEntity.notFound().build(); - } - - @Tag(name = "Generic Entities") - @DeleteMapping(value = "/{entityName}/{entityUrn}") - @Operation(summary = "Delete an entity") - public void deleteEntity( - @PathVariable("entityName") String entityName, @PathVariable("entityUrn") String entityUrn) { - - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); - Urn urn = UrnUtils.getUrn(entityUrn); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, DELETE, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + DELETE + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("deleteEntity", entityName), - authorizationChain, - authentication, - true); - - entityService.deleteAspect(opContext, entityUrn, entitySpec.getKeyAspectName(), Map.of(), true); - } - - @Tag(name = "Generic Entities") - @PostMapping(value = "/{entityName}", produces = MediaType.APPLICATION_JSON_VALUE) - @Operation(summary = "Create a batch of entities.") - public ResponseEntity> createEntity( - @PathVariable("entityName") String entityName, - @RequestParam(value = "async", required = false, defaultValue = "true") Boolean async, - @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") - Boolean withSystemMetadata, - @RequestBody @Nonnull String jsonEntityList) - throws URISyntaxException, JsonProcessingException { - - Authentication authentication = AuthenticationContext.getAuthentication(); - - if (!AuthUtil.isAPIAuthorizedEntityType( - authentication, authorizationChain, CREATE, entityName)) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + CREATE + " entities."); - } - - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("createEntity", entityName), - authorizationChain, - authentication, - true); - - AspectsBatch batch = toMCPBatch(opContext, jsonEntityList, authentication.getActor()); - Set results = entityService.ingestProposal(opContext, batch, async); - - if (!async) { - return ResponseEntity.ok(toEntityListResponse(results, withSystemMetadata)); - } else { - return ResponseEntity.accepted().body(toEntityListResponse(results, withSystemMetadata)); - } - } - - @Tag(name = "Generic Aspects") - @DeleteMapping(value = "/{entityName}/{entityUrn}/{aspectName}") - @Operation(summary = "Delete an entity aspect.") - public void deleteAspect( - @PathVariable("entityName") String entityName, - @PathVariable("entityUrn") String entityUrn, - @PathVariable("aspectName") String aspectName) { - - Urn urn = UrnUtils.getUrn(entityUrn); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, DELETE, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + DELETE + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("deleteAspect", entityName), - authorizationChain, - authentication, - true); - - entityService.deleteAspect( - opContext, entityUrn, lookupAspectSpec(urn, aspectName).getName(), Map.of(), true); - } - - @Tag(name = "Generic Aspects") - @PostMapping( - value = "/{entityName}/{entityUrn}/{aspectName}", - produces = MediaType.APPLICATION_JSON_VALUE) - @Operation(summary = "Create an entity aspect.") - public ResponseEntity createAspect( - @PathVariable("entityName") String entityName, - @PathVariable("entityUrn") String entityUrn, - @PathVariable("aspectName") String aspectName, - @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") - Boolean withSystemMetadata, - @RequestParam(value = "createIfNotExists", required = false, defaultValue = "false") - Boolean createIfNotExists, - @RequestBody @Nonnull String jsonAspect) - throws URISyntaxException { - - Urn urn = UrnUtils.getUrn(entityUrn); - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); - Authentication authentication = AuthenticationContext.getAuthentication(); - - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, CREATE, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + CREATE + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("createAspect", entityName), - authorizationChain, - authentication, - true); - - AspectSpec aspectSpec = lookupAspectSpec(entitySpec, aspectName); - ChangeMCP upsert = - toUpsertItem( - opContext.getRetrieverContext().get().getAspectRetriever(), - urn, - aspectSpec, - createIfNotExists, - jsonAspect, - authentication.getActor()); - - List results = - entityService.ingestAspects( - opContext, - AspectsBatchImpl.builder() - .retrieverContext(opContext.getRetrieverContext().get()) - .items(List.of(upsert)) - .build(), - true, - true); - - return ResponseEntity.of( - results.stream() - .findFirst() - .map( - result -> - GenericEntity.builder() - .urn(result.getUrn().toString()) - .build( - objectMapper, - Map.of( - aspectName, - Pair.of( - result.getNewValue(), - withSystemMetadata ? result.getNewSystemMetadata() : null))))); - } - - @Tag(name = "Generic Aspects") - @PatchMapping( - value = "/{entityName}/{entityUrn}/{aspectName}", - consumes = "application/json-patch+json", - produces = MediaType.APPLICATION_JSON_VALUE) - @Operation(summary = "Patch an entity aspect. (Experimental)") - public ResponseEntity patchAspect( - @PathVariable("entityName") String entityName, - @PathVariable("entityUrn") String entityUrn, - @PathVariable("aspectName") String aspectName, - @RequestParam(value = "systemMetadata", required = false, defaultValue = "false") - Boolean withSystemMetadata, - @RequestBody @Nonnull GenericJsonPatch patch) - throws URISyntaxException, - NoSuchMethodException, - InvocationTargetException, - InstantiationException, - IllegalAccessException { - - Urn urn = UrnUtils.getUrn(entityUrn); - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); - Authentication authentication = AuthenticationContext.getAuthentication(); - if (!AuthUtil.isAPIAuthorizedEntityUrns( - authentication, authorizationChain, UPDATE, List.of(urn))) { - throw new UnauthorizedException( - authentication.getActor().toUrnStr() + " is unauthorized to " + UPDATE + " entities."); - } - OperationContext opContext = - OperationContext.asSession( - systemOperationContext, - RequestContext.builder().buildOpenapi("patchAspect", entityName), - authorizationChain, - authentication, - true); - - AspectSpec aspectSpec = lookupAspectSpec(entitySpec, aspectName); - RecordTemplate currentValue = entityService.getAspect(opContext, urn, aspectSpec.getName(), 0); - - GenericPatchTemplate genericPatchTemplate = - GenericPatchTemplate.builder() - .genericJsonPatch(patch) - .templateType(aspectSpec.getDataTemplateClass()) - .templateDefault( - aspectSpec.getDataTemplateClass().getDeclaredConstructor().newInstance()) - .build(); - ChangeMCP upsert = - toUpsertItem( - opContext.getRetrieverContext().get().getAspectRetriever(), - UrnUtils.getUrn(entityUrn), - aspectSpec, - currentValue, - genericPatchTemplate, - authentication.getActor()); - - List results = - entityService.ingestAspects( - opContext, - AspectsBatchImpl.builder() - .retrieverContext(opContext.getRetrieverContext().get()) - .items(List.of(upsert)) - .build(), - true, - true); - - return ResponseEntity.of( - results.stream() - .findFirst() - .map( - result -> - GenericEntity.builder() - .urn(result.getUrn().toString()) - .build( - objectMapper, - Map.of( - aspectSpec.getName(), - Pair.of( - result.getNewValue(), - withSystemMetadata ? result.getNewSystemMetadata() : null))))); - } - - private List toRecordTemplates( - @Nonnull OperationContext opContext, - SearchEntityArray searchEntities, - Set aspectNames, - boolean withSystemMetadata) - throws URISyntaxException { - return toRecordTemplates( - opContext, - searchEntities.stream().map(SearchEntity::getEntity).collect(Collectors.toList()), - aspectNames, - withSystemMetadata); - } - - private Boolean exists(@Nonnull OperationContext opContext, Urn urn, @Nullable String aspect) { - return aspect == null - ? entityService.exists(opContext, urn, true) - : entityService.exists(opContext, urn, aspect, true); - } - - private List toRecordTemplates( - @Nonnull OperationContext opContext, - List urns, - Set aspectNames, - boolean withSystemMetadata) - throws URISyntaxException { - if (urns.isEmpty()) { - return List.of(); - } else { - Set urnsSet = new HashSet<>(urns); - - Map> aspects = - entityService.getLatestEnvelopedAspects( - opContext, - urnsSet, - resolveAspectNames(urnsSet, aspectNames).stream() - .map(AspectSpec::getName) - .collect(Collectors.toSet())); - - return urns.stream() - .map( - u -> - GenericEntity.builder() - .urn(u.toString()) - .build( - objectMapper, - toAspectMap(u, aspects.getOrDefault(u, List.of()), withSystemMetadata))) - .collect(Collectors.toList()); - } - } - - private Set resolveAspectNames(Set urns, Set requestedAspectNames) { - if (requestedAspectNames.isEmpty()) { - return urns.stream() - .flatMap(u -> entityRegistry.getEntitySpec(u.getEntityType()).getAspectSpecs().stream()) - .collect(Collectors.toSet()); - } else { - // ensure key is always present - return Stream.concat( - urns.stream() - .flatMap( - urn -> - requestedAspectNames.stream() - .map(aspectName -> lookupAspectSpec(urn, aspectName))), - urns.stream() - .map(u -> entityRegistry.getEntitySpec(u.getEntityType()).getKeyAspectSpec())) - .collect(Collectors.toSet()); - } - } - - private Map> toAspectMap( - Urn urn, List aspects, boolean withSystemMetadata) { - return aspects.stream() - .map( - a -> - Map.entry( - a.getName(), - Pair.of( - toRecordTemplate(lookupAspectSpec(urn, a.getName()), a), - withSystemMetadata ? a.getSystemMetadata() : null))) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - } - - private AspectSpec lookupAspectSpec(Urn urn, String aspectName) { - return lookupAspectSpec(entityRegistry.getEntitySpec(urn.getEntityType()), aspectName); - } - - private RecordTemplate toRecordTemplate(AspectSpec aspectSpec, EnvelopedAspect envelopedAspect) { - return RecordUtils.toRecordTemplate( - aspectSpec.getDataTemplateClass(), envelopedAspect.getValue().data()); - } - - private ChangeMCP toUpsertItem( - @Nonnull AspectRetriever aspectRetriever, - Urn entityUrn, - AspectSpec aspectSpec, - Boolean createIfNotExists, - String jsonAspect, - Actor actor) - throws URISyntaxException { - 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), - GenericRecordUtils.JSON, - aspectSpec)) - .build(aspectRetriever); - } - - private ChangeMCP toUpsertItem( - @Nonnull AspectRetriever aspectRetriever, - @Nonnull Urn urn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate currentValue, - @Nonnull GenericPatchTemplate genericPatchTemplate, - @Nonnull Actor actor) { - return ChangeItemImpl.fromPatch( - urn, - aspectSpec, - currentValue, - genericPatchTemplate, - AuditStampUtils.createAuditStamp(actor.toUrnStr()), - aspectRetriever); - } - - private AspectsBatch toMCPBatch( + @Override + protected AspectsBatch toMCPBatch( @Nonnull OperationContext opContext, String entityArrayList, Actor actor) - throws JsonProcessingException, URISyntaxException { + throws JsonProcessingException { JsonNode entities = objectMapper.readTree(entityArrayList); List items = new LinkedList<>(); @@ -707,8 +128,14 @@ private AspectsBatch toMCPBatch( Iterator entityItr = entities.iterator(); while (entityItr.hasNext()) { JsonNode entity = entityItr.next(); + if (!entity.has("urn")) { + throw new IllegalArgumentException("Missing `urn` field"); + } Urn entityUrn = UrnUtils.getUrn(entity.get("urn").asText()); + if (!entity.has("aspects")) { + throw new IllegalArgumentException("Missing `aspects` field"); + } Iterator> aspectItr = entity.get("aspects").fields(); while (aspectItr.hasNext()) { Map.Entry aspect = aspectItr.next(); @@ -747,9 +174,71 @@ private AspectsBatch toMCPBatch( .build(); } - public List toEntityListResponse( + @Override + protected List buildEntityList( + @Nonnull OperationContext opContext, + List urns, + Set aspectNames, + boolean withSystemMetadata) + throws URISyntaxException { + if (urns.isEmpty()) { + return List.of(); + } else { + Set urnsSet = new HashSet<>(urns); + + Map> aspects = + entityService.getLatestEnvelopedAspects( + opContext, + urnsSet, + resolveAspectNames(urnsSet, aspectNames).stream() + .map(AspectSpec::getName) + .collect(Collectors.toSet())); + + return urns.stream() + .map( + u -> + GenericEntityV2.builder() + .urn(u.toString()) + .build( + objectMapper, + toAspectMap(u, aspects.getOrDefault(u, List.of()), withSystemMetadata))) + .collect(Collectors.toList()); + } + } + + @Override + protected GenericEntityV2 buildGenericEntity( + @Nonnull String aspectName, + @Nonnull UpdateAspectResult updateAspectResult, + boolean withSystemMetadata) { + return GenericEntityV2.builder() + .urn(updateAspectResult.getUrn().toString()) + .build( + objectMapper, + Map.of( + aspectName, + Pair.of( + updateAspectResult.getNewValue(), + withSystemMetadata ? updateAspectResult.getNewSystemMetadata() : null))); + } + + private List toRecordTemplates( + @Nonnull OperationContext opContext, + SearchEntityArray searchEntities, + Set aspectNames, + boolean withSystemMetadata) + throws URISyntaxException { + return buildEntityList( + opContext, + searchEntities.stream().map(SearchEntity::getEntity).collect(Collectors.toList()), + aspectNames, + withSystemMetadata); + } + + @Override + protected List buildEntityList( Set ingestResults, boolean withSystemMetadata) { - List responseList = new LinkedList<>(); + List responseList = new LinkedList<>(); Map> entityMap = ingestResults.stream().collect(Collectors.groupingBy(IngestResult::getUrn)); @@ -765,24 +254,10 @@ public List toEntityListResponse( withSystemMetadata ? ingest.getRequest().getSystemMetadata() : null))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); responseList.add( - GenericEntity.builder() + GenericEntityV2.builder() .urn(urnAspects.getKey().toString()) .build(objectMapper, aspectsMap)); } return responseList; } - - /** - * Case-insensitive fallback - * - * @return - */ - private static AspectSpec lookupAspectSpec(EntitySpec entitySpec, String aspectName) { - return entitySpec.getAspectSpec(aspectName) != null - ? entitySpec.getAspectSpec(aspectName) - : entitySpec.getAspectSpecs().stream() - .filter(aspec -> aspec.getName().toLowerCase().equals(aspectName)) - .findFirst() - .get(); - } } diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java index ac0b9dd8c03ef..3e46e10857fbd 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java @@ -18,8 +18,8 @@ import com.linkedin.metadata.query.filter.RelationshipFilter; import com.linkedin.metadata.search.utils.QueryUtils; import io.datahubproject.openapi.exception.UnauthorizedException; +import io.datahubproject.openapi.models.GenericScrollResult; import io.datahubproject.openapi.v2.models.GenericRelationship; -import io.datahubproject.openapi.v2.models.GenericScrollResult; import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.tags.Tag; import java.util.Arrays; diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java index 267122d71a57b..1c404006d97a4 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/TimeseriesController.java @@ -19,7 +19,7 @@ import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.metadata.context.RequestContext; import io.datahubproject.openapi.exception.UnauthorizedException; -import io.datahubproject.openapi.v2.models.GenericScrollResult; +import io.datahubproject.openapi.models.GenericScrollResult; import io.datahubproject.openapi.v2.models.GenericTimeseriesAspect; import io.swagger.v3.oas.annotations.tags.Tag; import java.net.URISyntaxException; diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java index 689efbf8bc6ec..20e917f1f452e 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java @@ -1,5 +1,43 @@ package io.datahubproject.openapi.v3.controller; +import com.datahub.authentication.Actor; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +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; +import com.linkedin.metadata.aspect.batch.AspectsBatch; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.entity.EntityApiUtils; +import com.linkedin.metadata.entity.IngestResult; +import com.linkedin.metadata.entity.UpdateAspectResult; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.search.SearchEntity; +import com.linkedin.metadata.search.SearchEntityArray; +import com.linkedin.metadata.utils.AuditStampUtils; +import com.linkedin.metadata.utils.GenericRecordUtils; +import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.openapi.controller.GenericEntitiesController; +import io.datahubproject.openapi.v3.models.GenericEntityScrollResultV3; +import io.datahubproject.openapi.v3.models.GenericEntityV3; +import java.net.URISyntaxException; +import java.nio.charset.StandardCharsets; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +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.web.bind.annotation.RequestMapping; @@ -9,4 +47,167 @@ @RequiredArgsConstructor @RequestMapping("/v3/entity") @Slf4j -public class EntityController extends io.datahubproject.openapi.v2.controller.EntityController {} +public class EntityController + extends GenericEntitiesController< + GenericEntityV3, GenericEntityScrollResultV3> { + + @Override + public GenericEntityScrollResultV3 buildScrollResult( + @Nonnull OperationContext opContext, + SearchEntityArray searchEntities, + Set aspectNames, + boolean withSystemMetadata, + @Nullable String scrollId) + throws URISyntaxException { + return GenericEntityScrollResultV3.builder() + .entities(toRecordTemplates(opContext, searchEntities, aspectNames, withSystemMetadata)) + .scrollId(scrollId) + .build(); + } + + @Override + protected List buildEntityList( + @Nonnull OperationContext opContext, + List urns, + Set aspectNames, + boolean withSystemMetadata) + throws URISyntaxException { + if (urns.isEmpty()) { + return List.of(); + } else { + Set urnsSet = new HashSet<>(urns); + + Map> aspects = + entityService.getLatestEnvelopedAspects( + opContext, + urnsSet, + resolveAspectNames(urnsSet, aspectNames).stream() + .map(AspectSpec::getName) + .collect(Collectors.toSet())); + + return urns.stream() + .map( + u -> + GenericEntityV3.builder() + .build( + objectMapper, + u, + toAspectMap(u, aspects.getOrDefault(u, List.of()), withSystemMetadata))) + .collect(Collectors.toList()); + } + } + + @Override + protected List buildEntityList( + Set ingestResults, boolean withSystemMetadata) { + List responseList = new LinkedList<>(); + + Map> entityMap = + ingestResults.stream().collect(Collectors.groupingBy(IngestResult::getUrn)); + for (Map.Entry> urnAspects : entityMap.entrySet()) { + Map> aspectsMap = + urnAspects.getValue().stream() + .map( + ingest -> + Map.entry( + ingest.getRequest().getAspectName(), + Pair.of( + ingest.getRequest().getRecordTemplate(), + withSystemMetadata ? ingest.getRequest().getSystemMetadata() : null))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + responseList.add( + GenericEntityV3.builder().build(objectMapper, urnAspects.getKey(), aspectsMap)); + } + return responseList; + } + + @Override + protected GenericEntityV3 buildGenericEntity( + @Nonnull String aspectName, + @Nonnull UpdateAspectResult updateAspectResult, + boolean withSystemMetadata) { + return GenericEntityV3.builder() + .build( + objectMapper, + updateAspectResult.getUrn(), + Map.of( + aspectName, + Pair.of( + updateAspectResult.getNewValue(), + withSystemMetadata ? updateAspectResult.getNewSystemMetadata() : null))); + } + + private List toRecordTemplates( + @Nonnull OperationContext opContext, + SearchEntityArray searchEntities, + Set aspectNames, + boolean withSystemMetadata) + throws URISyntaxException { + return buildEntityList( + opContext, + searchEntities.stream().map(SearchEntity::getEntity).collect(Collectors.toList()), + aspectNames, + withSystemMetadata); + } + + @Override + protected AspectsBatch toMCPBatch( + @Nonnull OperationContext opContext, String entityArrayList, Actor actor) + throws JsonProcessingException { + JsonNode entities = objectMapper.readTree(entityArrayList); + + List items = new LinkedList<>(); + if (entities.isArray()) { + Iterator entityItr = entities.iterator(); + while (entityItr.hasNext()) { + JsonNode entity = entityItr.next(); + if (!entity.has("urn")) { + throw new IllegalArgumentException("Missing `urn` field"); + } + Urn entityUrn = UrnUtils.getUrn(entity.get("urn").asText()); + + Iterator> aspectItr = entity.fields(); + while (aspectItr.hasNext()) { + Map.Entry aspect = aspectItr.next(); + + if ("urn".equals(aspect.getKey())) { + continue; + } + + AspectSpec aspectSpec = lookupAspectSpec(entityUrn, aspect.getKey()); + + if (aspectSpec != null) { + + SystemMetadata systemMetadata = null; + if (aspect.getValue().has("systemMetadata")) { + systemMetadata = + EntityApiUtils.parseSystemMetadata( + objectMapper.writeValueAsString(aspect.getValue().get("systemMetadata"))); + ((ObjectNode) aspect.getValue()).remove("systemMetadata"); + } + + ChangeItemImpl.ChangeItemImplBuilder builder = + ChangeItemImpl.builder() + .urn(entityUrn) + .aspectName(aspectSpec.getName()) + .auditStamp(AuditStampUtils.createAuditStamp(actor.toUrnStr())) + .systemMetadata(systemMetadata) + .recordTemplate( + GenericRecordUtils.deserializeAspect( + ByteString.copyString( + objectMapper.writeValueAsString(aspect.getValue()), + StandardCharsets.UTF_8), + GenericRecordUtils.JSON, + aspectSpec)); + + items.add(builder.build(opContext.getRetrieverContext().get().getAspectRetriever())); + } + } + } + } + return AspectsBatchImpl.builder() + .items(items) + .retrieverContext(opContext.getRetrieverContext().get()) + .build(); + } +}