From 46530f249895ae6e34e04a6c0e7be6d720827057 Mon Sep 17 00:00:00 2001 From: david-leifker <114954101+david-leifker@users.noreply.github.com> Date: Fri, 12 Jul 2024 11:35:11 -0500 Subject: [PATCH] feat(conditional-writes): misc updates and fixes (#10901) --- docs/advanced/mcp-mcl.md | 12 + docs/api/openapi/openapi-usage-guide.md | 224 +++++++++++++++++- .../validation/ConditionalWriteValidator.java | 6 +- .../ConditionalWriteValidatorTest.java | 100 +++++++- .../controller/GenericEntitiesController.java | 89 ++++--- .../v2/controller/EntityController.java | 5 +- .../v3/controller/EntityController.java | 13 +- .../v3/controller/EntityControllerTest.java | 205 ++++++++++++++++ 8 files changed, 594 insertions(+), 60 deletions(-) create mode 100644 metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java diff --git a/docs/advanced/mcp-mcl.md b/docs/advanced/mcp-mcl.md index 2b2d2885428b56..9efb9b794954da 100644 --- a/docs/advanced/mcp-mcl.md +++ b/docs/advanced/mcp-mcl.md @@ -183,6 +183,9 @@ A writer can provide a header with the expected `version` when initiating the re match the actual `version` stored in the database, the write will fail. This prevents overwriting an aspect that has been modified by another process. +Note: If the aspect doesn't exist yet, then the `version` is `-1`. A writer can use this `version` to only create +an aspect if it doesn't. Also see _Change Types: [`CREATE`, `CREATE_ENTITY`]_ section below. + #### If-Modified-Since / If-Unmodified-Since A writer may also specify time-based conditions using http header semantics. Similar to version based conditional writes @@ -194,3 +197,12 @@ A writer can specify that the aspect must NOT have been modified after a specifi `If-Modified-Since` A writer can specify that the aspect must have been modified after a specific time, following [If-Modified-Since](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/If-Modified-Since) http headers. + +#### Change Types: [`CREATE`, `CREATE_ENTITY`] + +Another form of conditional writes which considers the existence of an aspect or entity uses the following Change Types. + +`CREATE` - Create the aspect if it doesn't already exist. + +`CREATE_ENTITY` - Create the aspect if no aspects exist for the entity. + diff --git a/docs/api/openapi/openapi-usage-guide.md b/docs/api/openapi/openapi-usage-guide.md index 9c6d0bfa12fc18..59a331b724cfa7 100644 --- a/docs/api/openapi/openapi-usage-guide.md +++ b/docs/api/openapi/openapi-usage-guide.md @@ -588,4 +588,226 @@ public class Main { ### Conditional Writes All the create/POST endpoints for aspects support `headers` in the POST body to support batch APIs. See the docs in the -[MetadataChangeProposal](../../advanced/mcp-mcl.md) section for the use of these headers to support conditional writes semantics. \ No newline at end of file +[MetadataChangeProposal](../../advanced/mcp-mcl.md) section for the use of these headers to support conditional writes semantics. + +### Batch Get + +Batch get endpoints in the form of `/v3/entity/{entityName}/batchGet` exist for all entities. This endpoint allows +fetching entity and aspects in bulk. In combination with the `If-Version-Match` header it can also retrieve +a specific version of the aspects, however it defaults to the latest aspect version. Currently, this interface is limited +to returning a single version for each entity/aspect however different versions can be specified across entities. + +A few example queries are as follows: + +Example Request: + +Fetch the latest aspects for the given URNs with the url parameter `systemMetadata=true` in order to view the current +versions of the aspects. + +```json +[ + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_deleted,PROD)", + "globalTags": {}, + "datasetProperties": {} + }, + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_created,PROD)", + "globalTags": {}, + "datasetProperties": {} + } +] +``` + +Example Response: + +Notice that `systemMetadata` contains `"version": "1"` for each of the aspects that exist in the system. + +```json +[ + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_deleted,PROD)", + "datasetProperties": { + "value": { + "description": "table containing all the users deleted on a single day", + "customProperties": { + "encoding": "utf-8" + }, + "tags": [] + }, + "systemMetadata": { + "properties": { + "clientVersion": "1!0.0.0.dev0", + "clientId": "acryl-datahub" + }, + "version": "1", + "lastObserved": 1720781548776, + "lastRunId": "file-2024_07_12-05_52_28", + "runId": "file-2024_07_12-05_52_28" + } + } + }, + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_created,PROD)", + "datasetProperties": { + "value": { + "description": "table containing all the users created on a single day", + "customProperties": { + "encoding": "utf-8" + }, + "tags": [] + }, + "systemMetadata": { + "properties": { + "clientVersion": "1!0.0.0.dev0", + "clientId": "acryl-datahub" + }, + "version": "1", + "lastObserved": 1720781548773, + "lastRunId": "file-2024_07_12-05_52_28", + "runId": "file-2024_07_12-05_52_28" + } + }, + "globalTags": { + "value": { + "tags": [ + { + "tag": "urn:li:tag:NeedsDocumentation" + } + ] + }, + "systemMetadata": { + "properties": { + "appSource": "ui" + }, + "version": "1", + "lastObserved": 0, + "lastRunId": "no-run-id-provided", + "runId": "no-run-id-provided" + } + } + } +] +``` + +Next let's mutate `globalTags` for the second URN by adding a new tag. This will increment the version of +the `globalTags` aspect. The response will then look at like the following, notice the incremented +`"version": "2"` in `systemMetadata` for the `globalTags` aspect. Also notice that there are now 2 tags present, unlike +previously where only `urn:li:tag:NeedsDocumentation` was present. + +```json +[ + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_deleted,PROD)", + "datasetProperties": { + "value": { + "description": "table containing all the users deleted on a single day", + "customProperties": { + "encoding": "utf-8" + }, + "tags": [] + }, + "systemMetadata": { + "properties": { + "clientVersion": "1!0.0.0.dev0", + "clientId": "acryl-datahub" + }, + "version": "1", + "lastObserved": 1720781548776, + "lastRunId": "file-2024_07_12-05_52_28", + "runId": "file-2024_07_12-05_52_28" + } + } + }, + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_created,PROD)", + "datasetProperties": { + "value": { + "description": "table containing all the users created on a single day", + "customProperties": { + "encoding": "utf-8" + }, + "tags": [] + }, + "systemMetadata": { + "properties": { + "clientVersion": "1!0.0.0.dev0", + "clientId": "acryl-datahub" + }, + "version": "1", + "lastObserved": 1720781548773, + "lastRunId": "file-2024_07_12-05_52_28", + "runId": "file-2024_07_12-05_52_28" + } + }, + "globalTags": { + "value": { + "tags": [ + { + "tag": "urn:li:tag:NeedsDocumentation" + }, + { + "tag": "urn:li:tag:Legacy" + } + ] + }, + "systemMetadata": { + "properties": { + "appSource": "ui" + }, + "version": "2", + "lastObserved": 0, + "lastRunId": "no-run-id-provided", + "runId": "no-run-id-provided" + } + } + } +] +``` + +Next, we'll retrieve the previous version of the `globalTags` for the one aspect with a version 2 with the following query. +We can do this by populating the `headers` map with `If-Version-Match` to retrieve the previous version 1. + +Example Request: +```json +[ + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_created,PROD)", + "globalTags": { + "headers": { + "If-Version-Match": "1" + } + } + } +] +``` + +Example Response: + +The previous version `1` of the `globalTags` aspect is returned as expected with only the single tag. + +```json +[ + { + "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,fct_users_created,PROD)", + "globalTags": { + "value": { + "tags": [ + { + "tag": "urn:li:tag:NeedsDocumentation" + } + ] + }, + "systemMetadata": { + "properties": { + "appSource": "ui" + }, + "version": "1", + "lastObserved": 0, + "lastRunId": "no-run-id-provided", + "runId": "no-run-id-provided" + } + } + } +] +``` diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/ConditionalWriteValidator.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/ConditionalWriteValidator.java index 9927ca4c5a0985..810693c80fa13f 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/ConditionalWriteValidator.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/ConditionalWriteValidator.java @@ -42,7 +42,7 @@ @Getter @Accessors(chain = true) public class ConditionalWriteValidator extends AspectPayloadValidator { - public static final String DEFAULT_ASPECT_VERSION = "1"; + public static final String UNVERSIONED_ASPECT_VERSION = "-1"; public static final long DEFAULT_LAST_MODIFIED_TIME = Long.MIN_VALUE; public static final String HTTP_HEADER_IF_VERSION_MATCH = "If-Version-Match"; public static final Set CREATE_CHANGE_TYPES = @@ -130,7 +130,7 @@ private static Optional validateVersionPrecondition( switch (item.getChangeType()) { case CREATE: case CREATE_ENTITY: - actualAspectVersion = DEFAULT_ASPECT_VERSION; + actualAspectVersion = UNVERSIONED_ASPECT_VERSION; break; default: actualAspectVersion = @@ -143,7 +143,7 @@ private static Optional validateVersionPrecondition( return String.valueOf(Math.max(1, prevSystemAspect.getVersion())); } }) - .orElse(DEFAULT_ASPECT_VERSION); + .orElse(UNVERSIONED_ASPECT_VERSION); break; } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/ConditionalWriteValidatorTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/ConditionalWriteValidatorTest.java index b62a95d8ab91c2..652328781c3bb3 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/ConditionalWriteValidatorTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/ConditionalWriteValidatorTest.java @@ -3,6 +3,7 @@ import static com.linkedin.metadata.aspect.validation.ConditionalWriteValidator.HTTP_HEADER_IF_VERSION_MATCH; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.reset; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; @@ -76,7 +77,6 @@ public void testNextVersionSuccess() { final ChangeMCP testMCP; switch (changeType) { case RESTATE: - case DELETE: case CREATE_ENTITY: case CREATE: testMCP = @@ -90,7 +90,7 @@ public void testNextVersionSuccess() { .getAspectSpec("status")) .recordTemplate(new Status().setRemoved(false)) // Expected - .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "1")) + .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "-1")) .build(); break; default: @@ -137,7 +137,63 @@ public void testNoSystemMetadataNextVersionNextVersionSuccess() { for (ChangeType changeType : supportedChangeTypes) { final ChangeMCP testMCP; switch (changeType) { + case DELETE: + reset(mockRetrieverContext.getAspectRetriever()); + when(mockRetrieverContext + .getAspectRetriever() + .getLatestSystemAspects(eq(Map.of(testEntityUrn, Set.of("status"))))) + .thenReturn( + Map.of( + testEntityUrn, + Map.of( + "status", + TestSystemAspect.builder() + .systemMetadata(new SystemMetadata().setVersion("1")) + .build()))); + testMCP = + TestMCP.builder() + .changeType(changeType) + .urn(testEntityUrn) + .entitySpec(entityRegistry.getEntitySpec(testEntityUrn.getEntityType())) + .aspectSpec( + entityRegistry + .getEntitySpec(testEntityUrn.getEntityType()) + .getAspectSpec("status")) + .recordTemplate(new Status().setRemoved(false)) + // Expected (cannot delete non-existent -1) + .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "1")) + .build(); + break; + case CREATE: + case CREATE_ENTITY: + reset(mockRetrieverContext.getAspectRetriever()); + testMCP = + TestMCP.builder() + .changeType(changeType) + .urn(testEntityUrn) + .entitySpec(entityRegistry.getEntitySpec(testEntityUrn.getEntityType())) + .aspectSpec( + entityRegistry + .getEntitySpec(testEntityUrn.getEntityType()) + .getAspectSpec("status")) + .recordTemplate(new Status().setRemoved(false)) + // Expected + .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "-1")) + .previousSystemAspect( + TestSystemAspect.builder() + .urn(testEntityUrn) + .entitySpec(entityRegistry.getEntitySpec(testEntityUrn.getEntityType())) + .aspectSpec( + entityRegistry + .getEntitySpec(testEntityUrn.getEntityType()) + .getAspectSpec("status")) + // Missing previous system metadata, expect fallback to version + .version(0) + .build()) + .build(); + break; default: + reset(mockRetrieverContext.getAspectRetriever()); testMCP = TestMCP.builder() .changeType(changeType) @@ -169,7 +225,11 @@ public void testNoSystemMetadataNextVersionNextVersionSuccess() { test.validatePreCommit(List.of(testMCP), mockRetrieverContext) .collect(Collectors.toSet()); - assertEquals(Set.of(), exceptions, "Expected no exceptions for change type " + changeType); + assertEquals( + Set.of(), + exceptions, + String.format( + "Expected no exceptions for change type %s but found %s", changeType, exceptions)); } } @@ -179,6 +239,7 @@ public void testNoPreviousVersionsLookupSchemaMetadataNextVersionSuccess() { Urn testEntityUrn = UrnUtils.getUrn("urn:li:chart:(looker,baz1)"); // Prepare mock lookup based on version + reset(mockRetrieverContext.getAspectRetriever()); when(mockRetrieverContext .getAspectRetriever() .getLatestSystemAspects(eq(Map.of(testEntityUrn, Set.of("status"))))) @@ -208,7 +269,7 @@ public void testNoPreviousVersionsLookupSchemaMetadataNextVersionSuccess() { .getAspectSpec("status")) .recordTemplate(new Status().setRemoved(false)) // Expected is always 1 - .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "1")) + .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "-1")) .build(); break; default: @@ -270,7 +331,7 @@ public void testNoPreviousVersionsLookupVersionNextVersionSuccess() { .getAspectSpec("status")) .recordTemplate(new Status().setRemoved(false)) // Expected is always 1 - .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "1")) + .headers(Map.of(HTTP_HEADER_IF_VERSION_MATCH, "-1")) .build(); break; default: @@ -306,8 +367,23 @@ public void testNextVersionFail() { for (ChangeType changeType : supportedChangeTypes) { final ChangeMCP testMCP; switch (changeType) { - case RESTATE: case DELETE: + // allow lookup of previous value + when(mockRetrieverContext + .getAspectRetriever() + .getLatestSystemAspects(Map.of(testEntityUrn, Set.of("status")))) + .thenReturn( + Map.of( + testEntityUrn, + Map.of( + "status", + TestSystemAspect.builder() + .urn(testEntityUrn) + .version(3) + .recordTemplate(new Status().setRemoved(false)) + .build()))); + // fall through + case RESTATE: case CREATE_ENTITY: case CREATE: testMCP = @@ -362,17 +438,17 @@ public void testNextVersionFail() { break; case CREATE: case CREATE_ENTITY: - case DELETE: assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); assertEquals( exceptions.stream().findFirst().get().getMessage(), - "Expected version 2, actual version 1"); + "Expected version 2, actual version -1"); break; default: assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); assertEquals( exceptions.stream().findFirst().get().getMessage(), - "Expected version 2, actual version 3"); + "Expected version 2, actual version 3", + "for changeType:" + changeType); break; } } @@ -427,7 +503,7 @@ public void testNoSystemMetadataNextVersionNextVersionFail() { assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); assertEquals( exceptions.stream().findFirst().get().getMessage(), - "Expected version 2, actual version 1"); + "Expected version 2, actual version -1"); break; default: assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); @@ -507,7 +583,7 @@ public void testNoPreviousVersionsLookupSchemaMetadataNextVersionFail() { assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); assertEquals( exceptions.stream().findFirst().get().getMessage(), - "Expected version 2, actual version 1"); + "Expected version 2, actual version -1"); break; default: assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); @@ -586,7 +662,7 @@ public void testNoPreviousVersionsLookupVersionNextVersionFail() { assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); assertEquals( exceptions.stream().findFirst().get().getMessage(), - "Expected version 2, actual version 1"); + "Expected version 2, actual version -1"); break; default: assertEquals(exceptions.size(), 1, "Expected exception for change type " + changeType); diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java index 6a6230622f44ff..c91c8ac987e5c9 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java @@ -59,7 +59,6 @@ import java.nio.charset.StandardCharsets; import java.util.*; 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; @@ -118,7 +117,8 @@ protected List buildEntityList( Set aspectNames, boolean withSystemMetadata) throws URISyntaxException { - Map> versionMap = + + LinkedHashMap> versionMap = resolveAspectNames( urns.stream() .map( @@ -128,13 +128,21 @@ protected List buildEntityList( aspectNames.stream() .map(aspectName -> Map.entry(aspectName, 0L)) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); + .collect( + Collectors.toMap( + Map.Entry::getKey, + Map.Entry::getValue, + (a, b) -> { + throw new IllegalStateException("Duplicate key"); + }, + LinkedHashMap::new)), + 0L); return buildEntityVersionedAspectList(opContext, versionMap, withSystemMetadata); } protected abstract List buildEntityVersionedAspectList( @Nonnull OperationContext opContext, - Map> urnAspectVersions, + LinkedHashMap> urnAspectVersions, boolean withSystemMetadata) throws URISyntaxException; @@ -335,7 +343,9 @@ public ResponseEntity getAspect( } else { resultList = buildEntityVersionedAspectList( - opContext, Map.of(urn, Map.of(aspectName, version)), withSystemMetadata); + opContext, + new LinkedHashMap<>(Map.of(urn, Map.of(aspectName, version))), + withSystemMetadata); } return resultList.stream() @@ -642,47 +652,54 @@ protected Boolean exists( * @return updated map * @param map values */ - protected Map> resolveAspectNames( - Map> requestedAspectNames) { + protected LinkedHashMap> resolveAspectNames( + LinkedHashMap> requestedAspectNames, @Nonnull T defaultValue) { return requestedAspectNames.entrySet().stream() .map( entry -> { final Urn urn = entry.getKey(); - final Set requestedNames; - if (entry.getValue().isEmpty()) { - requestedNames = + if (entry.getValue().isEmpty() || entry.getValue().containsKey("")) { + // All aspects specified + Set allNames = entityRegistry.getEntitySpec(urn.getEntityType()).getAspectSpecs().stream() .map(AspectSpec::getName) .collect(Collectors.toSet()); + return Map.entry( + urn, + allNames.stream() + .map( + aspectName -> + Map.entry( + aspectName, entry.getValue().getOrDefault("", defaultValue))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); } else { - // add key aspect - requestedNames = - Stream.concat( - entry.getValue().keySet().stream(), - Stream.of( - entityRegistry - .getEntitySpec(urn.getEntityType()) - .getKeyAspectName())) - .collect(Collectors.toSet()); + final Map normalizedNames = + entry.getValue().keySet().stream() + .map( + requestAspectName -> + Map.entry( + requestAspectName, + lookupAspectSpec(urn, requestAspectName).getName())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + return Map.entry( + urn, + entry.getValue().entrySet().stream() + .filter(reqEntry -> normalizedNames.containsKey(reqEntry.getKey())) + .map( + reqEntry -> + Map.entry( + normalizedNames.get(reqEntry.getKey()), reqEntry.getValue())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); } - final Map normalizedNames = - requestedNames.stream() - .map( - requestAspectName -> - Map.entry( - requestAspectName, - lookupAspectSpec(urn, requestAspectName).getName())) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - return Map.entry( - urn, - entry.getValue().entrySet().stream() - .map( - reqEntry -> - Map.entry( - normalizedNames.get(reqEntry.getKey()), reqEntry.getValue())) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); }) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + .collect( + Collectors.toMap( + Map.Entry::getKey, + Map.Entry::getValue, + (a, b) -> { + throw new IllegalStateException("Duplicate key"); + }, + LinkedHashMap::new)); } protected Map> toAspectMap( 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 92a4cb2bd79f38..54a7724cadd345 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 @@ -45,6 +45,7 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -187,12 +188,12 @@ protected AspectsBatch toMCPBatch( @Override protected List buildEntityVersionedAspectList( @Nonnull OperationContext opContext, - Map> urnAspectVersions, + LinkedHashMap> urnAspectVersions, boolean withSystemMetadata) throws URISyntaxException { Map> aspects = entityService.getEnvelopedVersionedAspects( - opContext, resolveAspectNames(urnAspectVersions), true); + opContext, resolveAspectNames(urnAspectVersions, 0L), true); return urnAspectVersions.keySet().stream() .map( 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 84cc4858d5fe47..d6feb6cc460c97 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 @@ -44,6 +44,7 @@ import java.nio.charset.StandardCharsets; import java.util.HashMap; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -80,7 +81,7 @@ public ResponseEntity> getEntityBatch( @RequestBody @Nonnull String jsonEntityList) throws URISyntaxException, JsonProcessingException { - Map> requestMap = toEntityVersionRequest(jsonEntityList); + LinkedHashMap> requestMap = toEntityVersionRequest(jsonEntityList); Authentication authentication = AuthenticationContext.getAuthentication(); OperationContext opContext = @@ -125,7 +126,7 @@ public GenericEntityScrollResultV3 buildScrollResult( @Override protected List buildEntityVersionedAspectList( @Nonnull OperationContext opContext, - Map> urnAspectVersions, + LinkedHashMap> urnAspectVersions, boolean withSystemMetadata) throws URISyntaxException { if (urnAspectVersions.isEmpty()) { @@ -133,7 +134,7 @@ protected List buildEntityVersionedAspectList( } else { Map> aspects = entityService.getEnvelopedVersionedAspects( - opContext, resolveAspectNames(urnAspectVersions), false); + opContext, resolveAspectNames(urnAspectVersions, 0L), false); return urnAspectVersions.keySet().stream() .filter(urn -> aspects.containsKey(urn) && !aspects.get(urn).isEmpty()) @@ -198,11 +199,11 @@ private List toRecordTemplates( withSystemMetadata); } - private Map> toEntityVersionRequest(@Nonnull String entityArrayList) - throws JsonProcessingException, InvalidUrnException { + private LinkedHashMap> toEntityVersionRequest( + @Nonnull String entityArrayList) throws JsonProcessingException, InvalidUrnException { JsonNode entities = objectMapper.readTree(entityArrayList); - Map> items = new HashMap<>(); + LinkedHashMap> items = new LinkedHashMap<>(); if (entities.isArray()) { Iterator entityItr = entities.iterator(); while (entityItr.hasNext()) { diff --git a/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java b/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java new file mode 100644 index 00000000000000..0855ad6c2e4ff2 --- /dev/null +++ b/metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java @@ -0,0 +1,205 @@ +package io.datahubproject.openapi.v3.controller; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.nullable; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; +import static org.testng.Assert.assertNotNull; + +import com.datahub.authentication.Actor; +import com.datahub.authentication.ActorType; +import com.datahub.authentication.Authentication; +import com.datahub.authentication.AuthenticationContext; +import com.datahub.authorization.AuthorizationResult; +import com.datahub.authorization.AuthorizerChain; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.linkedin.common.Status; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.entity.Aspect; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.EntityServiceImpl; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.query.filter.Filter; +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.SearchUtil; +import io.datahubproject.metadata.context.OperationContext; +import io.datahubproject.openapi.config.SpringWebConfig; +import io.datahubproject.test.metadata.context.TestOperationContexts; +import java.util.List; +import java.util.Map; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc; +import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureWebMvc; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.test.context.TestConfiguration; +import org.springframework.boot.test.mock.mockito.MockBean; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.ComponentScan; +import org.springframework.context.annotation.Import; +import org.springframework.context.annotation.Primary; +import org.springframework.http.MediaType; +import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; +import org.springframework.test.web.servlet.MockMvc; +import org.springframework.test.web.servlet.request.MockMvcRequestBuilders; +import org.springframework.test.web.servlet.result.MockMvcResultMatchers; +import org.testng.annotations.Test; + +@SpringBootTest(classes = {SpringWebConfig.class}) +@ComponentScan(basePackages = {"io.datahubproject.openapi.v3.controller"}) +@Import({SpringWebConfig.class, EntityControllerTest.EntityControllerTestConfig.class}) +@AutoConfigureWebMvc +@AutoConfigureMockMvc +public class EntityControllerTest extends AbstractTestNGSpringContextTests { + @Autowired private EntityController entityController; + @Autowired private MockMvc mockMvc; + @Autowired private SearchService mockSearchService; + @Autowired private EntityService mockEntityService; + + @Test + public void initTest() { + assertNotNull(entityController); + } + + @Test + public void testSearchOrderPreserved() throws Exception { + List TEST_URNS = + List.of( + UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:testPlatform,1,PROD)"), + UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:testPlatform,2,PROD)"), + UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:testPlatform,3,PROD)")); + + // Mock scroll ascending/descending results + ScrollResult expectedResultAscending = + new ScrollResult() + .setEntities( + new SearchEntityArray( + List.of( + new SearchEntity().setEntity(TEST_URNS.get(0)), + new SearchEntity().setEntity(TEST_URNS.get(1)), + new SearchEntity().setEntity(TEST_URNS.get(2))))); + when(mockSearchService.scrollAcrossEntities( + any(OperationContext.class), + eq(List.of("dataset")), + anyString(), + nullable(Filter.class), + eq(SearchUtil.sortBy("urn", SortOrder.valueOf("ASCENDING"))), + nullable(String.class), + nullable(String.class), + anyInt())) + .thenReturn(expectedResultAscending); + ScrollResult expectedResultDescending = + new ScrollResult() + .setEntities( + new SearchEntityArray( + List.of( + new SearchEntity().setEntity(TEST_URNS.get(2)), + new SearchEntity().setEntity(TEST_URNS.get(1)), + new SearchEntity().setEntity(TEST_URNS.get(0))))); + when(mockSearchService.scrollAcrossEntities( + any(OperationContext.class), + eq(List.of("dataset")), + anyString(), + nullable(Filter.class), + eq(SearchUtil.sortBy("urn", SortOrder.valueOf("DESCENDING"))), + nullable(String.class), + nullable(String.class), + anyInt())) + .thenReturn(expectedResultDescending); + // Mock entity aspect + when(mockEntityService.getEnvelopedVersionedAspects( + any(OperationContext.class), anyMap(), eq(false))) + .thenReturn( + Map.of( + TEST_URNS.get(0), + List.of( + new EnvelopedAspect() + .setName("status") + .setValue(new Aspect(new Status().data()))), + TEST_URNS.get(1), + List.of( + new EnvelopedAspect() + .setName("status") + .setValue(new Aspect(new Status().data()))), + TEST_URNS.get(2), + List.of( + new EnvelopedAspect() + .setName("status") + .setValue(new Aspect(new Status().data()))))); + + // test ASCENDING + mockMvc + .perform( + MockMvcRequestBuilders.get("/v3/entity/dataset") + .param("sortOrder", "ASCENDING") + .accept(MediaType.APPLICATION_JSON)) + .andExpect(status().is2xxSuccessful()) + .andExpect( + MockMvcResultMatchers.jsonPath("$.entities[0].urn").value(TEST_URNS.get(0).toString())) + .andExpect( + MockMvcResultMatchers.jsonPath("$.entities[1].urn").value(TEST_URNS.get(1).toString())) + .andExpect( + MockMvcResultMatchers.jsonPath("$.entities[2].urn").value(TEST_URNS.get(2).toString())); + + // test DESCENDING + mockMvc + .perform( + MockMvcRequestBuilders.get("/v3/entity/dataset") + .accept(MediaType.APPLICATION_JSON) + .param("sortOrder", "DESCENDING")) + .andExpect(status().is2xxSuccessful()) + .andExpect( + MockMvcResultMatchers.jsonPath("$.entities[0].urn").value(TEST_URNS.get(2).toString())) + .andExpect( + MockMvcResultMatchers.jsonPath("$.entities[1].urn").value(TEST_URNS.get(1).toString())) + .andExpect( + MockMvcResultMatchers.jsonPath("$.entities[2].urn").value(TEST_URNS.get(0).toString())); + } + + @TestConfiguration + public static class EntityControllerTestConfig { + @MockBean public EntityServiceImpl entityService; + @MockBean public SearchService searchService; + + @Bean + public ObjectMapper objectMapper() { + return new ObjectMapper(); + } + + @Bean(name = "systemOperationContext") + public OperationContext systemOperationContext() { + return TestOperationContexts.systemContextNoSearchAuthorization(); + } + + @Bean("entityRegistry") + @Primary + public EntityRegistry entityRegistry( + @Qualifier("systemOperationContext") final OperationContext testOperationContext) { + return testOperationContext.getEntityRegistry(); + } + + @Bean + public AuthorizerChain authorizerChain() { + AuthorizerChain authorizerChain = mock(AuthorizerChain.class); + + Authentication authentication = mock(Authentication.class); + when(authentication.getActor()).thenReturn(new Actor(ActorType.USER, "datahub")); + when(authorizerChain.authorize(any())) + .thenReturn(new AuthorizationResult(null, AuthorizationResult.Type.ALLOW, "")); + AuthenticationContext.setAuthentication(authentication); + + return authorizerChain; + } + } +}