From a10b7344c1ea49a957345d466aae7f6129a8d3ce Mon Sep 17 00:00:00 2001 From: David Leifker Date: Mon, 29 Jan 2024 12:23:14 -0600 Subject: [PATCH] feat(structured-properties): soft delete --- build.gradle | 2 +- .../linkedin/datahub/graphql/TestUtils.java | 20 +- .../resolvers/tag/AddTagsResolverTest.java | 4 +- .../restorebackup/RestoreStorageStep.java | 5 +- .../steps/BuildIndicesPreStep.java | 58 +- .../DatahubUpgradeNoSchemaRegistryTest.java | 2 +- entity-registry/build.gradle | 1 + .../validation => }/AspectRetriever.java | 2 +- .../aspect/CachingAspectRetriever.java | 4 + .../linkedin/metadata/aspect/ReadItem.java | 77 ++ .../metadata/aspect/SystemAspect.java | 25 + .../metadata/aspect/batch/AspectsBatch.java | 105 ++- .../metadata/aspect/batch/BatchItem.java | 56 +- .../metadata/aspect/batch/ChangeMCP.java | 42 + .../batch/{MCLBatchItem.java => MCLItem.java} | 19 +- .../batch/{MCPBatchItem.java => MCPItem.java} | 8 +- .../batch/{PatchItem.java => PatchMCP.java} | 9 +- .../metadata/aspect/batch/SystemAspect.java | 25 - .../metadata/aspect/batch/UpsertItem.java | 19 - .../hooks/StructuredPropertiesSoftDelete.java | 38 + .../metadata/aspect/plugins/PluginSpec.java | 27 +- .../plugins/config/AspectPluginConfig.java | 12 +- .../aspect/plugins/hooks/MCLSideEffect.java | 30 +- .../aspect/plugins/hooks/MCPSideEffect.java | 29 +- .../aspect/plugins/hooks/MutationHook.java | 86 +- .../validation/AspectPayloadValidator.java | 77 +- .../validation/AspectValidationException.java | 53 +- .../ValidationExceptionCollection.java | 68 ++ .../PropertyDefinitionValidator.java | 187 +++-- .../StructuredPropertiesValidator.java | 307 ++++--- .../models/StructuredPropertyUtils.java | 157 +++- .../models/registry/EntityRegistry.java | 65 +- .../StructuredPropertiesSoftDeleteTest.java | 96 +++ .../metadata/aspect/plugins/PluginsTest.java | 60 +- .../plugins/hooks/MCLSideEffectTest.java | 18 +- .../plugins/hooks/MCPSideEffectTest.java | 17 +- .../plugins/hooks/MutationPluginTest.java | 27 +- .../validation/ValidatorPluginTest.java | 41 +- .../PropertyDefinitionValidatorTest.java | 171 +++- .../StructuredPropertiesValidatorTest.java | 285 ++++--- .../metadata/aspect/MockAspectRetriever.java | 71 ++ .../metadata/aspect/TestEntityRegistry.java | 31 + .../test/metadata/aspect/batch/TestMCP.java | 128 +++ .../java/com/linkedin/metadata/Constants.java | 2 + .../aspect/utils/DefaultAspectsUtil.java | 14 +- .../client/EntityClientAspectRetriever.java | 25 +- .../metadata/client/JavaEntityClient.java | 129 +-- .../metadata/entity/EntityAspect.java | 129 ++- .../metadata/entity/EntityServiceImpl.java | 750 ++++++++---------- .../linkedin/metadata/entity/EntityUtils.java | 287 +++++-- .../entity/cassandra/CassandraAspectDao.java | 2 +- .../cassandra/CassandraRetentionService.java | 4 +- .../metadata/entity/ebean/EbeanAspectDao.java | 22 +- .../entity/ebean/EbeanRetentionService.java | 4 +- .../entity/ebean/batch/AspectsBatchImpl.java | 73 +- ...sertBatchItem.java => ChangeItemImpl.java} | 126 +-- .../entity/ebean/batch/DeleteItemImpl.java | 139 ++++ ...MCLBatchItemImpl.java => MCLItemImpl.java} | 40 +- ...PatchBatchItem.java => PatchItemImpl.java} | 44 +- .../entity/validation/ValidationUtils.java | 50 +- .../elasticsearch/ElasticSearchService.java | 8 + .../elasticsearch/query/ESBrowseDAO.java | 34 +- .../elasticsearch/query/ESSearchDAO.java | 52 +- .../request/AggregationQueryBuilder.java | 13 +- .../request/AutocompleteRequestHandler.java | 15 +- .../query/request/SearchRequestHandler.java | 85 +- .../SearchDocumentTransformer.java | 2 +- .../metadata/search/utils/ESUtils.java | 11 +- .../service/UpdateIndicesService.java | 36 +- .../ElasticSearchTimeseriesAspectService.java | 111 +-- .../TimeseriesAspectIndexBuilders.java | 25 +- .../elastic/query/ESAggregatedStatsDAO.java | 26 +- .../metadata/AspectIngestionUtils.java | 25 +- .../aspect/utils/DefaultAspectsUtilTest.java | 6 +- .../CassandraAspectMigrationsDaoTest.java | 10 +- .../entity/CassandraEntityServiceTest.java | 9 +- .../entity/DeleteEntityServiceTest.java | 9 +- .../entity/EbeanAspectMigrationsDaoTest.java | 10 +- .../entity/EbeanEntityServiceTest.java | 43 +- .../metadata/entity/EntityServiceTest.java | 166 ++-- .../search/LineageServiceTestBase.java | 257 +++--- .../search/SearchServiceTestBase.java | 120 +-- .../metadata/search/TestEntityTestBase.java | 125 +-- .../TestEntityElasticSearchTest.java | 22 +- ...eseriesAspectServiceElasticSearchTest.java | 6 +- ...TimeseriesAspectServiceOpenSearchTest.java | 6 +- .../metadata/search/query/BrowseDAOTest.java | 41 +- .../search/query/SearchDAOTestBase.java | 46 +- .../request/AggregationQueryBuilderTest.java | 37 +- .../AutocompleteRequestHandlerTest.java | 5 +- .../request/SearchRequestHandlerTest.java | 26 +- .../CassandraTimelineServiceTest.java | 9 +- .../timeline/EbeanTimelineServiceTest.java | 9 +- .../TimeseriesAspectServiceTestBase.java | 262 +++--- .../TimeseriesAspectServiceUnitTest.java | 17 +- .../io/datahubproject/test/DataGenerator.java | 6 +- .../SampleDataFixtureConfiguration.java | 28 +- .../SearchLineageFixtureConfiguration.java | 42 +- .../config/SearchCommonTestConfiguration.java | 29 + .../kafka/MetadataChangeLogProcessor.java | 5 + .../spring/MCLSpringTestConfiguration.java | 12 +- .../CustomDataQualityRulesMCLSideEffect.java | 55 +- .../CustomDataQualityRulesMCPSideEffect.java | 33 +- .../hooks/CustomDataQualityRulesMutator.java | 55 +- .../CustomDataQualityRulesValidator.java | 101 +-- .../src/main/resources/entity-registry.yml | 11 +- .../user/NativeUserServiceTest.java | 2 +- .../src/main/resources/application.yml | 3 +- .../factory/entity/EntityServiceFactory.java | 7 +- .../entity/RetentionServiceFactory.java | 12 +- .../indices/UpdateIndicesServiceFactory.java | 63 +- .../entityclient/AspectRetrieverFactory.java | 27 + .../factory/graphql/GraphQLEngineFactory.java | 2 +- .../search/ElasticSearchServiceFactory.java | 4 +- .../IngestDataPlatformInstancesStep.java | 11 +- .../boot/steps/IngestDataPlatformsStep.java | 13 +- .../IngestDataPlatformInstancesStepTest.java | 6 +- .../v2/delegates/EntityApiDelegateImpl.java | 2 +- .../openapi/entities/EntitiesController.java | 4 +- .../entities/PlatformEntitiesController.java | 4 +- .../openapi/util/MappingUtil.java | 4 +- .../v2/controller/EntityController.java | 34 +- .../java/entities/EntitiesControllerTest.java | 8 +- .../src/test/java/mock/MockEntityService.java | 7 +- .../linkedin/entity/client/EntityClient.java | 4 + .../resources/entity/AspectResource.java | 10 +- .../resources/entity/AspectResourceTest.java | 63 +- .../mock/MockTimeseriesAspectService.java | 6 + .../metadata/entity/EntityService.java | 82 +- .../metadata/entity/RetentionService.java | 4 +- .../metadata/entity/SearchIndicesService.java | 11 + .../metadata/entity/UpdateAspectResult.java | 4 +- .../metadata/search/EntitySearchService.java | 8 + .../timeseries/TimeseriesAspectService.java | 8 + .../gms/servlet/ConfigSearchExport.java | 11 +- smoke-test/build.gradle | 26 + .../test_structured_properties.py | 292 +++++-- .../tokens/revokable_access_token_test.py | 2 + 138 files changed, 4553 insertions(+), 2656 deletions(-) rename entity-registry/src/main/java/com/linkedin/metadata/aspect/{plugins/validation => }/AspectRetriever.java (95%) create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/CachingAspectRetriever.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/ReadItem.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/SystemAspect.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java rename entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/{MCLBatchItem.java => MCLItem.java} (70%) rename entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/{MCPBatchItem.java => MCPItem.java} (83%) rename entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/{PatchItem.java => PatchMCP.java} (58%) delete mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/SystemAspect.java delete mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDelete.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/ValidationExceptionCollection.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDeleteTest.java create mode 100644 entity-registry/src/test/java/com/linkedin/test/metadata/aspect/MockAspectRetriever.java create mode 100644 entity-registry/src/test/java/com/linkedin/test/metadata/aspect/TestEntityRegistry.java create mode 100644 entity-registry/src/test/java/com/linkedin/test/metadata/aspect/batch/TestMCP.java rename metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/{MCPUpsertBatchItem.java => ChangeItemImpl.java} (63%) create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/DeleteItemImpl.java rename metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/{MCLBatchItemImpl.java => MCLItemImpl.java} (77%) rename metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/{MCPPatchBatchItem.java => PatchItemImpl.java} (84%) create mode 100644 metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/AspectRetrieverFactory.java create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/SearchIndicesService.java diff --git a/build.gradle b/build.gradle index ea81d26355027d..09bb927c831098 100644 --- a/build.gradle +++ b/build.gradle @@ -73,7 +73,7 @@ plugins { id 'com.gorylenko.gradle-git-properties' version '2.4.1' id 'com.github.johnrengelman.shadow' version '8.1.1' apply false id 'com.palantir.docker' version '0.35.0' apply false - id 'com.avast.gradle.docker-compose' version '0.17.5' + id 'com.avast.gradle.docker-compose' version '0.17.6' id "com.diffplug.spotless" version "6.23.3" // https://blog.ltgt.net/javax-jakarta-mess-and-gradle-solution/ // TODO id "org.gradlex.java-ecosystem-capabilities" version "1.0" diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java index b75530773c352f..ac07053e59d75a 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java @@ -13,7 +13,7 @@ import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.mxe.MetadataChangeProposal; @@ -22,14 +22,14 @@ public class TestUtils { - public static EntityService getMockEntityService() { + public static EntityService getMockEntityService() { PathSpecBasedSchemaAnnotationVisitor.class .getClassLoader() .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); EntityRegistry registry = new ConfigEntityRegistry(TestUtils.class.getResourceAsStream("/test-entity-registry.yaml")); - EntityService mockEntityService = - (EntityService) Mockito.mock(EntityService.class); + EntityService mockEntityService = + (EntityService) Mockito.mock(EntityService.class); Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry); return mockEntityService; } @@ -111,14 +111,14 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ } public static void verifyIngestProposal( - EntityService mockService, + EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { verifyIngestProposal(mockService, numberOfInvocations, List.of(proposal)); } public static void verifyIngestProposal( - EntityService mockService, + EntityService mockService, int numberOfInvocations, List proposals) { AspectsBatchImpl batch = @@ -128,7 +128,7 @@ public static void verifyIngestProposal( } public static void verifySingleIngestProposal( - EntityService mockService, + EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)) @@ -136,13 +136,13 @@ public static void verifySingleIngestProposal( } public static void verifyIngestProposal( - EntityService mockService, int numberOfInvocations) { + EntityService mockService, int numberOfInvocations) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)) .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false)); } public static void verifySingleIngestProposal( - EntityService mockService, int numberOfInvocations) { + EntityService mockService, int numberOfInvocations) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)) .ingestProposal( Mockito.any(MetadataChangeProposal.class), @@ -150,7 +150,7 @@ public static void verifySingleIngestProposal( Mockito.eq(false)); } - public static void verifyNoIngestProposal(EntityService mockService) { + public static void verifyNoIngestProposal(EntityService mockService) { Mockito.verify(mockService, Mockito.times(0)) .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/AddTagsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/AddTagsResolverTest.java index 1898753e5ae76f..b8c4ce21949373 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/AddTagsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/AddTagsResolverTest.java @@ -18,7 +18,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; @@ -221,7 +221,7 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = getMockEntityService(); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class) .when(mockService) diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreStorageStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreStorageStep.java index c756407832a36e..147acc9c1e0f33 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreStorageStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreStorageStep.java @@ -178,8 +178,9 @@ private void readerExecutable(ReaderWrapper reader, UpgradeContext context) { final RecordTemplate aspectRecord; try { aspectRecord = - EntityUtils.toAspectRecord( - entityName, aspectName, aspect.getMetadata(), _entityRegistry); + EntityUtils.toSystemAspect(aspect.toEntityAspect(), _entityService) + .get() + .getRecordTemplate(); } catch (Exception e) { context .report() diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesPreStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesPreStep.java index 894075417a3498..0695dbe4b1acb0 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesPreStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/elasticsearch/steps/BuildIndicesPreStep.java @@ -2,10 +2,13 @@ import static com.linkedin.datahub.upgrade.system.elasticsearch.util.IndexUtils.INDEX_BLOCKS_WRITE_SETTING; import static com.linkedin.datahub.upgrade.system.elasticsearch.util.IndexUtils.getAllReindexConfigs; +import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME; import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME; import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME; +import com.datahub.util.RecordUtils; import com.google.common.collect.ImmutableMap; +import com.linkedin.common.Status; import com.linkedin.datahub.upgrade.UpgradeContext; import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; @@ -14,14 +17,15 @@ import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; import com.linkedin.metadata.entity.AspectDao; -import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.elasticsearch.indexbuilder.ReindexConfig; import com.linkedin.metadata.shared.ElasticSearchIndexed; import com.linkedin.structured.StructuredPropertyDefinition; +import com.linkedin.util.Pair; import java.io.IOException; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; import lombok.RequiredArgsConstructor; @@ -54,24 +58,13 @@ public int retryCount() { public Function executable() { return (context) -> { try { - List reindexConfigs = - _configurationProvider.getStructuredProperties().isSystemUpdateEnabled() - ? getAllReindexConfigs( - _services, - _aspectDao - .streamAspects( - STRUCTURED_PROPERTY_ENTITY_NAME, - STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME) - .map( - entityAspect -> - EntityUtils.toAspectRecord( - STRUCTURED_PROPERTY_ENTITY_NAME, - STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME, - entityAspect.getMetadata(), - _entityRegistry)) - .map(recordTemplate -> (StructuredPropertyDefinition) recordTemplate) - .collect(Collectors.toSet())) - : getAllReindexConfigs(_services); + final List reindexConfigs; + if (_configurationProvider.getStructuredProperties().isSystemUpdateEnabled()) { + reindexConfigs = + getAllReindexConfigs(_services, getActiveStructuredPropertiesDefinitions(_aspectDao)); + } else { + reindexConfigs = getAllReindexConfigs(_services); + } // Get indices to update List indexConfigs = @@ -160,4 +153,31 @@ private boolean blockWrites(String indexName) throws InterruptedException, IOExc return ack; } + + private static Set getActiveStructuredPropertiesDefinitions( + AspectDao aspectDao) { + Set removedStructuredPropertyUrns = + aspectDao + .streamAspects(STRUCTURED_PROPERTY_ENTITY_NAME, STATUS_ASPECT_NAME) + .map( + entityAspect -> + Pair.of( + entityAspect.getUrn(), + RecordUtils.toRecordTemplate(Status.class, entityAspect.getMetadata()))) + .filter(status -> status.getSecond().isRemoved()) + .map(Pair::getFirst) + .collect(Collectors.toSet()); + + return aspectDao + .streamAspects(STRUCTURED_PROPERTY_ENTITY_NAME, STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME) + .map( + entityAspect -> + Pair.of( + entityAspect.getUrn(), + RecordUtils.toRecordTemplate( + StructuredPropertyDefinition.class, entityAspect.getMetadata()))) + .filter(definition -> !removedStructuredPropertyUrns.contains(definition.getKey())) + .map(Pair::getSecond) + .collect(Collectors.toSet()); + } } diff --git a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/DatahubUpgradeNoSchemaRegistryTest.java b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/DatahubUpgradeNoSchemaRegistryTest.java index 4c9e12c0ed1511..ed09a4a5aec2b9 100644 --- a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/DatahubUpgradeNoSchemaRegistryTest.java +++ b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/DatahubUpgradeNoSchemaRegistryTest.java @@ -49,7 +49,7 @@ public void testSystemUpdateInit() { @Test public void testSystemUpdateKafkaProducerOverride() { assertEquals(kafkaEventProducer, duheKafkaEventProducer); - assertEquals(entityService.get_producer(), duheKafkaEventProducer); + assertEquals(entityService.getProducer(), duheKafkaEventProducer); } @Test diff --git a/entity-registry/build.gradle b/entity-registry/build.gradle index 315a29e305b77c..66e4ad4b930e07 100644 --- a/entity-registry/build.gradle +++ b/entity-registry/build.gradle @@ -32,6 +32,7 @@ dependencies { testImplementation externalDependency.mockito testImplementation externalDependency.mockitoInline testCompileOnly externalDependency.lombok + testAnnotationProcessor externalDependency.lombok testImplementation externalDependency.classGraph } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectRetriever.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/AspectRetriever.java similarity index 95% rename from entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectRetriever.java rename to entity-registry/src/main/java/com/linkedin/metadata/aspect/AspectRetriever.java index 11cd2352025efe..2ef22483da1ca5 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectRetriever.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/AspectRetriever.java @@ -1,4 +1,4 @@ -package com.linkedin.metadata.aspect.plugins.validation; +package com.linkedin.metadata.aspect; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/CachingAspectRetriever.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/CachingAspectRetriever.java new file mode 100644 index 00000000000000..77e799f752455c --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/CachingAspectRetriever.java @@ -0,0 +1,4 @@ +package com.linkedin.metadata.aspect; + +/** Responses can be cached based on application.yaml caching configuration for the EntityClient */ +public interface CachingAspectRetriever extends AspectRetriever {} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/ReadItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/ReadItem.java new file mode 100644 index 00000000000000..342b5376d8a755 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/ReadItem.java @@ -0,0 +1,77 @@ +package com.linkedin.metadata.aspect; + +import com.linkedin.common.urn.Urn; +import com.linkedin.data.DataMap; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.mxe.SystemMetadata; +import java.lang.reflect.InvocationTargetException; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +public interface ReadItem { + /** + * The urn associated with the aspect + * + * @return + */ + @Nonnull + Urn getUrn(); + + /** + * Aspect's name + * + * @return the name + */ + @Nonnull + default String getAspectName() { + return getAspectSpec().getName(); + } + + @Nullable + RecordTemplate getRecordTemplate(); + + default T getAspect(Class clazz) { + return getAspect(clazz, getRecordTemplate()); + } + + static T getAspect(Class clazz, @Nullable RecordTemplate recordTemplate) { + if (recordTemplate != null) { + try { + return clazz.getConstructor(DataMap.class).newInstance(recordTemplate.data()); + } catch (InstantiationException + | IllegalAccessException + | InvocationTargetException + | NoSuchMethodException e) { + throw new RuntimeException(e); + } + } else { + return null; + } + } + + /** + * System information + * + * @return the system metadata + */ + @Nullable + SystemMetadata getSystemMetadata(); + + /** + * The entity's schema + * + * @return entity specification + */ + @Nonnull + EntitySpec getEntitySpec(); + + /** + * The aspect's schema + * + * @return aspect's specification + */ + @Nonnull + AspectSpec getAspectSpec(); +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/SystemAspect.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/SystemAspect.java new file mode 100644 index 00000000000000..e83414c8c23a85 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/SystemAspect.java @@ -0,0 +1,25 @@ +package com.linkedin.metadata.aspect; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.UrnUtils; +import java.sql.Timestamp; +import javax.annotation.Nonnull; + +/** + * An aspect along with system metadata and creation timestamp. Represents an aspect as stored in + * primary storage. + */ +public interface SystemAspect extends ReadItem { + long getVersion(); + + Timestamp getCreatedOn(); + + String getCreatedBy(); + + @Nonnull + default AuditStamp getAuditStamp() { + return new AuditStamp() + .setActor(UrnUtils.getUrn(getCreatedBy())) + .setTime(getCreatedOn().getTime()); + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java index 3d803d238b4f92..ddedc96b385779 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java @@ -1,6 +1,10 @@ package com.linkedin.metadata.aspect.batch; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.ReadItem; +import com.linkedin.metadata.aspect.SystemAspect; +import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; +import com.linkedin.metadata.aspect.plugins.validation.ValidationExceptionCollection; import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; import java.util.Collection; @@ -20,27 +24,108 @@ public interface AspectsBatch { Collection getItems(); + AspectRetriever getAspectRetriever(); + /** - * Returns MCP items. Can be patch, upsert, etc. + * Returns MCP items. Could be patch, upsert, etc. * * @return batch items */ - default Collection getMCPItems() { + default Collection getMCPItems() { return getItems().stream() - .filter(item -> item instanceof MCPBatchItem) - .map(item -> (MCPBatchItem) item) + .filter(item -> item instanceof MCPItem) + .map(item -> (MCPItem) item) .collect(Collectors.toList()); } - Pair>, List> toUpsertBatchItems( - Map> latestAspects, AspectRetriever aspectRetriever); + /** + * Convert patches to upserts, apply hooks at the aspect and batch level. + * + * @param latestAspects latest version in the database + * @return The new urn/aspectnames and the uniform upserts, possibly expanded/mutated by the + * various hooks + */ + Pair>, List> toUpsertBatchItems( + Map> latestAspects); + + /** + * Apply read mutations to batch + * + * @param items + */ + default void applyReadMutationHooks(Collection items) { + applyReadMutationHooks(items, getAspectRetriever()); + } + + static void applyReadMutationHooks(Collection items, AspectRetriever aspectRetriever) { + for (MutationHook mutationHook : aspectRetriever.getEntityRegistry().getAllMutationHooks()) { + mutationHook.applyReadMutation(items, aspectRetriever); + } + } + + /** + * Apply write mutations to batch + * + * @param changeMCPS + */ + default void applyWriteMutationHooks(Collection changeMCPS) { + applyWriteMutationHooks(changeMCPS, getAspectRetriever()); + } + + static void applyWriteMutationHooks( + Collection changeMCPS, AspectRetriever aspectRetriever) { + for (MutationHook mutationHook : aspectRetriever.getEntityRegistry().getAllMutationHooks()) { + mutationHook.applyWriteMutation(changeMCPS, aspectRetriever); + } + } + + default ValidationExceptionCollection validateProposed( + Collection mcpItems) { + return validateProposed(mcpItems, getAspectRetriever()); + } + + static ValidationExceptionCollection validateProposed( + Collection mcpItems, AspectRetriever aspectRetriever) { + ValidationExceptionCollection exceptions = ValidationExceptionCollection.newCollection(); + aspectRetriever.getEntityRegistry().getAllAspectPayloadValidators().stream() + .flatMap(validator -> validator.validateProposed(mcpItems, aspectRetriever)) + .forEach(exceptions::addException); + return exceptions; + } + + default ValidationExceptionCollection validatePreCommit(Collection changeMCPs) { + return validatePreCommit(changeMCPs, getAspectRetriever()); + } + + static ValidationExceptionCollection validatePreCommit( + Collection changeMCPs, AspectRetriever aspectRetriever) { + ValidationExceptionCollection exceptions = ValidationExceptionCollection.newCollection(); + aspectRetriever.getEntityRegistry().getAllAspectPayloadValidators().stream() + .flatMap(validator -> validator.validatePreCommit(changeMCPs, aspectRetriever)) + .forEach(exceptions::addException); + return exceptions; + } + + default Stream applyMCPSideEffects(Collection items) { + return applyMCPSideEffects(items, getAspectRetriever()); + } - default Stream applyMCPSideEffects( - List items, AspectRetriever aspectRetriever) { + static Stream applyMCPSideEffects( + Collection items, AspectRetriever aspectRetriever) { return aspectRetriever.getEntityRegistry().getAllMCPSideEffects().stream() .flatMap(mcpSideEffect -> mcpSideEffect.apply(items, aspectRetriever)); } + default Stream applyMCLSideEffects(Collection items) { + return applyMCLSideEffects(items, getAspectRetriever()); + } + + static Stream applyMCLSideEffects( + Collection items, AspectRetriever aspectRetriever) { + return aspectRetriever.getEntityRegistry().getAllMCLSideEffects().stream() + .flatMap(mclSideEffect -> mclSideEffect.apply(items, aspectRetriever)); + } + default boolean containsDuplicateAspects() { return getItems().stream() .map(i -> String.format("%s_%s", i.getClass().getName(), i.hashCode())) @@ -81,7 +166,7 @@ default Map> getNewUrnAspectsMap( .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } - default Map> merge( + static Map> merge( @Nonnull Map> a, @Nonnull Map> b) { return Stream.concat(a.entrySet().stream(), b.entrySet().stream()) .flatMap( diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java index 60033cd6919d60..a6dfbc277e12ec 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java @@ -1,45 +1,19 @@ package com.linkedin.metadata.aspect.batch; import com.linkedin.common.AuditStamp; -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.mxe.SystemMetadata; +import com.linkedin.metadata.aspect.ReadItem; import javax.annotation.Nonnull; import javax.annotation.Nullable; -public interface BatchItem { - /** - * The urn associated with the aspect - * - * @return - */ - Urn getUrn(); - - /** - * Aspect's name - * - * @return the name - */ - @Nonnull - default String getAspectName() { - return getAspectSpec().getName(); - } - - /** - * System information - * - * @return the system metadata - */ - SystemMetadata getSystemMetadata(); +public interface BatchItem extends ReadItem { /** * Timestamp and actor * * @return the audit information */ + @Nullable AuditStamp getAuditStamp(); /** @@ -49,28 +23,4 @@ default String getAspectName() { */ @Nonnull ChangeType getChangeType(); - - /** - * The entity's schema - * - * @return entity specification - */ - @Nonnull - EntitySpec getEntitySpec(); - - /** - * The aspect's schema - * - * @return aspect's specification - */ - @Nonnull - AspectSpec getAspectSpec(); - - /** - * The aspect's record template. Null when patch - * - * @return record template if it exists - */ - @Nullable - RecordTemplate getRecordTemplate(); } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java new file mode 100644 index 00000000000000..94e8bbab3ceeb6 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/ChangeMCP.java @@ -0,0 +1,42 @@ +package com.linkedin.metadata.aspect.batch; + +import com.linkedin.data.DataMap; +import com.linkedin.metadata.aspect.SystemAspect; +import java.lang.reflect.InvocationTargetException; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** + * A proposal to write data to the primary datastore which includes system metadata and other + * related data stored along with the aspect + */ +public interface ChangeMCP extends MCPItem { + @Nonnull + SystemAspect getSystemAspect(@Nullable Long nextAspectVersion); + + @Nullable + SystemAspect getPreviousSystemAspect(); + + void setPreviousSystemAspect(@Nullable SystemAspect previousSystemAspect); + + long getNextAspectVersion(); + + void setNextAspectVersion(long nextAspectVersion); + + default T getPreviousAspect(Class clazz) { + if (getPreviousSystemAspect() != null) { + try { + return clazz + .getConstructor(DataMap.class) + .newInstance(getPreviousSystemAspect().getRecordTemplate().data()); + } catch (InstantiationException + | IllegalAccessException + | InvocationTargetException + | NoSuchMethodException e) { + throw new RuntimeException(e); + } + } else { + return null; + } + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLItem.java similarity index 70% rename from entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java rename to entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLItem.java index 17a910b125a34f..9fd2d6c607342f 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLItem.java @@ -2,15 +2,17 @@ import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; +import com.linkedin.data.DataMap; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.SystemMetadata; +import java.lang.reflect.InvocationTargetException; import javax.annotation.Nonnull; import javax.annotation.Nullable; /** An item that represents a change that has been written to primary storage. */ -public interface MCLBatchItem extends BatchItem { +public interface MCLItem extends BatchItem { @Nonnull MetadataChangeLog getMetadataChangeLog(); @@ -42,6 +44,21 @@ default SystemMetadata getPreviousSystemMetadata() { @Nullable RecordTemplate getPreviousRecordTemplate(); + default T getPreviousAspect(Class clazz) { + if (getPreviousRecordTemplate() != null) { + try { + return clazz.getConstructor(DataMap.class).newInstance(getPreviousRecordTemplate().data()); + } catch (InstantiationException + | IllegalAccessException + | InvocationTargetException + | NoSuchMethodException e) { + throw new RuntimeException(e); + } + } else { + return null; + } + } + @Override @Nonnull default ChangeType getChangeType() { diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPBatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPItem.java similarity index 83% rename from entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPBatchItem.java rename to entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPItem.java index dd0d0ec68dac6c..8c25f3c4f44de6 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPBatchItem.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPItem.java @@ -7,10 +7,10 @@ import javax.annotation.Nullable; /** Represents a proposal to write to the primary data store which may be represented by an MCP */ -public abstract class MCPBatchItem implements BatchItem { +public interface MCPItem extends BatchItem { @Nullable - public abstract MetadataChangeProposal getMetadataChangeProposal(); + MetadataChangeProposal getMetadataChangeProposal(); /** * Validates that a change type is valid for the given aspect @@ -19,7 +19,7 @@ public abstract class MCPBatchItem implements BatchItem { * @param aspectSpec * @return */ - protected static boolean isValidChangeType(ChangeType changeType, AspectSpec aspectSpec) { + static boolean isValidChangeType(ChangeType changeType, AspectSpec aspectSpec) { if (aspectSpec.isTimeseries()) { // Timeseries aspects only support UPSERT return ChangeType.UPSERT.equals(changeType); @@ -32,7 +32,7 @@ protected static boolean isValidChangeType(ChangeType changeType, AspectSpec asp } } - protected static boolean supportsPatch(AspectSpec aspectSpec) { + static boolean supportsPatch(AspectSpec aspectSpec) { // Limit initial support to defined templates if (!AspectTemplateEngine.SUPPORTED_TEMPLATES.contains(aspectSpec.getName())) { // Prevent unexpected behavior for aspects that do not currently have 1st class patch support, diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchMCP.java similarity index 58% rename from entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchItem.java rename to entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchMCP.java index e9e30f7f2bd96f..f04133e9e1ff8b 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchItem.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchMCP.java @@ -2,12 +2,12 @@ import com.github.fge.jsonpatch.Patch; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.aspect.AspectRetriever; /** * A change proposal represented as a patch to an exiting stored object in the primary data store. */ -public abstract class PatchItem extends MCPBatchItem { +public interface PatchMCP extends MCPItem { /** * Convert a Patch to an Upsert @@ -15,8 +15,7 @@ public abstract class PatchItem extends MCPBatchItem { * @param recordTemplate the current value record template * @return the upsert */ - public abstract UpsertItem applyPatch( - RecordTemplate recordTemplate, AspectRetriever aspectRetriever); + ChangeMCP applyPatch(RecordTemplate recordTemplate, AspectRetriever aspectRetriever); - public abstract Patch getPatch(); + Patch getPatch(); } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/SystemAspect.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/SystemAspect.java deleted file mode 100644 index 88ac902ae52fed..00000000000000 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/SystemAspect.java +++ /dev/null @@ -1,25 +0,0 @@ -package com.linkedin.metadata.aspect.batch; - -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.mxe.SystemMetadata; -import java.sql.Timestamp; - -/** - * An aspect along with system metadata and creation timestamp. Represents an aspect as stored in - * primary storage. - */ -public interface SystemAspect { - Urn getUrn(); - - String getAspectName(); - - long getVersion(); - - RecordTemplate getRecordTemplate(EntityRegistry entityRegistry); - - SystemMetadata getSystemMetadata(); - - Timestamp getCreatedOn(); -} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java deleted file mode 100644 index c64105637dfcc6..00000000000000 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java +++ /dev/null @@ -1,19 +0,0 @@ -package com.linkedin.metadata.aspect.batch; - -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; - -/** - * A proposal to write data to the primary datastore which includes system metadata and other - * related data stored along with the aspect - */ -public abstract class UpsertItem extends MCPBatchItem { - public abstract SystemAspect toLatestEntityAspect(); - - public abstract void validatePreCommit( - @Nullable RecordTemplate previous, @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException; -} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDelete.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDelete.java new file mode 100644 index 00000000000000..98e90cfaa45cfc --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDelete.java @@ -0,0 +1,38 @@ +package com.linkedin.metadata.aspect.hooks; + +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.ReadItem; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; +import com.linkedin.metadata.models.StructuredPropertyUtils; +import com.linkedin.structured.StructuredProperties; +import com.linkedin.util.Pair; +import java.util.Collection; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nonnull; + +public class StructuredPropertiesSoftDelete extends MutationHook { + public StructuredPropertiesSoftDelete(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + @Override + protected Stream> readMutation( + @Nonnull Collection items, @Nonnull AspectRetriever aspectRetriever) { + Map entityStructuredPropertiesMap = + items.stream() + .filter(i -> i.getRecordTemplate() != null) + .map(i -> Pair.of(i.getUrn(), i.getAspect(StructuredProperties.class))) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue)); + + // Apply filter + Map mutatedEntityStructuredPropertiesMap = + StructuredPropertyUtils.filterSoftDelete(entityStructuredPropertiesMap, aspectRetriever); + + return items.stream() + .map(i -> Pair.of(i, mutatedEntityStructuredPropertiesMap.getOrDefault(i.getUrn(), false))); + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginSpec.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginSpec.java index d88b05ede84548..6937070a684e29 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginSpec.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginSpec.java @@ -4,7 +4,9 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.AllArgsConstructor; import lombok.EqualsAndHashCode; @@ -20,22 +22,34 @@ protected AspectPluginConfig getConfig() { } public boolean shouldApply( - @Nonnull ChangeType changeType, @Nonnull Urn entityUrn, @Nonnull AspectSpec aspectSpec) { + @Nullable ChangeType changeType, @Nonnull Urn entityUrn, @Nonnull AspectSpec aspectSpec) { return shouldApply(changeType, entityUrn.getEntityType(), aspectSpec); } public boolean shouldApply( - @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull AspectSpec aspectSpec) { + @Nullable ChangeType changeType, + @Nonnull EntitySpec entitySpec, + @Nonnull AspectSpec aspectSpec) { + return shouldApply(changeType, entitySpec.getName(), aspectSpec.getName()); + } + + public boolean shouldApply( + @Nullable ChangeType changeType, @Nonnull String entityName, @Nonnull AspectSpec aspectSpec) { return shouldApply(changeType, entityName, aspectSpec.getName()); } public boolean shouldApply( - @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { + @Nullable ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { return getConfig().isEnabled() && isChangeTypeSupported(changeType) && isEntityAspectSupported(entityName, aspectName); } + protected boolean isEntityAspectSupported( + @Nonnull EntitySpec entitySpec, @Nonnull AspectSpec aspectSpec) { + return isEntityAspectSupported(entitySpec.getName(), aspectSpec.getName()); + } + protected boolean isEntityAspectSupported( @Nonnull String entityName, @Nonnull String aspectName) { return (getConfig().getSupportedEntityAspectNames().stream() @@ -51,8 +65,9 @@ protected boolean isAspectSupported(@Nonnull String aspectName) { .anyMatch(supported -> supported.getAspectName().equals(aspectName)); } - protected boolean isChangeTypeSupported(@Nonnull ChangeType changeType) { - return getConfig().getSupportedOperations().stream() - .anyMatch(supported -> changeType.toString().equals(supported)); + protected boolean isChangeTypeSupported(@Nullable ChangeType changeType) { + return (changeType == null && getConfig().getSupportedOperations().isEmpty()) + || getConfig().getSupportedOperations().stream() + .anyMatch(supported -> supported.equalsIgnoreCase(String.valueOf(changeType))); } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/AspectPluginConfig.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/AspectPluginConfig.java index 059f133ad27760..00ebcf6b464911 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/AspectPluginConfig.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/AspectPluginConfig.java @@ -1,7 +1,10 @@ package com.linkedin.metadata.aspect.plugins.config; +import java.util.Collections; import java.util.List; +import java.util.Objects; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; @@ -15,7 +18,7 @@ public class AspectPluginConfig { @Nonnull private String className; private boolean enabled; - @Nonnull private List supportedOperations; + @Nullable private List supportedOperations; @Nonnull private List supportedEntityAspectNames; @Data @@ -27,6 +30,11 @@ public static class EntityAspectName { @Nonnull private String aspectName; } + @Nonnull + public List getSupportedOperations() { + return supportedOperations != null ? supportedOperations : Collections.emptyList(); + } + /** * Used to determine is an earlier plugin is disabled by a subsequent plugin * @@ -44,7 +52,7 @@ private boolean isEqualExcludingEnabled(Object o) { AspectPluginConfig that = (AspectPluginConfig) o; if (!className.equals(that.className)) return false; - if (!supportedOperations.equals(that.supportedOperations)) return false; + if (!Objects.equals(supportedOperations, that.supportedOperations)) return false; return supportedEntityAspectNames.equals(that.supportedEntityAspectNames); } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffect.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffect.java index a21f3cd2436de7..c60af636e424f1 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffect.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffect.java @@ -1,15 +1,18 @@ package com.linkedin.metadata.aspect.plugins.hooks; -import com.linkedin.metadata.aspect.batch.MCLBatchItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.MCLItem; import com.linkedin.metadata.aspect.plugins.PluginSpec; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import java.util.List; +import java.util.Collection; +import java.util.function.BiFunction; +import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nonnull; /** Given an MCL produce additional MCLs for writing */ -public abstract class MCLSideEffect extends PluginSpec { +public abstract class MCLSideEffect extends PluginSpec + implements BiFunction, AspectRetriever, Stream> { public MCLSideEffect(AspectPluginConfig aspectPluginConfig) { super(aspectPluginConfig); @@ -18,16 +21,19 @@ public MCLSideEffect(AspectPluginConfig aspectPluginConfig) { /** * Given a list of MCLs, output additional MCLs * - * @param input list + * @param batchItems list * @return additional upserts */ - public final Stream apply( - @Nonnull List input, @Nonnull AspectRetriever aspectRetriever) { - return input.stream() - .filter(item -> shouldApply(item.getChangeType(), item.getUrn(), item.getAspectSpec())) - .flatMap(i -> applyMCLSideEffect(i, aspectRetriever)); + @Override + public final Stream apply( + @Nonnull Collection batchItems, @Nonnull AspectRetriever aspectRetriever) { + return applyMCLSideEffect( + batchItems.stream() + .filter(item -> shouldApply(item.getChangeType(), item.getUrn(), item.getAspectSpec())) + .collect(Collectors.toList()), + aspectRetriever); } - protected abstract Stream applyMCLSideEffect( - @Nonnull MCLBatchItem input, @Nonnull AspectRetriever aspectRetriever); + protected abstract Stream applyMCLSideEffect( + @Nonnull Collection batchItems, @Nonnull AspectRetriever aspectRetriever); } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffect.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffect.java index 80cb405201c876..c346695c51d30f 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffect.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffect.java @@ -1,15 +1,18 @@ package com.linkedin.metadata.aspect.plugins.hooks; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.PluginSpec; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import java.util.List; +import java.util.Collection; +import java.util.function.BiFunction; +import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nonnull; /** Given an MCP produce additional MCPs to write */ -public abstract class MCPSideEffect extends PluginSpec { +public abstract class MCPSideEffect extends PluginSpec + implements BiFunction, AspectRetriever, Stream> { public MCPSideEffect(AspectPluginConfig aspectPluginConfig) { super(aspectPluginConfig); @@ -18,16 +21,18 @@ public MCPSideEffect(AspectPluginConfig aspectPluginConfig) { /** * Given the list of MCP upserts, output additional upserts * - * @param input list + * @param changeMCPS list * @return additional upserts */ - public final Stream apply( - List input, @Nonnull AspectRetriever aspectRetriever) { - return input.stream() - .filter(item -> shouldApply(item.getChangeType(), item.getUrn(), item.getAspectSpec())) - .flatMap(i -> applyMCPSideEffect(i, aspectRetriever)); + public final Stream apply( + Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { + return applyMCPSideEffect( + changeMCPS.stream() + .filter(item -> shouldApply(item.getChangeType(), item.getUrn(), item.getAspectSpec())) + .collect(Collectors.toList()), + aspectRetriever); } - protected abstract Stream applyMCPSideEffect( - UpsertItem input, @Nonnull AspectRetriever aspectRetriever); + protected abstract Stream applyMCPSideEffect( + Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever); } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MutationHook.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MutationHook.java index 730a494c03d7b9..1d6b4eeb617f5d 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MutationHook.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MutationHook.java @@ -1,16 +1,15 @@ package com.linkedin.metadata.aspect.plugins.hooks; -import com.linkedin.common.AuditStamp; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.ReadItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.PluginSpec; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; +import java.util.Collection; +import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; -import javax.annotation.Nullable; /** Applies changes to the RecordTemplate prior to write */ public abstract class MutationHook extends PluginSpec { @@ -20,49 +19,38 @@ public MutationHook(AspectPluginConfig aspectPluginConfig) { } /** - * Mutating hook + * Mutating hook, original objects are potentially modified. * - * @param changeType Type of change to mutate - * @param entitySpec Entity specification - * @param aspectSpec Aspect specification - * @param oldAspectValue old aspect vale if it exists - * @param newAspectValue the new aspect - * @param oldSystemMetadata old system metadata if it exists - * @param newSystemMetadata the new system metadata - * @param auditStamp the audit stamp + * @param changeMCPS input upsert items + * @param aspectRetriever aspect retriever + * @return all items, with a boolean to indicate mutation */ - public final void applyMutation( - @Nonnull final ChangeType changeType, - @Nonnull EntitySpec entitySpec, - @Nonnull final AspectSpec aspectSpec, - @Nullable final RecordTemplate oldAspectValue, - @Nullable final RecordTemplate newAspectValue, - @Nullable final SystemMetadata oldSystemMetadata, - @Nullable final SystemMetadata newSystemMetadata, - @Nonnull AuditStamp auditStamp, - @Nonnull AspectRetriever aspectRetriever) { - if (shouldApply(changeType, entitySpec.getName(), aspectSpec)) { - mutate( - changeType, - entitySpec, - aspectSpec, - oldAspectValue, - newAspectValue, - oldSystemMetadata, - newSystemMetadata, - auditStamp, - aspectRetriever); - } + public final Stream> applyWriteMutation( + @Nonnull Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { + return writeMutation( + changeMCPS.stream() + .filter(i -> shouldApply(i.getChangeType(), i.getEntitySpec(), i.getAspectSpec())) + .collect(Collectors.toList()), + aspectRetriever); } - protected abstract void mutate( - @Nonnull final ChangeType changeType, - @Nonnull EntitySpec entitySpec, - @Nonnull final AspectSpec aspectSpec, - @Nullable final RecordTemplate oldAspectValue, - @Nullable final RecordTemplate newAspectValue, - @Nullable final SystemMetadata oldSystemMetadata, - @Nullable final SystemMetadata newSystemMetadata, - @Nonnull AuditStamp auditStamp, - @Nonnull AspectRetriever aspectRetriever); + // Read mutation + public final Stream> applyReadMutation( + @Nonnull Collection items, @Nonnull AspectRetriever aspectRetriever) { + return readMutation( + items.stream() + .filter(i -> isEntityAspectSupported(i.getEntitySpec(), i.getAspectSpec())) + .collect(Collectors.toList()), + aspectRetriever); + } + + protected Stream> readMutation( + @Nonnull Collection items, @Nonnull AspectRetriever aspectRetriever) { + return items.stream().map(i -> Pair.of(i, false)); + } + + protected Stream> writeMutation( + @Nonnull Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { + return changeMCPS.stream().map(i -> Pair.of(i, false)); + } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectPayloadValidator.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectPayloadValidator.java index 656d017724571e..6e4bc45e563b90 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectPayloadValidator.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectPayloadValidator.java @@ -1,13 +1,14 @@ package com.linkedin.metadata.aspect.plugins.validation; -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.PluginSpec; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.models.AspectSpec; +import java.util.Collection; +import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; -import javax.annotation.Nullable; public abstract class AspectPayloadValidator extends PluginSpec { @@ -19,65 +20,35 @@ public AspectPayloadValidator(AspectPluginConfig aspectPluginConfig) { * Validate a proposal for the given change type for an aspect within the context of the given * entity's urn. * - * @param changeType The change type - * @param entityUrn The parent entity for the aspect - * @param aspectSpec The aspect's specification - * @param aspectPayload The aspect's payload * @return whether the aspect proposal is valid - * @throws AspectValidationException */ - public final void validateProposed( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nonnull RecordTemplate aspectPayload, - @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { - if (shouldApply(changeType, entityUrn, aspectSpec)) { - validateProposedAspect(changeType, entityUrn, aspectSpec, aspectPayload, aspectRetriever); - } + public final Stream validateProposed( + @Nonnull Collection mcpItems, @Nonnull AspectRetriever aspectRetriever) { + return validateProposedAspects( + mcpItems.stream() + .filter(i -> shouldApply(i.getChangeType(), i.getUrn(), i.getAspectSpec())) + .collect(Collectors.toList()), + aspectRetriever); } /** * Validate the proposed aspect as its about to be written with the context of the previous * version of the aspect (if it existed) * - * @param changeType The change type - * @param entityUrn The parent entity for the aspect - * @param aspectSpec The aspect's specification - * @param previousAspect The previous version of the aspect if it exists - * @param proposedAspect The new version of the aspect * @return whether the aspect proposal is valid - * @throws AspectValidationException */ - public final void validatePreCommit( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate previousAspect, - @Nonnull RecordTemplate proposedAspect, - AspectRetriever aspectRetriever) - throws AspectValidationException { - if (shouldApply(changeType, entityUrn, aspectSpec)) { - validatePreCommitAspect( - changeType, entityUrn, aspectSpec, previousAspect, proposedAspect, aspectRetriever); - } + public final Stream validatePreCommit( + @Nonnull Collection changeMCPs, AspectRetriever aspectRetriever) { + return validatePreCommitAspects( + changeMCPs.stream() + .filter(i -> shouldApply(i.getChangeType(), i.getUrn(), i.getAspectSpec())) + .collect(Collectors.toList()), + aspectRetriever); } - protected abstract void validateProposedAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nonnull RecordTemplate aspectPayload, - @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException; + protected abstract Stream validateProposedAspects( + @Nonnull Collection mcpItems, @Nonnull AspectRetriever aspectRetriever); - protected abstract void validatePreCommitAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate previousAspect, - @Nonnull RecordTemplate proposedAspect, - AspectRetriever aspectRetriever) - throws AspectValidationException; + protected abstract Stream validatePreCommitAspects( + @Nonnull Collection changeMCPs, AspectRetriever aspectRetriever); } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectValidationException.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectValidationException.java index f858bdcf141aeb..f83642c5eed9ec 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectValidationException.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectValidationException.java @@ -1,12 +1,59 @@ package com.linkedin.metadata.aspect.plugins.validation; +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.util.Pair; +import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + public class AspectValidationException extends Exception { - public AspectValidationException(String msg) { - super(msg); + public static AspectValidationException forItem(BatchItem item, String msg) { + return forItem(item, msg, null); + } + + public static AspectValidationException forItem(BatchItem item, String msg, Exception e) { + return new AspectValidationException(item.getUrn(), item.getAspectName(), msg, e); } - public AspectValidationException(String msg, Exception e) { + @Nonnull private final Urn entityUrn; + @Nonnull private final String aspectName; + @Nullable private final String msg; + + public AspectValidationException(@Nonnull Urn entityUrn, @Nonnull String aspectName, String msg) { + this(entityUrn, aspectName, msg, null); + } + + public AspectValidationException( + @Nonnull Urn entityUrn, @Nonnull String aspectName, @Nonnull String msg, Exception e) { super(msg, e); + this.entityUrn = entityUrn; + this.aspectName = aspectName; + this.msg = msg; + } + + public Pair getExceptionKey() { + return Pair.of(entityUrn, aspectName); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + AspectValidationException that = (AspectValidationException) o; + + if (!entityUrn.equals(that.entityUrn)) return false; + if (!aspectName.equals(that.aspectName)) return false; + return Objects.equals(msg, that.msg); + } + + @Override + public int hashCode() { + int result = entityUrn.hashCode(); + result = 31 * result + aspectName.hashCode(); + result = 31 * result + (msg != null ? msg.hashCode() : 0); + return result; } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/ValidationExceptionCollection.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/ValidationExceptionCollection.java new file mode 100644 index 00000000000000..559fa85cff04c4 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/ValidationExceptionCollection.java @@ -0,0 +1,68 @@ +package com.linkedin.metadata.aspect.plugins.validation; + +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.util.Pair; +import java.util.Collection; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +/** Used to store a collection of exceptions, keyed by the URN/AspectName pair */ +public class ValidationExceptionCollection + extends HashMap, Set> { + + public static ValidationExceptionCollection newCollection() { + return new ValidationExceptionCollection(); + } + + public void addException(AspectValidationException exception) { + super.computeIfAbsent(exception.getExceptionKey(), key -> new HashSet<>()).add(exception); + } + + public void addException(BatchItem item, String message) { + addException(item, message, null); + } + + public void addException(BatchItem item, String message, Exception ex) { + super.computeIfAbsent(Pair.of(item.getUrn(), item.getAspectName()), key -> new HashSet<>()) + .add(AspectValidationException.forItem(item, message, ex)); + } + + public Stream streamAllExceptions() { + return values().stream().flatMap(Collection::stream); + } + + public Collection successful(Collection items) { + return streamSuccessful(items.stream()).collect(Collectors.toList()); + } + + public Stream streamSuccessful(Stream items) { + return items.filter(i -> !this.containsKey(Pair.of(i.getUrn(), i.getAspectName()))); + } + + public Collection exceptions(Collection items) { + return streamExceptions(items.stream()).collect(Collectors.toList()); + } + + public Stream streamExceptions(Stream items) { + return items.filter(i -> this.containsKey(Pair.of(i.getUrn(), i.getAspectName()))); + } + + @Override + public String toString() { + return String.format( + "ValidationExceptionCollection{%s}", + entrySet().stream() + // sort by entity/aspect + .sorted(Comparator.comparing(p -> p.getKey().toString())) + .map( + e -> + String.format( + "EntityAspect:%s Exceptions: %s", e.getKey().toString(), e.getValue())) + .collect(Collectors.joining("; "))); + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/PropertyDefinitionValidator.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/PropertyDefinitionValidator.java index 5a4635da433ae4..17b66de79d113c 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/PropertyDefinitionValidator.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/PropertyDefinitionValidator.java @@ -1,22 +1,37 @@ package com.linkedin.metadata.aspect.validation; +import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME; import static com.linkedin.structured.PropertyCardinality.*; +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.Status; import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; +import com.linkedin.entity.Aspect; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; -import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.aspect.plugins.validation.ValidationExceptionCollection; +import com.linkedin.r2.RemoteInvocationException; import com.linkedin.structured.PrimitivePropertyValue; import com.linkedin.structured.PropertyValue; import com.linkedin.structured.StructuredPropertyDefinition; +import java.net.URISyntaxException; +import java.util.Collection; +import java.util.Collections; +import java.util.Map; +import java.util.Objects; +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; public class PropertyDefinitionValidator extends AspectPayloadValidator { @@ -24,68 +39,128 @@ public PropertyDefinitionValidator(AspectPluginConfig aspectPluginConfig) { super(aspectPluginConfig); } + /** + * Prevent deletion of the definition or key aspect (only soft delete) + * + * @param mcpItems + * @param aspectRetriever + * @return + */ @Override - protected void validateProposedAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nonnull RecordTemplate aspectPayload, - @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { - // No-op + protected Stream validateProposedAspects( + @Nonnull Collection mcpItems, @Nonnull AspectRetriever aspectRetriever) { + final String entityKeyAspect = + aspectRetriever + .getEntityRegistry() + .getEntitySpec(STRUCTURED_PROPERTY_ENTITY_NAME) + .getKeyAspectName(); + + return mcpItems.stream() + .filter(i -> ChangeType.DELETE.equals(i.getChangeType())) + .map( + i -> { + if (ImmutableSet.of(entityKeyAspect, STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME) + .contains(i.getAspectSpec().getName())) { + return AspectValidationException.forItem( + i, "Hard delete of Structured Property Definitions is not supported."); + } + return null; + }) + .filter(Objects::nonNull); } @Override - protected void validatePreCommitAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate previousAspect, - @Nonnull RecordTemplate proposedAspect, - AspectRetriever aspectRetriever) - throws AspectValidationException { - validate(previousAspect, proposedAspect); + protected Stream validatePreCommitAspects( + @Nonnull Collection changeMCPs, AspectRetriever aspectRetriever) { + return validateDefinitionUpserts( + changeMCPs.stream() + .filter( + i -> + ChangeType.UPSERT.equals(i.getChangeType()) + && STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME.equals(i.getAspectName())) + .collect(Collectors.toList()), + aspectRetriever); } - public static boolean validate( - @Nullable RecordTemplate previousAspect, @Nonnull RecordTemplate proposedAspect) - throws AspectValidationException { - if (previousAspect != null) { - StructuredPropertyDefinition previousDefinition = - (StructuredPropertyDefinition) previousAspect; - StructuredPropertyDefinition newDefinition = (StructuredPropertyDefinition) proposedAspect; - if (!newDefinition.getValueType().equals(previousDefinition.getValueType())) { - throw new AspectValidationException( - "Value type cannot be changed as this is a backwards incompatible change"); - } - if (newDefinition.getCardinality().equals(SINGLE) - && previousDefinition.getCardinality().equals(MULTIPLE)) { - throw new AspectValidationException( - "Property definition cardinality cannot be changed from MULTI to SINGLE"); - } - if (!newDefinition.getQualifiedName().equals(previousDefinition.getQualifiedName())) { - throw new AspectValidationException( - "Cannot change the fully qualified name of a Structured Property"); - } - // Assure new definition has only added allowed values, not removed them - if (newDefinition.getAllowedValues() != null) { - if (!previousDefinition.hasAllowedValues() - || previousDefinition.getAllowedValues() == null) { - throw new AspectValidationException( - "Cannot restrict values that were previously allowed"); + public static Stream validateDefinitionUpserts( + @Nonnull Collection changeMCPs, @Nonnull AspectRetriever aspectRetriever) { + + ValidationExceptionCollection exceptions = ValidationExceptionCollection.newCollection(); + + Set propertyUrns = changeMCPs.stream().map(ChangeMCP::getUrn).collect(Collectors.toSet()); + + // Batch fetch status aspects + Map> structuredPropertyAspects = + fetchPropertyStatusAspects(propertyUrns, aspectRetriever); + + for (ChangeMCP item : changeMCPs) { + // Prevent updates to the definition, if soft deleted property + softDeleteCheck( + item, + structuredPropertyAspects.getOrDefault(item.getUrn(), Collections.emptyMap()), + "Cannot mutate a soft deleted Structured Property Definition") + .ifPresent(exceptions::addException); + + if (item.getPreviousSystemAspect() != null) { + + StructuredPropertyDefinition previousDefinition = + item.getPreviousSystemAspect().getAspect(StructuredPropertyDefinition.class); + StructuredPropertyDefinition newDefinition = + item.getAspect(StructuredPropertyDefinition.class); + + if (!newDefinition.getValueType().equals(previousDefinition.getValueType())) { + exceptions.addException( + item, "Value type cannot be changed as this is a backwards incompatible change"); + } + if (newDefinition.getCardinality().equals(SINGLE) + && previousDefinition.getCardinality().equals(MULTIPLE)) { + exceptions.addException( + item, "Property definition cardinality cannot be changed from MULTI to SINGLE"); + } + if (!newDefinition.getQualifiedName().equals(previousDefinition.getQualifiedName())) { + exceptions.addException( + item, "Cannot change the fully qualified name of a Structured Property"); } - Set newAllowedValues = - newDefinition.getAllowedValues().stream() - .map(PropertyValue::getValue) - .collect(Collectors.toSet()); - for (PropertyValue value : previousDefinition.getAllowedValues()) { - if (!newAllowedValues.contains(value.getValue())) { - throw new AspectValidationException( - "Cannot restrict values that were previously allowed"); + // Assure new definition has only added allowed values, not removed them + if (newDefinition.getAllowedValues() != null) { + if (!previousDefinition.hasAllowedValues() + || previousDefinition.getAllowedValues() == null) { + exceptions.addException(item, "Cannot restrict values that were previously allowed"); + } else { + Set newAllowedValues = + newDefinition.getAllowedValues().stream() + .map(PropertyValue::getValue) + .collect(Collectors.toSet()); + for (PropertyValue value : previousDefinition.getAllowedValues()) { + if (!newAllowedValues.contains(value.getValue())) { + exceptions.addException( + item, "Cannot restrict values that were previously allowed"); + } + } } } } } - return true; + + return exceptions.streamAllExceptions(); + } + + private static Map> fetchPropertyStatusAspects( + Set structuredPropertyUrns, AspectRetriever aspectRetriever) { + try { + return aspectRetriever.getLatestAspectObjects( + structuredPropertyUrns, ImmutableSet.of(Constants.STATUS_ASPECT_NAME)); + } catch (RemoteInvocationException | URISyntaxException e) { + throw new RuntimeException(e); + } + } + + static Optional softDeleteCheck( + T item, @Nonnull Map structuredPropertyAspects, String message) { + Aspect aspect = structuredPropertyAspects.get(STATUS_ASPECT_NAME); + if (aspect != null && new Status(aspect.data()).isRemoved()) { + return Optional.of(AspectValidationException.forItem(item, message)); + } + return Optional.empty(); } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/StructuredPropertiesValidator.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/StructuredPropertiesValidator.java index efd95e0c2e3f12..7650273cefd17d 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/StructuredPropertiesValidator.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/validation/StructuredPropertiesValidator.java @@ -1,18 +1,25 @@ package com.linkedin.metadata.aspect.validation; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME; +import static com.linkedin.metadata.aspect.validation.PropertyDefinitionValidator.softDeleteCheck; + +import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.template.StringArray; import com.linkedin.data.template.StringArrayMap; import com.linkedin.entity.Aspect; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; -import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.aspect.plugins.validation.ValidationExceptionCollection; import com.linkedin.metadata.models.LogicalValueType; import com.linkedin.metadata.models.StructuredPropertyUtils; +import com.linkedin.r2.RemoteInvocationException; import com.linkedin.structured.PrimitivePropertyValue; import com.linkedin.structured.PrimitivePropertyValueArray; import com.linkedin.structured.PropertyCardinality; @@ -23,14 +30,17 @@ import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; -import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; /** A Validator for StructuredProperties Aspect that is attached to entities like Datasets, etc. */ @@ -67,106 +77,155 @@ public static LogicalValueType getLogicalValueType(Urn valueType) { } @Override - protected void validateProposedAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nonnull RecordTemplate aspectPayload, - @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { - validate(aspectPayload, aspectRetriever); + protected Stream validateProposedAspects( + @Nonnull Collection mcpItems, @Nonnull AspectRetriever aspectRetriever) { + return validateProposedUpserts( + mcpItems.stream() + .filter(i -> ChangeType.UPSERT.equals(i.getChangeType())) + .collect(Collectors.toList()), + aspectRetriever); } - public static boolean validate( - @Nonnull RecordTemplate aspectPayload, @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { - StructuredProperties structuredProperties = (StructuredProperties) aspectPayload; - log.warn("Validator called with {}", structuredProperties); - Map> structuredPropertiesMap = - structuredProperties.getProperties().stream() - .collect( - Collectors.groupingBy( - x -> x.getPropertyUrn(), - HashMap::new, - Collectors.toCollection(ArrayList::new))); - for (Map.Entry> entry : - structuredPropertiesMap.entrySet()) { - // There should only be one entry per structured property - List values = entry.getValue(); - if (values.size() > 1) { - throw new AspectValidationException( - "Property: " + entry.getKey() + " has multiple entries: " + values); - } - } + @Override + protected Stream validatePreCommitAspects( + @Nonnull Collection changeMCPs, AspectRetriever aspectRetriever) { + return Stream.empty(); + } - for (StructuredPropertyValueAssignment structuredPropertyValueAssignment : - structuredProperties.getProperties()) { - Urn propertyUrn = structuredPropertyValueAssignment.getPropertyUrn(); - String property = propertyUrn.toString(); - if (!propertyUrn.getEntityType().equals("structuredProperty")) { - throw new IllegalStateException( - "Unexpected entity type. Expected: structuredProperty Found: " - + propertyUrn.getEntityType()); - } - Aspect structuredPropertyDefinitionAspect = null; - try { - structuredPropertyDefinitionAspect = - aspectRetriever.getLatestAspectObject(propertyUrn, "propertyDefinition"); + public static Stream validateProposedUpserts( + @Nonnull Collection mcpItems, @Nonnull AspectRetriever aspectRetriever) { + + ValidationExceptionCollection exceptions = ValidationExceptionCollection.newCollection(); + + // Validate propertyUrns + Set validPropertyUrns = validateStructuredPropertyUrns(mcpItems, exceptions); + + // Fetch property aspects for further validation + Map> allStructuredPropertiesAspects = + fetchPropertyAspects(validPropertyUrns, aspectRetriever); + + // Validate assignments + for (BatchItem i : exceptions.successful(mcpItems)) { + for (StructuredPropertyValueAssignment structuredPropertyValueAssignment : + i.getAspect(StructuredProperties.class).getProperties()) { + + Urn propertyUrn = structuredPropertyValueAssignment.getPropertyUrn(); + Map propertyAspects = + allStructuredPropertiesAspects.getOrDefault(propertyUrn, Collections.emptyMap()); + // check definition soft delete + softDeleteCheck(i, propertyAspects, "Cannot apply a soft deleted Structured Property value") + .ifPresent(exceptions::addException); + + Aspect structuredPropertyDefinitionAspect = + propertyAspects.get(STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME); if (structuredPropertyDefinitionAspect == null) { - throw new AspectValidationException("Unexpected null value found."); + exceptions.addException(i, "Unexpected null value found."); } - } catch (Exception e) { - log.error("Could not fetch latest aspect. PropertyUrn: {}", propertyUrn, e); - throw new AspectValidationException("Could not fetch latest aspect: " + e.getMessage(), e); - } - StructuredPropertyDefinition structuredPropertyDefinition = - new StructuredPropertyDefinition(structuredPropertyDefinitionAspect.data()); - log.warn( - "Retrieved property definition for {}. {}", propertyUrn, structuredPropertyDefinition); - if (structuredPropertyDefinition != null) { - PrimitivePropertyValueArray values = structuredPropertyValueAssignment.getValues(); - // Check cardinality - if (structuredPropertyDefinition.getCardinality() == PropertyCardinality.SINGLE) { - if (values.size() > 1) { - throw new AspectValidationException( - "Property: " - + property - + " has cardinality 1, but multiple values were assigned: " - + values); + StructuredPropertyDefinition structuredPropertyDefinition = + new StructuredPropertyDefinition(structuredPropertyDefinitionAspect.data()); + log.warn( + "Retrieved property definition for {}. {}", propertyUrn, structuredPropertyDefinition); + if (structuredPropertyDefinition != null) { + PrimitivePropertyValueArray values = structuredPropertyValueAssignment.getValues(); + // Check cardinality + if (structuredPropertyDefinition.getCardinality() == PropertyCardinality.SINGLE) { + if (values.size() > 1) { + exceptions.addException( + i, + "Property: " + + propertyUrn + + " has cardinality 1, but multiple values were assigned: " + + values); + } + } + + // Check values + for (PrimitivePropertyValue value : values) { + validateType(i, propertyUrn, structuredPropertyDefinition, value) + .ifPresent(exceptions::addException); + validateAllowedValues(i, propertyUrn, structuredPropertyDefinition, value) + .ifPresent(exceptions::addException); } } - // Check values - for (PrimitivePropertyValue value : values) { - validateType(propertyUrn, structuredPropertyDefinition, value); - validateAllowedValues(propertyUrn, structuredPropertyDefinition, value); + } + } + + return exceptions.streamAllExceptions(); + } + + private static Set validateStructuredPropertyUrns( + Collection mcpItems, ValidationExceptionCollection exceptions) { + Set validPropertyUrns = new HashSet<>(); + + for (BatchItem i : exceptions.successful(mcpItems)) { + StructuredProperties structuredProperties = i.getAspect(StructuredProperties.class); + + log.warn("Validator called with {}", structuredProperties); + Map> structuredPropertiesMap = + structuredProperties.getProperties().stream() + .collect( + Collectors.groupingBy( + x -> x.getPropertyUrn(), + HashMap::new, + Collectors.toCollection(ArrayList::new))); + for (Map.Entry> entry : + structuredPropertiesMap.entrySet()) { + + // There should only be one entry per structured property + List values = entry.getValue(); + if (values.size() > 1) { + exceptions.addException( + i, "Property: " + entry.getKey() + " has multiple entries: " + values); + } else { + for (StructuredPropertyValueAssignment structuredPropertyValueAssignment : + structuredProperties.getProperties()) { + Urn propertyUrn = structuredPropertyValueAssignment.getPropertyUrn(); + + if (!propertyUrn.getEntityType().equals("structuredProperty")) { + exceptions.addException( + i, + "Unexpected entity type. Expected: structuredProperty Found: " + + propertyUrn.getEntityType()); + } else { + validPropertyUrns.add(propertyUrn); + } + } } } } - return true; + return validPropertyUrns; } - private static void validateAllowedValues( - Urn propertyUrn, StructuredPropertyDefinition definition, PrimitivePropertyValue value) - throws AspectValidationException { + private static Optional validateAllowedValues( + BatchItem item, + Urn propertyUrn, + StructuredPropertyDefinition definition, + PrimitivePropertyValue value) { if (definition.getAllowedValues() != null) { Set definedValues = definition.getAllowedValues().stream() .map(PropertyValue::getValue) .collect(Collectors.toSet()); if (definedValues.stream().noneMatch(definedPrimitive -> definedPrimitive.equals(value))) { - throw new AspectValidationException( - String.format( - "Property: %s, value: %s should be one of %s", propertyUrn, value, definedValues)); + return Optional.of( + AspectValidationException.forItem( + item, + String.format( + "Property: %s, value: %s should be one of %s", + propertyUrn, value, definedValues))); } } + return Optional.empty(); } - private static void validateType( - Urn propertyUrn, StructuredPropertyDefinition definition, PrimitivePropertyValue value) - throws AspectValidationException { + private static Optional validateType( + BatchItem item, + Urn propertyUrn, + StructuredPropertyDefinition definition, + PrimitivePropertyValue value) { Urn valueType = definition.getValueType(); LogicalValueType typeDefinition = getLogicalValueType(valueType); @@ -175,16 +234,24 @@ private static void validateType( log.debug( "Property definition demands a string value. {}, {}", value.isString(), value.isDouble()); if (value.getString() == null) { - throw new AspectValidationException( - "Property: " + propertyUrn.toString() + ", value: " + value + " should be a string"); + return Optional.of( + AspectValidationException.forItem( + item, + "Property: " + + propertyUrn.toString() + + ", value: " + + value + + " should be a string")); } else if (typeDefinition.equals(LogicalValueType.DATE)) { if (!StructuredPropertyUtils.isValidDate(value)) { - throw new AspectValidationException( - "Property: " - + propertyUrn.toString() - + ", value: " - + value - + " should be a date with format YYYY-MM-DD"); + return Optional.of( + AspectValidationException.forItem( + item, + "Property: " + + propertyUrn.toString() + + ", value: " + + value + + " should be a date with format YYYY-MM-DD")); } } else if (typeDefinition.equals(LogicalValueType.URN)) { StringArrayMap valueTypeQualifier = definition.getTypeQualifier(); @@ -192,8 +259,11 @@ private static void validateType( try { typeValue = Urn.createFromString(value.getString()); } catch (URISyntaxException e) { - throw new AspectValidationException( - "Property: " + propertyUrn.toString() + ", value: " + value + " should be an urn", e); + return Optional.of( + AspectValidationException.forItem( + item, + "Property: " + propertyUrn.toString() + ", value: " + value + " should be an urn", + e)); } if (valueTypeQualifier != null) { if (valueTypeQualifier.containsKey("allowedTypes")) { @@ -216,13 +286,15 @@ private static void validateType( } } if (!matchedAny) { - throw new AspectValidationException( - "Property: " - + propertyUrn.toString() - + ", value: " - + value - + " is not of any supported urn types:" - + allowedTypes); + return Optional.of( + AspectValidationException.forItem( + item, + "Property: " + + propertyUrn.toString() + + ", value: " + + value + + " is not of any supported urn types:" + + allowedTypes)); } } } @@ -233,13 +305,25 @@ private static void validateType( Double doubleValue = value.getDouble() != null ? value.getDouble() : Double.parseDouble(value.getString()); } catch (NumberFormatException | NullPointerException e) { - throw new AspectValidationException( - "Property: " + propertyUrn.toString() + ", value: " + value + " should be a number"); + return Optional.of( + AspectValidationException.forItem( + item, + "Property: " + + propertyUrn.toString() + + ", value: " + + value + + " should be a number")); } } else { - throw new AspectValidationException( - "Validation support for type " + definition.getValueType() + " is not yet implemented."); + return Optional.of( + AspectValidationException.forItem( + item, + "Validation support for type " + + definition.getValueType() + + " is not yet implemented.")); } + + return Optional.empty(); } private static String getValueTypeId(@Nonnull final Urn valueType) { @@ -250,15 +334,18 @@ private static String getValueTypeId(@Nonnull final Urn valueType) { return valueTypeId; } - @Override - protected void validatePreCommitAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate previousAspect, - @Nonnull RecordTemplate proposedAspect, - AspectRetriever aspectRetriever) - throws AspectValidationException { - // No-op + private static Map> fetchPropertyAspects( + Set structuredPropertyUrns, AspectRetriever aspectRetriever) { + if (structuredPropertyUrns.isEmpty()) { + return Collections.emptyMap(); + } else { + try { + return aspectRetriever.getLatestAspectObjects( + structuredPropertyUrns, + ImmutableSet.of(Constants.STATUS_ASPECT_NAME, STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME)); + } catch (RemoteInvocationException | URISyntaxException e) { + throw new RuntimeException(e); + } + } } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/StructuredPropertyUtils.java b/entity-registry/src/main/java/com/linkedin/metadata/models/StructuredPropertyUtils.java index a8711429421f3b..6c720f6c83ffa5 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/StructuredPropertyUtils.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/StructuredPropertyUtils.java @@ -1,8 +1,34 @@ package com.linkedin.metadata.models; +import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_MAPPING_FIELD; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX; + +import com.google.common.collect.ImmutableSet; +import com.linkedin.common.Status; +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.query.filter.ConjunctiveCriterion; +import com.linkedin.metadata.query.filter.Criterion; +import com.linkedin.metadata.query.filter.Filter; +import com.linkedin.r2.RemoteInvocationException; import com.linkedin.structured.PrimitivePropertyValue; +import com.linkedin.structured.StructuredProperties; +import com.linkedin.structured.StructuredPropertyValueAssignment; +import com.linkedin.structured.StructuredPropertyValueAssignmentArray; +import com.linkedin.util.Pair; +import java.net.URISyntaxException; import java.sql.Date; import java.time.format.DateTimeParseException; +import java.util.Collection; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; public class StructuredPropertyUtils { @@ -18,9 +44,62 @@ private StructuredPropertyUtils() {} * @param fullyQualifiedName The original fully qualified name of the property * @return The sanitized version that can be used as a field name */ - public static String sanitizeStructuredPropertyFQN(String fullyQualifiedName) { - String sanitizedName = fullyQualifiedName.replace('.', '_').replace(' ', '_'); - return sanitizedName; + public static String sanitizeStructuredPropertyFQN(@Nonnull String fullyQualifiedName) { + if (fullyQualifiedName.contains(" ")) { + throw new IllegalArgumentException( + "Fully qualified structured property name cannot contain spaces"); + } + return fullyQualifiedName.replace('.', '_'); + } + + public static void validateStructuredPropertyFQN( + @Nonnull Collection fullyQualifiedNames, @Nonnull AspectRetriever aspectRetriever) { + Set structuredPropertyUrns = + fullyQualifiedNames.stream() + .map(StructuredPropertyUtils::toURNFromFieldName) + .collect(Collectors.toSet()); + Set removedUrns = getRemovedUrns(structuredPropertyUrns, aspectRetriever); + if (!removedUrns.isEmpty()) { + throw new IllegalArgumentException( + String.format("Cannot filter on deleted Structured Property %s", removedUrns)); + } + } + + public static Urn toURNFromFieldName(@Nonnull String fieldName) { + return UrnUtils.getUrn( + String.join(":", "urn:li", STRUCTURED_PROPERTY_ENTITY_NAME, fieldName.replace('_', '.'))); + } + + public static void validateFilter( + @Nullable Filter filter, @Nonnull AspectRetriever aspectRetriever) { + + if (filter == null) { + return; + } + + Set fieldNames = new HashSet<>(); + + if (filter.getCriteria() != null) { + for (Criterion c : filter.getCriteria()) { + if (c.getField().startsWith(STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX)) { + fieldNames.add(c.getField().substring(STRUCTURED_PROPERTY_MAPPING_FIELD.length() + 1)); + } + } + } + + if (filter.getOr() != null) { + for (ConjunctiveCriterion cc : filter.getOr()) { + for (Criterion c : cc.getAnd()) { + if (c.getField().startsWith(STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX)) { + fieldNames.add(c.getField().substring(STRUCTURED_PROPERTY_MAPPING_FIELD.length() + 1)); + } + } + } + } + + if (!fieldNames.isEmpty()) { + validateStructuredPropertyFQN(fieldNames, aspectRetriever); + } } public static Date toDate(PrimitivePropertyValue value) throws DateTimeParseException { @@ -42,4 +121,76 @@ public static boolean isValidDate(PrimitivePropertyValue value) { } return date.compareTo(MIN_DATE) >= 0 && date.compareTo(MAX_DATE) <= 0; } + + private static Set getRemovedUrns(Set urns, AspectRetriever aspectRetriever) { + try { + return aspectRetriever + .getLatestAspectObjects(urns, ImmutableSet.of(STATUS_ASPECT_NAME)) + .entrySet() + .stream() + .filter( + entry -> + entry.getValue().containsKey(STATUS_ASPECT_NAME) + && new Status(entry.getValue().get(STATUS_ASPECT_NAME).data()).isRemoved()) + .map(Map.Entry::getKey) + .collect(Collectors.toSet()); + } catch (RemoteInvocationException | URISyntaxException e) { + throw new RuntimeException(e); + } + } + + /** + * Given a collection of structured properties, return the structured properties with soft deleted + * assignments removed + * + * @param properties collection of structured properties + * @param aspectRetriever typically entity service or entity client + entity registry + * @return structured properties object without value assignments for deleted structured + * properties and whether values were filtered + */ + public static Map filterSoftDelete( + Map properties, AspectRetriever aspectRetriever) { + final Set structuredPropertiesUrns = + properties.values().stream() + .flatMap(structuredProperties -> structuredProperties.getProperties().stream()) + .map(StructuredPropertyValueAssignment::getPropertyUrn) + .collect(Collectors.toSet()); + + final Set removedUrns = getRemovedUrns(structuredPropertiesUrns, aspectRetriever); + + return properties.entrySet().stream() + .map( + entry -> + Pair.of( + entry.getKey(), filterSoftDelete(entry.getValue(), removedUrns).getSecond())) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue)); + } + + private static Pair filterSoftDelete( + StructuredProperties structuredProperties, Set softDeletedPropertyUrns) { + + Pair filtered = + filterValueAssignment(structuredProperties.getProperties(), softDeletedPropertyUrns); + + if (filtered.getSecond()) { + return Pair.of(structuredProperties.setProperties(filtered.getFirst()), true); + } else { + return Pair.of(structuredProperties, false); + } + } + + private static Pair filterValueAssignment( + StructuredPropertyValueAssignmentArray in, Set softDeletedPropertyUrns) { + if (in.stream().noneMatch(p -> softDeletedPropertyUrns.contains(p.getPropertyUrn()))) { + return Pair.of(in, false); + } else { + return Pair.of( + new StructuredPropertyValueAssignmentArray( + in.stream() + .filter( + assignment -> !softDeletedPropertyUrns.contains(assignment.getPropertyUrn())) + .collect(Collectors.toSet())), + true); + } + } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/EntityRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/EntityRegistry.java index c2aa1fab6c2c0f..f4d9926f13ae66 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/EntityRegistry.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/EntityRegistry.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.models.registry; -import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.aspect.patch.template.AspectTemplateEngine; import com.linkedin.metadata.aspect.plugins.PluginFactory; import com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffect; @@ -13,7 +12,6 @@ import com.linkedin.metadata.models.EventSpec; import java.util.List; import java.util.Map; -import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -77,45 +75,24 @@ default String getIdentifier() { AspectTemplateEngine getAspectTemplateEngine(); /** - * Returns applicable {@link AspectPayloadValidator} implementations given the change type and - * entity/aspect information. + * Prefer {@link com.linkedin.metadata.aspect.batch.AspectsBatch} instead of using this method + * directly. * - * @param changeType The type of change to be validated - * @param entityName The entity name - * @param aspectName The aspect name * @return List of validator implementations */ - @Nonnull - default List getAspectPayloadValidators( - @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { - return getAllAspectPayloadValidators().stream() - .filter( - aspectPayloadValidator -> - aspectPayloadValidator.shouldApply(changeType, entityName, aspectName)) - .collect(Collectors.toList()); - } - @Nonnull default List getAllAspectPayloadValidators() { return getPluginFactory().getAspectPayloadValidators(); } /** - * Return mutation hooks for {@link com.linkedin.data.template.RecordTemplate} + * Returns mutation hooks. + * + *

Prefer {@link com.linkedin.metadata.aspect.batch.AspectsBatch} instead of using this method + * directly. * - * @param changeType The type of change - * @param entityName The entity name - * @param aspectName The aspect name - * @return Mutation hooks + * @return list of mutation hooks. */ - @Nonnull - default List getMutationHooks( - @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { - return getAllMutationHooks().stream() - .filter(mutationHook -> mutationHook.shouldApply(changeType, entityName, aspectName)) - .collect(Collectors.toList()); - } - @Nonnull default List getAllMutationHooks() { return getPluginFactory().getMutationHooks(); @@ -125,19 +102,11 @@ default List getAllMutationHooks() { * Returns the side effects to apply to {@link com.linkedin.mxe.MetadataChangeProposal}. Side * effects can generate one or more additional MCPs during write operations. * - * @param changeType The type of change - * @param entityName The entity name - * @param aspectName The aspect name + *

Prefer {@link com.linkedin.metadata.aspect.batch.AspectsBatch} instead of using this method + * directly. + * * @return MCP side effects */ - @Nonnull - default List getMCPSideEffects( - @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { - return getAllMCPSideEffects().stream() - .filter(mcpSideEffect -> mcpSideEffect.shouldApply(changeType, entityName, aspectName)) - .collect(Collectors.toList()); - } - @Nonnull default List getAllMCPSideEffects() { return getPluginFactory().getMcpSideEffects(); @@ -147,19 +116,11 @@ default List getAllMCPSideEffects() { * Returns the side effects to apply to {@link com.linkedin.mxe.MetadataChangeLog}. Side effects * can generate one or more additional MCLs during write operations. * - * @param changeType The type of change - * @param entityName The entity name - * @param aspectName The aspect name + *

Prefer {@link com.linkedin.metadata.aspect.batch.AspectsBatch} instead of using this method + * directly. + * * @return MCL side effects */ - @Nonnull - default List getMCLSideEffects( - @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { - return getAllMCLSideEffects().stream() - .filter(mclSideEffect -> mclSideEffect.shouldApply(changeType, entityName, aspectName)) - .collect(Collectors.toList()); - } - @Nonnull default List getAllMCLSideEffects() { return getPluginFactory().getMclSideEffects(); diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDeleteTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDeleteTest.java new file mode 100644 index 00000000000000..6f4149b1031256 --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/StructuredPropertiesSoftDeleteTest.java @@ -0,0 +1,96 @@ +package com.linkedin.metadata.aspect.hooks; + +import static com.linkedin.metadata.Constants.DATASET_ENTITY_NAME; +import static org.testng.Assert.assertEquals; + +import com.linkedin.common.Status; +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.structured.PrimitivePropertyValue; +import com.linkedin.structured.PrimitivePropertyValueArray; +import com.linkedin.structured.StructuredProperties; +import com.linkedin.structured.StructuredPropertyDefinition; +import com.linkedin.structured.StructuredPropertyValueAssignment; +import com.linkedin.structured.StructuredPropertyValueAssignmentArray; +import com.linkedin.test.metadata.aspect.MockAspectRetriever; +import com.linkedin.test.metadata.aspect.TestEntityRegistry; +import com.linkedin.test.metadata.aspect.batch.TestMCP; +import java.net.URISyntaxException; +import java.util.List; +import java.util.Map; +import org.testng.annotations.Test; + +public class StructuredPropertiesSoftDeleteTest { + + private static final EntityRegistry TEST_REGISTRY = new TestEntityRegistry(); + + @Test + public void testSoftDeleteFilter() throws URISyntaxException, CloneNotSupportedException { + Urn propertyUrnA = + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime"); + StructuredPropertyDefinition stringPropertyDefA = + new StructuredPropertyDefinition() + .setValueType(Urn.createFromString("urn:li:type:datahub.string")); + StructuredPropertyValueAssignment assignmentA = + new StructuredPropertyValueAssignment() + .setPropertyUrn(propertyUrnA) + .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(0.0))); + + Urn propertyUrnB = + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTimeDeleted"); + StructuredPropertyDefinition stringPropertyDefB = + new StructuredPropertyDefinition() + .setValueType(Urn.createFromString("urn:li:type:datahub.string")); + StructuredPropertyValueAssignment assignmentB = + new StructuredPropertyValueAssignment() + .setPropertyUrn(propertyUrnB) + .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(0.0))); + + StructuredPropertiesSoftDelete testHook = + new StructuredPropertiesSoftDelete( + AspectPluginConfig.builder() + .enabled(true) + .className(StructuredPropertiesSoftDelete.class.getName()) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName(DATASET_ENTITY_NAME) + .aspectName(Constants.STRUCTURED_PROPERTIES_ASPECT_NAME) + .build())) + .build()); + + StructuredProperties expectedAllValues = new StructuredProperties(); + expectedAllValues.setProperties( + new StructuredPropertyValueAssignmentArray(assignmentA, assignmentB)); + + StructuredProperties test = expectedAllValues.copy(); + testHook.readMutation( + TestMCP.ofOneBatchItemDatasetUrn(test, TEST_REGISTRY), + new MockAspectRetriever( + Map.of( + propertyUrnA, + List.of(stringPropertyDefA), + propertyUrnB, + List.of(stringPropertyDefB)))); + assertEquals( + test.getProperties().size(), + 2, + "Expected all values because all definitions are NOT soft deleted"); + + StructuredProperties expectedOneValue = new StructuredProperties(); + expectedOneValue.setProperties(new StructuredPropertyValueAssignmentArray(assignmentA)); + test = expectedAllValues.copy(); + testHook.readMutation( + TestMCP.ofOneBatchItemDatasetUrn(test, TEST_REGISTRY), + new MockAspectRetriever( + Map.of( + propertyUrnA, + List.of(stringPropertyDefA), + propertyUrnB, + List.of(stringPropertyDefB, new Status().setRemoved(true))))); + assertEquals( + test.getProperties().size(), 1, "Expected 1 value because 1 definition is soft deleted"); + } +} diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java index f801ce7bf1ffe6..a9f903f4b7017d 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java @@ -61,27 +61,37 @@ public void testConfigEntityRegistry() throws FileNotFoundException { assertNotNull(eventSpec.getPegasusSchema()); assertEquals( - configEntityRegistry - .getAspectPayloadValidators(ChangeType.UPSERT, "chart", "status") - .size(), + configEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "chart", "status")) + .count(), 2); assertEquals( - configEntityRegistry - .getAspectPayloadValidators(ChangeType.DELETE, "chart", "status") - .size(), + configEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "chart", "status")) + .count(), 0); assertEquals( - configEntityRegistry.getMCPSideEffects(ChangeType.UPSERT, "dataset", "datasetKey").size(), + configEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "dataset", "datasetKey")) + .count(), 1); assertEquals( - configEntityRegistry.getMCPSideEffects(ChangeType.DELETE, "dataset", "datasetKey").size(), + configEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "dataset", "datasetKey")) + .count(), 0); assertEquals( - configEntityRegistry.getMutationHooks(ChangeType.UPSERT, "*", "schemaMetadata").size(), 1); + configEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "*", "schemaMetadata")) + .count(), + 1); assertEquals( - configEntityRegistry.getMutationHooks(ChangeType.DELETE, "*", "schemaMetadata").size(), 0); + configEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "*", "schemaMetadata")) + .count(), + 0); } @Test @@ -123,27 +133,37 @@ public void testMergedEntityRegistry() throws EntityRegistryException { assertNotNull(eventSpec.getPegasusSchema()); assertEquals( - mergedEntityRegistry - .getAspectPayloadValidators(ChangeType.UPSERT, "chart", "status") - .size(), + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "chart", "status")) + .count(), 2); assertEquals( - mergedEntityRegistry - .getAspectPayloadValidators(ChangeType.DELETE, "chart", "status") - .size(), + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "chart", "status")) + .count(), 1); assertEquals( - mergedEntityRegistry.getMCPSideEffects(ChangeType.UPSERT, "dataset", "datasetKey").size(), + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "dataset", "datasetKey")) + .count(), 2); assertEquals( - mergedEntityRegistry.getMCPSideEffects(ChangeType.DELETE, "dataset", "datasetKey").size(), + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "dataset", "datasetKey")) + .count(), 1); assertEquals( - mergedEntityRegistry.getMutationHooks(ChangeType.UPSERT, "*", "schemaMetadata").size(), 2); + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "*", "schemaMetadata")) + .count(), + 2); assertEquals( - mergedEntityRegistry.getMutationHooks(ChangeType.DELETE, "*", "schemaMetadata").size(), 1); + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "*", "schemaMetadata")) + .count(), + 1); } @Test diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffectTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffectTest.java index 8ee5ff4f998206..ac2397a6aaa335 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffectTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffectTest.java @@ -5,11 +5,13 @@ import com.datahub.test.TestEntityProfile; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.batch.MCLBatchItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.MCLItem; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import java.util.Collection; import java.util.List; +import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nonnull; import org.testng.annotations.BeforeTest; @@ -32,7 +34,10 @@ public void testCustomMCLSideEffect() { TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); List mclSideEffects = - configEntityRegistry.getMCLSideEffects(ChangeType.UPSERT, "chart", "chartInfo"); + configEntityRegistry.getAllMCLSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "chart", "chartInfo")) + .collect(Collectors.toList()); + assertEquals( mclSideEffects, List.of( @@ -52,15 +57,14 @@ public void testCustomMCLSideEffect() { } public static class TestMCLSideEffect extends MCLSideEffect { - public TestMCLSideEffect(AspectPluginConfig aspectPluginConfig) { super(aspectPluginConfig); } @Override - protected Stream applyMCLSideEffect( - @Nonnull MCLBatchItem input, @Nonnull AspectRetriever aspectRetriever) { - return Stream.of(input); + protected Stream applyMCLSideEffect( + @Nonnull Collection batchItems, @Nonnull AspectRetriever aspectRetriever) { + return null; } } } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffectTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffectTest.java index 8522e8facf3e08..e3499861d61986 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffectTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffectTest.java @@ -5,11 +5,13 @@ import com.datahub.test.TestEntityProfile; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import java.util.Collection; import java.util.List; +import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nonnull; import org.testng.annotations.BeforeTest; @@ -32,7 +34,10 @@ public void testCustomMCPSideEffect() { TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); List mcpSideEffects = - configEntityRegistry.getMCPSideEffects(ChangeType.UPSERT, "dataset", "datasetKey"); + configEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "dataset", "datasetKey")) + .collect(Collectors.toList()); + assertEquals( mcpSideEffects, List.of( @@ -58,9 +63,9 @@ public TestMCPSideEffect(AspectPluginConfig aspectPluginConfig) { } @Override - protected Stream applyMCPSideEffect( - UpsertItem input, @Nonnull AspectRetriever aspectRetriever) { - return Stream.of(input); + protected Stream applyMCPSideEffect( + Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { + return changeMCPS.stream(); } } } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MutationPluginTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MutationPluginTest.java index 5094fd7fdd443d..16ea003582b180 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MutationPluginTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MutationPluginTest.java @@ -3,19 +3,12 @@ import static org.testng.Assert.assertEquals; import com.datahub.test.TestEntityProfile; -import com.linkedin.common.AuditStamp; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; -import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; -import com.linkedin.mxe.SystemMetadata; import java.util.List; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; +import java.util.stream.Collectors; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -36,7 +29,10 @@ public void testCustomMutator() { TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); List mutators = - configEntityRegistry.getMutationHooks(ChangeType.UPSERT, "*", "schemaMetadata"); + configEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "*", "schemaMetadata")) + .collect(Collectors.toList()); + assertEquals( mutators, List.of( @@ -56,21 +52,8 @@ public void testCustomMutator() { } public static class TestMutator extends MutationHook { - public TestMutator(AspectPluginConfig aspectPluginConfig) { super(aspectPluginConfig); } - - @Override - protected void mutate( - @Nonnull ChangeType changeType, - @Nonnull EntitySpec entitySpec, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate oldAspectValue, - @Nullable RecordTemplate newAspectValue, - @Nullable SystemMetadata oldSystemMetadata, - @Nullable SystemMetadata newSystemMetadata, - @Nonnull AuditStamp auditStamp, - @Nonnull AspectRetriever aspectRetriever) {} } } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/validation/ValidatorPluginTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/validation/ValidatorPluginTest.java index eb132836be4656..10dbeafa822945 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/validation/ValidatorPluginTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/validation/ValidatorPluginTest.java @@ -3,16 +3,18 @@ import static org.testng.Assert.assertEquals; import com.datahub.test.TestEntityProfile; -import com.linkedin.common.urn.Urn; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; -import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import java.util.Collection; import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; -import javax.annotation.Nullable; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -33,7 +35,10 @@ public void testCustomValidator() { TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); List validators = - configEntityRegistry.getAspectPayloadValidators(ChangeType.UPSERT, "chart", "status"); + configEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "chart", "status")) + .collect(Collectors.toList()); + assertEquals( validators, List.of( @@ -72,26 +77,16 @@ public TestValidator(AspectPluginConfig config) { } @Override - protected void validateProposedAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nonnull RecordTemplate aspectPayload, - AspectRetriever aspectRetriever) - throws AspectValidationException { - if (entityUrn.toString().contains("dataset")) { - throw new AspectValidationException("test error"); - } + protected Stream validateProposedAspects( + @Nonnull Collection mcpItems, + @Nonnull AspectRetriever aspectRetriever) { + return mcpItems.stream().map(i -> AspectValidationException.forItem(i, "test error")); } @Override - protected void validatePreCommitAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate previousAspect, - @Nonnull RecordTemplate proposedAspect, - AspectRetriever aspectRetriever) - throws AspectValidationException {} + protected Stream validatePreCommitAspects( + @Nonnull Collection changeMCPs, AspectRetriever aspectRetriever) { + return Stream.empty(); + } } } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/PropertyDefinitionValidatorTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/PropertyDefinitionValidatorTest.java index 96e9fceb4a05d8..38ba87cfaae80c 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/PropertyDefinitionValidatorTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/PropertyDefinitionValidatorTest.java @@ -1,20 +1,49 @@ package com.linkedin.metadata.aspect.validators; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME; +import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.*; import com.linkedin.common.UrnArray; import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; import com.linkedin.metadata.aspect.validation.PropertyDefinitionValidator; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.structured.PrimitivePropertyValue; import com.linkedin.structured.PropertyCardinality; import com.linkedin.structured.PropertyValue; import com.linkedin.structured.PropertyValueArray; import com.linkedin.structured.StructuredPropertyDefinition; +import com.linkedin.structured.StructuredPropertyKey; +import com.linkedin.test.metadata.aspect.TestEntityRegistry; +import com.linkedin.test.metadata.aspect.batch.TestMCP; import java.net.URISyntaxException; +import java.util.List; +import java.util.Set; +import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; public class PropertyDefinitionValidatorTest { + + private EntityRegistry entityRegistry; + private Urn testPropertyUrn; + private AspectRetriever mockAspectRetriever; + + @BeforeTest + public void init() { + entityRegistry = new TestEntityRegistry(); + testPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); + mockAspectRetriever = mock(AspectRetriever.class); + when(mockAspectRetriever.getEntityRegistry()).thenReturn(entityRegistry); + } + @Test public void testValidatePreCommitNoPrevious() throws URISyntaxException, AspectValidationException { @@ -28,7 +57,11 @@ public void testValidatePreCommitNoPrevious() newProperty.setQualifiedName("prop3"); newProperty.setCardinality(PropertyCardinality.MULTIPLE); newProperty.setValueType(Urn.createFromString("urn:li:logicalType:STRING")); - assertTrue(PropertyDefinitionValidator.validate(null, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, newProperty, entityRegistry), mockAspectRetriever) + .count(), + 0); } @Test @@ -46,7 +79,12 @@ public void testCanChangeSingleToMultiple() oldProperty.setValueType(Urn.createFromString("urn:li:logicalType:STRING")); StructuredPropertyDefinition newProperty = oldProperty.copy(); newProperty.setCardinality(PropertyCardinality.MULTIPLE); - assertTrue(PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 0); } @Test @@ -64,9 +102,12 @@ public void testCannotChangeMultipleToSingle() oldProperty.setValueType(Urn.createFromString("urn:li:logicalType:STRING")); StructuredPropertyDefinition newProperty = oldProperty.copy(); newProperty.setCardinality(PropertyCardinality.SINGLE); - assertThrows( - AspectValidationException.class, - () -> PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 1); } @Test @@ -83,9 +124,12 @@ public void testCannotChangeValueType() throws URISyntaxException, CloneNotSuppo oldProperty.setValueType(Urn.createFromString("urn:li:logicalType:STRING")); StructuredPropertyDefinition newProperty = oldProperty.copy(); newProperty.setValueType(Urn.createFromString("urn:li:logicalType:NUMBER")); - assertThrows( - AspectValidationException.class, - () -> PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 1); } @Test @@ -103,7 +147,12 @@ public void testCanChangeDisplayName() oldProperty.setValueType(Urn.createFromString("urn:li:logicalType:STRING")); StructuredPropertyDefinition newProperty = oldProperty.copy(); newProperty.setDisplayName("newProp"); - assertTrue(PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 0); } @Test @@ -121,9 +170,12 @@ public void testCannotChangeFullyQualifiedName() oldProperty.setValueType(Urn.createFromString("urn:li:logicalType:STRING")); StructuredPropertyDefinition newProperty = oldProperty.copy(); newProperty.setQualifiedName("newProp"); - assertThrows( - AspectValidationException.class, - () -> PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 1); } @Test @@ -144,17 +196,23 @@ public void testCannotChangeRestrictAllowedValues() PropertyValue allowedValue = new PropertyValue().setValue(PrimitivePropertyValue.create(1.0)).setDescription("hello"); newProperty.setAllowedValues(new PropertyValueArray(allowedValue)); - assertThrows( - AspectValidationException.class, - () -> PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 1); // Remove allowed values from constraint case PropertyValue oldAllowedValue = new PropertyValue().setValue(PrimitivePropertyValue.create(3.0)).setDescription("hello"); oldProperty.setAllowedValues((new PropertyValueArray(allowedValue, oldAllowedValue))); - assertThrows( - AspectValidationException.class, - () -> PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 1); } @Test @@ -175,13 +233,23 @@ public void testCanExpandAllowedValues() PropertyValue allowedValue = new PropertyValue().setValue(PrimitivePropertyValue.create(1.0)).setDescription("hello"); oldProperty.setAllowedValues(new PropertyValueArray(allowedValue)); - assertTrue(PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 0); // Add allowed values to constraint case PropertyValue newAllowedValue = new PropertyValue().setValue(PrimitivePropertyValue.create(3.0)).setDescription("hello"); newProperty.setAllowedValues((new PropertyValueArray(allowedValue, newAllowedValue))); - assertTrue(PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 0); } @Test @@ -207,6 +275,67 @@ public void testCanChangeAllowedValueDescriptions() .setValue(PrimitivePropertyValue.create(1.0)) .setDescription("hello there"); newProperty.setAllowedValues(new PropertyValueArray(newAllowedValue)); - assertTrue(PropertyDefinitionValidator.validate(oldProperty, newProperty)); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(testPropertyUrn, oldProperty, newProperty, entityRegistry), + mockAspectRetriever) + .count(), + 0); + } + + @Test + public void testHardDeleteBlock() { + PropertyDefinitionValidator test = + new PropertyDefinitionValidator( + AspectPluginConfig.builder() + .enabled(true) + .className(PropertyDefinitionValidator.class.getName()) + .supportedOperations(List.of("DELETE")) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName(STRUCTURED_PROPERTY_ENTITY_NAME) + .aspectName(Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME) + .build(), + AspectPluginConfig.EntityAspectName.builder() + .entityName(STRUCTURED_PROPERTY_ENTITY_NAME) + .aspectName("structuredPropertyKey") + .build())) + .build()); + + assertEquals( + test.validateProposed( + Set.of( + TestMCP.builder() + .changeType(ChangeType.DELETE) + .urn(UrnUtils.getUrn("urn:li:structuredProperty:foo.bar")) + .entitySpec(entityRegistry.getEntitySpec("structuredProperty")) + .aspectSpec( + entityRegistry + .getEntitySpec(STRUCTURED_PROPERTY_ENTITY_NAME) + .getKeyAspectSpec()) + .recordTemplate(new StructuredPropertyKey()) + .build()), + mockAspectRetriever) + .count(), + 1); + + assertEquals( + test.validateProposed( + Set.of( + TestMCP.builder() + .changeType(ChangeType.DELETE) + .urn(UrnUtils.getUrn("urn:li:structuredProperty:foo.bar")) + .entitySpec(entityRegistry.getEntitySpec("structuredProperty")) + .aspectSpec( + entityRegistry + .getEntitySpec(STRUCTURED_PROPERTY_ENTITY_NAME) + .getAspectSpecMap() + .get(STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME)) + .recordTemplate(new StructuredPropertyDefinition()) + .build()), + mockAspectRetriever) + .count(), + 1); } } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/StructuredPropertiesValidatorTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/StructuredPropertiesValidatorTest.java index 450b299b48b34f..5d63d8c8ba5e7b 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/StructuredPropertiesValidatorTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/validators/StructuredPropertiesValidatorTest.java @@ -1,12 +1,11 @@ package com.linkedin.metadata.aspect.validators; +import static org.testng.Assert.assertEquals; + +import com.linkedin.common.Status; import com.linkedin.common.urn.Urn; -import com.linkedin.entity.Aspect; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; import com.linkedin.metadata.aspect.validation.StructuredPropertiesValidator; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.r2.RemoteInvocationException; import com.linkedin.structured.PrimitivePropertyValue; import com.linkedin.structured.PrimitivePropertyValueArray; import com.linkedin.structured.PropertyValue; @@ -15,42 +14,23 @@ import com.linkedin.structured.StructuredPropertyDefinition; import com.linkedin.structured.StructuredPropertyValueAssignment; import com.linkedin.structured.StructuredPropertyValueAssignmentArray; +import com.linkedin.test.metadata.aspect.MockAspectRetriever; +import com.linkedin.test.metadata.aspect.TestEntityRegistry; +import com.linkedin.test.metadata.aspect.batch.TestMCP; import java.net.URISyntaxException; import java.util.List; -import java.util.Map; -import java.util.Set; -import javax.annotation.Nonnull; import org.testng.Assert; import org.testng.annotations.Test; public class StructuredPropertiesValidatorTest { - static class MockAspectRetriever implements AspectRetriever { - StructuredPropertyDefinition _propertyDefinition; - - MockAspectRetriever(StructuredPropertyDefinition defToReturn) { - this._propertyDefinition = defToReturn; - } - - @Nonnull - @Override - public Map> getLatestAspectObjects( - Set urns, Set aspectNames) - throws RemoteInvocationException, URISyntaxException { - return Map.of( - urns.stream().findFirst().get(), - Map.of(aspectNames.stream().findFirst().get(), new Aspect(_propertyDefinition.data()))); - } - - @Nonnull - @Override - public EntityRegistry getEntityRegistry() { - return null; - } - } + private static final EntityRegistry TEST_REGISTRY = new TestEntityRegistry(); @Test public void testValidateAspectNumberUpsert() throws URISyntaxException { + Urn propertyUrn = + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime"); + StructuredPropertyDefinition numberPropertyDef = new StructuredPropertyDefinition() .setValueType(Urn.createFromString("urn:li:type:datahub.number")) @@ -61,40 +41,38 @@ public void testValidateAspectNumberUpsert() throws URISyntaxException { new PropertyValue().setValue(PrimitivePropertyValue.create(60.0)), new PropertyValue().setValue(PrimitivePropertyValue.create(90.0))))); - try { - StructuredPropertyValueAssignment assignment = - new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) - .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(30.0))); - StructuredProperties numberPayload = - new StructuredProperties() - .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); - - boolean isValid = - StructuredPropertiesValidator.validate( - numberPayload, new MockAspectRetriever(numberPropertyDef)); - Assert.assertTrue(isValid); - } catch (AspectValidationException e) { - throw new RuntimeException(e); - } - - try { - StructuredPropertyValueAssignment assignment = - new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) - .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(0.0))); - StructuredProperties numberPayload = - new StructuredProperties() - .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); - - StructuredPropertiesValidator.validate( - numberPayload, new MockAspectRetriever(numberPropertyDef)); - Assert.fail("Should have raised exception for disallowed value 0.0"); - } catch (AspectValidationException e) { - Assert.assertTrue(e.getMessage().contains("{double=0.0} should be one of [{")); - } + StructuredPropertyValueAssignment assignment = + new StructuredPropertyValueAssignment() + .setPropertyUrn(propertyUrn) + .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(30.0))); + StructuredProperties numberPayload = + new StructuredProperties() + .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); + + boolean isValid = + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(numberPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, numberPropertyDef)) + .count() + == 0; + Assert.assertTrue(isValid); + + assignment = + new StructuredPropertyValueAssignment() + .setPropertyUrn( + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) + .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(0.0))); + numberPayload = + new StructuredProperties() + .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); + + assertEquals( + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(numberPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, numberPropertyDef)) + .count(), + 1, + "Should have raised exception for disallowed value 0.0"); // Assign string value to number property StructuredPropertyValueAssignment stringAssignment = @@ -105,17 +83,21 @@ public void testValidateAspectNumberUpsert() throws URISyntaxException { StructuredProperties stringPayload = new StructuredProperties() .setProperties(new StructuredPropertyValueAssignmentArray(stringAssignment)); - try { - StructuredPropertiesValidator.validate( - stringPayload, new MockAspectRetriever(numberPropertyDef)); - Assert.fail("Should have raised exception for mis-matched types"); - } catch (AspectValidationException e) { - Assert.assertTrue(e.getMessage().contains("should be a number")); - } + + assertEquals( + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(stringPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, numberPropertyDef)) + .count(), + 2, + "Should have raised exception for mis-matched types `string` vs `number` && `hello` is not a valid value of [90.0, 30.0, 60.0]"); } @Test public void testValidateAspectDateUpsert() throws URISyntaxException { + Urn propertyUrn = + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime"); + // Assign string value StructuredPropertyValueAssignment stringAssignment = new StructuredPropertyValueAssignment() @@ -130,41 +112,43 @@ public void testValidateAspectDateUpsert() throws URISyntaxException { StructuredPropertyDefinition datePropertyDef = new StructuredPropertyDefinition() .setValueType(Urn.createFromString("urn:li:type:datahub.date")); - try { - StructuredPropertiesValidator.validate( - stringPayload, new MockAspectRetriever(datePropertyDef)); - Assert.fail("Should have raised exception for mis-matched types"); - } catch (AspectValidationException e) { - Assert.assertTrue(e.getMessage().contains("should be a date with format")); - } + + assertEquals( + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(stringPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, datePropertyDef)) + .count(), + 1, + "Should have raised exception for mis-matched types"); // Assign valid date StructuredPropertyValueAssignment dateAssignment = new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) + .setPropertyUrn(propertyUrn) .setValues( new PrimitivePropertyValueArray(PrimitivePropertyValue.create("2023-10-24"))); StructuredProperties datePayload = new StructuredProperties() .setProperties(new StructuredPropertyValueAssignmentArray(dateAssignment)); - try { - boolean isValid = - StructuredPropertiesValidator.validate( - datePayload, new MockAspectRetriever(datePropertyDef)); - Assert.assertTrue(isValid); - } catch (AspectValidationException e) { - throw new RuntimeException(e); - } + + boolean isValid = + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(datePayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, datePropertyDef)) + .count() + == 0; + Assert.assertTrue(isValid); } @Test public void testValidateAspectStringUpsert() throws URISyntaxException { + Urn propertyUrn = + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime"); + // Assign string value StructuredPropertyValueAssignment stringAssignment = new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) + .setPropertyUrn(propertyUrn) .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create("hello"))); StructuredProperties stringPayload = new StructuredProperties() @@ -173,8 +157,7 @@ public void testValidateAspectStringUpsert() throws URISyntaxException { // Assign date StructuredPropertyValueAssignment dateAssignment = new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) + .setPropertyUrn(propertyUrn) .setValues( new PrimitivePropertyValueArray(PrimitivePropertyValue.create("2023-10-24"))); StructuredProperties datePayload = @@ -184,8 +167,7 @@ public void testValidateAspectStringUpsert() throws URISyntaxException { // Assign number StructuredPropertyValueAssignment assignment = new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) + .setPropertyUrn(propertyUrn) .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(30.0))); StructuredProperties numberPayload = new StructuredProperties() @@ -202,45 +184,88 @@ public void testValidateAspectStringUpsert() throws URISyntaxException { .setValue(PrimitivePropertyValue.create("2023-10-24"))))); // Valid strings (both the date value and "hello" are valid) - try { - boolean isValid = - StructuredPropertiesValidator.validate( - stringPayload, new MockAspectRetriever(stringPropertyDef)); - Assert.assertTrue(isValid); - isValid = - StructuredPropertiesValidator.validate( - datePayload, new MockAspectRetriever(stringPropertyDef)); - Assert.assertTrue(isValid); - } catch (AspectValidationException e) { - throw new RuntimeException(e); - } + + boolean isValid = + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(stringPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, stringPropertyDef)) + .count() + == 0; + Assert.assertTrue(isValid); + isValid = + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(datePayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, stringPropertyDef)) + .count() + == 0; + Assert.assertTrue(isValid); // Invalid: assign a number to the string property - try { - StructuredPropertiesValidator.validate( - numberPayload, new MockAspectRetriever(stringPropertyDef)); - Assert.fail("Should have raised exception for mis-matched types"); - } catch (AspectValidationException e) { - Assert.assertTrue(e.getMessage().contains("should be a string")); - } + assertEquals( + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(numberPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, stringPropertyDef)) + .count(), + 2, + "Should have raised exception for mis-matched types. The double 30.0 is not a `string` && not one of the allowed types `2023-10-24` or `hello`"); // Invalid allowedValue - try { - assignment = - new StructuredPropertyValueAssignment() - .setPropertyUrn( - Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime")) - .setValues( - new PrimitivePropertyValueArray(PrimitivePropertyValue.create("not hello"))); - stringPayload = - new StructuredProperties() - .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); - - StructuredPropertiesValidator.validate( - stringPayload, new MockAspectRetriever(stringPropertyDef)); - Assert.fail("Should have raised exception for disallowed value `not hello`"); - } catch (AspectValidationException e) { - Assert.assertTrue(e.getMessage().contains("{string=not hello} should be one of [{")); - } + + assignment = + new StructuredPropertyValueAssignment() + .setPropertyUrn(propertyUrn) + .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create("not hello"))); + stringPayload = + new StructuredProperties() + .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); + + assertEquals( + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(stringPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, stringPropertyDef)) + .count(), + 1, + "Should have raised exception for disallowed value `not hello`"); + } + + @Test + public void testValidateSoftDeletedUpsert() throws URISyntaxException { + Urn propertyUrn = + Urn.createFromString("urn:li:structuredProperty:io.acryl.privacy.retentionTime"); + + StructuredPropertyDefinition numberPropertyDef = + new StructuredPropertyDefinition() + .setValueType(Urn.createFromString("urn:li:type:datahub.number")) + .setAllowedValues( + new PropertyValueArray( + List.of( + new PropertyValue().setValue(PrimitivePropertyValue.create(30.0)), + new PropertyValue().setValue(PrimitivePropertyValue.create(60.0)), + new PropertyValue().setValue(PrimitivePropertyValue.create(90.0))))); + + StructuredPropertyValueAssignment assignment = + new StructuredPropertyValueAssignment() + .setPropertyUrn(propertyUrn) + .setValues(new PrimitivePropertyValueArray(PrimitivePropertyValue.create(30.0))); + StructuredProperties numberPayload = + new StructuredProperties() + .setProperties(new StructuredPropertyValueAssignmentArray(assignment)); + + boolean isValid = + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(numberPayload, TEST_REGISTRY), + new MockAspectRetriever(propertyUrn, numberPropertyDef)) + .count() + == 0; + Assert.assertTrue(isValid); + + assertEquals( + StructuredPropertiesValidator.validateProposedUpserts( + TestMCP.ofOneUpsertItemDatasetUrn(numberPayload, TEST_REGISTRY), + new MockAspectRetriever( + propertyUrn, numberPropertyDef, new Status().setRemoved(true))) + .count(), + 1, + "Should have raised exception for soft deleted definition"); } } diff --git a/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/MockAspectRetriever.java b/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/MockAspectRetriever.java new file mode 100644 index 00000000000000..b98d78bf6ff933 --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/MockAspectRetriever.java @@ -0,0 +1,71 @@ +package com.linkedin.test.metadata.aspect; + +import static org.mockito.Mockito.mock; + +import com.linkedin.common.Status; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.DataMap; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.entity.Aspect; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.r2.RemoteInvocationException; +import com.linkedin.structured.StructuredPropertyDefinition; +import com.linkedin.util.Pair; +import java.net.URISyntaxException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; + +public class MockAspectRetriever implements AspectRetriever { + private final Map> data; + + public MockAspectRetriever(@Nonnull Map> data) { + this.data = + new HashMap<>( + data.entrySet().stream() + .map( + entry -> + Pair.of( + entry.getKey(), + entry.getValue().stream() + .map( + rt -> { + String aspectName = + ((DataMap) rt.schema().getProperties().get("Aspect")) + .get("name") + .toString(); + return Pair.of(aspectName, new Aspect(rt.data())); + }) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue)))) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue))); + } + + public MockAspectRetriever( + Urn propertyUrn, StructuredPropertyDefinition definition, Status status) { + this(Map.of(propertyUrn, List.of(definition, status))); + } + + public MockAspectRetriever(Urn propertyUrn, StructuredPropertyDefinition definition) { + this(Map.of(propertyUrn, List.of(definition))); + } + + @Nonnull + @Override + public Map> getLatestAspectObjects( + Set urns, Set aspectNames) throws RemoteInvocationException, URISyntaxException { + return urns.stream() + .filter(data::containsKey) + .map(urn -> Pair.of(urn, data.get(urn))) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue)); + } + + @Nonnull + @Override + public EntityRegistry getEntityRegistry() { + return mock(EntityRegistry.class); + } +} diff --git a/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/TestEntityRegistry.java b/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/TestEntityRegistry.java new file mode 100644 index 00000000000000..cad3b8c730e4bd --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/TestEntityRegistry.java @@ -0,0 +1,31 @@ +package com.linkedin.test.metadata.aspect; + +import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import java.util.Map; + +public class TestEntityRegistry extends ConfigEntityRegistry { + + static { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + public TestEntityRegistry() { + super(TestEntityRegistry.class.getClassLoader().getResourceAsStream("entity-registry.yml")); + } + + public static String getAspectName(T aspect) { + Map schemaProps = aspect.schema().getProperties(); + if (schemaProps != null && schemaProps.containsKey("Aspect")) { + Object aspectProps = schemaProps.get("Aspect"); + if (aspectProps instanceof Map aspectMap) { + return (String) aspectMap.get("name"); + } + } + + throw new IllegalStateException("Cannot determine aspect name"); + } +} diff --git a/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/batch/TestMCP.java b/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/batch/TestMCP.java new file mode 100644 index 00000000000000..20d01dc55934a6 --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/test/metadata/aspect/batch/TestMCP.java @@ -0,0 +1,128 @@ +package com.linkedin.test.metadata.aspect.batch; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.ReadItem; +import com.linkedin.metadata.aspect.SystemAspect; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; +import com.linkedin.test.metadata.aspect.TestEntityRegistry; +import java.net.URISyntaxException; +import java.util.Collection; +import java.util.Set; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Builder; +import lombok.Getter; +import lombok.Setter; + +@Builder +@Getter +public class TestMCP implements ChangeMCP { + private static final String TEST_DATASET_URN = + "urn:li:dataset:(urn:li:dataPlatform:datahub,Test,PROD)"; + + public static Collection ofOneBatchItem( + Urn urn, T aspect, EntityRegistry entityRegistry) { + return Set.of( + TestMCP.builder() + .urn(urn) + .entitySpec(entityRegistry.getEntitySpec(urn.getEntityType())) + .aspectSpec( + entityRegistry.getAspectSpecs().get(TestEntityRegistry.getAspectName(aspect))) + .recordTemplate(aspect) + .build()); + } + + public static Collection ofOneBatchItemDatasetUrn( + T aspect, EntityRegistry entityRegistry) { + try { + return ofOneBatchItem(Urn.createFromString(TEST_DATASET_URN), aspect, entityRegistry); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + + public static Set ofOneUpsertItem( + Urn urn, T aspect, EntityRegistry entityRegistry) { + return Set.of( + TestMCP.builder() + .urn(urn) + .entitySpec(entityRegistry.getEntitySpec(urn.getEntityType())) + .aspectSpec( + entityRegistry.getAspectSpecs().get(TestEntityRegistry.getAspectName(aspect))) + .recordTemplate(aspect) + .build()); + } + + public static Set ofOneUpsertItemDatasetUrn( + T aspect, EntityRegistry entityRegistry) { + try { + return ofOneUpsertItem(Urn.createFromString(TEST_DATASET_URN), aspect, entityRegistry); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + + public static Set ofOneMCP( + Urn urn, T newAspect, EntityRegistry entityRegistry) { + return ofOneMCP(urn, null, newAspect, entityRegistry); + } + + public static Set ofOneMCP( + Urn urn, @Nullable T oldAspect, T newAspect, EntityRegistry entityRegistry) { + + SystemAspect mockNewSystemAspect = mock(SystemAspect.class); + when(mockNewSystemAspect.getRecordTemplate()).thenReturn(newAspect); + when(mockNewSystemAspect.getAspect(any(Class.class))) + .thenAnswer(args -> ReadItem.getAspect(args.getArgument(0), newAspect)); + + SystemAspect mockOldSystemAspect = null; + if (oldAspect != null) { + mockOldSystemAspect = mock(SystemAspect.class); + when(mockOldSystemAspect.getRecordTemplate()).thenReturn(oldAspect); + when(mockOldSystemAspect.getAspect(any(Class.class))) + .thenAnswer(args -> ReadItem.getAspect(args.getArgument(0), oldAspect)); + } + + return Set.of( + TestMCP.builder() + .urn(urn) + .entitySpec(entityRegistry.getEntitySpec(urn.getEntityType())) + .aspectSpec( + entityRegistry.getAspectSpecs().get(TestEntityRegistry.getAspectName(newAspect))) + .recordTemplate(newAspect) + .systemAspect(mockNewSystemAspect) + .previousSystemAspect(mockOldSystemAspect) + .build()); + } + + private Urn urn; + private RecordTemplate recordTemplate; + private SystemMetadata systemMetadata; + private AuditStamp auditStamp; + private ChangeType changeType; + @Nonnull private final EntitySpec entitySpec; + @Nonnull private final AspectSpec aspectSpec; + private SystemAspect systemAspect; + private MetadataChangeProposal metadataChangeProposal; + @Setter private SystemAspect previousSystemAspect; + @Setter private long nextAspectVersion; + + @Nonnull + @Override + public SystemAspect getSystemAspect(@Nullable Long nextAspectVersion) { + return null; + } +} diff --git a/li-utils/src/main/java/com/linkedin/metadata/Constants.java b/li-utils/src/main/java/com/linkedin/metadata/Constants.java index 39a17612aa4b3a..0df1a959f9f603 100644 --- a/li-utils/src/main/java/com/linkedin/metadata/Constants.java +++ b/li-utils/src/main/java/com/linkedin/metadata/Constants.java @@ -16,6 +16,8 @@ public class Constants { public static final String ENTITY_TYPE_URN_PREFIX = "urn:li:entityType:"; public static final String DATA_TYPE_URN_PREFIX = "urn:li:dataType:"; public static final String STRUCTURED_PROPERTY_MAPPING_FIELD = "structuredProperties"; + public static final String STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX = + STRUCTURED_PROPERTY_MAPPING_FIELD + "."; // !!!!!!! IMPORTANT !!!!!!! // This effectively sets the max aspect size to 16 MB. Used in deserialization of messages. diff --git a/metadata-io/src/main/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtil.java b/metadata-io/src/main/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtil.java index a3711afb753dc8..544afc32a52e78 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtil.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtil.java @@ -19,10 +19,10 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.aspect.batch.BatchItem; -import com.linkedin.metadata.aspect.batch.MCPBatchItem; +import com.linkedin.metadata.aspect.batch.MCPItem; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.EntityUtils; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.utils.DataPlatformInstanceUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -49,10 +49,10 @@ private DefaultAspectsUtil() {} public static final Set SUPPORTED_TYPES = Set.of(ChangeType.UPSERT, ChangeType.CREATE, ChangeType.PATCH); - public static List getAdditionalChanges( + public static List getAdditionalChanges( @Nonnull AspectsBatch batch, @Nonnull EntityService entityService, boolean browsePathV2) { - Map> itemsByUrn = + Map> itemsByUrn = batch.getMCPItems().stream() .filter(item -> SUPPORTED_TYPES.contains(item.getChangeType())) .collect(Collectors.groupingBy(BatchItem::getUrn)); @@ -79,13 +79,13 @@ public static List getAdditionalChanges( RecordTemplate entityKeyAspect = defaultAspects.get(0).getSecond(); // pick the first item as a template (use entity information) - MCPBatchItem templateItem = aspectsEntry.getValue().get(0); + MCPItem templateItem = aspectsEntry.getValue().get(0); // generate default aspects (including key aspect, always upserts) return defaultAspects.stream() .map( entry -> - MCPUpsertBatchItem.MCPUpsertBatchItemBuilder.build( + ChangeItemImpl.ChangeItemImplBuilder.build( getProposalFromAspect( entry.getKey(), entry.getValue(), entityKeyAspect, templateItem), templateItem.getAuditStamp(), @@ -280,7 +280,7 @@ private static MetadataChangeProposal getProposalFromAspect( String aspectName, RecordTemplate aspect, RecordTemplate entityKeyAspect, - MCPBatchItem templateItem) { + MCPItem templateItem) { MetadataChangeProposal proposal = new MetadataChangeProposal(); GenericAspect genericAspect = GenericRecordUtils.serializeAspect(aspect); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java b/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java index 974406c0be0df1..0fcb765b340cf2 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java @@ -3,7 +3,7 @@ import com.linkedin.common.urn.Urn; import com.linkedin.entity.Aspect; import com.linkedin.entity.client.SystemEntityClient; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.aspect.CachingAspectRetriever; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.r2.RemoteInvocationException; import java.net.URISyntaxException; @@ -11,14 +11,29 @@ import java.util.Set; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import javax.annotation.PostConstruct; import lombok.Builder; import lombok.Getter; +import lombok.RequiredArgsConstructor; +import org.springframework.stereotype.Component; @Builder -public class EntityClientAspectRetriever implements AspectRetriever { +@Component +@RequiredArgsConstructor +public class EntityClientAspectRetriever implements CachingAspectRetriever { @Getter private final EntityRegistry entityRegistry; private final SystemEntityClient entityClient; + /** + * Preventing a circular dependency. Once constructed the AspectRetriever is injected into a few + * of the services which rely on the AspectRetriever when using the Java EntityClient. The Java + * EntityClient depends on services which in turn depend on the AspectRetriever + */ + @PostConstruct + public void postConstruct() { + entityClient.postConstruct(this); + } + @Nullable @Override public Aspect getLatestAspectObject(@Nonnull Urn urn, @Nonnull String aspectName) @@ -30,6 +45,10 @@ public Aspect getLatestAspectObject(@Nonnull Urn urn, @Nonnull String aspectName @Override public Map> getLatestAspectObjects( Set urns, Set aspectNames) throws RemoteInvocationException, URISyntaxException { - return entityClient.getLatestAspects(urns, aspectNames); + if (urns.isEmpty() || aspectNames.isEmpty()) { + return Map.of(); + } else { + return entityClient.getLatestAspects(urns, aspectNames); + } } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java index 15de029340a3c7..fed6379f921045 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java @@ -20,6 +20,7 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.aspect.EnvelopedAspectArray; import com.linkedin.metadata.aspect.VersionedAspect; @@ -83,15 +84,26 @@ public class JavaEntityClient implements EntityClient { private final Clock _clock = Clock.systemUTC(); - private final EntityService _entityService; - private final DeleteEntityService _deleteEntityService; - private final EntitySearchService _entitySearchService; - private final CachingEntitySearchService _cachingEntitySearchService; - private final SearchService _searchService; - private final LineageSearchService _lineageSearchService; - private final TimeseriesAspectService _timeseriesAspectService; + private final EntityService entityService; + private final DeleteEntityService deleteEntityService; + private final EntitySearchService entitySearchService; + private final CachingEntitySearchService cachingEntitySearchService; + private final SearchService searchService; + private final LineageSearchService lineageSearchService; + private final TimeseriesAspectService timeseriesAspectService; private final RollbackService rollbackService; - private final EventProducer _eventProducer; + private final EventProducer eventProducer; + + /** + * Preventing a circular dependency. Once constructed the AspectRetriever is injected into a few + * of the services which rely on the AspectRetriever when using the Java EntityClient. The Java + * EntityClient depends on services which in turn depend on the AspectRetriever + */ + @Override + public void postConstruct(AspectRetriever aspectRetriever) { + entitySearchService.postConstruct(aspectRetriever); + timeseriesAspectService.postConstruct(aspectRetriever); + } @Nullable public EntityResponse getV2( @@ -101,13 +113,13 @@ public EntityResponse getV2( @Nonnull final Authentication authentication) throws RemoteInvocationException, URISyntaxException { final Set projectedAspects = - aspectNames == null ? _entityService.getEntityAspectNames(entityName) : aspectNames; - return _entityService.getEntityV2(entityName, urn, projectedAspects); + aspectNames == null ? entityService.getEntityAspectNames(entityName) : aspectNames; + return entityService.getEntityV2(entityName, urn, projectedAspects); } @Nonnull public Entity get(@Nonnull final Urn urn, @Nonnull final Authentication authentication) { - return _entityService.getEntity(urn, ImmutableSet.of()); + return entityService.getEntity(urn, ImmutableSet.of()); } @Nonnull @@ -119,8 +131,8 @@ public Map batchGetV2( @Nonnull Authentication authentication) throws RemoteInvocationException, URISyntaxException { final Set projectedAspects = - aspectNames == null ? _entityService.getEntityAspectNames(entityName) : aspectNames; - return _entityService.getEntitiesV2(entityName, urns, projectedAspects); + aspectNames == null ? entityService.getEntityAspectNames(entityName) : aspectNames; + return entityService.getEntitiesV2(entityName, urns, projectedAspects); } @Nonnull @@ -131,14 +143,14 @@ public Map batchGetVersionedV2( @Nonnull final Authentication authentication) throws RemoteInvocationException, URISyntaxException { final Set projectedAspects = - aspectNames == null ? _entityService.getEntityAspectNames(entityName) : aspectNames; - return _entityService.getEntitiesVersionedV2(versionedUrns, projectedAspects); + aspectNames == null ? entityService.getEntityAspectNames(entityName) : aspectNames; + return entityService.getEntitiesVersionedV2(versionedUrns, projectedAspects); } @Nonnull public Map batchGet( @Nonnull final Set urns, @Nonnull final Authentication authentication) { - return _entityService.getEntities(urns, ImmutableSet.of()); + return entityService.getEntities(urns, ImmutableSet.of()); } /** @@ -160,7 +172,7 @@ public AutoCompleteResult autoComplete( @Nullable String field, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _cachingEntitySearchService.autoComplete( + return cachingEntitySearchService.autoComplete( entityType, query, field, filterOrDefaultEmptyFilter(requestFilters), limit, null); } @@ -181,7 +193,7 @@ public AutoCompleteResult autoComplete( @Nonnull int limit, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _cachingEntitySearchService.autoComplete( + return cachingEntitySearchService.autoComplete( entityType, query, "", filterOrDefaultEmptyFilter(requestFilters), limit, null); } @@ -205,9 +217,9 @@ public BrowseResult browse( @Nonnull final Authentication authentication) throws RemoteInvocationException { return ValidationUtils.validateBrowseResult( - _cachingEntitySearchService.browse( + cachingEntitySearchService.browse( entityType, path, newFilter(requestFilters), start, limit, null), - _entityService); + entityService); } /** @@ -232,8 +244,7 @@ public BrowseResultV2 browseV2( @Nonnull Authentication authentication, @Nullable SearchFlags searchFlags) { // TODO: cache browseV2 results - return _entitySearchService.browseV2( - entityName, path, filter, input, start, count, searchFlags); + return entitySearchService.browseV2(entityName, path, filter, input, start, count, searchFlags); } /** @@ -258,7 +269,7 @@ public BrowseResultV2 browseV2( @Nonnull Authentication authentication, @Nullable SearchFlags searchFlags) { // TODO: cache browseV2 results - return _entitySearchService.browseV2( + return entitySearchService.browseV2( entityNames, path, filter, input, start, count, searchFlags); } @@ -270,7 +281,7 @@ public void update(@Nonnull final Entity entity, @Nonnull final Authentication a AuditStamp auditStamp = new AuditStamp(); auditStamp.setActor(Urn.createFromString(authentication.getActor().toUrnStr())); auditStamp.setTime(Clock.systemUTC().millis()); - _entityService.ingestEntity(entity, auditStamp); + entityService.ingestEntity(entity, auditStamp); } @SneakyThrows @@ -289,7 +300,7 @@ public void updateWithSystemMetadata( auditStamp.setActor(Urn.createFromString(authentication.getActor().toUrnStr())); auditStamp.setTime(Clock.systemUTC().millis()); - _entityService.ingestEntity(entity, auditStamp, systemMetadata); + entityService.ingestEntity(entity, auditStamp, systemMetadata); tryIndexRunId( com.datahub.util.ModelUtils.getUrnFromSnapshotUnion(entity.getValue()), systemMetadata); } @@ -302,7 +313,7 @@ public void batchUpdate( AuditStamp auditStamp = new AuditStamp(); auditStamp.setActor(Urn.createFromString(authentication.getActor().toUrnStr())); auditStamp.setTime(Clock.systemUTC().millis()); - _entityService.ingestEntities( + entityService.ingestEntities( entities.stream().collect(Collectors.toList()), auditStamp, ImmutableList.of()); } @@ -331,9 +342,9 @@ public SearchResult search( throws RemoteInvocationException { return ValidationUtils.validateSearchResult( - _entitySearchService.search( + entitySearchService.search( List.of(entity), input, newFilter(requestFilters), null, start, count, searchFlags), - _entityService); + entityService); } /** @@ -358,8 +369,8 @@ public ListResult list( throws RemoteInvocationException { return ValidationUtils.validateListResult( toListResult( - _entitySearchService.filter(entity, newFilter(requestFilters), null, start, count)), - _entityService); + entitySearchService.filter(entity, newFilter(requestFilters), null, start, count)), + entityService); } /** @@ -386,9 +397,9 @@ public SearchResult search( @Nullable SearchFlags searchFlags) throws RemoteInvocationException { return ValidationUtils.validateSearchResult( - _entitySearchService.search( + entitySearchService.search( List.of(entity), input, filter, sortCriterion, start, count, searchFlags), - _entityService); + entityService); } @Nonnull @@ -434,9 +445,9 @@ public SearchResult searchAcrossEntities( final SearchFlags finalFlags = searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true); return ValidationUtils.validateSearchResult( - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( entities, input, filter, sortCriterion, start, count, finalFlags, facets), - _entityService); + entityService); } @Nonnull @@ -454,9 +465,9 @@ public ScrollResult scrollAcrossEntities( final SearchFlags finalFlags = searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true); return ValidationUtils.validateScrollResult( - _searchService.scrollAcrossEntities( + searchService.scrollAcrossEntities( entities, input, filter, null, scrollId, keepAlive, count, finalFlags), - _entityService); + entityService); } @Nonnull @@ -475,7 +486,7 @@ public LineageSearchResult searchAcrossLineage( @Nonnull final Authentication authentication) throws RemoteInvocationException { return ValidationUtils.validateLineageSearchResult( - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( sourceUrn, direction, entities, @@ -488,7 +499,7 @@ public LineageSearchResult searchAcrossLineage( null, null, searchFlags), - _entityService); + entityService); } @Nonnull @@ -509,7 +520,7 @@ public LineageSearchResult searchAcrossLineage( @Nonnull final Authentication authentication) throws RemoteInvocationException { return ValidationUtils.validateLineageSearchResult( - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( sourceUrn, direction, entities, @@ -522,7 +533,7 @@ public LineageSearchResult searchAcrossLineage( startTimeMillis, endTimeMillis, searchFlags), - _entityService); + entityService); } @Nonnull @@ -546,7 +557,7 @@ public LineageScrollResult scrollAcrossLineage( final SearchFlags finalFlags = searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true).setSkipCache(true); return ValidationUtils.validateLineageScrollResult( - _lineageSearchService.scrollAcrossLineage( + lineageSearchService.scrollAcrossLineage( sourceUrn, direction, entities, @@ -560,7 +571,7 @@ public LineageScrollResult scrollAcrossLineage( startTimeMillis, endTimeMillis, finalFlags), - _entityService); + entityService); } /** @@ -573,19 +584,19 @@ public LineageScrollResult scrollAcrossLineage( @Nonnull public StringArray getBrowsePaths(@Nonnull Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return new StringArray(_entitySearchService.getBrowsePaths(urn.getEntityType(), urn)); + return new StringArray(entitySearchService.getBrowsePaths(urn.getEntityType(), urn)); } public void setWritable(boolean canWrite, @Nonnull final Authentication authentication) throws RemoteInvocationException { - _entityService.setWritable(canWrite); + entityService.setWritable(canWrite); } @Nonnull public Map batchGetTotalEntityCount( @Nonnull List entityNames, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _searchService.docCountPerEntity(entityNames); + return searchService.docCountPerEntity(entityNames); } /** List all urns existing for a particular Entity type. */ @@ -595,19 +606,19 @@ public ListUrnsResult listUrns( final int count, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _entityService.listUrns(entityName, start, count); + return entityService.listUrns(entityName, start, count); } /** Hard delete an entity with a particular urn. */ public void deleteEntity(@Nonnull final Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException { - _entityService.deleteUrn(urn); + entityService.deleteUrn(urn); } @Override public void deleteEntityReferences(@Nonnull Urn urn, @Nonnull Authentication authentication) throws RemoteInvocationException { - withRetry(() -> _deleteEntityService.deleteReferencesTo(urn, false), "deleteEntityReferences"); + withRetry(() -> deleteEntityService.deleteReferencesTo(urn, false), "deleteEntityReferences"); } @Nonnull @@ -621,13 +632,13 @@ public SearchResult filter( @Nonnull final Authentication authentication) throws RemoteInvocationException { return ValidationUtils.validateSearchResult( - _entitySearchService.filter(entity, filter, sortCriterion, start, count), _entityService); + entitySearchService.filter(entity, filter, sortCriterion, start, count), entityService); } @Override public boolean exists(@Nonnull Urn urn, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _entityService.exists(urn, true); + return entityService.exists(urn, true); } @SneakyThrows @@ -638,7 +649,7 @@ public VersionedAspect getAspect( @Nonnull Long version, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); + return entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); } @SneakyThrows @@ -649,7 +660,7 @@ public VersionedAspect getAspectOrNull( @Nonnull Long version, @Nonnull final Authentication authentication) throws RemoteInvocationException { - return _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); + return entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); } @SneakyThrows @@ -682,7 +693,7 @@ public List getTimeseriesAspectValues( } response.setValues( new EnvelopedAspectArray( - _timeseriesAspectService.getAspectValues( + timeseriesAspectService.getAspectValues( Urn.createFromString(urn), entity, aspect, @@ -711,10 +722,10 @@ public String ingestProposal( AspectsBatch batch = AspectsBatchImpl.builder() - .mcps(List.of(metadataChangeProposal), auditStamp, _entityService) + .mcps(List.of(metadataChangeProposal), auditStamp, entityService) .build(); - IngestResult one = _entityService.ingestProposal(batch, async).stream().findFirst().get(); + IngestResult one = entityService.ingestProposal(batch, async).stream().findFirst().get(); Urn urn = one.getUrn(); tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata()); @@ -731,7 +742,7 @@ public Optional getVersionedAspect( @Nonnull final Authentication authentication) throws RemoteInvocationException { VersionedAspect entity = - _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); + entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); if (entity != null && entity.hasAspect()) { DataMap rawAspect = ((DataMap) entity.data().get("aspect")); if (rawAspect.containsKey(aspectClass.getCanonicalName())) { @@ -750,7 +761,7 @@ public DataMap getRawAspect( @Nonnull Authentication authentication) throws RemoteInvocationException { VersionedAspect entity = - _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); + entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version); if (entity == null) { return null; } @@ -770,7 +781,7 @@ public void producePlatformEvent( @Nonnull PlatformEvent event, @Nonnull Authentication authentication) throws Exception { - _eventProducer.producePlatformEvent(name, key, event); + eventProducer.producePlatformEvent(name, key, event); } @Override @@ -782,7 +793,7 @@ public void rollbackIngestion( private void tryIndexRunId(Urn entityUrn, @Nullable SystemMetadata systemMetadata) { if (systemMetadata != null && systemMetadata.hasRunId()) { - _entitySearchService.appendRunId( + entitySearchService.appendRunId( entityUrn.getEntityType(), entityUrn, systemMetadata.getRunId()); } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspect.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspect.java index d72586e289ea78..ae1b3007ed647e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspect.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityAspect.java @@ -1,19 +1,27 @@ package com.linkedin.metadata.entity; +import static com.linkedin.metadata.entity.EntityUtils.parseSystemMetadata; + +import com.datahub.util.RecordUtils; import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.aspect.batch.SystemAspect; -import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.entity.AspectType; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.metadata.aspect.SystemAspect; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; import com.linkedin.mxe.SystemMetadata; -import java.net.URISyntaxException; import java.sql.Timestamp; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.AllArgsConstructor; +import lombok.Builder; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.NoArgsConstructor; import lombok.Setter; +import lombok.extern.slf4j.Slf4j; /** * This is an internal representation of an entity aspect record {@link EntityServiceImpl} and @@ -21,6 +29,7 @@ * own aspect record implementations, they cary implementation details that should not leak outside. * Therefore, this is the type to use in public {@link AspectDao} methods. */ +@Slf4j @Getter @Setter @NoArgsConstructor @@ -44,60 +53,52 @@ public class EntityAspect { private String createdFor; - public EntityAspectIdentifier toAspectIdentifier() { + public EntityAspectIdentifier getAspectIdentifier() { return new EntityAspectIdentifier(getUrn(), getAspect(), getVersion()); } - @Nonnull - public SystemAspect asSystemAspect() { - return EntitySystemAspect.from(this); - } - /** * Provide a typed EntityAspect without breaking the existing public contract with generic types. */ + @Builder @Getter - @AllArgsConstructor @EqualsAndHashCode public static class EntitySystemAspect implements SystemAspect { - - @Nullable - public static EntitySystemAspect from(EntityAspect entityAspect) { - return entityAspect != null ? new EntitySystemAspect(entityAspect) : null; - } - @Nonnull private final EntityAspect entityAspect; + @Nonnull private final Urn urn; - @Nonnull - public Urn getUrn() { - try { - return Urn.createFromString(entityAspect.getUrn()); - } catch (URISyntaxException e) { - throw new RuntimeException(e); - } - } + /** Note that read mutations depend on the mutability of recordTemplate */ + @Nullable private final RecordTemplate recordTemplate; + + @Nonnull private final EntitySpec entitySpec; + @Nonnull private final AspectSpec aspectSpec; @Nonnull public String getUrnRaw() { return entityAspect.getUrn(); } - @Override - public SystemMetadata getSystemMetadata() { - return EntityUtils.parseSystemMetadata(entityAspect.getSystemMetadata()); - } - @Nullable public String getSystemMetadataRaw() { return entityAspect.getSystemMetadata(); } + public String getMetadataRaw() { + return entityAspect.getMetadata(); + } + @Override public Timestamp getCreatedOn() { return entityAspect.getCreatedOn(); } @Override + public String getCreatedBy() { + return entityAspect.getCreatedBy(); + } + + @Override + @Nonnull public String getAspectName() { return entityAspect.aspect; } @@ -107,14 +108,72 @@ public long getVersion() { return entityAspect.getVersion(); } - @Override - public RecordTemplate getRecordTemplate(EntityRegistry entityRegistry) { - return EntityUtils.toAspectRecord( - getUrn().getEntityType(), getAspectName(), entityAspect.getMetadata(), entityRegistry); + @Nullable + public SystemMetadata getSystemMetadata() { + return parseSystemMetadata(getSystemMetadataRaw()); + } + + public EntityAspectIdentifier getAspectIdentifier() { + return entityAspect.getAspectIdentifier(); } - public EntityAspect asRaw() { - return entityAspect; + /** + * Convert to enveloped aspect + * + * @return enveloped aspect + */ + public EnvelopedAspect toEnvelopedAspects() { + // Now turn it into an EnvelopedAspect + final com.linkedin.entity.Aspect aspect = + new com.linkedin.entity.Aspect(getRecordTemplate().data()); + + final EnvelopedAspect envelopedAspect = new EnvelopedAspect(); + envelopedAspect.setName(getAspectName()); + envelopedAspect.setVersion(getVersion()); + + // TODO: I think we can assume this here, adding as it's a required field so object mapping + // barfs when trying to access it, + // since nowhere else is using it should be safe for now at least + envelopedAspect.setType(AspectType.VERSIONED); + envelopedAspect.setValue(aspect); + + try { + if (getSystemMetadata() != null) { + envelopedAspect.setSystemMetadata(getSystemMetadata()); + } + } catch (Exception e) { + log.warn( + "Exception encountered when setting system metadata on enveloped aspect {}. Error: {}", + envelopedAspect.getName(), + e.toString()); + } + + envelopedAspect.setCreated(getAuditStamp()); + + return envelopedAspect; + } + + public static class EntitySystemAspectBuilder { + + private EntityAspect.EntitySystemAspect build() { + return null; + } + + public EntityAspect.EntitySystemAspect build( + @Nonnull EntitySpec entitySpec, + @Nonnull AspectSpec aspectSpec, + @Nonnull EntityAspect entityAspect) { + this.entityAspect = entityAspect; + this.urn = UrnUtils.getUrn(entityAspect.getUrn()); + this.aspectSpec = aspectSpec; + if (entityAspect.getMetadata() != null) { + this.recordTemplate = + RecordUtils.toRecordTemplate( + aspectSpec.getDataTemplateClass(), entityAspect.getMetadata()); + } + + return new EntitySystemAspect(entityAspect, urn, recordTemplate, entitySpec, aspectSpec); + } } } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java index eec5c6120886dd..ed3a78ceddba4a 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java @@ -14,7 +14,6 @@ import com.codahale.metrics.Timer; import com.datahub.util.RecordUtils; -import com.datahub.util.exception.ModelConversionException; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; @@ -34,30 +33,30 @@ import com.linkedin.data.template.SetMode; import com.linkedin.data.template.StringMap; import com.linkedin.data.template.UnionTemplate; -import com.linkedin.entity.AspectType; import com.linkedin.entity.Entity; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; -import com.linkedin.entity.EnvelopedAspectMap; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.Aspect; +import com.linkedin.metadata.aspect.SystemAspect; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.aspect.batch.BatchItem; -import com.linkedin.metadata.aspect.batch.MCPBatchItem; -import com.linkedin.metadata.aspect.batch.SystemAspect; -import com.linkedin.metadata.aspect.batch.UpsertItem; -import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.batch.MCPItem; +import com.linkedin.metadata.aspect.plugins.validation.ValidationExceptionCollection; import com.linkedin.metadata.aspect.utils.DefaultAspectsUtil; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; +import com.linkedin.metadata.entity.ebean.batch.DeleteItemImpl; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; +import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -65,9 +64,7 @@ import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; import com.linkedin.metadata.run.AspectRowSummary; -import com.linkedin.metadata.service.UpdateIndicesService; import com.linkedin.metadata.snapshot.Snapshot; -import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.metadata.utils.PegasusUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; @@ -138,7 +135,7 @@ * class. */ @Slf4j -public class EntityServiceImpl implements EntityService { +public class EntityServiceImpl implements EntityService { /** * As described above, the latest version of an aspect should always take the value 0, with @@ -146,15 +143,15 @@ public class EntityServiceImpl implements EntityService { */ private static final int DEFAULT_MAX_TRANSACTION_RETRY = 3; - protected final AspectDao _aspectDao; + protected final AspectDao aspectDao; - @VisibleForTesting @Getter private final EventProducer _producer; - private final EntityRegistry _entityRegistry; - private final Map> _entityToValidAspects; - private RetentionService _retentionService; - private final Boolean _alwaysEmitChangeLog; - @Getter private final UpdateIndicesService _updateIndicesService; - private final PreProcessHooks _preProcessHooks; + @VisibleForTesting @Getter private final EventProducer producer; + private final EntityRegistry entityRegistry; + private final Map> entityToValidAspects; + private RetentionService retentionService; + private final Boolean alwaysEmitChangeLog; + @Nullable @Getter private SearchIndicesService updateIndicesService; + private final PreProcessHooks preProcessHooks; protected static final int MAX_KEYS_PER_QUERY = 500; private final Integer ebeanMaxTransactionRetry; @@ -165,7 +162,6 @@ public EntityServiceImpl( @Nonnull final EventProducer producer, @Nonnull final EntityRegistry entityRegistry, final boolean alwaysEmitChangeLog, - @Nullable final UpdateIndicesService updateIndicesService, final PreProcessHooks preProcessHooks, final boolean enableBrowsePathV2) { this( @@ -173,7 +169,6 @@ public EntityServiceImpl( producer, entityRegistry, alwaysEmitChangeLog, - updateIndicesService, preProcessHooks, DEFAULT_MAX_TRANSACTION_RETRY, enableBrowsePathV2); @@ -184,25 +179,27 @@ public EntityServiceImpl( @Nonnull final EventProducer producer, @Nonnull final EntityRegistry entityRegistry, final boolean alwaysEmitChangeLog, - @Nullable final UpdateIndicesService updateIndicesService, final PreProcessHooks preProcessHooks, @Nullable final Integer retry, final boolean enableBrowseV2) { - _aspectDao = aspectDao; - _producer = producer; - _entityRegistry = entityRegistry; - _entityToValidAspects = buildEntityToValidAspects(entityRegistry); - _alwaysEmitChangeLog = alwaysEmitChangeLog; - _updateIndicesService = updateIndicesService; - if (_updateIndicesService != null) { - _updateIndicesService.initializeAspectRetriever(this); - } - _preProcessHooks = preProcessHooks; + this.aspectDao = aspectDao; + this.producer = producer; + this.entityRegistry = entityRegistry; + entityToValidAspects = EntityUtils.buildEntityToValidAspects(entityRegistry); + this.alwaysEmitChangeLog = alwaysEmitChangeLog; + this.preProcessHooks = preProcessHooks; ebeanMaxTransactionRetry = retry != null ? retry : DEFAULT_MAX_TRANSACTION_RETRY; this.enableBrowseV2 = enableBrowseV2; } + public void setUpdateIndicesService(@Nullable SearchIndicesService updateIndicesService) { + this.updateIndicesService = updateIndicesService; + if (this.updateIndicesService != null) { + this.updateIndicesService.initializeAspectRetriever(this); + } + } + @Override public RecordTemplate getLatestAspect(@Nonnull Urn urn, @Nonnull String aspectName) { log.debug("Invoked getLatestAspect with urn {}, aspect {}", urn, aspectName); @@ -236,24 +233,22 @@ public Map> getLatestAspects( .keySet() .forEach( key -> { - final RecordTemplate keyAspect = EntityUtils.buildKeyAspect(_entityRegistry, key); + final RecordTemplate keyAspect = EntityUtils.buildKeyAspect(entityRegistry, key); urnToAspects.get(key).add(keyAspect); }); - batchGetResults.forEach( - (key, aspectEntry) -> { - final Urn urn = toUrn(key.getUrn()); - final String aspectName = key.getAspect(); - // for now, don't add the key aspect here- we have already added it above - if (aspectName.equals(getKeyAspectName(urn))) { - return; - } + List systemAspects = EntityUtils.toSystemAspects(batchGetResults.values(), this); - final RecordTemplate aspectRecord = - aspectEntry.asSystemAspect().getRecordTemplate(getEntityRegistry()); - urnToAspects.putIfAbsent(urn, new ArrayList<>()); - urnToAspects.get(urn).add(aspectRecord); - }); + systemAspects.stream() + // for now, don't add the key aspect here we have already added it above + .filter( + systemAspect -> + !getKeyAspectName(systemAspect.getUrn()).equals(systemAspect.getAspectName())) + .forEach( + systemAspect -> + urnToAspects + .computeIfAbsent(systemAspect.getUrn(), u -> new ArrayList<>()) + .add(systemAspect.getRecordTemplate())); return urnToAspects; } @@ -265,15 +260,10 @@ public Map getLatestAspectsForUrn( Map batchGetResults = getLatestAspect(new HashSet<>(Arrays.asList(urn)), aspectNames); - final Map result = new HashMap<>(); - batchGetResults.forEach( - (key, aspectEntry) -> { - final String aspectName = key.getAspect(); - final RecordTemplate aspectRecord = - aspectEntry.asSystemAspect().getRecordTemplate(getEntityRegistry()); - result.put(aspectName, aspectRecord); - }); - return result; + return EntityUtils.toSystemAspects(batchGetResults.values(), this).stream() + .map( + systemAspect -> Pair.of(systemAspect.getAspectName(), systemAspect.getRecordTemplate())) + .collect(Collectors.toMap(Pair::getKey, Pair::getValue)); } /** @@ -291,7 +281,12 @@ public Map getLatestAspectsForUrn( @Nullable @Override public RecordTemplate getAspect( - @Nonnull final Urn urn, @Nonnull final String aspectName, @Nonnull long version) { + @Nonnull final Urn urn, @Nonnull final String aspectName, long version) { + return getAspectVersionPair(urn, aspectName, version).getFirst(); + } + + public Pair getAspectVersionPair( + @Nonnull final Urn urn, @Nonnull final String aspectName, long version) { log.debug( "Invoked getAspect with urn: {}, aspectName: {}, version: {}", urn, aspectName, version); @@ -299,14 +294,13 @@ public RecordTemplate getAspect( version = calculateVersionNumber(urn, aspectName, version); final EntityAspectIdentifier primaryKey = new EntityAspectIdentifier(urn.toString(), aspectName, version); - final Optional maybeAspect = - Optional.ofNullable(_aspectDao.getAspect(primaryKey)); - return maybeAspect - .map( - aspect -> - EntityUtils.toAspectRecord( - urn, aspectName, aspect.getMetadata(), getEntityRegistry())) - .orElse(null); + final Optional maybeAspect = Optional.ofNullable(aspectDao.getAspect(primaryKey)); + + return Pair.of( + EntityUtils.toSystemAspect(maybeAspect.orElse(null), this) + .map(SystemAspect::getRecordTemplate) + .orElse(null), + version); } /** @@ -347,7 +341,8 @@ public Map getEntitiesV2( return getLatestEnvelopedAspects(urns, aspectNames).entrySet().stream() .collect( Collectors.toMap( - Map.Entry::getKey, entry -> toEntityResponse(entry.getKey(), entry.getValue()))); + Map.Entry::getKey, + entry -> EntityUtils.toEntityResponse(entry.getKey(), entry.getValue()))); } /** @@ -366,7 +361,8 @@ public Map getEntitiesVersionedV2( return getVersionedEnvelopedAspects(versionedUrns, aspectNames).entrySet().stream() .collect( Collectors.toMap( - Map.Entry::getKey, entry -> toEntityResponse(entry.getKey(), entry.getValue()))); + Map.Entry::getKey, + entry -> EntityUtils.toEntityResponse(entry.getKey(), entry.getValue()))); } /** @@ -473,7 +469,7 @@ private Map> getCorrespondingAspects( for (Urn urn : urns) { List aspects = urnToAspects.getOrDefault(urn.toString(), Collections.emptyList()); - EnvelopedAspect keyAspect = getKeyEnvelopedAspect(urn); + EnvelopedAspect keyAspect = EntityUtils.getKeyEnvelopedAspect(urn, entityRegistry); // Add key aspect if it does not exist in the returned aspects if (aspects.isEmpty() || aspects.stream().noneMatch(aspect -> keyAspect.getName().equals(aspect.getName()))) { @@ -520,29 +516,16 @@ public VersionedAspect getVersionedAspect( VersionedAspect result = new VersionedAspect(); - version = calculateVersionNumber(urn, aspectName, version); - - final EntityAspectIdentifier primaryKey = - new EntityAspectIdentifier(urn.toString(), aspectName, version); - final Optional maybeAspect = - Optional.ofNullable(_aspectDao.getAspect(primaryKey)); - RecordTemplate aspectRecord = - maybeAspect - .map( - aspect -> - EntityUtils.toAspectRecord( - urn, aspectName, aspect.getMetadata(), getEntityRegistry())) - .orElse(null); - - if (aspectRecord == null) { + Pair aspectRecord = getAspectVersionPair(urn, aspectName, version); + if (aspectRecord.getFirst() == null) { return null; } Aspect resultAspect = new Aspect(); - RecordUtils.setSelectedRecordTemplateInUnion(resultAspect, aspectRecord); + RecordUtils.setSelectedRecordTemplateInUnion(resultAspect, aspectRecord.getFirst()); result.setAspect(resultAspect); - result.setVersion(version); + result.setVersion(aspectRecord.getSecond()); return result; } @@ -575,20 +558,22 @@ public ListResult listLatestAspects( count); final ListResult aspectMetadataList = - _aspectDao.listLatestAspectMetadata(entityName, aspectName, start, count); + aspectDao.listLatestAspectMetadata(entityName, aspectName, start, count); - final List aspects = new ArrayList<>(); + List entityAspects = new ArrayList<>(); for (int i = 0; i < aspectMetadataList.getValues().size(); i++) { - aspects.add( - EntityUtils.toAspectRecord( - aspectMetadataList.getMetadata().getExtraInfos().get(i).getUrn(), - aspectName, - aspectMetadataList.getValues().get(i), - getEntityRegistry())); + EntityAspect entityAspect = new EntityAspect(); + entityAspect.setUrn( + aspectMetadataList.getMetadata().getExtraInfos().get(i).getUrn().toString()); + entityAspect.setAspect(aspectName); + entityAspect.setMetadata(aspectMetadataList.getValues().get(i)); + entityAspects.add(entityAspect); } return new ListResult<>( - aspects, + EntityUtils.toSystemAspects(entityAspects, this).stream() + .map(SystemAspect::getRecordTemplate) + .collect(Collectors.toList()), aspectMetadataList.getMetadata(), aspectMetadataList.getNextStart(), aspectMetadataList.isHasNext(), @@ -612,11 +597,11 @@ public List ingestAspects( List> pairList, @Nonnull final AuditStamp auditStamp, SystemMetadata systemMetadata) { - List items = + List items = pairList.stream() .map( pair -> - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(pair.getKey()) .recordTemplate(pair.getValue()) @@ -624,7 +609,8 @@ public List ingestAspects( .auditStamp(auditStamp) .build(this)) .collect(Collectors.toList()); - return ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + return ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(this).items(items).build(), true, true); } /** @@ -643,7 +629,8 @@ public List ingestAspects( // Generate additional items as needed items.addAll(DefaultAspectsUtil.getAdditionalChanges(aspectsBatch, this, enableBrowseV2)); - AspectsBatch withDefaults = AspectsBatchImpl.builder().items(items).build(); + AspectsBatch withDefaults = + AspectsBatchImpl.builder().aspectRetriever(this).items(items).build(); Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectsToLocalDB").time(); @@ -672,73 +659,84 @@ private List ingestAspectsToLocalDB( log.warn(String.format("Batch contains duplicates: %s", aspectsBatch)); } - return _aspectDao + return aspectDao .runInTransactionWithRetry( (tx) -> { // Read before write is unfortunate, however batch it final Map> urnAspects = aspectsBatch.getUrnAspectsMap(); // read #1 final Map> latestAspects = - toSystemEntityAspects(_aspectDao.getLatestAspects(urnAspects)); + EntityUtils.toSystemAspects(aspectDao.getLatestAspects(urnAspects), this); // read #2 final Map> nextVersions = - _aspectDao.getNextVersions(urnAspects); + aspectDao.getNextVersions(urnAspects); // 1. Convert patches to full upserts // 2. Run any entity/aspect level hooks - Pair>, List> updatedItems = - aspectsBatch.toUpsertBatchItems(latestAspects, this); + Pair>, List> updatedItems = + aspectsBatch.toUpsertBatchItems(latestAspects); // Fetch additional information if needed final Map> updatedLatestAspects; final Map> updatedNextVersions; if (!updatedItems.getFirst().isEmpty()) { Map> newLatestAspects = - toSystemEntityAspects(_aspectDao.getLatestAspects(updatedItems.getFirst())); + EntityUtils.toSystemAspects( + aspectDao.getLatestAspects(updatedItems.getFirst()), this); Map> newNextVersions = - _aspectDao.getNextVersions(updatedItems.getFirst()); + aspectDao.getNextVersions(updatedItems.getFirst()); // merge - updatedLatestAspects = aspectsBatch.merge(latestAspects, newLatestAspects); - updatedNextVersions = aspectsBatch.merge(nextVersions, newNextVersions); + updatedLatestAspects = AspectsBatch.merge(latestAspects, newLatestAspects); + updatedNextVersions = AspectsBatch.merge(nextVersions, newNextVersions); } else { updatedLatestAspects = latestAspects; updatedNextVersions = nextVersions; } - // do final pre-commit checks with previous aspect value - updatedItems - .getSecond() - .forEach( - item -> { - SystemAspect previousAspect = + // Add previous version to each upsert + List changeMCPs = + updatedItems.getSecond().stream() + .peek( + changeMCP -> { + String urnStr = changeMCP.getUrn().toString(); + long nextVersion = + updatedNextVersions + .getOrDefault(urnStr, Map.of()) + .getOrDefault(changeMCP.getAspectName(), 0L); + + changeMCP.setPreviousSystemAspect( + updatedLatestAspects + .getOrDefault(urnStr, Map.of()) + .getOrDefault(changeMCP.getAspectName(), null)); + + changeMCP.setNextAspectVersion(nextVersion); + + // support inner-batch upserts updatedLatestAspects - .getOrDefault(item.getUrn().toString(), Map.of()) - .get(item.getAspectSpec().getName()); - try { - item.validatePreCommit( - previousAspect == null - ? null - : previousAspect.getRecordTemplate(_entityRegistry), - this); - } catch (AspectValidationException e) { - throw new RuntimeException(e); - } - }); + .computeIfAbsent(urnStr, key -> new HashMap<>()) + .put( + changeMCP.getAspectName(), + changeMCP.getSystemAspect(nextVersion)); + updatedNextVersions + .computeIfAbsent(urnStr, key -> new HashMap<>()) + .put(changeMCP.getAspectName(), nextVersion + 1); + }) + .collect(Collectors.toList()); + + // do final pre-commit checks with previous aspect value + ValidationExceptionCollection exceptions = + AspectsBatch.validatePreCommit(changeMCPs, this); + if (!exceptions.isEmpty()) { + throw new ValidationException(exceptions.toString()); + } // Database Upsert results List upsertResults = - updatedItems.getSecond().stream() + changeMCPs.stream() .map( item -> { - final String urnStr = item.getUrn().toString(); - final SystemAspect latest = - updatedLatestAspects - .getOrDefault(urnStr, Map.of()) - .get(item.getAspectName()); - final long nextVersion = - updatedNextVersions - .getOrDefault(urnStr, Map.of()) - .getOrDefault(item.getAspectName(), 0L); + final EntityAspect.EntitySystemAspect latest = + (EntityAspect.EntitySystemAspect) item.getPreviousSystemAspect(); final UpdateAspectResult result; if (overwrite || latest == null) { @@ -750,26 +748,17 @@ private List ingestAspectsToLocalDB( item.getRecordTemplate(), item.getAuditStamp(), item.getSystemMetadata(), - latest == null - ? null - : ((EntityAspect.EntitySystemAspect) latest).asRaw(), - nextVersion) + latest == null ? null : latest, + item.getNextAspectVersion()) .toBuilder() .request(item) .build(); - // support inner-batch upserts - latestAspects - .computeIfAbsent(urnStr, key -> new HashMap<>()) - .put(item.getAspectName(), item.toLatestEntityAspect()); - nextVersions - .computeIfAbsent(urnStr, key -> new HashMap<>()) - .put(item.getAspectName(), nextVersion + 1); } else { - RecordTemplate oldValue = latest.getRecordTemplate(_entityRegistry); + RecordTemplate oldValue = latest.getRecordTemplate(); SystemMetadata oldMetadata = latest.getSystemMetadata(); result = - UpdateAspectResult.builder() + UpdateAspectResult.builder() .urn(item.getUrn()) .request(item) .oldValue(oldValue) @@ -792,7 +781,7 @@ private List ingestAspectsToLocalDB( } // Retention optimization and tx - if (_retentionService != null) { + if (retentionService != null) { List retentionBatch = upsertResults.stream() // Only consider retention when there was a previous version @@ -810,7 +799,7 @@ private List ingestAspectsToLocalDB( // value return oldAspect != newAspect && oldAspect != null - && _retentionService != null; + && retentionService != null; }) .map( result -> @@ -820,7 +809,7 @@ private List ingestAspectsToLocalDB( .maxVersion(Optional.of(result.getMaxVersion())) .build()) .collect(Collectors.toList()); - _retentionService.applyRetentionWithPolicyDefaults(retentionBatch); + retentionService.applyRetentionWithPolicyDefaults(retentionBatch); } else { log.warn("Retention service is missing!"); } @@ -834,25 +823,6 @@ private List ingestAspectsToLocalDB( .collect(Collectors.toList()); } - /** - * Convert EntityAspect to EntitySystemAspect - * - * @param latestAspects latest aspect map - * @return map with converted values - */ - private static Map> toSystemEntityAspects( - Map> latestAspects) { - return latestAspects.entrySet().stream() - .map( - e -> - Map.entry( - e.getKey(), - e.getValue().entrySet().stream() - .map(e2 -> Map.entry(e2.getKey(), e2.getValue().asSystemAspect())) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - } - @Nonnull private List emitMCL(List sqlResults, boolean emitMCL) { List withEmitMCL = @@ -924,13 +894,14 @@ public RecordTemplate ingestAspectIfNotPresent( AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() .one( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(urn) .aspectName(aspectName) .recordTemplate(newValue) .systemMetadata(systemMetadata) .auditStamp(auditStamp) - .build(this)) + .build(this), + this) .build(); List ingested = ingestAspects(aspectsBatch, true, false); @@ -1007,17 +978,20 @@ private Stream ingestTimeseriesProposal( .filter(item -> item.getAspectSpec().isTimeseries()) .collect(Collectors.toList()); - List defaultAspects = + List defaultAspects = DefaultAspectsUtil.getAdditionalChanges( - AspectsBatchImpl.builder().items(timeseriesItems).build(), this, enableBrowseV2); - ingestProposalSync(AspectsBatchImpl.builder().items(defaultAspects).build()); + AspectsBatchImpl.builder().aspectRetriever(this).items(timeseriesItems).build(), + this, + enableBrowseV2); + ingestProposalSync( + AspectsBatchImpl.builder().aspectRetriever(this).items(defaultAspects).build()); } // Emit timeseries MCLs - List, Boolean>>>> timeseriesResults = + List, Boolean>>>> timeseriesResults = aspectsBatch.getItems().stream() .filter(item -> item.getAspectSpec().isTimeseries()) - .map(item -> (MCPUpsertBatchItem) item) + .map(item -> (ChangeItemImpl) item) .map( item -> Pair.of( @@ -1047,7 +1021,7 @@ private Stream ingestTimeseriesProposal( } }); - MCPUpsertBatchItem request = result.getFirst(); + ChangeItemImpl request = result.getFirst(); return IngestResult.builder() .urn(request.getUrn()) .request(request) @@ -1065,7 +1039,7 @@ private Stream ingestTimeseriesProposal( * @return produced items to the MCP topic */ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { - List nonTimeseries = + List nonTimeseries = aspectsBatch.getMCPItems().stream() .filter(item -> !item.getAspectSpec().isTimeseries()) .collect(Collectors.toList()); @@ -1075,7 +1049,7 @@ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { .map( item -> // When async is turned on, we write to proposal log and return without waiting - _producer.produceMetadataChangeProposal( + producer.produceMetadataChangeProposal( item.getUrn(), item.getMetadataChangeProposal())) .filter(Objects::nonNull) .collect(Collectors.toList()); @@ -1084,7 +1058,7 @@ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { return nonTimeseries.stream() .map( item -> - IngestResult.builder() + IngestResult.builder() .urn(item.getUrn()) .request(item) .publishedMCP(true) @@ -1104,13 +1078,14 @@ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { private Stream ingestProposalSync(AspectsBatch aspectsBatch) { AspectsBatchImpl nonTimeseries = AspectsBatchImpl.builder() + .aspectRetriever(this) .items( aspectsBatch.getItems().stream() .filter(item -> !item.getAspectSpec().isTimeseries()) .collect(Collectors.toList())) .build(); - List unsupported = + List unsupported = nonTimeseries.getMCPItems().stream() .filter( item -> @@ -1130,7 +1105,7 @@ private Stream ingestProposalSync(AspectsBatch aspectsBatch) { return upsertResults.stream() .map( result -> { - UpsertItem item = result.getRequest(); + ChangeMCP item = result.getRequest(); return IngestResult.builder() .urn(item.getUrn()) @@ -1160,18 +1135,20 @@ public String batchApplyRetention( args.attemptWithVersion = attemptWithVersion; args.aspectName = aspectName; args.urn = urn; - BulkApplyRetentionResult result = _retentionService.batchApplyRetentionEntities(args); + BulkApplyRetentionResult result = retentionService.batchApplyRetentionEntities(args); return result.toString(); } private boolean preprocessEvent(MetadataChangeLog metadataChangeLog) { - if (_preProcessHooks.isUiEnabled()) { + if (preProcessHooks.isUiEnabled()) { if (metadataChangeLog.getSystemMetadata() != null) { if (metadataChangeLog.getSystemMetadata().getProperties() != null) { if (UI_SOURCE.equals( metadataChangeLog.getSystemMetadata().getProperties().get(APP_SOURCE))) { // Pre-process the update indices hook for UI updates to avoid perceived lag from Kafka - _updateIndicesService.handleChangeEvent(metadataChangeLog); + if (updateIndicesService != null) { + updateIndicesService.handleChangeEvent(metadataChangeLog); + } return true; } } @@ -1182,7 +1159,7 @@ private boolean preprocessEvent(MetadataChangeLog metadataChangeLog) { @Override public Integer getCountAspect(@Nonnull String aspectName, @Nullable String urnLike) { - return _aspectDao.countAspect(aspectName, urnLike); + return aspectDao.countAspect(aspectName, urnLike); } @Nonnull @@ -1198,7 +1175,7 @@ public RestoreIndicesResult restoreIndices( "Reading rows %s through %s from the aspects table started.", args.start, args.start + args.batchSize)); long startTime = System.currentTimeMillis(); - PagedList rows = _aspectDao.getPagedAspects(args); + PagedList rows = aspectDao.getPagedAspects(args); result.timeSqlQueryMs = System.currentTimeMillis() - startTime; startTime = System.currentTimeMillis(); logger.accept( @@ -1208,19 +1185,23 @@ public RestoreIndicesResult restoreIndices( LinkedList> futures = new LinkedList<>(); - for (EbeanAspectV2 aspect : rows != null ? rows.getList() : List.of()) { + List systemAspects = + EntityUtils.toSystemAspectFromEbeanAspects( + rows != null ? rows.getList() : List.of(), this); + + for (SystemAspect aspect : systemAspects) { // 1. Extract an Entity type from the entity Urn result.timeGetRowMs = System.currentTimeMillis() - startTime; startTime = System.currentTimeMillis(); Urn urn; try { - urn = Urn.createFromString(aspect.getKey().getUrn()); + urn = aspect.getUrn(); result.lastUrn = urn.toString(); } catch (Exception e) { logger.accept( String.format( "Failed to bind Urn with value %s into Urn object: %s. Ignoring row.", - aspect.getKey().getUrn(), e)); + aspect.getUrn(), e)); ignored = ignored + 1; continue; } @@ -1231,7 +1212,7 @@ public RestoreIndicesResult restoreIndices( final String entityName = urn.getEntityType(); final EntitySpec entitySpec; try { - entitySpec = _entityRegistry.getEntitySpec(entityName); + entitySpec = entityRegistry.getEntitySpec(entityName); } catch (Exception e) { logger.accept( String.format( @@ -1242,7 +1223,7 @@ public RestoreIndicesResult restoreIndices( } result.timeEntityRegistryCheckMs += System.currentTimeMillis() - startTime; startTime = System.currentTimeMillis(); - final String aspectName = aspect.getKey().getAspect(); + final String aspectName = aspect.getAspectName(); result.lastAspect = aspectName; // 3. Verify that the aspect is a valid aspect associated with the entity @@ -1261,14 +1242,12 @@ public RestoreIndicesResult restoreIndices( // 4. Create record from json aspect final RecordTemplate aspectRecord; try { - aspectRecord = - EntityUtils.toAspectRecord( - entityName, aspectName, aspect.getMetadata(), _entityRegistry); + aspectRecord = aspect.getRecordTemplate(); } catch (Exception e) { logger.accept( String.format( - "Failed to deserialize row %s for entity %s, aspect %s: %s. Ignoring row.", - aspect.getMetadata(), entityName, aspectName, e)); + "Failed to deserialize for entity %s, aspect %s: %s. Ignoring row.", + entityName, aspectName, e)); ignored = ignored + 1; continue; } @@ -1276,8 +1255,8 @@ public RestoreIndicesResult restoreIndices( startTime = System.currentTimeMillis(); // Force indexing to skip diff mode and fix error states - SystemMetadata latestSystemMetadata = - EntityUtils.parseSystemMetadata(aspect.getSystemMetadata()); + SystemMetadata latestSystemMetadata = aspect.getSystemMetadata(); + StringMap properties = latestSystemMetadata.getProperties() != null ? latestSystemMetadata.getProperties() @@ -1343,7 +1322,7 @@ public ListUrnsResult listUrns( final String keyAspectName = getEntityRegistry().getEntitySpec(entityName).getKeyAspectSpec().getName(); final ListResult keyAspectList = - _aspectDao.listUrns(entityName, keyAspectName, start, count); + aspectDao.listUrns(entityName, keyAspectName, start, count); final ListUrnsResult result = new ListUrnsResult(); result.setStart(start); @@ -1393,7 +1372,8 @@ public Map getEntities( return Collections.emptyMap(); } return getSnapshotUnions(urns, aspectNames).entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> toEntity(entry.getValue()))); + .collect( + Collectors.toMap(Map.Entry::getKey, entry -> EntityUtils.toEntity(entry.getValue()))); } @Override @@ -1401,7 +1381,7 @@ public Pair, Boolean> alwaysProduceMCLAsync( @Nonnull final Urn urn, @Nonnull final AspectSpec aspectSpec, @Nonnull final MetadataChangeLog metadataChangeLog) { - Future future = _producer.produceMetadataChangeLog(urn, aspectSpec, metadataChangeLog); + Future future = producer.produceMetadataChangeLog(urn, aspectSpec, metadataChangeLog); return Pair.of(future, preprocessEvent(metadataChangeLog)); } @@ -1442,7 +1422,7 @@ public Optional, Boolean>> conditionallyProduceMCLAsync( AuditStamp auditStamp, AspectSpec aspectSpec) { boolean isNoOp = oldAspect == newAspect; - if (!isNoOp || _alwaysEmitChangeLog || shouldAspectEmitChangeLog(aspectSpec)) { + if (!isNoOp || alwaysEmitChangeLog || shouldAspectEmitChangeLog(aspectSpec)) { log.debug( "Producing MetadataChangeLog for ingested aspect {}, urn {}", aspectSpec.getName(), @@ -1475,7 +1455,7 @@ public Optional, Boolean>> conditionallyProduceMCLAsync( } private UpdateAspectResult conditionallyProduceMCLAsync(UpdateAspectResult result) { - UpsertItem request = result.getRequest(); + ChangeMCP request = result.getRequest(); Optional, Boolean>> emissionStatus = conditionallyProduceMCLAsync( result.getOldValue(), @@ -1537,7 +1517,9 @@ public void ingestEntity( protected Map getSnapshotUnions( @Nonnull final Set urns, @Nonnull final Set aspectNames) { return getSnapshotRecords(urns, aspectNames).entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> toSnapshotUnion(entry.getValue()))); + .collect( + Collectors.toMap( + Map.Entry::getKey, entry -> EntityUtils.toSnapshotUnion(entry.getValue()))); } @Nonnull @@ -1582,11 +1564,12 @@ private void ingestSnapshotUnion( AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() + .aspectRetriever(this) .items( aspectRecordsToIngest.stream() .map( pair -> - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(urn) .aspectName(pair.getKey()) .recordTemplate(pair.getValue()) @@ -1606,38 +1589,28 @@ public AspectSpec getKeyAspectSpec(@Nonnull final Urn urn) { @Override public AspectSpec getKeyAspectSpec(@Nonnull final String entityName) { - final EntitySpec spec = _entityRegistry.getEntitySpec(entityName); + final EntitySpec spec = entityRegistry.getEntitySpec(entityName); return spec.getKeyAspectSpec(); } @Override public Optional getAspectSpec( @Nonnull final String entityName, @Nonnull final String aspectName) { - final EntitySpec entitySpec = _entityRegistry.getEntitySpec(entityName); + final EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); return Optional.ofNullable(entitySpec.getAspectSpec(aspectName)); } @Override public String getKeyAspectName(@Nonnull final Urn urn) { - final EntitySpec spec = _entityRegistry.getEntitySpec(urnToEntityName(urn)); + final EntitySpec spec = entityRegistry.getEntitySpec(urnToEntityName(urn)); final AspectSpec keySpec = spec.getKeyAspectSpec(); return keySpec.getName(); } - protected Entity toEntity(@Nonnull final Snapshot snapshot) { - return new Entity().setValue(snapshot); - } - - protected Snapshot toSnapshotUnion(@Nonnull final RecordTemplate snapshotRecord) { - final Snapshot snapshot = new Snapshot(); - RecordUtils.setSelectedRecordTemplateInUnion(snapshot, snapshotRecord); - return snapshot; - } - protected RecordTemplate toSnapshotRecord( @Nonnull final Urn urn, @Nonnull final List aspectUnionTemplates) { final String entityName = urnToEntityName(urn); - final EntitySpec entitySpec = _entityRegistry.getEntitySpec(entityName); + final EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); return com.datahub.util.ModelUtils.newSnapshot( getDataTemplateClassFromSchema(entitySpec.getSnapshotSchema(), RecordTemplate.class), urn, @@ -1646,7 +1619,7 @@ protected RecordTemplate toSnapshotRecord( protected UnionTemplate toAspectUnion( @Nonnull final Urn urn, @Nonnull final RecordTemplate aspectRecord) { - final EntitySpec entitySpec = _entityRegistry.getEntitySpec(urnToEntityName(urn)); + final EntitySpec entitySpec = entityRegistry.getEntitySpec(urnToEntityName(urn)); final TyperefDataSchema aspectSchema = entitySpec.getAspectTyperefSchema(); if (aspectSchema == null) { throw new RuntimeException( @@ -1659,49 +1632,15 @@ protected UnionTemplate toAspectUnion( aspectRecord); } - protected Urn toUrn(final String urnStr) { - try { - return Urn.createFromString(urnStr); - } catch (URISyntaxException e) { - log.error("Failed to convert urn string {} into Urn object", urnStr); - throw new ModelConversionException( - String.format("Failed to convert urn string %s into Urn object ", urnStr), e); - } - } - - private EntityResponse toEntityResponse( - final Urn urn, final List envelopedAspects) { - final EntityResponse response = new EntityResponse(); - response.setUrn(urn); - response.setEntityName(urnToEntityName(urn)); - response.setAspects( - new EnvelopedAspectMap( - envelopedAspects.stream() - .collect(Collectors.toMap(EnvelopedAspect::getName, aspect -> aspect)))); - return response; - } - - private static Map> buildEntityToValidAspects( - final EntityRegistry entityRegistry) { - return entityRegistry.getEntitySpecs().values().stream() - .collect( - Collectors.toMap( - EntitySpec::getName, - entry -> - entry.getAspectSpecs().stream() - .map(AspectSpec::getName) - .collect(Collectors.toSet()))); - } - @Override @Nonnull public EntityRegistry getEntityRegistry() { - return _entityRegistry; + return entityRegistry; } @Override - public void setRetentionService(RetentionService retentionService) { - _retentionService = retentionService; + public void setRetentionService(RetentionService retentionService) { + this.retentionService = retentionService; } protected Set getEntityAspectNames(final Urn entityUrn) { @@ -1710,13 +1649,13 @@ protected Set getEntityAspectNames(final Urn entityUrn) { @Override public Set getEntityAspectNames(final String entityName) { - return _entityToValidAspects.get(entityName); + return entityToValidAspects.get(entityName); } @Override public void setWritable(boolean canWrite) { log.debug("Setting writable to {}", canWrite); - _aspectDao.setWritable(canWrite); + aspectDao.setWritable(canWrite); } @Override @@ -1797,7 +1736,7 @@ public RollbackRunResult deleteUrn(Urn urn) { EntityAspect latestKey = null; try { - latestKey = _aspectDao.getLatestAspect(urn.toString(), keyAspectName); + latestKey = aspectDao.getLatestAspect(urn.toString(), keyAspectName); } catch (EntityNotFoundException e) { log.warn("Entity to delete does not exist. {}", urn.toString()); } @@ -1805,7 +1744,8 @@ public RollbackRunResult deleteUrn(Urn urn) { return new RollbackRunResult(removedAspects, rowsDeletedFromEntityDeletion); } - SystemMetadata latestKeySystemMetadata = latestKey.asSystemAspect().getSystemMetadata(); + SystemMetadata latestKeySystemMetadata = + EntityUtils.toSystemAspect(latestKey, this).map(SystemAspect::getSystemMetadata).get(); RollbackResult result = deleteAspect( urn.toString(), @@ -1850,29 +1790,26 @@ public RollbackRunResult deleteUrn(Urn urn) { return new RollbackRunResult(removedAspects, rowsDeletedFromEntityDeletion); } - /** - * Returns a set of urns of entities that exist (has materialized aspects). - * - * @param urns the list of urns of the entities to check - * @param includeSoftDeleted whether to consider soft delete - * @return a set of urns of entities that exist. - */ @Override - public Set exists(@Nonnull final Collection urns, boolean includeSoftDeleted) { - + public Set exists( + @Nonnull final Collection urns, + @Nullable String aspectName, + boolean includeSoftDeleted) { final Set dbKeys = urns.stream() .map( urn -> new EntityAspectIdentifier( urn.toString(), - _entityRegistry - .getEntitySpec(urn.getEntityType()) - .getKeyAspectSpec() - .getName(), + aspectName == null + ? entityRegistry + .getEntitySpec(urn.getEntityType()) + .getKeyAspectSpec() + .getName() + : aspectName, ASPECT_LATEST_VERSION)) .collect(Collectors.toSet()); - final Map aspects = _aspectDao.batchGet(dbKeys); + final Map aspects = aspectDao.batchGet(dbKeys); final Set existingUrnStrings = aspects.values().stream() .filter(aspect -> aspect != null) @@ -1901,37 +1838,43 @@ public Set exists(@Nonnull final Collection urns, boolean includeSoftD } } - @Override - public Boolean exists(Urn urn, String aspectName) { - EntityAspectIdentifier dbKey = - new EntityAspectIdentifier(urn.toString(), aspectName, ASPECT_LATEST_VERSION); - Map aspects = _aspectDao.batchGet(Set.of(dbKey)); - return aspects.values().stream().anyMatch(Objects::nonNull); - } - @Nullable @Override public RollbackResult deleteAspect( String urn, String aspectName, @Nonnull Map conditions, boolean hardDelete) { + final AuditStamp auditStamp = + new AuditStamp() + .setActor(UrnUtils.getUrn(Constants.SYSTEM_ACTOR)) + .setTime(System.currentTimeMillis()); + // Validate pre-conditions before running queries - Urn entityUrn; - EntitySpec entitySpec; - try { - entityUrn = Urn.createFromString(urn); - String entityName = PegasusUtils.urnToEntityName(entityUrn); - entitySpec = getEntityRegistry().getEntitySpec(entityName); - } catch (URISyntaxException uriSyntaxException) { - // don't expect this to happen, so raising RuntimeException here - throw new RuntimeException(String.format("Failed to extract urn from %s", urn)); + Urn entityUrn = UrnUtils.getUrn(urn); + + // Runs simple validations + MCPItem deleteItem = + DeleteItemImpl.builder() + .urn(entityUrn) + .aspectName(aspectName) + .auditStamp(auditStamp) + .build(this); + + // Delete validation hooks + ValidationExceptionCollection exceptions = + AspectsBatch.validateProposed(List.of(deleteItem), this); + if (!exceptions.isEmpty()) { + throw new ValidationException(exceptions.toString()); } final RollbackResult result = - _aspectDao.runInTransactionWithRetry( + aspectDao.runInTransactionWithRetry( (tx) -> { Integer additionalRowsDeleted = 0; // 1. Fetch the latest existing version of the aspect. - final EntityAspect latest = _aspectDao.getLatestAspect(urn, aspectName); + final EntityAspect.EntitySystemAspect latest = + (EntityAspect.EntitySystemAspect) + EntityUtils.toSystemAspect(aspectDao.getLatestAspect(urn, aspectName), this) + .orElse(null); // 1.1 If no latest exists, skip this aspect if (latest == null) { @@ -1939,63 +1882,86 @@ public RollbackResult deleteAspect( } // 2. Compare the match conditions, if they don't match, ignore. - SystemMetadata latestSystemMetadata = latest.asSystemAspect().getSystemMetadata(); + SystemMetadata latestSystemMetadata = latest.getSystemMetadata(); if (!filterMatch(latestSystemMetadata, conditions)) { return null; } - String latestMetadata = latest.getMetadata(); // 3. Check if this is a key aspect Boolean isKeyAspect = getKeyAspectName(entityUrn).equals(aspectName); // 4. Fetch all preceding aspects, that match List aspectsToDelete = new ArrayList<>(); - long maxVersion = _aspectDao.getMaxVersion(urn, aspectName); - EntityAspect survivingAspect = null; + long maxVersion = aspectDao.getMaxVersion(urn, aspectName); + EntityAspect.EntitySystemAspect survivingAspect = null; String previousMetadata = null; boolean filterMatch = true; while (maxVersion > 0 && filterMatch) { - EntityAspect candidateAspect = _aspectDao.getAspect(urn, aspectName, maxVersion); + EntityAspect.EntitySystemAspect candidateAspect = + (EntityAspect.EntitySystemAspect) + EntityUtils.toSystemAspect( + aspectDao.getAspect(urn, aspectName, maxVersion), this) + .orElse(null); SystemMetadata previousSysMetadata = - candidateAspect != null - ? candidateAspect.asSystemAspect().getSystemMetadata() - : null; + candidateAspect != null ? candidateAspect.getSystemMetadata() : null; filterMatch = previousSysMetadata != null && filterMatch(previousSysMetadata, conditions); if (filterMatch) { - aspectsToDelete.add(candidateAspect); + aspectsToDelete.add(candidateAspect.getEntityAspect()); maxVersion = maxVersion - 1; } else { survivingAspect = candidateAspect; - previousMetadata = survivingAspect.getMetadata(); + previousMetadata = survivingAspect.getMetadataRaw(); } } - // 5. Apply deletes and fix up latest row + // Delete validation hooks + ValidationExceptionCollection preCommitExceptions = + AspectsBatch.validatePreCommit( + aspectsToDelete.stream() + .map( + toDelete -> + DeleteItemImpl.builder() + .urn(UrnUtils.getUrn(toDelete.getUrn())) + .aspectName(toDelete.getAspect()) + .auditStamp(auditStamp) + .build(this)) + .collect(Collectors.toList()), + this); + if (!preCommitExceptions.isEmpty()) { + throw new ValidationException(preCommitExceptions.toString()); + } - aspectsToDelete.forEach(aspect -> _aspectDao.deleteAspect(tx, aspect)); + // 5. Apply deletes and fix up latest row + aspectsToDelete.forEach(aspect -> aspectDao.deleteAspect(tx, aspect)); if (survivingAspect != null) { // if there was a surviving aspect, copy its information into the latest row // eBean does not like us updating a pkey column (version) for the surviving aspect // as a result we copy information from survivingAspect to latest and delete // survivingAspect - latest.setMetadata(survivingAspect.getMetadata()); - latest.setSystemMetadata(survivingAspect.getSystemMetadata()); - latest.setCreatedOn(survivingAspect.getCreatedOn()); - latest.setCreatedBy(survivingAspect.getCreatedBy()); - latest.setCreatedFor(survivingAspect.getCreatedFor()); - _aspectDao.saveAspect(tx, latest, false); + latest + .getEntityAspect() + .setMetadata(survivingAspect.getEntityAspect().getMetadata()); + latest + .getEntityAspect() + .setSystemMetadata(survivingAspect.getEntityAspect().getSystemMetadata()); + latest.getEntityAspect().setCreatedOn(survivingAspect.getCreatedOn()); + latest.getEntityAspect().setCreatedBy(survivingAspect.getCreatedBy()); + latest + .getEntityAspect() + .setCreatedFor(survivingAspect.getEntityAspect().getCreatedFor()); + aspectDao.saveAspect(tx, latest.getEntityAspect(), false); // metrics - _aspectDao.incrementWriteMetrics( - aspectName, 1, latest.getAspect().getBytes(StandardCharsets.UTF_8).length); - _aspectDao.deleteAspect(tx, survivingAspect); + aspectDao.incrementWriteMetrics( + aspectName, 1, latest.getMetadataRaw().getBytes(StandardCharsets.UTF_8).length); + aspectDao.deleteAspect(tx, survivingAspect.getEntityAspect()); } else { if (isKeyAspect) { if (hardDelete) { // If this is the key aspect, delete the entity entirely. - additionalRowsDeleted = _aspectDao.deleteUrn(tx, urn); - } else if (entitySpec.hasAspect(Constants.STATUS_ASPECT_NAME)) { + additionalRowsDeleted = aspectDao.deleteUrn(tx, urn); + } else if (deleteItem.getEntitySpec().hasAspect(Constants.STATUS_ASPECT_NAME)) { // soft delete by setting status.removed=true (if applicable) final Status statusAspect = new Status(); statusAspect.setRemoved(true); @@ -2006,38 +1972,21 @@ public RollbackResult deleteAspect( gmce.setEntityType(entityUrn.getEntityType()); gmce.setAspectName(Constants.STATUS_ASPECT_NAME); gmce.setAspect(GenericRecordUtils.serializeAspect(statusAspect)); - final AuditStamp auditStamp = - new AuditStamp() - .setActor(UrnUtils.getUrn(Constants.SYSTEM_ACTOR)) - .setTime(System.currentTimeMillis()); this.ingestProposal(gmce, auditStamp, false); } } else { // Else, only delete the specific aspect. - _aspectDao.deleteAspect(tx, latest); + aspectDao.deleteAspect(tx, latest.getEntityAspect()); } } // 6. Emit the Update try { final RecordTemplate latestValue = - latest == null - ? null - : EntityUtils.toAspectRecord( - entitySpec.getName(), - latest.getAspect(), - latestMetadata, - getEntityRegistry()); - + latest == null ? null : latest.getRecordTemplate(); final RecordTemplate previousValue = - survivingAspect == null - ? null - : EntityUtils.toAspectRecord( - entitySpec.getName(), - survivingAspect.getAspect(), - previousMetadata, - getEntityRegistry()); + survivingAspect == null ? null : latest.getRecordTemplate(); final Urn urnObj = Urn.createFromString(urn); // We are not deleting key aspect if hardDelete has not been set so do not return a @@ -2048,13 +1997,11 @@ public RollbackResult deleteAspect( return new RollbackResult( urnObj, urnObj.getEntityType(), - latest.getAspect(), + latest.getAspectName(), latestValue, previousValue, latestSystemMetadata, - previousValue == null - ? null - : survivingAspect.asSystemAspect().getSystemMetadata(), + previousValue == null ? null : survivingAspect.getSystemMetadata(), survivingAspect == null ? ChangeType.DELETE : ChangeType.UPSERT, isKeyAspect, additionalRowsDeleted); @@ -2128,7 +2075,7 @@ private Map getLatestAspect( Map batchGetResults = new HashMap<>(); Iterators.partition(dbKeys.iterator(), MAX_KEYS_PER_QUERY) .forEachRemaining( - batch -> batchGetResults.putAll(_aspectDao.batchGet(ImmutableSet.copyOf(batch)))); + batch -> batchGetResults.putAll(aspectDao.batchGet(ImmutableSet.copyOf(batch)))); return batchGetResults; } @@ -2140,83 +2087,24 @@ private Map getLatestAspect( private long calculateVersionNumber( @Nonnull final Urn urn, @Nonnull final String aspectName, @Nonnull long version) { if (version < 0) { - return _aspectDao.getMaxVersion(urn.toString(), aspectName) + version + 1; + return aspectDao.getMaxVersion(urn.toString(), aspectName) + version + 1; } return version; } private Map getEnvelopedAspects( final Set dbKeys) { - final Map result = new HashMap<>(); - final Map dbEntries = _aspectDao.batchGet(dbKeys); - - for (EntityAspectIdentifier currKey : dbKeys) { - - final EntityAspect currAspectEntry = dbEntries.get(currKey); - - if (currAspectEntry == null) { - // No aspect found. - continue; - } + final Map dbEntries = aspectDao.batchGet(dbKeys); - result.put(currKey, toEnvelopedAspect(currAspectEntry)); - } - return result; - } - - private static EnvelopedAspect toEnvelopedAspect(EntityAspect entityAspect) { - // Aspect found. Now turn it into an EnvelopedAspect - final com.linkedin.entity.Aspect aspect = - RecordUtils.toRecordTemplate(com.linkedin.entity.Aspect.class, entityAspect.getMetadata()); - final EnvelopedAspect envelopedAspect = new EnvelopedAspect(); - envelopedAspect.setName(entityAspect.getAspect()); - envelopedAspect.setVersion(entityAspect.getVersion()); - // TODO: I think we can assume this here, adding as it's a required field so object mapping - // barfs when trying to access it, - // since nowhere else is using it should be safe for now at least - envelopedAspect.setType(AspectType.VERSIONED); - envelopedAspect.setValue(aspect); + List envelopedAspects = EntityUtils.toSystemAspects(dbEntries.values(), this); - try { - if (entityAspect.getSystemMetadata() != null) { - final SystemMetadata systemMetadata = entityAspect.asSystemAspect().getSystemMetadata(); - envelopedAspect.setSystemMetadata(systemMetadata); - } - } catch (Exception e) { - log.warn( - "Exception encountered when setting system metadata on enveloped aspect {}. Error: {}", - envelopedAspect.getName(), - e.toString()); - } - - envelopedAspect.setCreated( - new AuditStamp() - .setActor(UrnUtils.getUrn(entityAspect.getCreatedBy())) - .setTime(entityAspect.getCreatedOn().getTime())); - - return envelopedAspect; - } - - private EnvelopedAspect getKeyEnvelopedAspect(final Urn urn) { - final EntitySpec spec = getEntityRegistry().getEntitySpec(PegasusUtils.urnToEntityName(urn)); - final AspectSpec keySpec = spec.getKeyAspectSpec(); - final com.linkedin.entity.Aspect aspect = - new com.linkedin.entity.Aspect(EntityKeyUtils.convertUrnToEntityKey(urn, keySpec).data()); - - final EnvelopedAspect envelopedAspect = new EnvelopedAspect(); - envelopedAspect.setName(keySpec.getName()); - envelopedAspect.setVersion(ASPECT_LATEST_VERSION); - envelopedAspect.setValue(aspect); - // TODO: I think we can assume this here, adding as it's a required field so object mapping - // barfs when trying to access it, - // since nowhere else is using it should be safe for now at least - envelopedAspect.setType(AspectType.VERSIONED); - envelopedAspect.setCreated( - new AuditStamp() - .setActor(UrnUtils.getUrn(SYSTEM_ACTOR)) - .setTime(System.currentTimeMillis())); - - return envelopedAspect; + return envelopedAspects.stream() + .collect( + Collectors.toMap( + systemAspect -> + ((EntityAspect.EntitySystemAspect) systemAspect).getAspectIdentifier(), + systemAspect -> + ((EntityAspect.EntitySystemAspect) systemAspect).toEnvelopedAspects())); } @Nonnull @@ -2227,7 +2115,7 @@ private UpdateAspectResult ingestAspectToLocalDB( @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp, @Nonnull final SystemMetadata providedSystemMetadata, - @Nullable final EntityAspect latest, + @Nullable final EntityAspect.EntitySystemAspect latest, @Nonnull final Long nextVersion) { // Set the "last run id" to be the run id provided with the new system metadata. This will be @@ -2237,34 +2125,30 @@ private UpdateAspectResult ingestAspectToLocalDB( providedSystemMetadata.getRunId(GetMode.NULL), SetMode.IGNORE_NULL); // 2. Compare the latest existing and new. - final RecordTemplate oldValue = - latest == null - ? null - : EntityUtils.toAspectRecord( - urn, aspectName, latest.getMetadata(), getEntityRegistry()); + final RecordTemplate oldValue = latest == null ? null : latest.getRecordTemplate(); // 3. If there is no difference between existing and new, we just update // the lastObserved in system metadata. RunId should stay as the original runId if (oldValue != null && DataTemplateUtil.areEqual(oldValue, newValue)) { - SystemMetadata latestSystemMetadata = latest.asSystemAspect().getSystemMetadata(); + SystemMetadata latestSystemMetadata = latest.getSystemMetadata(); latestSystemMetadata.setLastObserved(providedSystemMetadata.getLastObserved()); latestSystemMetadata.setLastRunId( providedSystemMetadata.getLastRunId(GetMode.NULL), SetMode.IGNORE_NULL); - latest.setSystemMetadata(RecordUtils.toJsonString(latestSystemMetadata)); + latest.getEntityAspect().setSystemMetadata(RecordUtils.toJsonString(latestSystemMetadata)); log.info("Ingesting aspect with name {}, urn {}", aspectName, urn); - _aspectDao.saveAspect(tx, latest, false); + aspectDao.saveAspect(tx, latest.getEntityAspect(), false); // metrics - _aspectDao.incrementWriteMetrics( - aspectName, 1, latest.getAspect().getBytes(StandardCharsets.UTF_8).length); + aspectDao.incrementWriteMetrics( + aspectName, 1, latest.getMetadataRaw().getBytes(StandardCharsets.UTF_8).length); return UpdateAspectResult.builder() .urn(urn) .oldValue(oldValue) .newValue(oldValue) - .oldSystemMetadata(latest.asSystemAspect().getSystemMetadata()) + .oldSystemMetadata(latest.getSystemMetadata()) .newSystemMetadata(latestSystemMetadata) .operation(MetadataAuditOperation.UPDATE) .auditStamp(auditStamp) @@ -2276,15 +2160,15 @@ private UpdateAspectResult ingestAspectToLocalDB( log.debug("Ingesting aspect with name {}, urn {}", aspectName, urn); String newValueStr = EntityUtils.toJsonAspect(newValue); long versionOfOld = - _aspectDao.saveLatestAspect( + aspectDao.saveLatestAspect( tx, urn.toString(), aspectName, latest == null ? null : EntityUtils.toJsonAspect(oldValue), latest == null ? null : latest.getCreatedBy(), - latest == null ? null : latest.getCreatedFor(), + latest == null ? null : latest.getEntityAspect().getCreatedFor(), latest == null ? null : latest.getCreatedOn(), - latest == null ? null : latest.getSystemMetadata(), + latest == null ? null : latest.getSystemMetadataRaw(), newValueStr, auditStamp.getActor().toString(), auditStamp.hasImpersonator() ? auditStamp.getImpersonator().toString() : null, @@ -2293,14 +2177,14 @@ private UpdateAspectResult ingestAspectToLocalDB( nextVersion); // metrics - _aspectDao.incrementWriteMetrics( + aspectDao.incrementWriteMetrics( aspectName, 1, newValueStr.getBytes(StandardCharsets.UTF_8).length); return UpdateAspectResult.builder() .urn(urn) .oldValue(oldValue) .newValue(newValue) - .oldSystemMetadata(latest == null ? null : latest.asSystemAspect().getSystemMetadata()) + .oldSystemMetadata(latest == null ? null : latest.getSystemMetadata()) .newSystemMetadata(providedSystemMetadata) .operation(MetadataAuditOperation.UPDATE) .auditStamp(auditStamp) @@ -2318,7 +2202,11 @@ private static boolean shouldAspectEmitChangeLog(@Nonnull final AspectSpec aspec @Override public Map> getLatestAspectObjects( Set urns, Set aspectNames) throws RemoteInvocationException, URISyntaxException { - String entityName = urns.stream().findFirst().map(Urn::getEntityType).get(); - return entityResponseToAspectMap(getEntitiesV2(entityName, urns, aspectNames)); + if (urns.isEmpty() || aspectNames.isEmpty()) { + return Map.of(); + } else { + String entityName = urns.stream().findFirst().map(Urn::getEntityType).get(); + return entityResponseToAspectMap(getEntitiesV2(entityName, urns, aspectNames)); + } } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java index f353e5142755d1..fecf246b31c02b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityUtils.java @@ -7,21 +7,38 @@ import com.google.common.base.Preconditions; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; -import com.linkedin.data.schema.RecordDataSchema; +import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.template.RecordTemplate; +import com.linkedin.entity.AspectType; +import com.linkedin.entity.Entity; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.ReadItem; +import com.linkedin.metadata.aspect.SystemAspect; +import com.linkedin.metadata.aspect.batch.AspectsBatch; +import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.validation.EntityRegistryUrnValidator; import com.linkedin.metadata.entity.validation.RecordTemplateValidator; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.snapshot.Snapshot; import com.linkedin.metadata.utils.EntityKeyUtils; +import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.metadata.utils.PegasusUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; import java.net.URISyntaxException; -import java.net.URLEncoder; +import java.util.Collection; import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; @@ -31,9 +48,6 @@ public class EntityUtils { private EntityUtils() {} - public static final int URN_NUM_BYTES_LIMIT = 512; - public static final String URN_DELIMITER_SEPARATOR = "␟"; - @Nonnull public static String toJsonAspect(@Nonnull final RecordTemplate aspectRecord) { return RecordUtils.toJsonString(aspectRecord); @@ -101,45 +115,205 @@ public static RecordTemplate getAspectFromEntity( } } + public static RecordTemplate buildKeyAspect( + @Nonnull EntityRegistry entityRegistry, @Nonnull final Urn urn) { + final EntitySpec spec = entityRegistry.getEntitySpec(urnToEntityName(urn)); + final AspectSpec keySpec = spec.getKeyAspectSpec(); + return EntityKeyUtils.convertUrnToEntityKey(urn, keySpec); + } + + static Entity toEntity(@Nonnull final Snapshot snapshot) { + return new Entity().setValue(snapshot); + } + + static Snapshot toSnapshotUnion(@Nonnull final RecordTemplate snapshotRecord) { + final Snapshot snapshot = new Snapshot(); + RecordUtils.setSelectedRecordTemplateInUnion(snapshot, snapshotRecord); + return snapshot; + } + + static EnvelopedAspect getKeyEnvelopedAspect(final Urn urn, final EntityRegistry entityRegistry) { + final EntitySpec spec = entityRegistry.getEntitySpec(PegasusUtils.urnToEntityName(urn)); + final AspectSpec keySpec = spec.getKeyAspectSpec(); + final com.linkedin.entity.Aspect aspect = + new com.linkedin.entity.Aspect(EntityKeyUtils.convertUrnToEntityKey(urn, keySpec).data()); + + final EnvelopedAspect envelopedAspect = new EnvelopedAspect(); + envelopedAspect.setName(keySpec.getName()); + envelopedAspect.setVersion(ASPECT_LATEST_VERSION); + envelopedAspect.setValue(aspect); + // TODO: I think we can assume this here, adding as it's a required field so object mapping + // barfs when trying to access it, + // since nowhere else is using it should be safe for now at least + envelopedAspect.setType(AspectType.VERSIONED); + envelopedAspect.setCreated( + new AuditStamp() + .setActor(UrnUtils.getUrn(SYSTEM_ACTOR)) + .setTime(System.currentTimeMillis())); + + return envelopedAspect; + } + + static EntityResponse toEntityResponse( + final Urn urn, final List envelopedAspects) { + final EntityResponse response = new EntityResponse(); + response.setUrn(urn); + response.setEntityName(urnToEntityName(urn)); + response.setAspects( + new EnvelopedAspectMap( + envelopedAspects.stream() + .collect(Collectors.toMap(EnvelopedAspect::getName, aspect -> aspect)))); + return response; + } + + static Map> buildEntityToValidAspects(final EntityRegistry entityRegistry) { + return entityRegistry.getEntitySpecs().values().stream() + .collect( + Collectors.toMap( + EntitySpec::getName, + entry -> + entry.getAspectSpecs().stream() + .map(AspectSpec::getName) + .collect(Collectors.toSet()))); + } + + /** + * Prefer batched interfaces + * + * @param entityAspect optional entity aspect + * @param aspectRetriever + * @return + */ + public static Optional toSystemAspect( + @Nullable EntityAspect entityAspect, @Nonnull AspectRetriever aspectRetriever) { + return Optional.ofNullable(entityAspect) + .map(aspect -> EntityUtils.toSystemAspects(List.of(aspect), aspectRetriever)) + .filter(systemAspects -> !systemAspects.isEmpty()) + .map(systemAspects -> systemAspects.get(0)); + } + + /** + * Given a `Map>` from the database representation, + * translate that into our java classes + * + * @param rawAspects `Map>` + * @param aspectRetriever used for read mutations + * @return the java map for the given database object map + */ + @Nonnull + public static Map> toSystemAspects( + @Nonnull Map> rawAspects, + @Nonnull AspectRetriever aspectRetriever) { + List systemAspects = + toSystemAspects( + rawAspects.values().stream() + .flatMap(m -> m.values().stream()) + .collect(Collectors.toList()), + aspectRetriever); + + // map the list into the desired shape + return systemAspects.stream() + .collect(Collectors.groupingBy(SystemAspect::getUrn)) + .entrySet() + .stream() + .map( + entry -> + Pair.of( + entry.getKey(), + entry.getValue().stream() + .collect(Collectors.groupingBy(SystemAspect::getAspectName)))) + .collect( + Collectors.toMap( + p -> p.getFirst().toString(), + p -> + p.getSecond().entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().get(0))))); + } + @Nonnull - public static RecordTemplate toAspectRecord( - @Nonnull final Urn entityUrn, - @Nonnull final String aspectName, - @Nonnull final String jsonAspect, - @Nonnull final EntityRegistry entityRegistry) { - return toAspectRecord( - PegasusUtils.urnToEntityName(entityUrn), aspectName, jsonAspect, entityRegistry); + public static List toSystemAspectFromEbeanAspects( + @Nonnull Collection rawAspects, @Nonnull AspectRetriever aspectRetriever) { + return toSystemAspects( + rawAspects.stream().map(EbeanAspectV2::toEntityAspect).collect(Collectors.toList()), + aspectRetriever); } /** - * @param entityName - * @param aspectName - * @param jsonAspect - * @param entityRegistry - * @return a RecordTemplate which has been validated, validation errors are logged as warnings + * Convert EntityAspect to EntitySystemAspect + * + *

This should be the 1 point that all conversions from database representations to java + * objects happens since we need to enforce read mutations happen. + * + * @param rawAspects raw aspects to convert + * @return map converted aspects */ - public static RecordTemplate toAspectRecord( - @Nonnull final String entityName, - @Nonnull final String aspectName, - @Nonnull final String jsonAspect, - @Nonnull final EntityRegistry entityRegistry) { - final EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); - final AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); - // TODO: aspectSpec can be null here - Preconditions.checkState( - aspectSpec != null, String.format("Aspect %s could not be found", aspectName)); - final RecordDataSchema aspectSchema = aspectSpec.getPegasusSchema(); - RecordTemplate aspectRecord = - RecordUtils.toRecordTemplate(aspectSpec.getDataTemplateClass(), jsonAspect); - RecordTemplateValidator.validate( - aspectRecord, - validationFailure -> { - log.warn(String.format("Failed to validate record %s against its schema.", aspectRecord)); + @Nonnull + public static List toSystemAspects( + @Nonnull Collection rawAspects, @Nonnull AspectRetriever aspectRetriever) { + EntityRegistry entityRegistry = aspectRetriever.getEntityRegistry(); + + // Build + List systemAspects = + rawAspects.stream() + .map( + raw -> { + Urn urn = UrnUtils.getUrn(raw.getUrn()); + AspectSpec aspectSpec = + entityRegistry + .getEntitySpec(urn.getEntityType()) + .getAspectSpec(raw.getAspect()); + + // TODO: aspectSpec can be null here + Preconditions.checkState( + aspectSpec != null, + String.format("Aspect %s could not be found", raw.getAspect())); + + return EntityAspect.EntitySystemAspect.builder() + .build(entityRegistry.getEntitySpec(urn.getEntityType()), aspectSpec, raw); + }) + .collect(Collectors.toList()); + + // Read Mutate + Map, List> grouped = + systemAspects.stream() + .collect( + Collectors.groupingBy(item -> Pair.of(item.getEntitySpec(), item.getAspectSpec()))); + + grouped.forEach( + (key, value) -> { + AspectsBatch.applyReadMutationHooks(value, aspectRetriever); }); - return aspectRecord; + + // Read Validate + systemAspects.forEach( + systemAspect -> + RecordTemplateValidator.validate( + systemAspect.getRecordTemplate(), + validationFailure -> + log.warn( + String.format( + "Failed to validate record %s against its schema.", + systemAspect.getRecordTemplate())))); + + // TODO consider applying write validation plugins + + return systemAspects; + } + + public static MetadataChangeProposal buildMCP( + Urn entityUrn, String aspectName, ChangeType changeType, @Nullable T aspect) { + MetadataChangeProposal proposal = new MetadataChangeProposal(); + proposal.setEntityUrn(entityUrn); + proposal.setChangeType(changeType); + proposal.setEntityType(entityUrn.getEntityType()); + proposal.setAspectName(aspectName); + if (aspect != null) { + proposal.setAspect(GenericRecordUtils.serializeAspect(aspect)); + } + return proposal; } - public static SystemMetadata parseSystemMetadata(String jsonSystemMetadata) { + static SystemMetadata parseSystemMetadata(String jsonSystemMetadata) { if (jsonSystemMetadata == null || jsonSystemMetadata.equals("")) { SystemMetadata response = new SystemMetadata(); response.setRunId(DEFAULT_RUN_ID); @@ -148,43 +322,4 @@ public static SystemMetadata parseSystemMetadata(String jsonSystemMetadata) { } return RecordUtils.toRecordTemplate(SystemMetadata.class, jsonSystemMetadata); } - - public static RecordTemplate buildKeyAspect( - @Nonnull EntityRegistry entityRegistry, @Nonnull final Urn urn) { - final EntitySpec spec = entityRegistry.getEntitySpec(urnToEntityName(urn)); - final AspectSpec keySpec = spec.getKeyAspectSpec(); - return EntityKeyUtils.convertUrnToEntityKey(urn, keySpec); - } - - public static void validateUrn(@Nonnull EntityRegistry entityRegistry, @Nonnull final Urn urn) { - EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(entityRegistry); - validator.setCurrentEntitySpec(entityRegistry.getEntitySpec(urn.getEntityType())); - RecordTemplateValidator.validate( - EntityUtils.buildKeyAspect(entityRegistry, urn), - validationResult -> { - throw new IllegalArgumentException( - "Invalid urn: " + urn + "\n Cause: " + validationResult.getMessages()); - }, - validator); - - if (urn.toString().trim().length() != urn.toString().length()) { - throw new IllegalArgumentException( - "Error: cannot provide an URN with leading or trailing whitespace"); - } - if (URLEncoder.encode(urn.toString()).length() > URN_NUM_BYTES_LIMIT) { - throw new IllegalArgumentException( - "Error: cannot provide an URN longer than " - + Integer.toString(URN_NUM_BYTES_LIMIT) - + " bytes (when URL encoded)"); - } - if (urn.toString().contains(URN_DELIMITER_SEPARATOR)) { - throw new IllegalArgumentException( - "Error: URN cannot contain " + URN_DELIMITER_SEPARATOR + " character"); - } - try { - Urn.createFromString(urn.toString()); - } catch (URISyntaxException e) { - throw new IllegalArgumentException(e); - } - } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java index f37f63913abe45..c1e76e7c678363 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java @@ -201,7 +201,7 @@ public Map batchGet( return keys.stream() .map(this::getAspect) .filter(Objects::nonNull) - .collect(Collectors.toMap(EntityAspect::toAspectIdentifier, aspect -> aspect)); + .collect(Collectors.toMap(EntityAspect::getAspectIdentifier, aspect -> aspect)); } @Override diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraRetentionService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraRetentionService.java index 4d9d2b3c416b7b..91e31975298771 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraRetentionService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraRetentionService.java @@ -17,7 +17,7 @@ import com.linkedin.common.urn.Urn; import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.batch.AspectsBatch; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.EntityAspectIdentifier; import com.linkedin.metadata.entity.EntityService; @@ -45,7 +45,7 @@ @Slf4j @RequiredArgsConstructor -public class CassandraRetentionService extends RetentionService { +public class CassandraRetentionService extends RetentionService { private final EntityService _entityService; private final CqlSession _cqlSession; private final int _batchSize; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java index 3342d4632f642e..23d443c10b71fc 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java @@ -10,8 +10,9 @@ import com.google.common.cache.LoadingCache; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.batch.AspectsBatch; -import com.linkedin.metadata.aspect.batch.MCPBatchItem; +import com.linkedin.metadata.aspect.batch.MCPItem; import com.linkedin.metadata.config.EbeanConfiguration; import com.linkedin.metadata.entity.AspectDao; import com.linkedin.metadata.entity.AspectMigrationsDao; @@ -637,13 +638,14 @@ public List runInTransactionWithRetry( Set urnsWithKeyAspects = batch.getMCPItems().stream() .filter(i -> i.getEntitySpec().getKeyAspectSpec().equals(i.getAspectSpec())) - .map(MCPBatchItem::getUrn) + .map(MCPItem::getUrn) .collect(Collectors.toSet()); if (!urnsWithKeyAspects.isEmpty()) { // Split into batches by urn with key aspect, remaining aspects in the pair's second - Pair, AspectsBatch> splitBatches = splitByUrn(batch, urnsWithKeyAspects); + Pair, AspectsBatch> splitBatches = + splitByUrn(batch, urnsWithKeyAspects, batch.getAspectRetriever()); // Run non-key aspect `other` batch per normal if (!splitBatches.getSecond().getItems().isEmpty()) { @@ -902,12 +904,13 @@ private static String buildMetricName( * @return separated batches */ private static Pair, AspectsBatch> splitByUrn( - AspectsBatch batch, Set urns) { - Map> itemsByUrn = - batch.getMCPItems().stream().collect(Collectors.groupingBy(MCPBatchItem::getUrn)); + AspectsBatch batch, Set urns, AspectRetriever aspectRetriever) { + Map> itemsByUrn = + batch.getMCPItems().stream().collect(Collectors.groupingBy(MCPItem::getUrn)); AspectsBatch other = AspectsBatchImpl.builder() + .aspectRetriever(aspectRetriever) .items( itemsByUrn.entrySet().stream() .filter(entry -> !urns.contains(entry.getKey())) @@ -917,7 +920,12 @@ private static Pair, AspectsBatch> splitByUrn( List nonEmptyBatches = urns.stream() - .map(urn -> AspectsBatchImpl.builder().items(itemsByUrn.get(urn)).build()) + .map( + urn -> + AspectsBatchImpl.builder() + .aspectRetriever(aspectRetriever) + .items(itemsByUrn.get(urn)) + .build()) .filter(b -> !b.getItems().isEmpty()) .collect(Collectors.toList()); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanRetentionService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanRetentionService.java index eba550714766b8..250a81d9c8edcf 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanRetentionService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanRetentionService.java @@ -5,7 +5,7 @@ import com.linkedin.common.urn.Urn; import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.batch.AspectsBatch; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RetentionService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; @@ -40,7 +40,7 @@ @Slf4j @RequiredArgsConstructor -public class EbeanRetentionService extends RetentionService { +public class EbeanRetentionService extends RetentionService { private final EntityService _entityService; private final Database _server; private final int _batchSize; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java index 1718bd835dc31f..3edb55f265dc13 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java @@ -3,13 +3,13 @@ import com.linkedin.common.AuditStamp; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.SystemAspect; import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.aspect.batch.BatchItem; -import com.linkedin.metadata.aspect.batch.SystemAspect; -import com.linkedin.metadata.aspect.batch.UpsertItem; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.plugins.validation.ValidationExceptionCollection; import com.linkedin.mxe.MetadataChangeProposal; -import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; import java.util.Collection; import java.util.LinkedList; @@ -18,6 +18,7 @@ import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; +import javax.annotation.Nonnull; import lombok.Builder; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -27,7 +28,8 @@ @Builder(toBuilder = true) public class AspectsBatchImpl implements AspectsBatch { - private final Collection items; + @Nonnull private final Collection items; + @Nonnull private final AspectRetriever aspectRetriever; /** * Convert patches to upserts, apply hooks at the aspect and batch level. @@ -37,10 +39,11 @@ public class AspectsBatchImpl implements AspectsBatch { * various hooks */ @Override - public Pair>, List> toUpsertBatchItems( - final Map> latestAspects, AspectRetriever aspectRetriever) { + public Pair>, List> toUpsertBatchItems( + final Map> latestAspects) { - LinkedList upsertBatchItems = + // Convert patches to upserts if needed + LinkedList upsertBatchItems = items.stream() .map( item -> { @@ -49,35 +52,26 @@ public Pair>, List> toUpsertBatchItems( final SystemAspect latest = latestAspects.getOrDefault(urnStr, Map.of()).get(item.getAspectName()); - final MCPUpsertBatchItem upsertItem; - if (item instanceof MCPUpsertBatchItem) { - upsertItem = (MCPUpsertBatchItem) item; + final ChangeItemImpl upsertItem; + if (item instanceof ChangeItemImpl) { + upsertItem = (ChangeItemImpl) item; } else { // patch to upsert - MCPPatchBatchItem patchBatchItem = (MCPPatchBatchItem) item; + PatchItemImpl patchBatchItem = (PatchItemImpl) item; final RecordTemplate currentValue = - latest != null - ? latest.getRecordTemplate(aspectRetriever.getEntityRegistry()) - : null; + latest != null ? latest.getRecordTemplate() : null; upsertItem = patchBatchItem.applyPatch(currentValue, aspectRetriever); } - // Apply hooks - final SystemMetadata oldSystemMetadata = - latest != null ? latest.getSystemMetadata() : null; - final RecordTemplate oldAspectValue = - latest != null - ? latest.getRecordTemplate(aspectRetriever.getEntityRegistry()) - : null; - upsertItem.applyMutationHooks(oldAspectValue, oldSystemMetadata, aspectRetriever); - return upsertItem; }) .collect(Collectors.toCollection(LinkedList::new)); - LinkedList newItems = - applyMCPSideEffects(upsertBatchItems, aspectRetriever) - .collect(Collectors.toCollection(LinkedList::new)); + // Apply write hooks before side effects + applyWriteMutationHooks(upsertBatchItems); + + LinkedList newItems = + applyMCPSideEffects(upsertBatchItems).collect(Collectors.toCollection(LinkedList::new)); Map> newUrnAspectNames = getNewUrnAspectsMap(getUrnAspectsMap(), newItems); upsertBatchItems.addAll(newItems); @@ -91,28 +85,41 @@ public static class AspectsBatchImplBuilder { * @param data aspect data * @return builder */ - public AspectsBatchImplBuilder one(BatchItem data) { - this.items = List.of(data); + public AspectsBatchImplBuilder one(BatchItem data, AspectRetriever aspectRetriever) { + aspectRetriever(aspectRetriever); + items(List.of(data)); return this; } public AspectsBatchImplBuilder mcps( List mcps, AuditStamp auditStamp, AspectRetriever aspectRetriever) { - this.items = + + aspectRetriever(aspectRetriever); + items( mcps.stream() .map( mcp -> { if (mcp.getChangeType().equals(ChangeType.PATCH)) { - return MCPPatchBatchItem.MCPPatchBatchItemBuilder.build( + return PatchItemImpl.PatchItemImplBuilder.build( mcp, auditStamp, aspectRetriever.getEntityRegistry()); } else { - return MCPUpsertBatchItem.MCPUpsertBatchItemBuilder.build( + return ChangeItemImpl.ChangeItemImplBuilder.build( mcp, auditStamp, aspectRetriever); } }) - .collect(Collectors.toList()); + .collect(Collectors.toList())); return this; } + + public AspectsBatchImpl build() { + ValidationExceptionCollection exceptions = + AspectsBatch.validateProposed(this.items, this.aspectRetriever); + if (!exceptions.isEmpty()) { + throw new IllegalArgumentException("Failed to validate MCP due to: " + exceptions); + } + + return new AspectsBatchImpl(this.items, this.aspectRetriever); + } } @Override diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPUpsertBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/ChangeItemImpl.java similarity index 63% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPUpsertBatchItem.java rename to metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/ChangeItemImpl.java index b9d5f24e7ce084..b2e3363547dd02 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPUpsertBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/ChangeItemImpl.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.entity.ebean.batch; -import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; import static com.linkedin.metadata.entity.AspectUtils.validateAspect; import com.datahub.util.exception.ModelConversionException; @@ -9,13 +8,11 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.batch.SystemAspect; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.SystemAspect; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.batch.MCPItem; import com.linkedin.metadata.aspect.patch.template.common.GenericPatchTemplate; -import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; -import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.metadata.entity.validation.ValidationUtils; @@ -33,26 +30,24 @@ import javax.annotation.Nullable; import lombok.Builder; import lombok.Getter; +import lombok.Setter; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; @Slf4j @Getter @Builder(toBuilder = true) -public class MCPUpsertBatchItem extends UpsertItem { +public class ChangeItemImpl implements ChangeMCP { - public static MCPUpsertBatchItem fromPatch( + public static ChangeItemImpl fromPatch( @Nonnull Urn urn, @Nonnull AspectSpec aspectSpec, @Nullable RecordTemplate recordTemplate, GenericPatchTemplate genericPatchTemplate, @Nonnull AuditStamp auditStamp, AspectRetriever aspectRetriever) { - MCPUpsertBatchItem.MCPUpsertBatchItemBuilder builder = - MCPUpsertBatchItem.builder() - .urn(urn) - .auditStamp(auditStamp) - .aspectName(aspectSpec.getName()); + ChangeItemImplBuilder builder = + ChangeItemImpl.builder().urn(urn).auditStamp(auditStamp).aspectName(aspectSpec.getName()); RecordTemplate currentValue = recordTemplate != null ? recordTemplate : genericPatchTemplate.getDefault(); @@ -84,76 +79,45 @@ public static MCPUpsertBatchItem fromPatch( @Nonnull private final EntitySpec entitySpec; @Nonnull private final AspectSpec aspectSpec; + @Setter @Nullable private SystemAspect previousSystemAspect; + @Setter private long nextAspectVersion; + @Nonnull @Override public ChangeType getChangeType() { return ChangeType.UPSERT; } - public void applyMutationHooks( - @Nullable RecordTemplate oldAspectValue, - @Nullable SystemMetadata oldSystemMetadata, - @Nonnull AspectRetriever aspectRetriever) { - // add audit stamp/system meta if needed - for (MutationHook mutationHook : - aspectRetriever - .getEntityRegistry() - .getMutationHooks(getChangeType(), entitySpec.getName(), aspectSpec.getName())) { - mutationHook.applyMutation( - getChangeType(), - entitySpec, - aspectSpec, - oldAspectValue, - recordTemplate, - oldSystemMetadata, - systemMetadata, - auditStamp, - aspectRetriever); - } - } - - @Override - public SystemAspect toLatestEntityAspect() { - EntityAspect latest = new EntityAspect(); - latest.setAspect(getAspectName()); - latest.setMetadata(EntityUtils.toJsonAspect(getRecordTemplate())); - latest.setUrn(getUrn().toString()); - latest.setVersion(ASPECT_LATEST_VERSION); - latest.setCreatedOn(new Timestamp(auditStamp.getTime())); - latest.setCreatedBy(auditStamp.getActor().toString()); - return latest.asSystemAspect(); - } - + @Nonnull @Override - public void validatePreCommit( - @Nullable RecordTemplate previous, @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { - - for (AspectPayloadValidator validator : - aspectRetriever - .getEntityRegistry() - .getAspectPayloadValidators( - getChangeType(), entitySpec.getName(), aspectSpec.getName())) { - validator.validatePreCommit( - getChangeType(), urn, getAspectSpec(), previous, this.recordTemplate, aspectRetriever); - } + public SystemAspect getSystemAspect(@Nullable Long version) { + EntityAspect entityAspect = new EntityAspect(); + entityAspect.setAspect(getAspectName()); + entityAspect.setMetadata(EntityUtils.toJsonAspect(getRecordTemplate())); + entityAspect.setUrn(getUrn().toString()); + entityAspect.setVersion(version == null ? getNextAspectVersion() : version); + entityAspect.setCreatedOn(new Timestamp(getAuditStamp().getTime())); + entityAspect.setCreatedBy(getAuditStamp().getActor().toString()); + entityAspect.setSystemMetadata(EntityUtils.toJsonAspect(getSystemMetadata())); + return EntityAspect.EntitySystemAspect.builder() + .build(getEntitySpec(), getAspectSpec(), entityAspect); } - public static class MCPUpsertBatchItemBuilder { + public static class ChangeItemImplBuilder { // Ensure use of other builders - private MCPUpsertBatchItem build() { + private ChangeItemImpl build() { return null; } - public MCPUpsertBatchItemBuilder systemMetadata(SystemMetadata systemMetadata) { + public ChangeItemImplBuilder systemMetadata(SystemMetadata systemMetadata) { this.systemMetadata = SystemMetadataUtils.generateSystemMetadataIfEmpty(systemMetadata); return this; } @SneakyThrows - public MCPUpsertBatchItem build(AspectRetriever aspectRetriever) { - EntityUtils.validateUrn(aspectRetriever.getEntityRegistry(), this.urn); + public ChangeItemImpl build(AspectRetriever aspectRetriever) { + ValidationUtils.validateUrn(aspectRetriever.getEntityRegistry(), this.urn); log.debug("entity type = {}", this.urn.getEntityType()); entitySpec(aspectRetriever.getEntityRegistry().getEntitySpec(this.urn.getEntityType())); @@ -163,14 +127,9 @@ public MCPUpsertBatchItem build(AspectRetriever aspectRetriever) { log.debug("aspect spec = {}", this.aspectSpec); ValidationUtils.validateRecordTemplate( - ChangeType.UPSERT, - this.entitySpec, - this.aspectSpec, - this.urn, - this.recordTemplate, - aspectRetriever); + this.entitySpec, this.urn, this.recordTemplate, aspectRetriever); - return new MCPUpsertBatchItem( + return new ChangeItemImpl( this.urn, this.aspectName, this.recordTemplate, @@ -178,10 +137,12 @@ public MCPUpsertBatchItem build(AspectRetriever aspectRetriever) { this.auditStamp, this.metadataChangeProposal, this.entitySpec, - this.aspectSpec); + this.aspectSpec, + this.previousSystemAspect, + this.nextAspectVersion); } - public static MCPUpsertBatchItem build( + public static ChangeItemImpl build( MetadataChangeProposal mcp, AuditStamp auditStamp, AspectRetriever aspectRetriever) { if (!mcp.getChangeType().equals(ChangeType.UPSERT)) { throw new IllegalArgumentException( @@ -193,7 +154,7 @@ public static MCPUpsertBatchItem build( aspectRetriever.getEntityRegistry().getEntitySpec(mcp.getEntityType()); AspectSpec aspectSpec = validateAspect(mcp, entitySpec); - if (!isValidChangeType(ChangeType.UPSERT, aspectSpec)) { + if (!MCPItem.isValidChangeType(ChangeType.UPSERT, aspectSpec)) { throw new UnsupportedOperationException( "ChangeType not supported: " + mcp.getChangeType() @@ -206,7 +167,7 @@ public static MCPUpsertBatchItem build( urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } - return MCPUpsertBatchItem.builder() + return ChangeItemImpl.builder() .urn(urn) .aspectName(mcp.getAspectName()) .systemMetadata( @@ -217,16 +178,6 @@ public static MCPUpsertBatchItem build( .build(aspectRetriever); } - private MCPUpsertBatchItemBuilder entitySpec(EntitySpec entitySpec) { - this.entitySpec = entitySpec; - return this; - } - - private MCPUpsertBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { - this.aspectSpec = aspectSpec; - return this; - } - private static RecordTemplate convertToRecordTemplate( MetadataChangeProposal mcp, AspectSpec aspectSpec) { RecordTemplate aspect; @@ -241,7 +192,6 @@ private static RecordTemplate convertToRecordTemplate( "Could not deserialize %s for aspect %s", mcp.getAspect().getValue(), mcp.getAspectName())); } - log.debug("aspect = {}", aspect); return aspect; } } @@ -254,7 +204,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { return false; } - MCPUpsertBatchItem that = (MCPUpsertBatchItem) o; + ChangeItemImpl that = (ChangeItemImpl) o; return urn.equals(that.urn) && aspectName.equals(that.aspectName) && Objects.equals(systemMetadata, that.systemMetadata) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/DeleteItemImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/DeleteItemImpl.java new file mode 100644 index 00000000000000..0ab854198a2828 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/DeleteItemImpl.java @@ -0,0 +1,139 @@ +package com.linkedin.metadata.entity.ebean.batch; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.SystemAspect; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.entity.EntityAspect; +import com.linkedin.metadata.entity.EntityUtils; +import com.linkedin.metadata.entity.validation.ValidationUtils; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; +import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Builder; +import lombok.Getter; +import lombok.Setter; +import lombok.SneakyThrows; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class DeleteItemImpl implements ChangeMCP { + + // urn an urn associated with the new aspect + @Nonnull private final Urn urn; + + // aspectName name of the aspect being inserted + @Nonnull private final String aspectName; + + @Nonnull private final AuditStamp auditStamp; + + // derived + @Nonnull private final EntitySpec entitySpec; + @Nonnull private final AspectSpec aspectSpec; + + @Setter @Nullable private SystemAspect previousSystemAspect; + + @Nonnull + @Override + public ChangeType getChangeType() { + return ChangeType.DELETE; + } + + @Nullable + @Override + public RecordTemplate getRecordTemplate() { + return null; + } + + @Nullable + @Override + public SystemMetadata getSystemMetadata() { + return null; + } + + @Nullable + @Override + public MetadataChangeProposal getMetadataChangeProposal() { + return EntityUtils.buildMCP(getUrn(), aspectName, getChangeType(), null); + } + + @Nonnull + @Override + public SystemAspect getSystemAspect(@Nullable Long nextAspectVersion) { + EntityAspect entityAspect = new EntityAspect(); + entityAspect.setAspect(getAspectName()); + entityAspect.setUrn(getUrn().toString()); + entityAspect.setVersion(0); + return EntityAspect.EntitySystemAspect.builder() + .build(getEntitySpec(), getAspectSpec(), entityAspect); + } + + @Override + public long getNextAspectVersion() { + return 0; + } + + @Override + public void setNextAspectVersion(long nextAspectVersion) { + throw new IllegalStateException("Next aspect version is always zero"); + } + + public static class DeleteItemImplBuilder { + + // Ensure use of other builders + private DeleteItemImpl build() { + return null; + } + + @SneakyThrows + public DeleteItemImpl build(AspectRetriever aspectRetriever) { + ValidationUtils.validateUrn(aspectRetriever.getEntityRegistry(), this.urn); + log.debug("entity type = {}", this.urn.getEntityType()); + + entitySpec(aspectRetriever.getEntityRegistry().getEntitySpec(this.urn.getEntityType())); + log.debug("entity spec = {}", this.entitySpec); + + aspectSpec(ValidationUtils.validate(this.entitySpec, this.aspectName)); + log.debug("aspect spec = {}", this.aspectSpec); + + return new DeleteItemImpl( + this.urn, + this.aspectName, + this.auditStamp, + this.entitySpec, + this.aspectSpec, + this.previousSystemAspect); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + DeleteItemImpl that = (DeleteItemImpl) o; + return urn.equals(that.urn) && aspectName.equals(that.aspectName); + } + + @Override + public int hashCode() { + return Objects.hash(urn, aspectName); + } + + @Override + public String toString() { + return "UpsertBatchItem{" + "urn=" + urn + ", aspectName='" + aspectName + '\'' + '}'; + } +} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLBatchItemImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLItemImpl.java similarity index 77% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLBatchItemImpl.java rename to metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLItemImpl.java index a2ed2eb18fe6a3..6efc1e78b543c1 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLBatchItemImpl.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLItemImpl.java @@ -1,14 +1,12 @@ package com.linkedin.metadata.entity.ebean.batch; -import static com.linkedin.metadata.entity.AspectUtils.validateAspect; - import com.datahub.util.exception.ModelConversionException; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.batch.MCLBatchItem; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.entity.EntityUtils; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.MCLItem; +import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.validation.ValidationUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -26,7 +24,7 @@ @Slf4j @Getter @Builder(toBuilder = true) -public class MCLBatchItemImpl implements MCLBatchItem { +public class MCLItemImpl implements MCLItem { @Nonnull private final MetadataChangeLog metadataChangeLog; @@ -38,19 +36,18 @@ public class MCLBatchItemImpl implements MCLBatchItem { private final EntitySpec entitySpec; private final AspectSpec aspectSpec; - public static class MCLBatchItemImplBuilder { + public static class MCLItemImplBuilder { // Ensure use of other builders - private MCLBatchItemImpl build() { + private MCLItemImpl build() { return null; } - public MCLBatchItemImpl build( - MetadataChangeLog metadataChangeLog, AspectRetriever aspectRetriever) { - return MCLBatchItemImpl.builder().metadataChangeLog(metadataChangeLog).build(aspectRetriever); + public MCLItemImpl build(MetadataChangeLog metadataChangeLog, AspectRetriever aspectRetriever) { + return MCLItemImpl.builder().metadataChangeLog(metadataChangeLog).build(aspectRetriever); } - public MCLBatchItemImpl build(AspectRetriever aspectRetriever) { + public MCLItemImpl build(AspectRetriever aspectRetriever) { EntityRegistry entityRegistry = aspectRetriever.getEntityRegistry(); log.debug("entity type = {}", this.metadataChangeLog.getEntityType()); @@ -58,7 +55,7 @@ public MCLBatchItemImpl build(AspectRetriever aspectRetriever) { aspectRetriever .getEntityRegistry() .getEntitySpec(this.metadataChangeLog.getEntityType())); - aspectSpec(validateAspect(this.metadataChangeLog, this.entitySpec)); + aspectSpec(AspectUtils.validateAspect(this.metadataChangeLog, this.entitySpec)); Urn urn = this.metadataChangeLog.getEntityUrn(); if (urn == null) { @@ -66,7 +63,7 @@ public MCLBatchItemImpl build(AspectRetriever aspectRetriever) { EntityKeyUtils.getUrnFromLog( this.metadataChangeLog, this.entitySpec.getKeyAspectSpec()); } - EntityUtils.validateUrn(entityRegistry, urn); + ValidationUtils.validateUrn(entityRegistry, urn); log.debug("entity type = {}", urn.getEntityType()); entitySpec(entityRegistry.getEntitySpec(urn.getEntityType())); @@ -80,14 +77,9 @@ public MCLBatchItemImpl build(AspectRetriever aspectRetriever) { // validate new ValidationUtils.validateRecordTemplate( - this.metadataChangeLog.getChangeType(), - this.entitySpec, - this.aspectSpec, - urn, - aspects.getFirst(), - aspectRetriever); + this.entitySpec, urn, aspects.getFirst(), aspectRetriever); - return new MCLBatchItemImpl( + return new MCLItemImpl( this.metadataChangeLog, aspects.getFirst(), aspects.getSecond(), @@ -95,12 +87,12 @@ public MCLBatchItemImpl build(AspectRetriever aspectRetriever) { this.aspectSpec); } - private MCLBatchItemImplBuilder entitySpec(EntitySpec entitySpec) { + private MCLItemImplBuilder entitySpec(EntitySpec entitySpec) { this.entitySpec = entitySpec; return this; } - private MCLBatchItemImplBuilder aspectSpec(AspectSpec aspectSpec) { + private MCLItemImplBuilder aspectSpec(AspectSpec aspectSpec) { this.aspectSpec = aspectSpec; return this; } @@ -150,7 +142,7 @@ public boolean equals(Object o) { return false; } - MCLBatchItemImpl that = (MCLBatchItemImpl) o; + MCLItemImpl that = (MCLItemImpl) o; return metadataChangeLog.equals(that.metadataChangeLog); } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPPatchBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/PatchItemImpl.java similarity index 84% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPPatchBatchItem.java rename to metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/PatchItemImpl.java index d0cb2a4cc59b8a..cf9c3978e3a374 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPPatchBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/PatchItemImpl.java @@ -15,10 +15,10 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.batch.PatchItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.MCPItem; +import com.linkedin.metadata.aspect.batch.PatchMCP; import com.linkedin.metadata.aspect.patch.template.AspectTemplateEngine; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.metadata.entity.validation.ValidationUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -39,7 +39,7 @@ @Slf4j @Getter @Builder(toBuilder = true) -public class MCPPatchBatchItem extends PatchItem { +public class PatchItemImpl implements PatchMCP { private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); static { @@ -79,10 +79,9 @@ public RecordTemplate getRecordTemplate() { return null; } - public MCPUpsertBatchItem applyPatch( - RecordTemplate recordTemplate, AspectRetriever aspectRetriever) { - MCPUpsertBatchItem.MCPUpsertBatchItemBuilder builder = - MCPUpsertBatchItem.builder() + public ChangeItemImpl applyPatch(RecordTemplate recordTemplate, AspectRetriever aspectRetriever) { + ChangeItemImpl.ChangeItemImplBuilder builder = + ChangeItemImpl.builder() .urn(getUrn()) .aspectName(getAspectName()) .metadataChangeProposal(getMetadataChangeProposal()) @@ -116,16 +115,15 @@ public MCPUpsertBatchItem applyPatch( return builder.build(aspectRetriever); } - public static class MCPPatchBatchItemBuilder { + public static class PatchItemImplBuilder { - public MCPPatchBatchItem.MCPPatchBatchItemBuilder systemMetadata( - SystemMetadata systemMetadata) { + public PatchItemImpl.PatchItemImplBuilder systemMetadata(SystemMetadata systemMetadata) { this.systemMetadata = SystemMetadataUtils.generateSystemMetadataIfEmpty(systemMetadata); return this; } - public MCPPatchBatchItem build(EntityRegistry entityRegistry) { - EntityUtils.validateUrn(entityRegistry, this.urn); + public PatchItemImpl build(EntityRegistry entityRegistry) { + ValidationUtils.validateUrn(entityRegistry, this.urn); log.debug("entity type = {}", this.urn.getEntityType()); entitySpec(entityRegistry.getEntitySpec(this.urn.getEntityType())); @@ -139,7 +137,7 @@ public MCPPatchBatchItem build(EntityRegistry entityRegistry) { String.format("Missing patch to apply. Aspect: %s", this.aspectSpec.getName())); } - return new MCPPatchBatchItem( + return new PatchItemImpl( this.urn, this.aspectName, SystemMetadataUtils.generateSystemMetadataIfEmpty(this.systemMetadata), @@ -150,13 +148,13 @@ public MCPPatchBatchItem build(EntityRegistry entityRegistry) { this.aspectSpec); } - public static MCPPatchBatchItem build( + public static PatchItemImpl build( MetadataChangeProposal mcp, AuditStamp auditStamp, EntityRegistry entityRegistry) { log.debug("entity type = {}", mcp.getEntityType()); EntitySpec entitySpec = entityRegistry.getEntitySpec(mcp.getEntityType()); AspectSpec aspectSpec = validateAspect(mcp, entitySpec); - if (!PatchItem.isValidChangeType(ChangeType.PATCH, aspectSpec)) { + if (!MCPItem.isValidChangeType(ChangeType.PATCH, aspectSpec)) { throw new UnsupportedOperationException( "ChangeType not supported: " + mcp.getChangeType() @@ -169,7 +167,7 @@ public static MCPPatchBatchItem build( urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } - return MCPPatchBatchItem.builder() + return PatchItemImpl.builder() .urn(urn) .aspectName(mcp.getAspectName()) .systemMetadata( @@ -180,16 +178,6 @@ public static MCPPatchBatchItem build( .build(entityRegistry); } - private MCPPatchBatchItemBuilder entitySpec(EntitySpec entitySpec) { - this.entitySpec = entitySpec; - return this; - } - - private MCPPatchBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { - this.aspectSpec = aspectSpec; - return this; - } - private static Patch convertToJsonPatch(MetadataChangeProposal mcp) { JsonNode json; try { @@ -209,7 +197,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { return false; } - MCPPatchBatchItem that = (MCPPatchBatchItem) o; + PatchItemImpl that = (PatchItemImpl) o; return urn.equals(that.urn) && aspectName.equals(that.aspectName) && Objects.equals(systemMetadata, that.systemMetadata) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java index 947f0116b587c6..16942a02b0e4a3 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java @@ -3,14 +3,13 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.schema.validation.ValidationResult; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; +import java.net.URISyntaxException; +import java.net.URLEncoder; import java.util.function.Consumer; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -18,6 +17,8 @@ @Slf4j public class ValidationUtils { + public static final int URN_NUM_BYTES_LIMIT = 512; + public static final String URN_DELIMITER_SEPARATOR = "␟"; /** * Validates a {@link RecordTemplate} and throws {@link @@ -66,9 +67,7 @@ public static AspectSpec validate(EntitySpec entitySpec, String aspectName) { } public static void validateRecordTemplate( - ChangeType changeType, EntitySpec entitySpec, - AspectSpec aspectSpec, Urn urn, @Nullable RecordTemplate aspect, @Nonnull AspectRetriever aspectRetriever) { @@ -89,17 +88,38 @@ public static void validateRecordTemplate( if (aspect != null) { RecordTemplateValidator.validate(aspect, resultFunction, validator); + } + } - for (AspectPayloadValidator aspectValidator : - entityRegistry.getAspectPayloadValidators( - changeType, entitySpec.getName(), aspectSpec.getName())) { - try { - aspectValidator.validateProposed(changeType, urn, aspectSpec, aspect, aspectRetriever); - } catch (AspectValidationException e) { + public static void validateUrn(@Nonnull EntityRegistry entityRegistry, @Nonnull final Urn urn) { + EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(entityRegistry); + validator.setCurrentEntitySpec(entityRegistry.getEntitySpec(urn.getEntityType())); + RecordTemplateValidator.validate( + EntityUtils.buildKeyAspect(entityRegistry, urn), + validationResult -> { throw new IllegalArgumentException( - "Failed to validate aspect due to: " + e.getMessage(), e); - } - } + "Invalid urn: " + urn + "\n Cause: " + validationResult.getMessages()); + }, + validator); + + if (urn.toString().trim().length() != urn.toString().length()) { + throw new IllegalArgumentException( + "Error: cannot provide an URN with leading or trailing whitespace"); + } + if (URLEncoder.encode(urn.toString()).length() > URN_NUM_BYTES_LIMIT) { + throw new IllegalArgumentException( + "Error: cannot provide an URN longer than " + + Integer.toString(URN_NUM_BYTES_LIMIT) + + " bytes (when URL encoded)"); + } + if (urn.toString().contains(URN_DELIMITER_SEPARATOR)) { + throw new IllegalArgumentException( + "Error: URN cannot contain " + URN_DELIMITER_SEPARATOR + " character"); + } + try { + Urn.createFromString(urn.toString()); + } catch (URISyntaxException e) { + throw new IllegalArgumentException(e); } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java index c20c16e0ea7d1e..4c1d8da7c4a364 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.search.elasticsearch; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.browse.BrowseResultV2; import com.linkedin.metadata.query.AutoCompleteResult; @@ -40,6 +41,13 @@ public class ElasticSearchService implements EntitySearchService, ElasticSearchI private final ESBrowseDAO esBrowseDAO; private final ESWriteDAO esWriteDAO; + @Override + public ElasticSearchService postConstruct(AspectRetriever aspectRetriever) { + esSearchDAO.setAspectRetriever(aspectRetriever); + esBrowseDAO.setAspectRetriever(aspectRetriever); + return this; + } + @Override public void configure() { indexBuilders.reindexAll(); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java index b8085885200892..dd1c09853114d1 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java @@ -7,6 +7,7 @@ import com.google.common.annotations.VisibleForTesting; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.browse.BrowseResultEntity; import com.linkedin.metadata.browse.BrowseResultEntityArray; @@ -20,7 +21,6 @@ import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.annotation.SearchableAnnotation; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.search.elasticsearch.query.request.SearchRequestHandler; @@ -41,7 +41,9 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.RequiredArgsConstructor; +import lombok.Setter; import lombok.Value; +import lombok.experimental.Accessors; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang.StringUtils; import org.opensearch.action.search.SearchRequest; @@ -62,9 +64,10 @@ @Slf4j @RequiredArgsConstructor +@Accessors(chain = true) public class ESBrowseDAO { - private final EntityRegistry entityRegistry; + @Setter private AspectRetriever aspectRetriever; private final RestHighLevelClient client; private final IndexConvention indexConvention; @Nonnull private final SearchConfiguration searchConfiguration; @@ -118,7 +121,8 @@ public BrowseResult browse( try { final String indexName = - indexConvention.getIndexName(entityRegistry.getEntitySpec(entityName)); + indexConvention.getIndexName( + aspectRetriever.getEntityRegistry().getEntitySpec(entityName)); final SearchResponse groupsResponse; try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esGroupSearch").time()) { @@ -373,7 +377,8 @@ private static int getPathDepth(@Nonnull String path) { */ @Nonnull public List getBrowsePaths(@Nonnull String entityName, @Nonnull Urn urn) { - final String indexName = indexConvention.getIndexName(entityRegistry.getEntitySpec(entityName)); + final String indexName = + indexConvention.getIndexName(aspectRetriever.getEntityRegistry().getEntitySpec(entityName)); final SearchRequest searchRequest = new SearchRequest(indexName); searchRequest.source( new SearchSourceBuilder().query(QueryBuilders.termQuery(URN, urn.toString()))); @@ -478,7 +483,8 @@ private SearchRequest constructGroupsSearchRequestV2( @Nullable Filter filter, @Nonnull String input, @Nullable SearchFlags searchFlags) { - final String indexName = indexConvention.getIndexName(entityRegistry.getEntitySpec(entityName)); + final String indexName = + indexConvention.getIndexName(aspectRetriever.getEntityRegistry().getEntitySpec(entityName)); final SearchRequest searchRequest = new SearchRequest(indexName); final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.size(0); @@ -503,7 +509,9 @@ private SearchRequest constructGroupsSearchRequestBrowseAcrossEntities( @Nullable SearchFlags searchFlags) { List entitySpecs = - entities.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList()); + entities.stream() + .map(name -> aspectRetriever.getEntityRegistry().getEntitySpec(name)) + .collect(Collectors.toList()); String[] indexArray = entities.stream().map(indexConvention::getEntityIndexName).toArray(String[]::new); @@ -553,9 +561,10 @@ private QueryBuilder buildQueryStringV2( final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery(); - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec(entityName); QueryBuilder query = - SearchRequestHandler.getBuilder(entitySpec, searchConfiguration, customSearchConfiguration) + SearchRequestHandler.getBuilder( + entitySpec, searchConfiguration, customSearchConfiguration, aspectRetriever) .getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext())); queryBuilder.must(query); @@ -568,7 +577,8 @@ private QueryBuilder buildQueryStringV2( queryBuilder.filter(QueryBuilders.rangeQuery(BROWSE_PATH_V2_DEPTH).gt(browseDepthVal)); queryBuilder.filter( - SearchRequestHandler.getFilterQuery(filter, entitySpec.getSearchableFieldTypes())); + SearchRequestHandler.getFilterQuery( + filter, entitySpec.getSearchableFieldTypes(), aspectRetriever)); return queryBuilder; } @@ -587,7 +597,8 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities( final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery(); QueryBuilder query = - SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration) + SearchRequestHandler.getBuilder( + entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext())); queryBuilder.must(query); @@ -608,7 +619,8 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities( set1.addAll(set2); return set1; })); - queryBuilder.filter(SearchRequestHandler.getFilterQuery(filter, searchableFields)); + queryBuilder.filter( + SearchRequestHandler.getFilterQuery(filter, searchableFields, aspectRetriever)); return queryBuilder; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java index 76153a8d2adb3f..688bf85d4cba70 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java @@ -9,10 +9,10 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.google.common.annotations.VisibleForTesting; import com.linkedin.data.template.LongMap; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.AutoCompleteResult; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; @@ -41,6 +41,8 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.RequiredArgsConstructor; +import lombok.Setter; +import lombok.experimental.Accessors; import lombok.extern.slf4j.Slf4j; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; @@ -60,6 +62,7 @@ /** A search DAO for Elasticsearch backend. */ @Slf4j @RequiredArgsConstructor +@Accessors(chain = true) public class ESSearchDAO { private static final NamedXContentRegistry X_CONTENT_REGISTRY; @@ -68,7 +71,7 @@ public class ESSearchDAO { X_CONTENT_REGISTRY = new NamedXContentRegistry(searchModule.getNamedXContents()); } - private final EntityRegistry entityRegistry; + @Setter private AspectRetriever aspectRetriever; private final RestHighLevelClient client; private final IndexConvention indexConvention; private final boolean pointInTimeCreationEnabled; @@ -77,10 +80,12 @@ public class ESSearchDAO { @Nullable private final CustomSearchConfiguration customSearchConfiguration; public long docCount(@Nonnull String entityName) { - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec(entityName); CountRequest countRequest = new CountRequest(indexConvention.getIndexName(entitySpec)) - .query(SearchRequestHandler.getFilterQuery(null, entitySpec.getSearchableFieldTypes())); + .query( + SearchRequestHandler.getFilterQuery( + null, entitySpec.getSearchableFieldTypes(), aspectRetriever)); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "docCount").time()) { return client.count(countRequest, RequestOptions.DEFAULT).getCount(); } catch (IOException e) { @@ -105,7 +110,7 @@ private SearchResult executeAndExtract( // extract results, validated against document model as well return transformIndexIntoEntityName( SearchRequestHandler.getBuilder( - entitySpec, searchConfiguration, customSearchConfiguration) + entitySpec, searchConfiguration, customSearchConfiguration, aspectRetriever) .extractResult(searchResponse, filter, from, size)); } catch (Exception e) { log.error("Search query failed", e); @@ -189,7 +194,7 @@ private ScrollResult executeAndExtract( // extract results, validated against document model as well return transformIndexIntoEntityName( SearchRequestHandler.getBuilder( - entitySpecs, searchConfiguration, customSearchConfiguration) + entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .extractScrollResult( searchResponse, filter, scrollId, keepAlive, size, supportsPointInTime())); } catch (Exception e) { @@ -226,11 +231,14 @@ public SearchResult search( final String finalInput = input.isEmpty() ? "*" : input; Timer.Context searchRequestTimer = MetricUtils.timer(this.getClass(), "searchRequest").time(); List entitySpecs = - entityNames.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList()); + entityNames.stream() + .map(name -> aspectRetriever.getEntityRegistry().getEntitySpec(name)) + .collect(Collectors.toList()); Filter transformedFilters = transformFilterForEntities(postFilters, indexConvention); // Step 1: construct the query final SearchRequest searchRequest = - SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration) + SearchRequestHandler.getBuilder( + entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getSearchRequest( finalInput, transformedFilters, sortCriterion, from, size, searchFlags, facets); searchRequest.indices( @@ -258,10 +266,11 @@ public SearchResult filter( @Nullable SortCriterion sortCriterion, int from, int size) { - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec(entityName); Filter transformedFilters = transformFilterForEntities(filters, indexConvention); final SearchRequest searchRequest = - SearchRequestHandler.getBuilder(entitySpec, searchConfiguration, customSearchConfiguration) + SearchRequestHandler.getBuilder( + entitySpec, searchConfiguration, customSearchConfiguration, aspectRetriever) .getFilterRequest(transformedFilters, sortCriterion, from, size); searchRequest.indices(indexConvention.getIndexName(entitySpec)); @@ -288,8 +297,9 @@ public AutoCompleteResult autoComplete( @Nullable Filter requestParams, int limit) { try { - EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); - AutocompleteRequestHandler builder = AutocompleteRequestHandler.getBuilder(entitySpec); + EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec(entityName); + AutocompleteRequestHandler builder = + AutocompleteRequestHandler.getBuilder(entitySpec, aspectRetriever); SearchRequest req = builder.getSearchRequest( query, field, transformFilterForEntities(requestParams, indexConvention), limit); @@ -319,13 +329,16 @@ public Map aggregateByValue( int limit) { List entitySpecs; if (entityNames == null || entityNames.isEmpty()) { - entitySpecs = QueryUtils.getQueryByDefaultEntitySpecs(entityRegistry); + entitySpecs = QueryUtils.getQueryByDefaultEntitySpecs(aspectRetriever.getEntityRegistry()); } else { entitySpecs = - entityNames.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList()); + entityNames.stream() + .map(name -> aspectRetriever.getEntityRegistry().getEntitySpec(name)) + .collect(Collectors.toList()); } final SearchRequest searchRequest = - SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration) + SearchRequestHandler.getBuilder( + entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getAggregationRequest( field, transformFilterForEntities(requestParams, indexConvention), limit); if (entityNames == null) { @@ -334,7 +347,7 @@ public Map aggregateByValue( } else { Stream stream = entityNames.stream() - .map(entityRegistry::getEntitySpec) + .map(name -> aspectRetriever.getEntityRegistry().getEntitySpec(name)) .map(indexConvention::getIndexName); searchRequest.indices(stream.toArray(String[]::new)); } @@ -379,7 +392,9 @@ public ScrollResult scroll( entities.stream().map(indexConvention::getEntityIndexName).toArray(String[]::new); Timer.Context scrollRequestTimer = MetricUtils.timer(this.getClass(), "scrollRequest").time(); List entitySpecs = - entities.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList()); + entities.stream() + .map(name -> aspectRetriever.getEntityRegistry().getEntitySpec(name)) + .collect(Collectors.toList()); String pitId = null; Object[] sort = null; if (scrollId != null) { @@ -399,7 +414,8 @@ public ScrollResult scroll( Filter transformedFilters = transformFilterForEntities(postFilters, indexConvention); // Step 1: construct the query final SearchRequest searchRequest = - SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration) + SearchRequestHandler.getBuilder( + entitySpecs, searchConfiguration, customSearchConfiguration, aspectRetriever) .getSearchRequest( finalInput, transformedFilters, diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java index 887d4b22f37e24..fb3b51930370c4 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java @@ -5,6 +5,7 @@ import static com.linkedin.metadata.utils.SearchUtil.*; import com.linkedin.data.template.LongMap; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.StructuredPropertyUtils; @@ -47,6 +48,7 @@ public class AggregationQueryBuilder { private static final String URN_FILTER = "urn"; + private final AspectRetriever aspectRetriever; private final SearchConfiguration configs; private final Set defaultFacetFields; private final Set allFacetFields; @@ -56,7 +58,8 @@ public class AggregationQueryBuilder { public AggregationQueryBuilder( @Nonnull final SearchConfiguration configs, - @Nonnull Map> entitySearchAnnotations) { + @Nonnull Map> entitySearchAnnotations, + @Nonnull AspectRetriever aspectRetriever) { this.configs = Objects.requireNonNull(configs, "configs must not be null"); this.entitySearchAnnotations = entitySearchAnnotations; @@ -66,6 +69,7 @@ public AggregationQueryBuilder( .collect(Collectors.toList()); this.defaultFacetFields = getDefaultFacetFields(annotations); this.allFacetFields = getAllFacetFields(annotations); + this.aspectRetriever = aspectRetriever; } /** Get the set of default aggregations, across all facets. */ @@ -130,11 +134,12 @@ private AggregationBuilder facetToAggregationBuilder(final String inputFacet) { AggregationBuilder lastAggBuilder = null; for (int i = facets.size() - 1; i >= 0; i--) { String facet = facets.get(i); - if (facet.startsWith(STRUCTURED_PROPERTY_MAPPING_FIELD + ".")) { + if (facet.startsWith(STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX)) { String structPropFqn = facet.substring(STRUCTURED_PROPERTY_MAPPING_FIELD.length() + 1); + StructuredPropertyUtils.validateStructuredPropertyFQN( + Set.of(structPropFqn), aspectRetriever); facet = - STRUCTURED_PROPERTY_MAPPING_FIELD - + "." + STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX + StructuredPropertyUtils.sanitizeStructuredPropertyFQN(structPropFqn); } AggregationBuilder aggBuilder; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java index 38350322478741..de35d53bcde49b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java @@ -5,6 +5,7 @@ import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.StringArray; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.SearchableFieldSpec; import com.linkedin.metadata.models.annotation.SearchableAnnotation; @@ -46,7 +47,10 @@ public class AutocompleteRequestHandler { private static final Map AUTOCOMPLETE_QUERY_BUILDER_BY_ENTITY_NAME = new ConcurrentHashMap<>(); - public AutocompleteRequestHandler(@Nonnull EntitySpec entitySpec) { + private final AspectRetriever aspectRetriever; + + public AutocompleteRequestHandler( + @Nonnull EntitySpec entitySpec, @Nonnull AspectRetriever aspectRetriever) { List fieldSpecs = entitySpec.getSearchableFieldSpecs(); _defaultAutocompleteFields = Stream.concat( @@ -70,11 +74,13 @@ public AutocompleteRequestHandler(@Nonnull EntitySpec entitySpec) { set1.addAll(set2); return set1; })); + this.aspectRetriever = aspectRetriever; } - public static AutocompleteRequestHandler getBuilder(@Nonnull EntitySpec entitySpec) { + public static AutocompleteRequestHandler getBuilder( + @Nonnull EntitySpec entitySpec, @Nonnull AspectRetriever aspectRetriever) { return AUTOCOMPLETE_QUERY_BUILDER_BY_ENTITY_NAME.computeIfAbsent( - entitySpec, k -> new AutocompleteRequestHandler(entitySpec)); + entitySpec, k -> new AutocompleteRequestHandler(entitySpec, aspectRetriever)); } public SearchRequest getSearchRequest( @@ -83,7 +89,8 @@ public SearchRequest getSearchRequest( SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.size(limit); searchSourceBuilder.query(getQuery(input, field)); - searchSourceBuilder.postFilter(ESUtils.buildFilterQuery(filter, false, searchableFieldTypes)); + searchSourceBuilder.postFilter( + ESUtils.buildFilterQuery(filter, false, searchableFieldTypes, aspectRetriever)); searchSourceBuilder.highlighter(getHighlights(field)); searchRequest.source(searchSourceBuilder); return searchRequest; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java index 3ac05ed122cd70..402ccc08fee716 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java @@ -9,6 +9,7 @@ import com.google.common.collect.ImmutableMap; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.DoubleMap; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; import com.linkedin.metadata.models.EntitySpec; @@ -72,40 +73,45 @@ public class SearchRequestHandler { .setSkipHighlighting(false); private static final Map, SearchRequestHandler> REQUEST_HANDLER_BY_ENTITY_NAME = new ConcurrentHashMap<>(); - private final List _entitySpecs; - private final Set _defaultQueryFieldNames; - private final HighlightBuilder _highlights; + private final List entitySpecs; + private final Set defaultQueryFieldNames; + private final HighlightBuilder highlights; - private final SearchConfiguration _configs; - private final SearchQueryBuilder _searchQueryBuilder; - private final AggregationQueryBuilder _aggregationQueryBuilder; + private final SearchConfiguration configs; + private final SearchQueryBuilder searchQueryBuilder; + private final AggregationQueryBuilder aggregationQueryBuilder; private final Map> searchableFieldTypes; + private final AspectRetriever aspectRetriever; + private SearchRequestHandler( @Nonnull EntitySpec entitySpec, @Nonnull SearchConfiguration configs, - @Nullable CustomSearchConfiguration customSearchConfiguration) { - this(ImmutableList.of(entitySpec), configs, customSearchConfiguration); + @Nullable CustomSearchConfiguration customSearchConfiguration, + @Nonnull AspectRetriever aspectRetriever) { + this(ImmutableList.of(entitySpec), configs, customSearchConfiguration, aspectRetriever); } private SearchRequestHandler( @Nonnull List entitySpecs, @Nonnull SearchConfiguration configs, - @Nullable CustomSearchConfiguration customSearchConfiguration) { - _entitySpecs = entitySpecs; + @Nullable CustomSearchConfiguration customSearchConfiguration, + @Nonnull AspectRetriever aspectRetriever) { + this.entitySpecs = entitySpecs; Map> entitySearchAnnotations = getSearchableAnnotations(); List annotations = entitySearchAnnotations.values().stream() .flatMap(List::stream) .collect(Collectors.toList()); - _defaultQueryFieldNames = getDefaultQueryFieldNames(annotations); - _highlights = getHighlights(); - _searchQueryBuilder = new SearchQueryBuilder(configs, customSearchConfiguration); - _aggregationQueryBuilder = new AggregationQueryBuilder(configs, entitySearchAnnotations); - _configs = configs; + defaultQueryFieldNames = getDefaultQueryFieldNames(annotations); + highlights = getHighlights(); + searchQueryBuilder = new SearchQueryBuilder(configs, customSearchConfiguration); + aggregationQueryBuilder = + new AggregationQueryBuilder(configs, entitySearchAnnotations, aspectRetriever); + this.configs = configs; searchableFieldTypes = - _entitySpecs.stream() + this.entitySpecs.stream() .flatMap(entitySpec -> entitySpec.getSearchableFieldTypes().entrySet().stream()) .collect( Collectors.toMap( @@ -115,28 +121,35 @@ private SearchRequestHandler( set1.addAll(set2); return set1; })); + this.aspectRetriever = aspectRetriever; } public static SearchRequestHandler getBuilder( @Nonnull EntitySpec entitySpec, @Nonnull SearchConfiguration configs, - @Nullable CustomSearchConfiguration customSearchConfiguration) { + @Nullable CustomSearchConfiguration customSearchConfiguration, + @Nonnull AspectRetriever aspectRetriever) { return REQUEST_HANDLER_BY_ENTITY_NAME.computeIfAbsent( ImmutableList.of(entitySpec), - k -> new SearchRequestHandler(entitySpec, configs, customSearchConfiguration)); + k -> + new SearchRequestHandler( + entitySpec, configs, customSearchConfiguration, aspectRetriever)); } public static SearchRequestHandler getBuilder( @Nonnull List entitySpecs, @Nonnull SearchConfiguration configs, - @Nullable CustomSearchConfiguration customSearchConfiguration) { + @Nullable CustomSearchConfiguration customSearchConfiguration, + @Nonnull AspectRetriever aspectRetriever) { return REQUEST_HANDLER_BY_ENTITY_NAME.computeIfAbsent( ImmutableList.copyOf(entitySpecs), - k -> new SearchRequestHandler(entitySpecs, configs, customSearchConfiguration)); + k -> + new SearchRequestHandler( + entitySpecs, configs, customSearchConfiguration, aspectRetriever)); } private Map> getSearchableAnnotations() { - return _entitySpecs.stream() + return entitySpecs.stream() .map( spec -> Pair.of( @@ -158,13 +171,15 @@ private Set getDefaultQueryFieldNames(List annotat } public BoolQueryBuilder getFilterQuery(@Nullable Filter filter) { - return getFilterQuery(filter, searchableFieldTypes); + return getFilterQuery(filter, searchableFieldTypes, aspectRetriever); } public static BoolQueryBuilder getFilterQuery( @Nullable Filter filter, - Map> searchableFieldTypes) { - BoolQueryBuilder filterQuery = ESUtils.buildFilterQuery(filter, false, searchableFieldTypes); + Map> searchableFieldTypes, + @Nonnull AspectRetriever aspectRetriever) { + BoolQueryBuilder filterQuery = + ESUtils.buildFilterQuery(filter, false, searchableFieldTypes, aspectRetriever); return filterSoftDeletedByDefault(filter, filterQuery); } @@ -209,12 +224,12 @@ public SearchRequest getSearchRequest( .must(getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext()))) .filter(filterQuery)); if (Boolean.FALSE.equals(finalSearchFlags.isSkipAggregates())) { - _aggregationQueryBuilder.getAggregations(facets).forEach(searchSourceBuilder::aggregation); + aggregationQueryBuilder.getAggregations(facets).forEach(searchSourceBuilder::aggregation); } if (Boolean.FALSE.equals(finalSearchFlags.isSkipHighlighting())) { - searchSourceBuilder.highlighter(_highlights); + searchSourceBuilder.highlighter(highlights); } - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, entitySpecs); if (Boolean.TRUE.equals(finalSearchFlags.isGetSuggestions())) { ESUtils.buildNameSuggestions(searchSourceBuilder, input); @@ -265,12 +280,12 @@ public SearchRequest getSearchRequest( .must(getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext()))) .filter(filterQuery)); if (Boolean.FALSE.equals(finalSearchFlags.isSkipAggregates())) { - _aggregationQueryBuilder.getAggregations().forEach(searchSourceBuilder::aggregation); + aggregationQueryBuilder.getAggregations().forEach(searchSourceBuilder::aggregation); } if (Boolean.FALSE.equals(finalSearchFlags.isSkipHighlighting())) { - searchSourceBuilder.highlighter(_highlights); + searchSourceBuilder.highlighter(highlights); } - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, entitySpecs); searchRequest.source(searchSourceBuilder); log.debug("Search request is: " + searchRequest); searchRequest.indicesOptions(null); @@ -297,7 +312,7 @@ public SearchRequest getFilterRequest( final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.query(filterQuery); searchSourceBuilder.from(from).size(size); - ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, _entitySpecs); + ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion, entitySpecs); searchRequest.source(searchSourceBuilder); return searchRequest; @@ -328,7 +343,7 @@ public SearchRequest getAggregationRequest( } public QueryBuilder getQuery(@Nonnull String query, boolean fulltext) { - return _searchQueryBuilder.buildQuery(_entitySpecs, query, fulltext); + return searchQueryBuilder.buildQuery(entitySpecs, query, fulltext); } @VisibleForTesting @@ -340,7 +355,7 @@ public HighlightBuilder getHighlights() { highlightBuilder.postTags(""); // Check for each field name and any subfields - _defaultQueryFieldNames.stream() + defaultQueryFieldNames.stream() .flatMap(fieldName -> Stream.of(fieldName, fieldName + ".*")) .distinct() .forEach(highlightBuilder::field); @@ -445,7 +460,7 @@ private List extractMatchedFields(@Nonnull SearchHit hit) { @Nonnull private Optional getFieldName(String matchedField) { - return _defaultQueryFieldNames.stream().filter(matchedField::startsWith).findFirst(); + return defaultQueryFieldNames.stream().filter(matchedField::startsWith).findFirst(); } private Map extractFeatures(@Nonnull SearchHit searchHit) { @@ -498,7 +513,7 @@ private SearchResultMetadata extractSearchResultMetadata( new SearchResultMetadata().setAggregations(new AggregationMetadataArray()); final List aggregationMetadataList = - _aggregationQueryBuilder.extractAggregationMetadata(searchResponse, filter); + aggregationQueryBuilder.extractAggregationMetadata(searchResponse, filter); searchResultMetadata.setAggregations(new AggregationMetadataArray(aggregationMetadataList)); final List searchSuggestions = extractSearchSuggestions(searchResponse); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java b/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java index d52a80d685fd5b..75c3d23d26c667 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java @@ -11,7 +11,7 @@ import com.linkedin.data.schema.DataSchema; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.Aspect; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.validation.StructuredPropertiesValidator; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java index 86d411e9b5b928..492dc53783904b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java @@ -7,6 +7,7 @@ import static com.linkedin.metadata.search.utils.SearchUtils.isUrn; import com.google.common.collect.ImmutableList; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.SearchableFieldSpec; import com.linkedin.metadata.models.StructuredPropertyUtils; @@ -132,11 +133,15 @@ private ESUtils() {} public static BoolQueryBuilder buildFilterQuery( @Nullable Filter filter, boolean isTimeseries, - final Map> searchableFieldTypes) { + final Map> searchableFieldTypes, + @Nonnull AspectRetriever aspectRetriever) { BoolQueryBuilder finalQueryBuilder = QueryBuilders.boolQuery(); if (filter == null) { return finalQueryBuilder; } + + StructuredPropertyUtils.validateFilter(filter, aspectRetriever); + if (filter.getOr() != null) { // If caller is using the new Filters API, build boolean query from that. filter @@ -387,11 +392,11 @@ public static String escapeReservedCharacters(@Nonnull String input) { public static String toFacetField(@Nonnull final String filterField) { String fieldName = filterField; if (fieldName.startsWith(STRUCTURED_PROPERTY_MAPPING_FIELD + ".")) { + String fqn = fieldName.substring(STRUCTURED_PROPERTY_MAPPING_FIELD.length() + 1); fieldName = STRUCTURED_PROPERTY_MAPPING_FIELD + "." - + StructuredPropertyUtils.sanitizeStructuredPropertyFQN( - fieldName.substring(STRUCTURED_PROPERTY_MAPPING_FIELD.length() + 1)); + + StructuredPropertyUtils.sanitizeStructuredPropertyFQN(fqn); } return fieldName.replace(ESUtils.KEYWORD_SUFFIX, ""); } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/service/UpdateIndicesService.java b/metadata-io/src/main/java/com/linkedin/metadata/service/UpdateIndicesService.java index 3c73d1acab5c25..52f0d680ff4ba1 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/service/UpdateIndicesService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/service/UpdateIndicesService.java @@ -20,9 +20,11 @@ import com.linkedin.dataset.UpstreamLineage; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.aspect.batch.MCLBatchItem; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.entity.ebean.batch.MCLBatchItemImpl; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.AspectsBatch; +import com.linkedin.metadata.aspect.batch.MCLItem; +import com.linkedin.metadata.entity.SearchIndicesService; +import com.linkedin.metadata.entity.ebean.batch.MCLItemImpl; import com.linkedin.metadata.graph.Edge; import com.linkedin.metadata.graph.GraphIndexUtils; import com.linkedin.metadata.graph.GraphService; @@ -67,7 +69,7 @@ import org.springframework.beans.factory.annotation.Value; @Slf4j -public class UpdateIndicesService { +public class UpdateIndicesService implements SearchIndicesService { private static final String DOWNSTREAM_OF = "DownstreamOf"; private final GraphService _graphService; @@ -120,24 +122,21 @@ public UpdateIndicesService( _entityIndexBuilders = entityIndexBuilders; } + @Override public void handleChangeEvent(@Nonnull final MetadataChangeLog event) { try { - MCLBatchItemImpl batch = MCLBatchItemImpl.builder().build(event, aspectRetriever); + MCLItemImpl batch = MCLItemImpl.builder().build(event, aspectRetriever); - Stream sideEffects = - _entityRegistry - .getMCLSideEffects( - event.getChangeType(), event.getEntityType(), event.getAspectName()) - .stream() - .flatMap(mclSideEffect -> mclSideEffect.apply(List.of(batch), aspectRetriever)); + Stream sideEffects = + AspectsBatch.applyMCLSideEffects(List.of(batch), aspectRetriever); - for (MCLBatchItem mclBatchItem : + for (MCLItem mclItem : Stream.concat(Stream.of(batch), sideEffects).collect(Collectors.toList())) { - MetadataChangeLog hookEvent = mclBatchItem.getMetadataChangeLog(); + MetadataChangeLog hookEvent = mclItem.getMetadataChangeLog(); if (UPDATE_CHANGE_TYPES.contains(hookEvent.getChangeType())) { - handleUpdateChangeEvent(mclBatchItem); + handleUpdateChangeEvent(mclItem); } else if (hookEvent.getChangeType() == ChangeType.DELETE) { - handleDeleteChangeEvent(mclBatchItem); + handleDeleteChangeEvent(mclItem); } } } catch (IOException e) { @@ -154,7 +153,7 @@ public void handleChangeEvent(@Nonnull final MetadataChangeLog event) { * * @param event the change event to be processed. */ - private void handleUpdateChangeEvent(@Nonnull final MCLBatchItem event) throws IOException { + private void handleUpdateChangeEvent(@Nonnull final MCLItem event) throws IOException { final EntitySpec entitySpec = event.getEntitySpec(); final AspectSpec aspectSpec = event.getAspectSpec(); @@ -251,7 +250,7 @@ public void updateIndexMappings( * * @param event the change event to be processed. */ - private void handleDeleteChangeEvent(@Nonnull final MCLBatchItem event) { + private void handleDeleteChangeEvent(@Nonnull final MCLItem event) { final EntitySpec entitySpec = event.getEntitySpec(); final Urn urn = event.getUrn(); @@ -696,7 +695,8 @@ private EntitySpec getEventEntitySpec(@Nonnull final MetadataChangeLog event) { * * @param aspectRetriever aspect Retriever */ - public void initializeAspectRetriever(AspectRetriever aspectRetriever) { + @Override + public void initializeAspectRetriever(@Nonnull AspectRetriever aspectRetriever) { this.aspectRetriever = aspectRetriever; this._entityRegistry = aspectRetriever.getEntityRegistry(); this._searchDocumentTransformer.setAspectRetriever(aspectRetriever); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java index cb06dc75c70bc9..e9ace7bf449ef7 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/ElasticSearchTimeseriesAspectService.java @@ -11,6 +11,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.common.urn.Urn; import com.linkedin.data.ByteString; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -95,13 +96,14 @@ public class ElasticSearchTimeseriesAspectService private static final Integer DEFAULT_LIMIT = 10000; - private final IndexConvention _indexConvention; - private final ESBulkProcessor _bulkProcessor; - private final int _numRetries; - private final TimeseriesAspectIndexBuilders _indexBuilders; - private final RestHighLevelClient _searchClient; - private final ESAggregatedStatsDAO _esAggregatedStatsDAO; - private final EntityRegistry _entityRegistry; + private final IndexConvention indexConvention; + private final ESBulkProcessor bulkProcessor; + private final int numRetries; + private final TimeseriesAspectIndexBuilders indexBuilders; + private final RestHighLevelClient searchClient; + private final ESAggregatedStatsDAO esAggregatedStatsDAO; + private final EntityRegistry entityRegistry; + private AspectRetriever aspectRetriever; public ElasticSearchTimeseriesAspectService( @Nonnull RestHighLevelClient searchClient, @@ -110,14 +112,21 @@ public ElasticSearchTimeseriesAspectService( @Nonnull EntityRegistry entityRegistry, @Nonnull ESBulkProcessor bulkProcessor, int numRetries) { - _indexConvention = indexConvention; - _indexBuilders = indexBuilders; - _searchClient = searchClient; - _bulkProcessor = bulkProcessor; - _entityRegistry = entityRegistry; - _numRetries = numRetries; - - _esAggregatedStatsDAO = new ESAggregatedStatsDAO(indexConvention, searchClient, entityRegistry); + this.indexConvention = indexConvention; + this.indexBuilders = indexBuilders; + this.searchClient = searchClient; + this.bulkProcessor = bulkProcessor; + this.entityRegistry = entityRegistry; + this.numRetries = numRetries; + + esAggregatedStatsDAO = new ESAggregatedStatsDAO(indexConvention, searchClient, entityRegistry); + } + + @Override + public ElasticSearchTimeseriesAspectService postConstruct(AspectRetriever aspectRetriever) { + this.aspectRetriever = aspectRetriever; + esAggregatedStatsDAO.setAspectRetriever(aspectRetriever); + return this; } private static EnvelopedAspect parseDocument(@Nonnull SearchHit doc) { @@ -209,24 +218,24 @@ private static Pair toEnvAspectGener @Override public void configure() { - _indexBuilders.reindexAll(); + indexBuilders.reindexAll(); } @Override public List buildReindexConfigs() { - return _indexBuilders.buildReindexConfigs(); + return indexBuilders.buildReindexConfigs(); } @Override public List buildReindexConfigsWithAllStructProps( Collection properties) throws IOException { - return _indexBuilders.buildReindexConfigsWithAllStructProps(properties); + return indexBuilders.buildReindexConfigsWithAllStructProps(properties); } public String reindexAsync( String index, @Nullable QueryBuilder filterQuery, BatchWriteOperationsOptions options) throws Exception { - return _indexBuilders.reindexAsync(index, filterQuery, options); + return indexBuilders.reindexAsync(index, filterQuery, options); } @Override @@ -240,23 +249,23 @@ public void upsertDocument( @Nonnull String aspectName, @Nonnull String docId, @Nonnull JsonNode document) { - String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); final UpdateRequest updateRequest = new UpdateRequest(indexName, docId) .detectNoop(false) .docAsUpsert(true) .doc(document.toString(), XContentType.JSON) - .retryOnConflict(_numRetries); - _bulkProcessor.add(updateRequest); + .retryOnConflict(numRetries); + bulkProcessor.add(updateRequest); } @Override public List getIndexSizes() { List res = new ArrayList<>(); try { - String indicesPattern = _indexConvention.getAllTimeseriesAspectIndicesPattern(); + String indicesPattern = indexConvention.getAllTimeseriesAspectIndicesPattern(); Response r = - _searchClient + searchClient .getLowLevelClient() .performRequest(new Request("GET", "/" + indicesPattern + "/_stats")); JsonNode body = new ObjectMapper().readTree(r.getEntity().getContent()); @@ -267,7 +276,7 @@ public List getIndexSizes() { TimeseriesIndexSizeResult elemResult = new TimeseriesIndexSizeResult(); elemResult.setIndexName(entry.getKey()); Optional> indexEntityAndAspect = - _indexConvention.getEntityAndAspectName(entry.getKey()); + indexConvention.getEntityAndAspectName(entry.getKey()); if (indexEntityAndAspect.isPresent()) { elemResult.setEntityName(indexEntityAndAspect.get().getFirst()); elemResult.setAspectName(indexEntityAndAspect.get().getSecond()); @@ -289,19 +298,20 @@ public long countByFilter( @Nonnull final String entityName, @Nonnull final String aspectName, @Nullable final Filter filter) { - final String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + final String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); final BoolQueryBuilder filterQueryBuilder = QueryBuilders.boolQuery() .must( ESUtils.buildFilterQuery( filter, true, - _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes())); + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(), + aspectRetriever)); CountRequest countRequest = new CountRequest(); countRequest.query(filterQueryBuilder); countRequest.indices(indexName); try { - CountResponse resp = _searchClient.count(countRequest, RequestOptions.DEFAULT); + CountResponse resp = searchClient.count(countRequest, RequestOptions.DEFAULT); return resp.getCount(); } catch (IOException e) { log.error("Count query failed:", e); @@ -320,10 +330,10 @@ public List getAspectValues( @Nullable final Filter filter, @Nullable final SortCriterion sort) { Map> searchableFieldTypes = - _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(); + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(); final BoolQueryBuilder filterQueryBuilder = QueryBuilders.boolQuery() - .must(ESUtils.buildFilterQuery(filter, true, searchableFieldTypes)); + .must(ESUtils.buildFilterQuery(filter, true, searchableFieldTypes, aspectRetriever)); filterQueryBuilder.must(QueryBuilders.matchQuery("urn", urn.toString())); // NOTE: We are interested only in the un-exploded rows as only they carry the `event` payload. filterQueryBuilder.mustNot(QueryBuilders.termQuery(MappingsBuilder.IS_EXPLODED_FIELD, true)); @@ -363,7 +373,7 @@ public List getAspectValues( final SearchRequest searchRequest = new SearchRequest(); searchRequest.source(searchSourceBuilder); - String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); searchRequest.indices(indexName); log.debug("Search request is: " + searchRequest); @@ -371,7 +381,7 @@ public List getAspectValues( try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "searchAspectValues_search").time()) { final SearchResponse searchResponse = - _searchClient.search(searchRequest, RequestOptions.DEFAULT); + searchClient.search(searchRequest, RequestOptions.DEFAULT); hits = searchResponse.getHits(); } catch (Exception e) { log.error("Search query failed:", e); @@ -390,7 +400,7 @@ public GenericTable getAggregatedStats( @Nonnull AggregationSpec[] aggregationSpecs, @Nullable Filter filter, @Nullable GroupingBucket[] groupingBuckets) { - return _esAggregatedStatsDAO.getAggregatedStats( + return esAggregatedStatsDAO.getAggregatedStats( entityName, aspectName, aggregationSpecs, filter, groupingBuckets); } @@ -410,13 +420,16 @@ public GenericTable getAggregatedStats( @Override public DeleteAspectValuesResult deleteAspectValues( @Nonnull String entityName, @Nonnull String aspectName, @Nonnull Filter filter) { - final String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + final String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); final BoolQueryBuilder filterQueryBuilder = ESUtils.buildFilterQuery( - filter, true, _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes()); + filter, + true, + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(), + aspectRetriever); final Optional result = - _bulkProcessor + bulkProcessor .deleteByQuery( filterQueryBuilder, false, DEFAULT_LIMIT, TimeValue.timeValueMinutes(10), indexName) .map( @@ -438,17 +451,20 @@ public String deleteAspectValuesAsync( @Nonnull String aspectName, @Nonnull Filter filter, @Nonnull BatchWriteOperationsOptions options) { - final String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + final String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); final BoolQueryBuilder filterQueryBuilder = ESUtils.buildFilterQuery( - filter, true, _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes()); + filter, + true, + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(), + aspectRetriever); final int batchSize = options.getBatchSize() > 0 ? options.getBatchSize() : DEFAULT_LIMIT; TimeValue timeout = options.getTimeoutSeconds() > 0 ? TimeValue.timeValueSeconds(options.getTimeoutSeconds()) : null; final Optional result = - _bulkProcessor.deleteByQueryAsync(filterQueryBuilder, false, batchSize, timeout, indexName); + bulkProcessor.deleteByQueryAsync(filterQueryBuilder, false, batchSize, timeout, indexName); if (result.isPresent()) { return result.get().getTask(); @@ -464,10 +480,13 @@ public String reindexAsync( @Nonnull String aspectName, @Nonnull Filter filter, @Nonnull BatchWriteOperationsOptions options) { - final String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + final String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); final BoolQueryBuilder filterQueryBuilder = ESUtils.buildFilterQuery( - filter, true, _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes()); + filter, + true, + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(), + aspectRetriever); try { return this.reindexAsync(indexName, filterQueryBuilder, options); } catch (Exception e) { @@ -484,7 +503,7 @@ public DeleteAspectValuesResult rollbackTimeseriesAspects(@Nonnull String runId) Filter filter = QueryUtils.newFilter("runId", runId); // Delete the timeseries aspects across all entities with the runId. - for (Map.Entry entry : _entityRegistry.getEntitySpecs().entrySet()) { + for (Map.Entry entry : entityRegistry.getEntitySpecs().entrySet()) { for (AspectSpec aspectSpec : entry.getValue().getAspectSpecs()) { if (aspectSpec.isTimeseries()) { DeleteAspectValuesResult result = @@ -517,10 +536,10 @@ public TimeseriesScrollResult scrollAspects( @Nullable Long endTimeMillis) { Map> searchableFieldTypes = - _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(); + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(); final BoolQueryBuilder filterQueryBuilder = QueryBuilders.boolQuery() - .filter(ESUtils.buildFilterQuery(filter, true, searchableFieldTypes)); + .filter(ESUtils.buildFilterQuery(filter, true, searchableFieldTypes, aspectRetriever)); if (startTimeMillis != null) { Criterion startTimeCriterion = @@ -583,11 +602,11 @@ private SearchResponse executeScrollSearchQuery( searchRequest.source(searchSourceBuilder); ESUtils.setSearchAfter(searchSourceBuilder, sort, null, null); - searchRequest.indices(_indexConvention.getTimeseriesAspectIndexName(entityName, aspectName)); + searchRequest.indices(indexConvention.getTimeseriesAspectIndexName(entityName, aspectName)); try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "scrollAspects_search").time()) { - return _searchClient.search(searchRequest, RequestOptions.DEFAULT); + return searchClient.search(searchRequest, RequestOptions.DEFAULT); } catch (Exception e) { log.error("Search query failed", e); throw new ESQueryException("Search query failed:", e); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/TimeseriesAspectIndexBuilders.java b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/TimeseriesAspectIndexBuilders.java index 6437bbc390d829..b59cd3a647d71c 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/TimeseriesAspectIndexBuilders.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/indexbuilder/TimeseriesAspectIndexBuilders.java @@ -15,6 +15,7 @@ import java.util.List; import java.util.Optional; import java.util.stream.Collectors; +import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -23,15 +24,15 @@ @Slf4j @RequiredArgsConstructor public class TimeseriesAspectIndexBuilders implements ElasticSearchIndexed { - private final ESIndexBuilder _indexBuilder; - private final EntityRegistry _entityRegistry; - private final IndexConvention _indexConvention; + @Nonnull private final ESIndexBuilder indexBuilder; + @Nonnull private final EntityRegistry entityRegistry; + @Nonnull private final IndexConvention indexConvention; @Override public void reindexAll() { for (ReindexConfig config : buildReindexConfigs()) { try { - _indexBuilder.buildIndex(config); + indexBuilder.buildIndex(config); } catch (IOException e) { throw new RuntimeException(e); } @@ -41,13 +42,13 @@ public void reindexAll() { public String reindexAsync( String index, @Nullable QueryBuilder filterQuery, BatchWriteOperationsOptions options) throws Exception { - Optional> entityAndAspect = _indexConvention.getEntityAndAspectName(index); + Optional> entityAndAspect = indexConvention.getEntityAndAspectName(index); if (entityAndAspect.isEmpty()) { throw new IllegalArgumentException("Could not extract entity and aspect from index " + index); } String entityName = entityAndAspect.get().getFirst(); String aspectName = entityAndAspect.get().getSecond(); - EntitySpec entitySpec = _entityRegistry.getEntitySpec(entityName); + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); for (String aspect : entitySpec.getAspectSpecMap().keySet()) { if (aspect.toLowerCase().equals(aspectName)) { aspectName = aspect; @@ -59,17 +60,17 @@ public String reindexAsync( String.format("Could not find aspect %s of entity %s", aspectName, entityName)); } ReindexConfig config = - _indexBuilder.buildReindexState( + indexBuilder.buildReindexState( index, MappingsBuilder.getMappings( - _entityRegistry.getEntitySpec(entityName).getAspectSpec(aspectName)), + entityRegistry.getEntitySpec(entityName).getAspectSpec(aspectName)), Collections.emptyMap()); - return _indexBuilder.reindexInPlaceAsync(index, filterQuery, options, config); + return indexBuilder.reindexInPlaceAsync(index, filterQuery, options, config); } @Override public List buildReindexConfigs() { - return _entityRegistry.getEntitySpecs().values().stream() + return entityRegistry.getEntitySpecs().values().stream() .flatMap( entitySpec -> entitySpec.getAspectSpecs().stream() @@ -78,8 +79,8 @@ public List buildReindexConfigs() { .map( pair -> { try { - return _indexBuilder.buildReindexState( - _indexConvention.getTimeseriesAspectIndexName( + return indexBuilder.buildReindexState( + indexConvention.getTimeseriesAspectIndexName( pair.getFirst().getName(), pair.getSecond().getName()), MappingsBuilder.getMappings(pair.getSecond()), Collections.emptyMap()); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/query/ESAggregatedStatsDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/query/ESAggregatedStatsDAO.java index 580888e54b7007..1324aebb80006d 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/query/ESAggregatedStatsDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeseries/elastic/query/ESAggregatedStatsDAO.java @@ -5,6 +5,7 @@ import com.linkedin.data.schema.DataSchema; import com.linkedin.data.template.StringArray; import com.linkedin.data.template.StringArrayArray; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.TimeseriesFieldCollectionSpec; @@ -29,6 +30,7 @@ import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import lombok.Setter; import lombok.extern.slf4j.Slf4j; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; @@ -61,17 +63,18 @@ public class ESAggregatedStatsDAO { ES_AGGREGATION_PREFIX + ES_MAX_AGGREGATION_PREFIX + ES_FIELD_TIMESTAMP; private static final int MAX_TERM_BUCKETS = 24 * 60; // minutes in a day. - private final IndexConvention _indexConvention; - private final RestHighLevelClient _searchClient; - private final EntityRegistry _entityRegistry; + private final IndexConvention indexConvention; + private final RestHighLevelClient searchClient; + private final EntityRegistry entityRegistry; + @Setter private AspectRetriever aspectRetriever; public ESAggregatedStatsDAO( @Nonnull IndexConvention indexConvention, @Nonnull RestHighLevelClient searchClient, @Nonnull EntityRegistry entityRegistry) { - _indexConvention = indexConvention; - _searchClient = searchClient; - _entityRegistry = entityRegistry; + this.indexConvention = indexConvention; + this.searchClient = searchClient; + this.entityRegistry = entityRegistry; } private static String toEsAggName(final String aggName) { @@ -353,7 +356,7 @@ private static String extractAggregationValue( private AspectSpec getTimeseriesAspectSpec( @Nonnull String entityName, @Nonnull String aspectName) { - EntitySpec entitySpec = _entityRegistry.getEntitySpec(entityName); + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); if (aspectSpec == null) { new IllegalArgumentException( @@ -379,7 +382,10 @@ public GenericTable getAggregatedStats( // Setup the filter query builder using the input filter provided. final BoolQueryBuilder filterQueryBuilder = ESUtils.buildFilterQuery( - filter, true, _entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes()); + filter, + true, + entityRegistry.getEntitySpec(entityName).getSearchableFieldTypes(), + aspectRetriever); AspectSpec aspectSpec = getTimeseriesAspectSpec(entityName, aspectName); // Build and attach the grouping aggregations @@ -402,14 +408,14 @@ public GenericTable getAggregatedStats( final SearchRequest searchRequest = new SearchRequest(); searchRequest.source(searchSourceBuilder); - final String indexName = _indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); + final String indexName = indexConvention.getTimeseriesAspectIndexName(entityName, aspectName); searchRequest.indices(indexName); log.debug("Search request is: " + searchRequest); try { final SearchResponse searchResponse = - _searchClient.search(searchRequest, RequestOptions.DEFAULT); + searchClient.search(searchRequest, RequestOptions.DEFAULT); return generateResponseFromElastic( searchResponse, groupingBuckets, aggregationSpecs, aspectSpec); } catch (Exception e) { diff --git a/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java b/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java index 72bbc794171ff9..84d084e14d54d0 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java @@ -6,7 +6,7 @@ import com.linkedin.identity.CorpUserInfo; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.key.CorpUserKey; import java.util.HashMap; import java.util.LinkedList; @@ -26,16 +26,16 @@ public static Map ingestCorpUserKeyAspects( @Nonnull public static Map ingestCorpUserKeyAspects( - EntityService entityService, int aspectCount, int startIndex) { + EntityService entityService, int aspectCount, int startIndex) { String aspectName = AspectGenerationUtils.getAspectName(new CorpUserKey()); Map aspects = new HashMap<>(); - List items = new LinkedList<>(); + List items = new LinkedList<>(); for (int i = startIndex; i < startIndex + aspectCount; i++) { Urn urn = UrnUtils.getUrn(String.format("urn:li:corpuser:tester%d", i)); CorpUserKey aspect = AspectGenerationUtils.createCorpUserKey(urn); aspects.put(urn, aspect); items.add( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(urn) .aspectName(aspectName) .recordTemplate(aspect) @@ -43,7 +43,8 @@ public static Map ingestCorpUserKeyAspects( .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService)); } - entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + entityService.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(entityService).items(items).build(), true, true); return aspects; } @@ -58,14 +59,14 @@ public static Map ingestCorpUserInfoAspects( @Nonnull final EntityService entityService, int aspectCount, int startIndex) { String aspectName = AspectGenerationUtils.getAspectName(new CorpUserInfo()); Map aspects = new HashMap<>(); - List items = new LinkedList<>(); + List items = new LinkedList<>(); for (int i = startIndex; i < startIndex + aspectCount; i++) { Urn urn = UrnUtils.getUrn(String.format("urn:li:corpuser:tester%d", i)); String email = String.format("email%d@test.com", i); CorpUserInfo aspect = AspectGenerationUtils.createCorpUserInfo(email); aspects.put(urn, aspect); items.add( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(urn) .aspectName(aspectName) .recordTemplate(aspect) @@ -73,7 +74,8 @@ public static Map ingestCorpUserInfoAspects( .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService)); } - entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + entityService.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(entityService).items(items).build(), true, true); return aspects; } @@ -88,7 +90,7 @@ public static Map ingestChartInfoAspects( @Nonnull final EntityService entityService, int aspectCount, int startIndex) { String aspectName = AspectGenerationUtils.getAspectName(new ChartInfo()); Map aspects = new HashMap<>(); - List items = new LinkedList<>(); + List items = new LinkedList<>(); for (int i = startIndex; i < startIndex + aspectCount; i++) { Urn urn = UrnUtils.getUrn(String.format("urn:li:chart:(looker,test%d)", i)); String title = String.format("Test Title %d", i); @@ -96,7 +98,7 @@ public static Map ingestChartInfoAspects( ChartInfo aspect = AspectGenerationUtils.createChartInfo(title, description); aspects.put(urn, aspect); items.add( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(urn) .aspectName(aspectName) .recordTemplate(aspect) @@ -104,7 +106,8 @@ public static Map ingestChartInfoAspects( .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService)); } - entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + entityService.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(entityService).items(items).build(), true, true); return aspects; } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtilTest.java b/metadata-io/src/test/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtilTest.java index 308832a9c63ef0..c38e14711fe966 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtilTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtilTest.java @@ -8,7 +8,7 @@ import com.linkedin.common.urn.DatasetUrn; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.EbeanTestUtils; -import com.linkedin.metadata.aspect.batch.MCPBatchItem; +import com.linkedin.metadata.aspect.batch.MCPItem; import com.linkedin.metadata.aspect.patch.builder.DatasetPropertiesPatchBuilder; import com.linkedin.metadata.config.EbeanConfiguration; import com.linkedin.metadata.config.PreProcessHooks; @@ -50,7 +50,7 @@ public void testAdditionalChanges() { preProcessHooks.setUiEnabled(true); EntityServiceImpl entityServiceImpl = new EntityServiceImpl( - aspectDao, mockProducer, _testEntityRegistry, true, null, preProcessHooks, false); + aspectDao, mockProducer, _testEntityRegistry, true, preProcessHooks, false); MetadataChangeProposal proposal1 = new DatasetPropertiesPatchBuilder() @@ -71,7 +71,7 @@ public void testAdditionalChanges() { entityServiceImpl, false) .stream() - .map(MCPBatchItem::getMetadataChangeProposal) + .map(MCPItem::getMetadataChangeProposal) .collect(Collectors.toList()); // proposals for key aspect, browsePath, browsePathV2, dataPlatformInstance Assert.assertEquals(proposalList.size(), 4); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraAspectMigrationsDaoTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraAspectMigrationsDaoTest.java index d191ea2b9fa971..660fb1af47be4e 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraAspectMigrationsDaoTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraAspectMigrationsDaoTest.java @@ -48,14 +48,8 @@ private void configureComponents() { PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); _entityServiceImpl = - new EntityServiceImpl( - dao, - _mockProducer, - _testEntityRegistry, - true, - _mockUpdateIndicesService, - preProcessHooks, - true); + new EntityServiceImpl(dao, _mockProducer, _testEntityRegistry, true, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); _retentionService = new CassandraRetentionService(_entityServiceImpl, session, 1000); _entityServiceImpl.setRetentionService(_retentionService); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraEntityServiceTest.java index 8d30fb02915c70..e1dd9eb21e78be 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraEntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/CassandraEntityServiceTest.java @@ -72,13 +72,8 @@ private void configureComponents() { preProcessHooks.setUiEnabled(true); _entityServiceImpl = new EntityServiceImpl( - _aspectDao, - _mockProducer, - _testEntityRegistry, - false, - _mockUpdateIndicesService, - preProcessHooks, - true); + _aspectDao, _mockProducer, _testEntityRegistry, false, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); _retentionService = new CassandraRetentionService(_entityServiceImpl, session, 1000); _entityServiceImpl.setRetentionService(_retentionService); } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java index 42fa2acb542375..390cbc7392b2ea 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java @@ -54,13 +54,8 @@ public DeleteEntityServiceTest() { preProcessHooks.setUiEnabled(true); _entityServiceImpl = new EntityServiceImpl( - _aspectDao, - mock(EventProducer.class), - _entityRegistry, - true, - _mockUpdateIndicesService, - preProcessHooks, - true); + _aspectDao, mock(EventProducer.class), _entityRegistry, true, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); _deleteEntityService = new DeleteEntityService(_entityServiceImpl, _graphService); } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java index d241fb3b9581b4..75a4a76192a207 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java @@ -39,14 +39,8 @@ public void setupTest() { PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); _entityServiceImpl = - new EntityServiceImpl( - dao, - _mockProducer, - _testEntityRegistry, - true, - _mockUpdateIndicesService, - preProcessHooks, - true); + new EntityServiceImpl(dao, _mockProducer, _testEntityRegistry, true, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); _retentionService = new EbeanRetentionService(_entityServiceImpl, server, 1000); _entityServiceImpl.setRetentionService(_retentionService); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java index 1e2cf4d4255d2e..586f6d3b79a8f4 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java @@ -18,7 +18,7 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectDao; import com.linkedin.metadata.entity.ebean.EbeanRetentionService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; import com.linkedin.metadata.models.registry.EntityRegistryException; @@ -71,13 +71,8 @@ public void setupTest() { preProcessHooks.setUiEnabled(true); _entityServiceImpl = new EntityServiceImpl( - _aspectDao, - _mockProducer, - _testEntityRegistry, - false, - _mockUpdateIndicesService, - preProcessHooks, - true); + _aspectDao, _mockProducer, _testEntityRegistry, false, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); _retentionService = new EbeanRetentionService(_entityServiceImpl, server, 1000); _entityServiceImpl.setRetentionService(_retentionService); } @@ -118,30 +113,33 @@ public void testIngestListLatestAspects() throws AssertionError { // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn2) .aspectName(aspectName) .recordTemplate(writeAspect2) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn3) .aspectName(aspectName) .recordTemplate(writeAspect3) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // List aspects ListResult batch1 = @@ -187,30 +185,33 @@ public void testIngestListUrns() throws AssertionError { // Ingest CorpUserInfo Aspect #3 RecordTemplate writeAspect3 = AspectGenerationUtils.createCorpUserKey(entityUrn3); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn2) .aspectName(aspectName) .recordTemplate(writeAspect2) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn3) .aspectName(aspectName) .recordTemplate(writeAspect3) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // List aspects urns ListUrnsResult batch1 = _entityServiceImpl.listUrns(entityUrn1.getEntityType(), 0, 2); @@ -284,7 +285,7 @@ public void dataGeneratorThreadingTest() { @Test // ensure same thread as h2 public void multiThreadingTest() { DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl); - Database server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); + Database server = ((EbeanAspectDao) _entityServiceImpl.aspectDao).getServer(); // Add data List aspects = List.of("status", "globalTags", "glossaryTerms"); @@ -340,7 +341,7 @@ public void multiThreadingTest() { @Test public void singleThreadingTest() { DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl); - Database server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); + Database server = ((EbeanAspectDao) _entityServiceImpl.aspectDao).getServer(); // Add data List aspects = List.of("status", "globalTags", "glossaryTerms"); @@ -393,7 +394,7 @@ private static void executeThreadingTest( EntityServiceImpl entityService, List> testData, int threadCount) { - Database server = ((EbeanAspectDao) entityService._aspectDao).getServer(); + Database server = ((EbeanAspectDao) entityService.aspectDao).getServer(); server.sqlUpdate("truncate metadata_aspect_v2"); int count = diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java index 384b54c7a1c8d3..e325e23ef86070 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java @@ -39,8 +39,9 @@ import com.linkedin.metadata.aspect.CorpUserAspectArray; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; +import com.linkedin.metadata.entity.validation.ValidationUtils; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; import com.linkedin.metadata.models.AspectSpec; @@ -864,37 +865,40 @@ public void testRollbackAspect() throws AssertionError { CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn2) .aspectName(aspectName) .recordTemplate(writeAspect2) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn3) .aspectName(aspectName) .recordTemplate(writeAspect3) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1Overwrite) .systemMetadata(metadata2) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // this should no-op since this run has been overwritten AspectRowSummary rollbackOverwrittenAspect = new AspectRowSummary(); @@ -943,30 +947,33 @@ public void testRollbackKey() throws AssertionError { CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(keyAspectName) .recordTemplate(writeKey1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1Overwrite) .systemMetadata(metadata2) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // this should no-op since the key should have been written in the furst run AspectRowSummary rollbackKeyWithWrongRunId = new AspectRowSummary(); @@ -1023,44 +1030,47 @@ public void testRollbackUrn() throws AssertionError { CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(keyAspectName) .recordTemplate(writeKey1) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn2) .aspectName(aspectName) .recordTemplate(writeAspect2) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn3) .aspectName(aspectName) .recordTemplate(writeAspect3) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn1) .aspectName(aspectName) .recordTemplate(writeAspect1Overwrite) .systemMetadata(metadata2) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // this should no-op since the key should have been written in the furst run AspectRowSummary rollbackKeyWithWrongRunId = new AspectRowSummary(); @@ -1090,16 +1100,19 @@ public void testIngestGetLatestAspect() throws AssertionError { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // Validate retrieval of CorpUserInfo Aspect #1 RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1132,14 +1145,17 @@ public void testIngestGetLatestAspect() throws AssertionError { items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect2) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata2) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // Validate retrieval of CorpUserInfo Aspect #2 RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1176,16 +1192,19 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1) .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // Validate retrieval of CorpUserInfo Aspect #1 EnvelopedAspect readAspect1 = @@ -1198,14 +1217,17 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect2) .systemMetadata(metadata2) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // Validate retrieval of CorpUserInfo Aspect #2 EnvelopedAspect readAspect2 = @@ -1250,16 +1272,19 @@ public void testIngestSameAspect() throws AssertionError { SystemMetadata metadata3 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-123", "run-456"); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // Validate retrieval of CorpUserInfo Aspect #1 RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1292,14 +1317,17 @@ public void testIngestSameAspect() throws AssertionError { items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect2) .systemMetadata(metadata2) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); // Validate retrieval of CorpUserInfo Aspect #2 RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1343,51 +1371,54 @@ public void testRetention() throws AssertionError { Status writeAspect2a = new Status().setRemoved(false); Status writeAspect2b = new Status().setRemoved(true); - List items = + List items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1a) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1b) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName2) .recordTemplate(writeAspect2) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName2) .recordTemplate(writeAspect2a) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName2) .recordTemplate(writeAspect2b) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName, 1), writeAspect1); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1412,21 +1443,24 @@ public void testRetention() throws AssertionError { items = List.of( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName) .recordTemplate(writeAspect1c) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl), - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectName2) .recordTemplate(writeAspect2c) .systemMetadata(metadata1) .auditStamp(TEST_AUDIT_STAMP) .build(_entityServiceImpl)); - _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityServiceImpl.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityServiceImpl).items(items).build(), + true, + true); assertNull(_entityServiceImpl.getAspect(entityUrn, aspectName, 1)); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1564,12 +1598,12 @@ public void testRestoreIndices() throws Exception { public void testValidateUrn() throws Exception { // Valid URN Urn validTestUrn = new Urn("li", "corpuser", new TupleKey("testKey")); - EntityUtils.validateUrn(_testEntityRegistry, validTestUrn); + ValidationUtils.validateUrn(_testEntityRegistry, validTestUrn); // URN with trailing whitespace Urn testUrnWithTrailingWhitespace = new Urn("li", "corpuser", new TupleKey("testKey ")); try { - EntityUtils.validateUrn(_testEntityRegistry, testUrnWithTrailingWhitespace); + ValidationUtils.validateUrn(_testEntityRegistry, testUrnWithTrailingWhitespace); Assert.fail("Should have raised IllegalArgumentException for URN with trailing whitespace"); } catch (IllegalArgumentException e) { assertEquals( @@ -1581,7 +1615,7 @@ public void testValidateUrn() throws Exception { Urn testUrnTooLong = new Urn("li", "corpuser", new TupleKey(stringTooLong)); try { - EntityUtils.validateUrn(_testEntityRegistry, testUrnTooLong); + ValidationUtils.validateUrn(_testEntityRegistry, testUrnTooLong); Assert.fail("Should have raised IllegalArgumentException for URN too long"); } catch (IllegalArgumentException e) { assertEquals( @@ -1600,9 +1634,9 @@ public void testValidateUrn() throws Exception { Urn testUrnSameLengthWhenEncoded = new Urn("li", "corpUser", new TupleKey(buildStringSameLengthWhenEncoded.toString())); // Same length when encoded should be allowed, the encoded one should not be - EntityUtils.validateUrn(_testEntityRegistry, testUrnSameLengthWhenEncoded); + ValidationUtils.validateUrn(_testEntityRegistry, testUrnSameLengthWhenEncoded); try { - EntityUtils.validateUrn(_testEntityRegistry, testUrnTooLongWhenEncoded); + ValidationUtils.validateUrn(_testEntityRegistry, testUrnTooLongWhenEncoded); Assert.fail("Should have raised IllegalArgumentException for URN too long"); } catch (IllegalArgumentException e) { assertEquals( @@ -1612,9 +1646,9 @@ public void testValidateUrn() throws Exception { // Urn containing disallowed character Urn testUrnSpecialCharValid = new Urn("li", "corpUser", new TupleKey("bob␇")); Urn testUrnSpecialCharInvalid = new Urn("li", "corpUser", new TupleKey("bob␟")); - EntityUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharValid); + ValidationUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharValid); try { - EntityUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharInvalid); + ValidationUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharInvalid); Assert.fail( "Should have raised IllegalArgumentException for URN containing the illegal char"); } catch (IllegalArgumentException e) { @@ -1623,7 +1657,7 @@ public void testValidateUrn() throws Exception { Urn urnWithMismatchedParens = new Urn("li", "corpuser", new TupleKey("test(Key")); try { - EntityUtils.validateUrn(_testEntityRegistry, urnWithMismatchedParens); + ValidationUtils.validateUrn(_testEntityRegistry, urnWithMismatchedParens); Assert.fail("Should have raised IllegalArgumentException for URN with mismatched parens"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("mismatched paren nesting")); @@ -1631,7 +1665,7 @@ public void testValidateUrn() throws Exception { Urn invalidType = new Urn("li", "fakeMadeUpType", new TupleKey("testKey")); try { - EntityUtils.validateUrn(_testEntityRegistry, invalidType); + ValidationUtils.validateUrn(_testEntityRegistry, invalidType); Assert.fail( "Should have raised IllegalArgumentException for URN with non-existent entity type"); } catch (IllegalArgumentException e) { @@ -1640,12 +1674,12 @@ public void testValidateUrn() throws Exception { Urn validFabricType = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "PROD")); - EntityUtils.validateUrn(_testEntityRegistry, validFabricType); + ValidationUtils.validateUrn(_testEntityRegistry, validFabricType); Urn invalidFabricType = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "prod")); try { - EntityUtils.validateUrn(_testEntityRegistry, invalidFabricType); + ValidationUtils.validateUrn(_testEntityRegistry, invalidFabricType); Assert.fail("Should have raised IllegalArgumentException for URN with invalid fabric type"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains(invalidFabricType.toString())); @@ -1654,7 +1688,7 @@ public void testValidateUrn() throws Exception { Urn urnEndingInComma = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "PROD", "")); try { - EntityUtils.validateUrn(_testEntityRegistry, urnEndingInComma); + ValidationUtils.validateUrn(_testEntityRegistry, urnEndingInComma); Assert.fail("Should have raised IllegalArgumentException for URN ending in comma"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains(urnEndingInComma.toString())); @@ -1750,12 +1784,9 @@ public void testStructuredPropertyIngestProposal() throws Exception { STRUCTURED_PROPERTY_ENTITY_NAME, STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME) .map( entityAspect -> - EntityUtils.toAspectRecord( - STRUCTURED_PROPERTY_ENTITY_NAME, - STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME, - entityAspect.getMetadata(), - _testEntityRegistry)) - .map(recordTemplate -> (StructuredPropertyDefinition) recordTemplate) + EntityUtils.toSystemAspect(entityAspect, _entityServiceImpl) + .get() + .getAspect(StructuredPropertyDefinition.class)) .collect(Collectors.toSet()); assertEquals(defs.size(), 1); assertEquals(defs, Set.of(structuredPropertyDefinition)); @@ -1826,12 +1857,9 @@ public void testStructuredPropertyIngestProposal() throws Exception { STRUCTURED_PROPERTY_ENTITY_NAME, STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME) .map( entityAspect -> - EntityUtils.toAspectRecord( - STRUCTURED_PROPERTY_ENTITY_NAME, - STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME, - entityAspect.getMetadata(), - _testEntityRegistry)) - .map(recordTemplate -> (StructuredPropertyDefinition) recordTemplate) + EntityUtils.toSystemAspect(entityAspect, _entityServiceImpl) + .get() + .getAspect(StructuredPropertyDefinition.class)) .collect(Collectors.toSet()); assertEquals(defs.size(), 2); assertEquals(defs, Set.of(secondDefinition, structuredPropertyDefinition)); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java index 079ec084625150..9588140bebd65f 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java @@ -31,6 +31,7 @@ import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; import com.linkedin.data.template.LongMap; import com.linkedin.metadata.TestEntityUtil; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.cache.EntityDocCountCacheConfiguration; import com.linkedin.metadata.config.cache.SearchLineageCacheConfiguration; import com.linkedin.metadata.config.search.SearchConfiguration; @@ -40,7 +41,6 @@ import com.linkedin.metadata.graph.LineageDirection; import com.linkedin.metadata.graph.LineageRelationship; import com.linkedin.metadata.graph.LineageRelationshipArray; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; @@ -63,6 +63,7 @@ import com.linkedin.metadata.search.utils.QueryUtils; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import com.linkedin.r2.RemoteInvocationException; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collections; @@ -102,14 +103,14 @@ public abstract class LineageServiceTestBase extends AbstractTestNGSpringContext @Nonnull protected abstract CustomSearchConfiguration getCustomSearchConfiguration(); - private EntityRegistry _entityRegistry; - private IndexConvention _indexConvention; - private SettingsBuilder _settingsBuilder; - private ElasticSearchService _elasticSearchService; - private GraphService _graphService; - private CacheManager _cacheManager; - private LineageSearchService _lineageSearchService; - private RestHighLevelClient _searchClientSpy; + private AspectRetriever aspectRetriever; + private IndexConvention indexConvention; + private SettingsBuilder settingsBuilder; + private ElasticSearchService elasticSearchService; + private GraphService graphService; + private CacheManager cacheManager; + private LineageSearchService lineageSearchService; + private RestHighLevelClient searchClientSpy; private static final String ENTITY_NAME = "testEntity"; private static final Urn TEST_URN = TestEntityUtil.getTestEntityUrn(); @@ -126,20 +127,23 @@ public void disableAssert() { } @BeforeClass - public void setup() { - _entityRegistry = new SnapshotEntityRegistry(new Snapshot()); - _indexConvention = new IndexConventionImpl("lineage_search_service_test"); - _settingsBuilder = new SettingsBuilder(null); - _elasticSearchService = buildEntitySearchService(); - _elasticSearchService.configure(); - _cacheManager = new ConcurrentMapCacheManager(); - _graphService = mock(GraphService.class); + public void setup() throws RemoteInvocationException, URISyntaxException { + aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()) + .thenReturn(new SnapshotEntityRegistry(new Snapshot())); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + indexConvention = new IndexConventionImpl("lineage_search_service_test"); + settingsBuilder = new SettingsBuilder(null); + elasticSearchService = buildEntitySearchService(); + elasticSearchService.configure(); + cacheManager = new ConcurrentMapCacheManager(); + graphService = mock(GraphService.class); resetService(true, false); } private void resetService(boolean withCache, boolean withLightingCache) { CachingEntitySearchService cachingEntitySearchService = - new CachingEntitySearchService(_cacheManager, _elasticSearchService, 100, true); + new CachingEntitySearchService(cacheManager, elasticSearchService, 100, true); EntityDocCountCacheConfiguration entityDocCountCacheConfiguration = new EntityDocCountCacheConfiguration(); entityDocCountCacheConfiguration.setTtlSeconds(600L); @@ -149,23 +153,25 @@ private void resetService(boolean withCache, boolean withLightingCache) { searchLineageCacheConfiguration.setTtlSeconds(600L); searchLineageCacheConfiguration.setLightningThreshold(withLightingCache ? -1 : 300); - _lineageSearchService = + lineageSearchService = spy( new LineageSearchService( new SearchService( new EntityDocCountCache( - _entityRegistry, _elasticSearchService, entityDocCountCacheConfiguration), + aspectRetriever.getEntityRegistry(), + elasticSearchService, + entityDocCountCacheConfiguration), cachingEntitySearchService, new SimpleRanker()), - _graphService, - _cacheManager.getCache("test"), + graphService, + cacheManager.getCache("test"), withCache, searchLineageCacheConfiguration)); } @BeforeMethod public void wipe() throws Exception { - _elasticSearchService.clear(); + elasticSearchService.clear(); clearCache(false); syncAfterWrite(getBulkProcessor()); } @@ -174,31 +180,38 @@ public void wipe() throws Exception { private ElasticSearchService buildEntitySearchService() { EntityIndexBuilders indexBuilders = new EntityIndexBuilders( - getIndexBuilder(), _entityRegistry, _indexConvention, _settingsBuilder); - _searchClientSpy = spy(getSearchClient()); + getIndexBuilder(), + aspectRetriever.getEntityRegistry(), + indexConvention, + settingsBuilder); + searchClientSpy = spy(getSearchClient()); ESSearchDAO searchDAO = new ESSearchDAO( - _entityRegistry, - _searchClientSpy, - _indexConvention, + searchClientSpy, + indexConvention, false, ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, getSearchConfiguration(), null); ESBrowseDAO browseDAO = new ESBrowseDAO( - _entityRegistry, - _searchClientSpy, - _indexConvention, + searchClientSpy, + indexConvention, getSearchConfiguration(), getCustomSearchConfiguration()); ESWriteDAO writeDAO = - new ESWriteDAO(_entityRegistry, _searchClientSpy, _indexConvention, getBulkProcessor(), 1); - return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); + new ESWriteDAO( + aspectRetriever.getEntityRegistry(), + searchClientSpy, + indexConvention, + getBulkProcessor(), + 1); + return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO) + .postConstruct(aspectRetriever); } private void clearCache(boolean withLightingCache) { - _cacheManager.getCacheNames().forEach(cache -> _cacheManager.getCache(cache).clear()); + cacheManager.getCacheNames().forEach(cache -> cacheManager.getCache(cache).clear()); resetService(true, withLightingCache); } @@ -212,7 +225,7 @@ private EntityLineageResult mockResult(List lineageRelation @Test public void testSearchService() throws Exception { - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -227,7 +240,7 @@ public void testSearchService() throws Exception { assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(false); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -254,10 +267,10 @@ public void testSearchService() throws Exception { document.set("keyPart1", JsonNodeFactory.instance.textNode("test")); document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride")); document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(getBulkProcessor()); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -271,7 +284,7 @@ public void testSearchService() throws Exception { assertEquals(searchResult.getEntities().size(), 0); clearCache(false); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -304,22 +317,22 @@ public void testSearchService() throws Exception { document2.set("keyPart1", JsonNodeFactory.instance.textNode("random")); document2.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride2")); document2.set("browsePaths", JsonNodeFactory.instance.textNode("/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(getBulkProcessor()); - Mockito.reset(_searchClientSpy); + Mockito.reset(searchClientSpy); searchResult = searchAcrossLineage(null, TEST1); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); // Verify that highlighting was turned off in the query ArgumentCaptor searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class); - Mockito.verify(_searchClientSpy, times(1)).search(searchRequestCaptor.capture(), any()); + Mockito.verify(searchClientSpy, times(1)).search(searchRequestCaptor.capture(), any()); SearchRequest capturedRequest = searchRequestCaptor.getValue(); assertNull(capturedRequest.source().highlighter()); clearCache(false); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -337,10 +350,10 @@ public void testSearchService() throws Exception { clearCache(false); // Test Cache Behavior - Mockito.reset(_graphService); + Mockito.reset(graphService); // Case 1: Use the maxHops in the cache. - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -354,7 +367,7 @@ public void testSearchService() throws Exception { new LineageRelationship().setDegree(3).setType("type").setEntity(urn)))); searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -369,7 +382,7 @@ public void testSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - Mockito.verify(_graphService, times(1)) + Mockito.verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -381,7 +394,7 @@ public void testSearchService() throws Exception { // Hit the cache on second attempt searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -395,7 +408,7 @@ public void testSearchService() throws Exception { null, new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - Mockito.verify(_graphService, times(1)) + Mockito.verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -406,7 +419,7 @@ public void testSearchService() throws Exception { eq(null)); // Case 2: Use the start and end time in the cache. - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -420,7 +433,7 @@ public void testSearchService() throws Exception { new LineageRelationship().setDegree(3).setType("type").setEntity(urn)))); searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -435,7 +448,7 @@ public void testSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - Mockito.verify(_graphService, times(1)) + Mockito.verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -447,7 +460,7 @@ public void testSearchService() throws Exception { // Hit the cache on second attempt searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -461,7 +474,7 @@ public void testSearchService() throws Exception { 1L, new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - Mockito.verify(_graphService, times(1)) + Mockito.verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -474,11 +487,11 @@ public void testSearchService() throws Exception { clearCache(false); // Cleanup - _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(getBulkProcessor()); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(), anyInt())) .thenReturn( mockResult( @@ -491,7 +504,7 @@ public void testSearchService() throws Exception { @Test public void testScrollAcrossLineage() throws Exception { - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -508,7 +521,7 @@ public void testScrollAcrossLineage() throws Exception { assertNull(scrollResult.getScrollId()); clearCache(false); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -537,10 +550,10 @@ public void testScrollAcrossLineage() throws Exception { document.set("keyPart1", JsonNodeFactory.instance.textNode("test")); document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride")); document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(getBulkProcessor()); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -555,7 +568,7 @@ public void testScrollAcrossLineage() throws Exception { assertNull(scrollResult.getScrollId()); clearCache(false); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -586,10 +599,10 @@ public void testScrollAcrossLineage() throws Exception { clearCache(false); // Cleanup - _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); syncAfterWrite(getBulkProcessor()); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(), anyInt())) .thenReturn( mockResult( @@ -611,7 +624,7 @@ public void testLightningSearchService() throws Exception { // Enable lightning resetService(true, true); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -624,7 +637,7 @@ public void testLightningSearchService() throws Exception { assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(true); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -646,10 +659,10 @@ public void testLightningSearchService() throws Exception { document.set("keyPart1", JsonNodeFactory.instance.textNode("test")); document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride")); document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(getBulkProcessor()); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -663,7 +676,7 @@ public void testLightningSearchService() throws Exception { assertEquals(searchResult.getEntities().size(), 0); clearCache(true); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -679,20 +692,20 @@ public void testLightningSearchService() throws Exception { assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); assertEquals(searchResult.getEntities().get(0).getDegree().intValue(), 1); - verify(_lineageSearchService, times(1)) + verify(lineageSearchService, times(1)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); searchResult = searchAcrossLineage(QueryUtils.newFilter("degree.keyword", "1"), testStar); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); assertEquals(searchResult.getEntities().get(0).getDegree().intValue(), 1); - verify(_lineageSearchService, times(2)) + verify(lineageSearchService, times(2)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); searchResult = searchAcrossLineage(QueryUtils.newFilter("degree.keyword", "2"), testStar); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); - verify(_lineageSearchService, times(3)) + verify(lineageSearchService, times(3)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); clearCache(true); // resets spy @@ -702,17 +715,17 @@ public void testLightningSearchService() throws Exception { document2.set("keyPart1", JsonNodeFactory.instance.textNode("random")); document2.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride2")); document2.set("browsePaths", JsonNodeFactory.instance.textNode("/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = searchAcrossLineage(null, testStar); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); - verify(_lineageSearchService, times(1)) + verify(lineageSearchService, times(1)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); clearCache(true); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -727,16 +740,16 @@ public void testLightningSearchService() throws Exception { searchResult = searchAcrossLineage(null, testStar); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().size(), 1); - verify(_lineageSearchService, times(1)) + verify(lineageSearchService, times(1)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); clearCache(true); // Test Cache Behavior - reset(_graphService); - reset(_lineageSearchService); + reset(graphService); + reset(lineageSearchService); // Case 1: Use the maxHops in the cache. - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -750,7 +763,7 @@ public void testLightningSearchService() throws Exception { new LineageRelationship().setDegree(3).setType("type").setEntity(urn)))); searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -765,7 +778,7 @@ public void testLightningSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - verify(_graphService, times(1)) + verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -774,12 +787,12 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(null), eq(null)); - verify(_lineageSearchService, times(1)) + verify(lineageSearchService, times(1)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); // Hit the cache on second attempt searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -793,7 +806,7 @@ public void testLightningSearchService() throws Exception { null, new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - verify(_graphService, times(1)) + verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -802,11 +815,11 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(null), eq(null)); - verify(_lineageSearchService, times(2)) + verify(lineageSearchService, times(2)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); // Case 2: Use the start and end time in the cache. - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), @@ -820,7 +833,7 @@ public void testLightningSearchService() throws Exception { new LineageRelationship().setDegree(3).setType("type").setEntity(urn)))); searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -835,7 +848,7 @@ public void testLightningSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - verify(_graphService, times(1)) + verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -844,12 +857,12 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(0L), eq(1L)); - verify(_lineageSearchService, times(3)) + verify(lineageSearchService, times(3)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); // Hit the cache on second attempt searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -863,7 +876,7 @@ public void testLightningSearchService() throws Exception { 1L, new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); - verify(_graphService, times(1)) + verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -872,17 +885,17 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(0L), eq(1L)); - verify(_lineageSearchService, times(4)) + verify(lineageSearchService, times(4)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); /* * Test filtering */ - reset(_lineageSearchService); + reset(lineageSearchService); // Entity searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(DATASET_ENTITY_NAME), @@ -897,12 +910,12 @@ public void testLightningSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); - verify(_lineageSearchService, times(1)) + verify(lineageSearchService, times(1)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); // Cached searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(DATASET_ENTITY_NAME), @@ -915,7 +928,7 @@ public void testLightningSearchService() throws Exception { null, null, new SearchFlags().setSkipCache(false)); - Mockito.verify(_graphService, times(1)) + Mockito.verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -924,7 +937,7 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(0L), eq(1L)); - verify(_lineageSearchService, times(2)) + verify(lineageSearchService, times(2)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); @@ -945,7 +958,7 @@ public void testLightningSearchService() throws Exception { Filter filter = new Filter().setOr(conCritArr); searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -960,12 +973,12 @@ public void testLightningSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); - verify(_lineageSearchService, times(3)) + verify(lineageSearchService, times(3)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); // Cached searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -978,7 +991,7 @@ public void testLightningSearchService() throws Exception { null, null, new SearchFlags().setSkipCache(false)); - verify(_graphService, times(1)) + verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -987,7 +1000,7 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(0L), eq(1L)); - verify(_lineageSearchService, times(4)) + verify(lineageSearchService, times(4)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); @@ -995,7 +1008,7 @@ public void testLightningSearchService() throws Exception { // Environment Filter originFilter = QueryUtils.newFilter("origin", "PROD"); searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -1010,12 +1023,12 @@ public void testLightningSearchService() throws Exception { new SearchFlags().setSkipCache(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); - verify(_lineageSearchService, times(5)) + verify(lineageSearchService, times(5)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); // Cached searchResult = - _lineageSearchService.searchAcrossLineage( + lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME), @@ -1028,7 +1041,7 @@ public void testLightningSearchService() throws Exception { null, null, new SearchFlags().setSkipCache(false)); - verify(_graphService, times(1)) + verify(graphService, times(1)) .getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), @@ -1037,7 +1050,7 @@ public void testLightningSearchService() throws Exception { eq(1000), eq(0L), eq(1L)); - verify(_lineageSearchService, times(6)) + verify(lineageSearchService, times(6)) .getLightningSearchResult(any(), any(), anyInt(), anyInt(), anySet()); assertEquals(searchResult.getNumEntities().intValue(), 0); assertEquals(searchResult.getEntities().size(), 0); @@ -1045,11 +1058,11 @@ public void testLightningSearchService() throws Exception { clearCache(true); // Cleanup - _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(getBulkProcessor()); - when(_graphService.getLineage( + when(graphService.getLineage( eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(), anyInt())) .thenReturn( mockResult( @@ -1089,7 +1102,7 @@ public void testLightningEnvFiltering() throws Exception { Set entityNames = Collections.emptySet(); LineageSearchResult lineageSearchResult = - _lineageSearchService.getLightningSearchResult( + lineageSearchService.getLightningSearchResult( lineageRelationships, filter, from, size, entityNames); assertEquals(lineageSearchResult.getNumEntities(), Integer.valueOf(500)); @@ -1131,7 +1144,7 @@ public void testLightningEnvFiltering() throws Exception { filter = new Filter().setOr(conCritArr); lineageSearchResult = - _lineageSearchService.getLightningSearchResult( + lineageSearchService.getLightningSearchResult( lineageRelationships, filter, from, size, entityNames); // assert that if the query has an env filter, it is applied correctly @@ -1168,7 +1181,7 @@ public void testLightningPagination() throws Exception { Set entityNames = Collections.emptySet(); LineageSearchResult lineageSearchResult = - _lineageSearchService.getLightningSearchResult( + lineageSearchService.getLightningSearchResult( lineageRelationships, filter, from, size, entityNames); assertEquals(lineageSearchResult.getNumEntities(), Integer.valueOf(500)); @@ -1181,7 +1194,7 @@ public void testLightningPagination() throws Exception { from = 50; size = 20; lineageSearchResult = - _lineageSearchService.getLightningSearchResult( + lineageSearchService.getLightningSearchResult( lineageRelationships, filter, from, size, entityNames); assertEquals(lineageSearchResult.getNumEntities(), Integer.valueOf(500)); @@ -1206,7 +1219,7 @@ public void testLightningPagination() throws Exception { size = 10; filter = new Filter().setOr(conCritArr); lineageSearchResult = - _lineageSearchService.getLightningSearchResult( + lineageSearchService.getLightningSearchResult( lineageRelationships, filter, from, size, entityNames); assertEquals(lineageSearchResult.getNumEntities(), Integer.valueOf(600)); @@ -1220,7 +1233,7 @@ public void testLightningPagination() throws Exception { from = 0; size = 10; lineageSearchResult = - _lineageSearchService.getLightningSearchResult( + lineageSearchService.getLightningSearchResult( lineageRelationships, null, from, size, entityNames); // Static Degree agg is the first element @@ -1270,7 +1283,7 @@ private LineageRelationship constructLineageRelationship(Urn urn) { // Convenience method to reduce spots where we're sending the same params private LineageSearchResult searchAcrossLineage(@Nullable Filter filter, @Nullable String input) { - return _lineageSearchService.searchAcrossLineage( + return lineageSearchService.searchAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -1287,7 +1300,7 @@ private LineageSearchResult searchAcrossLineage(@Nullable Filter filter, @Nullab private LineageScrollResult scrollAcrossLineage( @Nullable Filter filter, @Nullable String input, String scrollId, int size) { - return _lineageSearchService.scrollAcrossLineage( + return lineageSearchService.scrollAcrossLineage( TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), @@ -1326,8 +1339,7 @@ public void testCanDoLightning() throws Exception { int size = 10; Set entityNames = Collections.emptySet(); - Assert.assertTrue( - _lineageSearchService.canDoLightning(lineageRelationships, "*", filter, null)); + Assert.assertTrue(lineageSearchService.canDoLightning(lineageRelationships, "*", filter, null)); // Set up filters ConjunctiveCriterionArray conCritArr = new ConjunctiveCriterionArray(); @@ -1350,7 +1362,6 @@ public void testCanDoLightning() throws Exception { from = 500; size = 10; filter = new Filter().setOr(conCritArr); - Assert.assertTrue( - _lineageSearchService.canDoLightning(lineageRelationships, "*", filter, null)); + Assert.assertTrue(lineageSearchService.canDoLightning(lineageRelationships, "*", filter, null)); } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java index 71f35adabce368..d860776a316815 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java @@ -2,6 +2,9 @@ import static com.linkedin.metadata.Constants.ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH; import static io.datahubproject.test.search.SearchTestUtils.syncAfterWrite; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import com.datahub.test.Snapshot; @@ -11,10 +14,10 @@ import com.linkedin.common.urn.TestEntityUrn; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.StringArray; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.cache.EntityDocCountCacheConfiguration; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Condition; @@ -36,6 +39,9 @@ import com.linkedin.metadata.search.ranker.SimpleRanker; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import com.linkedin.r2.RemoteInvocationException; +import java.net.URISyntaxException; +import java.util.Map; import javax.annotation.Nonnull; import org.opensearch.client.RestHighLevelClient; import org.springframework.cache.CacheManager; @@ -62,44 +68,49 @@ public abstract class SearchServiceTestBase extends AbstractTestNGSpringContextT @Nonnull protected abstract CustomSearchConfiguration getCustomSearchConfiguration(); - private EntityRegistry _entityRegistry; - private IndexConvention _indexConvention; - private SettingsBuilder _settingsBuilder; - private ElasticSearchService _elasticSearchService; - private CacheManager _cacheManager; - private SearchService _searchService; + private AspectRetriever aspectRetriever; + private IndexConvention indexConvention; + private SettingsBuilder settingsBuilder; + private ElasticSearchService elasticSearchService; + private CacheManager cacheManager; + private SearchService searchService; private static final String ENTITY_NAME = "testEntity"; @BeforeClass - public void setup() { - _entityRegistry = new SnapshotEntityRegistry(new Snapshot()); - _indexConvention = new IndexConventionImpl("search_service_test"); - _settingsBuilder = new SettingsBuilder(null); - _elasticSearchService = buildEntitySearchService(); - _elasticSearchService.configure(); - _cacheManager = new ConcurrentMapCacheManager(); + public void setup() throws RemoteInvocationException, URISyntaxException { + aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()) + .thenReturn(new SnapshotEntityRegistry(new Snapshot())); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + indexConvention = new IndexConventionImpl("search_service_test"); + settingsBuilder = new SettingsBuilder(null); + elasticSearchService = buildEntitySearchService(); + elasticSearchService.configure(); + cacheManager = new ConcurrentMapCacheManager(); resetSearchService(); } private void resetSearchService() { CachingEntitySearchService cachingEntitySearchService = - new CachingEntitySearchService(_cacheManager, _elasticSearchService, 100, true); + new CachingEntitySearchService(cacheManager, elasticSearchService, 100, true); EntityDocCountCacheConfiguration entityDocCountCacheConfiguration = new EntityDocCountCacheConfiguration(); entityDocCountCacheConfiguration.setTtlSeconds(600L); - _searchService = + searchService = new SearchService( new EntityDocCountCache( - _entityRegistry, _elasticSearchService, entityDocCountCacheConfiguration), + aspectRetriever.getEntityRegistry(), + elasticSearchService, + entityDocCountCacheConfiguration), cachingEntitySearchService, new SimpleRanker()); } @BeforeMethod public void wipe() throws Exception { - _elasticSearchService.clear(); + elasticSearchService.clear(); syncAfterWrite(getBulkProcessor()); } @@ -107,37 +118,44 @@ public void wipe() throws Exception { private ElasticSearchService buildEntitySearchService() { EntityIndexBuilders indexBuilders = new EntityIndexBuilders( - getIndexBuilder(), _entityRegistry, _indexConvention, _settingsBuilder); + getIndexBuilder(), + aspectRetriever.getEntityRegistry(), + indexConvention, + settingsBuilder); ESSearchDAO searchDAO = new ESSearchDAO( - _entityRegistry, getSearchClient(), - _indexConvention, + indexConvention, false, ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, getSearchConfiguration(), null); ESBrowseDAO browseDAO = new ESBrowseDAO( - _entityRegistry, getSearchClient(), - _indexConvention, + indexConvention, getSearchConfiguration(), getCustomSearchConfiguration()); ESWriteDAO writeDAO = - new ESWriteDAO(_entityRegistry, getSearchClient(), _indexConvention, getBulkProcessor(), 1); - return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); + new ESWriteDAO( + aspectRetriever.getEntityRegistry(), + getSearchClient(), + indexConvention, + getBulkProcessor(), + 1); + return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO) + .postConstruct(aspectRetriever); } private void clearCache() { - _cacheManager.getCacheNames().forEach(cache -> _cacheManager.getCache(cache).clear()); + cacheManager.getCacheNames().forEach(cache -> cacheManager.getCache(cache).clear()); resetSearchService(); } @Test public void testSearchService() throws Exception { SearchResult searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(ENTITY_NAME), "test", null, @@ -147,7 +165,7 @@ public void testSearchService() throws Exception { new SearchFlags().setFulltext(true).setSkipCache(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); clearCache(); @@ -158,11 +176,11 @@ public void testSearchService() throws Exception { document.set("keyPart1", JsonNodeFactory.instance.textNode("test")); document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride")); document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); @@ -174,24 +192,24 @@ public void testSearchService() throws Exception { document2.set("keyPart1", JsonNodeFactory.instance.textNode("random")); document2.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride2")); document2.set("browsePaths", JsonNodeFactory.instance.textNode("/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "'test2'", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); clearCache(); - long docCount = _elasticSearchService.docCount(ENTITY_NAME); + long docCount = elasticSearchService.docCount(ENTITY_NAME); assertEquals(docCount, 2L); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "'test2'", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); } @@ -222,7 +240,7 @@ public void testAdvancedSearchOr() throws Exception { .setAnd(new CriterionArray(ImmutableList.of(subtypeCriterion))))); SearchResult searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(ENTITY_NAME), "test", filterWithCondition, @@ -242,7 +260,7 @@ public void testAdvancedSearchOr() throws Exception { document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); document.set("subtypes", JsonNodeFactory.instance.textNode("view")); document.set("platform", JsonNodeFactory.instance.textNode("snowflake")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); Urn urn2 = new TestEntityUrn("test", "testUrn", "VALUE_2"); ObjectNode document2 = JsonNodeFactory.instance.objectNode(); @@ -252,7 +270,7 @@ public void testAdvancedSearchOr() throws Exception { document2.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); document2.set("subtypes", JsonNodeFactory.instance.textNode("table")); document2.set("platform", JsonNodeFactory.instance.textNode("hive")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); Urn urn3 = new TestEntityUrn("test", "testUrn", "VALUE_3"); ObjectNode document3 = JsonNodeFactory.instance.objectNode(); @@ -262,12 +280,12 @@ public void testAdvancedSearchOr() throws Exception { document3.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); document3.set("subtypes", JsonNodeFactory.instance.textNode("table")); document3.set("platform", JsonNodeFactory.instance.textNode("snowflake")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document3.toString(), urn3.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document3.toString(), urn3.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "test", filterWithCondition, @@ -307,7 +325,7 @@ public void testAdvancedSearchSoftDelete() throws Exception { ImmutableList.of(filterCriterion, removedCriterion))))); SearchResult searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(ENTITY_NAME), "test", filterWithCondition, @@ -328,7 +346,7 @@ public void testAdvancedSearchSoftDelete() throws Exception { document.set("subtypes", JsonNodeFactory.instance.textNode("view")); document.set("platform", JsonNodeFactory.instance.textNode("hive")); document.set("removed", JsonNodeFactory.instance.booleanNode(true)); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); Urn urn2 = new TestEntityUrn("test", "testUrn", "VALUE_2"); ObjectNode document2 = JsonNodeFactory.instance.objectNode(); @@ -339,7 +357,7 @@ public void testAdvancedSearchSoftDelete() throws Exception { document2.set("subtypes", JsonNodeFactory.instance.textNode("table")); document2.set("platform", JsonNodeFactory.instance.textNode("hive")); document.set("removed", JsonNodeFactory.instance.booleanNode(false)); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); Urn urn3 = new TestEntityUrn("test", "testUrn", "VALUE_3"); ObjectNode document3 = JsonNodeFactory.instance.objectNode(); @@ -350,12 +368,12 @@ public void testAdvancedSearchSoftDelete() throws Exception { document3.set("subtypes", JsonNodeFactory.instance.textNode("table")); document3.set("platform", JsonNodeFactory.instance.textNode("snowflake")); document.set("removed", JsonNodeFactory.instance.booleanNode(false)); - _elasticSearchService.upsertDocument(ENTITY_NAME, document3.toString(), urn3.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document3.toString(), urn3.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "test", filterWithCondition, @@ -386,7 +404,7 @@ public void testAdvancedSearchNegated() throws Exception { .setAnd(new CriterionArray(ImmutableList.of(filterCriterion))))); SearchResult searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(ENTITY_NAME), "test", filterWithCondition, @@ -407,7 +425,7 @@ public void testAdvancedSearchNegated() throws Exception { document.set("subtypes", JsonNodeFactory.instance.textNode("view")); document.set("platform", JsonNodeFactory.instance.textNode("hive")); document.set("removed", JsonNodeFactory.instance.booleanNode(true)); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); Urn urn2 = new TestEntityUrn("test", "testUrn", "VALUE_2"); ObjectNode document2 = JsonNodeFactory.instance.objectNode(); @@ -418,7 +436,7 @@ public void testAdvancedSearchNegated() throws Exception { document2.set("subtypes", JsonNodeFactory.instance.textNode("table")); document2.set("platform", JsonNodeFactory.instance.textNode("hive")); document.set("removed", JsonNodeFactory.instance.booleanNode(false)); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); Urn urn3 = new TestEntityUrn("test", "testUrn", "VALUE_3"); ObjectNode document3 = JsonNodeFactory.instance.objectNode(); @@ -429,12 +447,12 @@ public void testAdvancedSearchNegated() throws Exception { document3.set("subtypes", JsonNodeFactory.instance.textNode("table")); document3.set("platform", JsonNodeFactory.instance.textNode("snowflake")); document.set("removed", JsonNodeFactory.instance.booleanNode(false)); - _elasticSearchService.upsertDocument(ENTITY_NAME, document3.toString(), urn3.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document3.toString(), urn3.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _searchService.searchAcrossEntities( + searchService.searchAcrossEntities( ImmutableList.of(), "test", filterWithCondition, diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java index b544faa061f0ed..40ccc8dfb5047e 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java @@ -4,18 +4,16 @@ import static io.datahubproject.test.search.SearchTestUtils.syncAfterWrite; import static org.testng.Assert.assertEquals; -import com.datahub.test.Snapshot; import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.common.collect.ImmutableList; import com.linkedin.common.urn.TestEntityUrn; import com.linkedin.common.urn.Urn; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; -import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.search.elasticsearch.ElasticSearchService; import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder; @@ -30,6 +28,8 @@ import java.util.List; import javax.annotation.Nonnull; import org.opensearch.client.RestHighLevelClient; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; import org.testcontainers.shaded.com.google.common.collect.ImmutableMap; import org.testng.annotations.BeforeClass; @@ -53,20 +53,22 @@ public abstract class TestEntityTestBase extends AbstractTestNGSpringContextTest @Nonnull protected abstract CustomSearchConfiguration getCustomSearchConfiguration(); - private EntityRegistry _entityRegistry; - private IndexConvention _indexConvention; - private SettingsBuilder _settingsBuilder; - private ElasticSearchService _elasticSearchService; + @Autowired + @Qualifier("snapshotRegistryAspectRetriever") + AspectRetriever aspectRetriever; + + private IndexConvention indexConvention; + private SettingsBuilder settingsBuilder; + private ElasticSearchService elasticSearchService; private static final String ENTITY_NAME = "testEntity"; @BeforeClass public void setup() { - _entityRegistry = new SnapshotEntityRegistry(new Snapshot()); - _indexConvention = new IndexConventionImpl("es_service_test"); - _settingsBuilder = new SettingsBuilder(null); - _elasticSearchService = buildService(); - _elasticSearchService.configure(); + indexConvention = new IndexConventionImpl("es_service_test"); + settingsBuilder = new SettingsBuilder(null); + elasticSearchService = buildService(); + elasticSearchService.configure(); } @BeforeClass @@ -78,46 +80,55 @@ public void disableAssert() { @BeforeMethod public void wipe() throws Exception { - _elasticSearchService.clear(); + elasticSearchService.clear(); } @Nonnull private ElasticSearchService buildService() { EntityIndexBuilders indexBuilders = new EntityIndexBuilders( - getIndexBuilder(), _entityRegistry, _indexConvention, _settingsBuilder); + getIndexBuilder(), + aspectRetriever.getEntityRegistry(), + indexConvention, + settingsBuilder); ESSearchDAO searchDAO = new ESSearchDAO( - _entityRegistry, getSearchClient(), - _indexConvention, + indexConvention, false, ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, getSearchConfiguration(), null); ESBrowseDAO browseDAO = new ESBrowseDAO( - _entityRegistry, getSearchClient(), - _indexConvention, + indexConvention, getSearchConfiguration(), getCustomSearchConfiguration()); ESWriteDAO writeDAO = - new ESWriteDAO(_entityRegistry, getSearchClient(), _indexConvention, getBulkProcessor(), 1); - return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); + new ESWriteDAO( + aspectRetriever.getEntityRegistry(), + getSearchClient(), + indexConvention, + getBulkProcessor(), + 1); + ElasticSearchService searchService = + new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); + searchService.postConstruct(aspectRetriever); + return searchService; } @Test public void testElasticSearchServiceStructuredQuery() throws Exception { SearchResult searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); - BrowseResult browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + BrowseResult browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 0); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 0); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 0); assertEquals( - _elasticSearchService + elasticSearchService .aggregateByValue(ImmutableList.of(ENTITY_NAME), "textField", null, 10) .size(), 0); @@ -129,16 +140,16 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride")); document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); document.set("foreignKey", JsonNodeFactory.instance.textNode("urn:li:tag:Node.Value")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "foreignKey:Node", null, @@ -148,15 +159,15 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); - browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 1); assertEquals(browseResult.getGroups().get(0).getName(), "a"); - browseResult = _elasticSearchService.browse(ENTITY_NAME, "/a", null, 0, 10); + browseResult = elasticSearchService.browse(ENTITY_NAME, "/a", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 1); assertEquals(browseResult.getGroups().get(0).getName(), "b"); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 1); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 1); assertEquals( - _elasticSearchService.aggregateByValue( + elasticSearchService.aggregateByValue( ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), ImmutableMap.of("textFieldOverride", 1L)); @@ -166,39 +177,39 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { document2.set("keyPart1", JsonNodeFactory.instance.textNode("random")); document2.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride2")); document2.set("browsePaths", JsonNodeFactory.instance.textNode("/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); - browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 2); assertEquals(browseResult.getGroups().get(0).getName(), "a"); assertEquals(browseResult.getGroups().get(1).getName(), "b"); - browseResult = _elasticSearchService.browse(ENTITY_NAME, "/a", null, 0, 10); + browseResult = elasticSearchService.browse(ENTITY_NAME, "/a", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 1); assertEquals(browseResult.getGroups().get(0).getName(), "b"); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 2); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 2); assertEquals( - _elasticSearchService.aggregateByValue( + elasticSearchService.aggregateByValue( ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), ImmutableMap.of("textFieldOverride", 1L, "textFieldOverride2", 1L)); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(false)); assertEquals(searchResult.getNumEntities().intValue(), 0); - browseResult = _elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); + browseResult = elasticSearchService.browse(ENTITY_NAME, "", null, 0, 10); assertEquals(browseResult.getMetadata().getTotalNumEntities().longValue(), 0); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 0); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 0); assertEquals( - _elasticSearchService + elasticSearchService .aggregateByValue(ImmutableList.of(ENTITY_NAME), "textField", null, 10) .size(), 0); @@ -207,7 +218,7 @@ public void testElasticSearchServiceStructuredQuery() throws Exception { @Test public void testElasticSearchServiceFulltext() throws Exception { SearchResult searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); @@ -218,18 +229,18 @@ public void testElasticSearchServiceFulltext() throws Exception { document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride")); document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c")); document.set("foreignKey", JsonNodeFactory.instance.textNode("urn:li:tag:Node.Value")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 1); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 1); assertEquals( - _elasticSearchService.aggregateByValue( + elasticSearchService.aggregateByValue( ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), ImmutableMap.of("textFieldOverride", 1L)); @@ -239,32 +250,32 @@ public void testElasticSearchServiceFulltext() throws Exception { document2.set("keyPart1", JsonNodeFactory.instance.textNode("random")); document2.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride2")); document2.set("browsePaths", JsonNodeFactory.instance.textNode("/b/c")); - _elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); + elasticSearchService.upsertDocument(ENTITY_NAME, document2.toString(), urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 1); assertEquals(searchResult.getEntities().get(0).getEntity(), urn2); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 2); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 2); assertEquals( - _elasticSearchService.aggregateByValue( + elasticSearchService.aggregateByValue( ImmutableList.of(ENTITY_NAME), "textFieldOverride", null, 10), ImmutableMap.of("textFieldOverride", 1L, "textFieldOverride2", 1L)); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); - _elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString()); + elasticSearchService.deleteDocument(ENTITY_NAME, urn2.toString()); syncAfterWrite(getBulkProcessor()); searchResult = - _elasticSearchService.search( + elasticSearchService.search( List.of(ENTITY_NAME), "test2", null, null, 0, 10, new SearchFlags().setFulltext(true)); assertEquals(searchResult.getNumEntities().intValue(), 0); - assertEquals(_elasticSearchService.docCount(ENTITY_NAME), 0); + assertEquals(elasticSearchService.docCount(ENTITY_NAME), 0); assertEquals( - _elasticSearchService + elasticSearchService .aggregateByValue(ImmutableList.of(ENTITY_NAME), "textField", null, 10) .size(), 0); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TestEntityElasticSearchTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TestEntityElasticSearchTest.java index 843da17fbd1321..5ad7b1218a5bf4 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TestEntityElasticSearchTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TestEntityElasticSearchTest.java @@ -21,44 +21,44 @@ }) public class TestEntityElasticSearchTest extends TestEntityTestBase { - @Autowired private RestHighLevelClient _searchClient; - @Autowired private ESBulkProcessor _bulkProcessor; - @Autowired private ESIndexBuilder _esIndexBuilder; - @Autowired private SearchConfiguration _searchConfiguration; - @Autowired private CustomSearchConfiguration _customSearchConfiguration; + @Autowired private RestHighLevelClient searchClient; + @Autowired private ESBulkProcessor bulkProcessor; + @Autowired private ESIndexBuilder esIndexBuilder; + @Autowired private SearchConfiguration searchConfiguration; + @Autowired private CustomSearchConfiguration customSearchConfiguration; @NotNull @Override protected RestHighLevelClient getSearchClient() { - return _searchClient; + return searchClient; } @NotNull @Override protected ESBulkProcessor getBulkProcessor() { - return _bulkProcessor; + return bulkProcessor; } @NotNull @Override protected ESIndexBuilder getIndexBuilder() { - return _esIndexBuilder; + return esIndexBuilder; } @NotNull @Override protected SearchConfiguration getSearchConfiguration() { - return _searchConfiguration; + return searchConfiguration; } @NotNull @Override protected CustomSearchConfiguration getCustomSearchConfiguration() { - return _customSearchConfiguration; + return customSearchConfiguration; } @Test public void initTest() { - AssertJUnit.assertNotNull(_searchClient); + AssertJUnit.assertNotNull(searchClient); } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TimeseriesAspectServiceElasticSearchTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TimeseriesAspectServiceElasticSearchTest.java index 6ebe42d0181e44..1f51d463a2963a 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TimeseriesAspectServiceElasticSearchTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/TimeseriesAspectServiceElasticSearchTest.java @@ -7,6 +7,7 @@ import org.jetbrains.annotations.NotNull; import org.opensearch.client.RestHighLevelClient; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.context.annotation.Import; import org.testng.AssertJUnit; import org.testng.annotations.Test; @@ -16,7 +17,10 @@ public class TimeseriesAspectServiceElasticSearchTest extends TimeseriesAspectSe @Autowired private RestHighLevelClient _searchClient; @Autowired private ESBulkProcessor _bulkProcessor; - @Autowired private ESIndexBuilder _esIndexBuilder; + + @Autowired + @Qualifier("searchIndexBuilder") + private ESIndexBuilder _esIndexBuilder; @NotNull @Override diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/TimeseriesAspectServiceOpenSearchTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/TimeseriesAspectServiceOpenSearchTest.java index 63dffa9c210045..16ac03415ee5c2 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/TimeseriesAspectServiceOpenSearchTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/opensearch/TimeseriesAspectServiceOpenSearchTest.java @@ -7,6 +7,7 @@ import org.jetbrains.annotations.NotNull; import org.opensearch.client.RestHighLevelClient; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.context.annotation.Import; import org.testng.AssertJUnit; import org.testng.annotations.Test; @@ -16,7 +17,10 @@ public class TimeseriesAspectServiceOpenSearchTest extends TimeseriesAspectServi @Autowired private RestHighLevelClient _searchClient; @Autowired private ESBulkProcessor _bulkProcessor; - @Autowired private ESIndexBuilder _esIndexBuilder; + + @Autowired + @Qualifier("searchIndexBuilder") + private ESIndexBuilder _esIndexBuilder; @NotNull @Override diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/BrowseDAOTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/BrowseDAOTest.java index a261b53f25c605..53eeb9dc8314c3 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/BrowseDAOTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/BrowseDAOTest.java @@ -7,11 +7,13 @@ import static org.testng.Assert.assertEquals; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration; import com.linkedin.metadata.entity.TestEntityRegistry; import com.linkedin.metadata.search.elasticsearch.query.ESBrowseDAO; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import com.linkedin.r2.RemoteInvocationException; import io.datahubproject.test.search.config.SearchCommonTestConfiguration; import java.net.URISyntaxException; import java.util.Collections; @@ -31,22 +33,25 @@ @Import(SearchCommonTestConfiguration.class) public class BrowseDAOTest extends AbstractTestNGSpringContextTests { - private RestHighLevelClient _mockClient; - private ESBrowseDAO _browseDAO; + private RestHighLevelClient mockClient; + private ESBrowseDAO browseDAO; - @Autowired private SearchConfiguration _searchConfiguration; - @Autowired private CustomSearchConfiguration _customSearchConfiguration; + @Autowired private SearchConfiguration searchConfiguration; + @Autowired private CustomSearchConfiguration customSearchConfiguration; @BeforeMethod - public void setup() { - _mockClient = mock(RestHighLevelClient.class); - _browseDAO = + public void setup() throws RemoteInvocationException, URISyntaxException { + mockClient = mock(RestHighLevelClient.class); + AspectRetriever aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()).thenReturn(new TestEntityRegistry()); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + browseDAO = new ESBrowseDAO( - new TestEntityRegistry(), - _mockClient, - new IndexConventionImpl("es_browse_dao_test"), - _searchConfiguration, - _customSearchConfiguration); + mockClient, + new IndexConventionImpl("es_browse_dao_test"), + searchConfiguration, + customSearchConfiguration) + .setAspectRetriever(aspectRetriever); } public static Urn makeUrn(Object id) { @@ -68,24 +73,24 @@ public void testGetBrowsePath() throws Exception { // Test when there is no search hit for getBrowsePaths when(mockSearchHits.getHits()).thenReturn(new SearchHit[0]); when(mockSearchResponse.getHits()).thenReturn(mockSearchHits); - when(_mockClient.search(any(), eq(RequestOptions.DEFAULT))).thenReturn(mockSearchResponse); - assertEquals(_browseDAO.getBrowsePaths("dataset", dummyUrn).size(), 0); + when(mockClient.search(any(), eq(RequestOptions.DEFAULT))).thenReturn(mockSearchResponse); + assertEquals(browseDAO.getBrowsePaths("dataset", dummyUrn).size(), 0); // Test the case of single search hit & browsePaths field doesn't exist sourceMap.remove("browse_paths"); when(mockSearchHit.getSourceAsMap()).thenReturn(sourceMap); when(mockSearchHits.getHits()).thenReturn(new SearchHit[] {mockSearchHit}); when(mockSearchResponse.getHits()).thenReturn(mockSearchHits); - when(_mockClient.search(any(), eq(RequestOptions.DEFAULT))).thenReturn(mockSearchResponse); - assertEquals(_browseDAO.getBrowsePaths("dataset", dummyUrn).size(), 0); + when(mockClient.search(any(), eq(RequestOptions.DEFAULT))).thenReturn(mockSearchResponse); + assertEquals(browseDAO.getBrowsePaths("dataset", dummyUrn).size(), 0); // Test the case of single search hit & browsePaths field exists sourceMap.put("browsePaths", Collections.singletonList("foo")); when(mockSearchHit.getSourceAsMap()).thenReturn(sourceMap); when(mockSearchHits.getHits()).thenReturn(new SearchHit[] {mockSearchHit}); when(mockSearchResponse.getHits()).thenReturn(mockSearchHits); - when(_mockClient.search(any(), eq(RequestOptions.DEFAULT))).thenReturn(mockSearchResponse); - List browsePaths = _browseDAO.getBrowsePaths("dataset", dummyUrn); + when(mockClient.search(any(), eq(RequestOptions.DEFAULT))).thenReturn(mockSearchResponse); + List browsePaths = browseDAO.getBrowsePaths("dataset", dummyUrn); assertEquals(browsePaths.size(), 1); assertEquals(browsePaths.get(0), "foo"); } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java index ba909dc3822c55..047efe7cf4313d 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/SearchDAOTestBase.java @@ -2,6 +2,9 @@ import static com.linkedin.metadata.Constants.ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH; import static com.linkedin.metadata.utils.SearchUtil.AGGREGATION_SEPARATOR_CHAR; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotEquals; import static org.testng.Assert.assertNotNull; @@ -11,8 +14,8 @@ import com.google.common.collect.ImmutableList; import com.linkedin.data.template.LongMap; import com.linkedin.data.template.StringArray; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; import com.linkedin.metadata.query.filter.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; @@ -29,12 +32,15 @@ import com.linkedin.metadata.search.elasticsearch.query.ESSearchDAO; import com.linkedin.metadata.utils.SearchUtil; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; +import com.linkedin.r2.RemoteInvocationException; +import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; import org.opensearch.client.RestHighLevelClient; import org.springframework.test.context.testng.AbstractTestNGSpringContextTests; +import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; public abstract class SearchDAOTestBase extends AbstractTestNGSpringContextTests { @@ -45,7 +51,15 @@ public abstract class SearchDAOTestBase extends AbstractTestNGSpringContextTests protected abstract IndexConvention getIndexConvention(); - EntityRegistry _entityRegistry = new SnapshotEntityRegistry(new Snapshot()); + protected AspectRetriever aspectRetriever; + + @BeforeClass + public void setup() throws RemoteInvocationException, URISyntaxException { + aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()) + .thenReturn(new SnapshotEntityRegistry(new Snapshot())); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + } @Test public void testTransformFilterForEntitiesNoChange() { @@ -219,13 +233,13 @@ public void testTransformFilterForEntitiesWithSomeChanges() { public void testTransformIndexIntoEntityNameSingle() { ESSearchDAO searchDAO = new ESSearchDAO( - _entityRegistry, - getSearchClient(), - getIndexConvention(), - false, - ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, - getSearchConfiguration(), - null); + getSearchClient(), + getIndexConvention(), + false, + ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, + getSearchConfiguration(), + null) + .setAspectRetriever(aspectRetriever); // Empty aggregations final SearchResultMetadata searchResultMetadata = new SearchResultMetadata().setAggregations(new AggregationMetadataArray()); @@ -302,13 +316,13 @@ public void testTransformIndexIntoEntityNameSingle() { public void testTransformIndexIntoEntityNameNested() { ESSearchDAO searchDAO = new ESSearchDAO( - _entityRegistry, - getSearchClient(), - getIndexConvention(), - false, - ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, - getSearchConfiguration(), - null); + getSearchClient(), + getIndexConvention(), + false, + ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, + getSearchConfiguration(), + null) + .setAspectRetriever(aspectRetriever); // One nested facet Map entityTypeMap = Map.of( diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AggregationQueryBuilderTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AggregationQueryBuilderTest.java index ed4c9db5db6430..4dd53775bbef7f 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AggregationQueryBuilderTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AggregationQueryBuilderTest.java @@ -1,27 +1,44 @@ package com.linkedin.metadata.search.query.request; import static com.linkedin.metadata.utils.SearchUtil.*; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.annotation.SearchableAnnotation; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.elasticsearch.query.request.AggregationQueryBuilder; +import com.linkedin.r2.RemoteInvocationException; +import java.net.URISyntaxException; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder; import org.testng.Assert; +import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; public class AggregationQueryBuilderTest { + private static AspectRetriever aspectRetriever; + + @BeforeClass + public static void setup() throws RemoteInvocationException, URISyntaxException { + aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()).thenReturn(mock(EntityRegistry.class)); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + } + @Test public void testGetDefaultAggregationsHasFields() { SearchableAnnotation annotation = @@ -46,7 +63,9 @@ public void testGetDefaultAggregationsHasFields() { AggregationQueryBuilder builder = new AggregationQueryBuilder( - config, ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation))); + config, + ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation)), + aspectRetriever); List aggs = builder.getAggregations(); @@ -78,7 +97,9 @@ public void testGetDefaultAggregationsFields() { AggregationQueryBuilder builder = new AggregationQueryBuilder( - config, ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation))); + config, + ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation)), + aspectRetriever); List aggs = builder.getAggregations(); @@ -127,7 +148,8 @@ public void testGetSpecificAggregationsHasFields() { AggregationQueryBuilder builder = new AggregationQueryBuilder( config, - ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation1, annotation2))); + ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation1, annotation2)), + aspectRetriever); // Case 1: Ask for fields that should exist. List aggs = @@ -148,7 +170,7 @@ public void testAggregateOverStructuredProperty() { AggregationQueryBuilder builder = new AggregationQueryBuilder( - config, ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of())); + config, ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of()), aspectRetriever); List aggs = builder.getAggregations(List.of("structuredProperties.ab.fgh.ten")); @@ -213,7 +235,8 @@ public void testAggregateOverFieldsAndStructProp() { AggregationQueryBuilder builder = new AggregationQueryBuilder( config, - ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation1, annotation2))); + ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation1, annotation2)), + aspectRetriever); // Aggregate over fields and structured properties List aggs = @@ -264,7 +287,9 @@ public void testMissingAggregation() { AggregationQueryBuilder builder = new AggregationQueryBuilder( - config, ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation))); + config, + ImmutableMap.of(mock(EntitySpec.class), ImmutableList.of(annotation)), + aspectRetriever); List aggs = builder.getAggregations(); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java index ab832eb1ac24fa..bb37fb3f3b206a 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java @@ -1,9 +1,11 @@ package com.linkedin.metadata.search.query.request; +import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; import com.linkedin.metadata.TestEntitySpecBuilder; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.search.elasticsearch.query.request.AutocompleteRequestHandler; import java.util.List; import java.util.Map; @@ -18,7 +20,8 @@ public class AutocompleteRequestHandlerTest { private AutocompleteRequestHandler handler = - AutocompleteRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec()); + AutocompleteRequestHandler.getBuilder( + TestEntitySpecBuilder.getSpec(), mock(AspectRetriever.class)); @Test public void testDefaultAutocompleteRequest() { diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java index 02c9ea800f0af3..37ad1df608ad58 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java @@ -6,12 +6,12 @@ import com.google.common.collect.ImmutableList; import com.linkedin.data.template.StringArray; import com.linkedin.metadata.TestEntitySpecBuilder; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.ExactMatchConfiguration; import com.linkedin.metadata.config.search.PartialConfiguration; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.config.search.WordGramConfiguration; 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.Condition; import com.linkedin.metadata.query.filter.ConjunctiveCriterion; @@ -49,7 +49,7 @@ @Import(SearchCommonTestConfiguration.class) public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests { - @Autowired private EntityRegistry entityRegistry; + @Autowired private AspectRetriever aspectRetriever; public static SearchConfiguration testQueryConfig; @@ -81,9 +81,9 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests { @Test public void testDatasetFieldsAndHighlights() { - EntitySpec entitySpec = entityRegistry.getEntitySpec("dataset"); + EntitySpec entitySpec = aspectRetriever.getEntityRegistry().getEntitySpec("dataset"); SearchRequestHandler datasetHandler = - SearchRequestHandler.getBuilder(entitySpec, testQueryConfig, null); + SearchRequestHandler.getBuilder(entitySpec, testQueryConfig, null, aspectRetriever); /* Ensure efficient query performance, we do not expect upstream/downstream/fineGrained lineage @@ -102,7 +102,8 @@ public void testDatasetFieldsAndHighlights() { @Test public void testSearchRequestHandlerHighlightingTurnedOff() { SearchRequestHandler requestHandler = - SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequestHandler.getBuilder( + TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); SearchRequest searchRequest = requestHandler.getSearchRequest( "testQuery", @@ -141,7 +142,8 @@ public void testSearchRequestHandlerHighlightingTurnedOff() { @Test public void testSearchRequestHandler() { SearchRequestHandler requestHandler = - SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequestHandler.getBuilder( + TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); SearchRequest searchRequest = requestHandler.getSearchRequest( "testQuery", null, null, 0, 10, new SearchFlags().setFulltext(false), null); @@ -196,7 +198,8 @@ public void testSearchRequestHandler() { @Test public void testAggregationsInSearch() { SearchRequestHandler requestHandler = - SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequestHandler.getBuilder( + TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); final String nestedAggString = String.format("_entityType%stextFieldOverride", AGGREGATION_SEPARATOR_CHAR); SearchRequest searchRequest = @@ -264,7 +267,8 @@ public void testAggregationsInSearch() { public void testFilteredSearch() { final SearchRequestHandler requestHandler = - SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequestHandler.getBuilder( + TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); final BoolQueryBuilder testQuery = constructFilterQuery(requestHandler, false); @@ -614,7 +618,8 @@ public void testBrowsePathQueryFilter() { Filter filter = new Filter(); filter.setOr(conjunctiveCriterionArray); - BoolQueryBuilder test = SearchRequestHandler.getFilterQuery(filter, new HashMap<>()); + BoolQueryBuilder test = + SearchRequestHandler.getFilterQuery(filter, new HashMap<>(), aspectRetriever); assertEquals(test.should().size(), 1); @@ -637,7 +642,8 @@ private BoolQueryBuilder getQuery(final Criterion filterCriterion) { .setAnd(new CriterionArray(ImmutableList.of(filterCriterion))))); final SearchRequestHandler requestHandler = - SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequestHandler.getBuilder( + TestEntitySpecBuilder.getSpec(), testQueryConfig, null, aspectRetriever); return (BoolQueryBuilder) requestHandler diff --git a/metadata-io/src/test/java/com/linkedin/metadata/timeline/CassandraTimelineServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/timeline/CassandraTimelineServiceTest.java index 552cb0b52994f9..00168aefab1ef1 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/timeline/CassandraTimelineServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/timeline/CassandraTimelineServiceTest.java @@ -56,13 +56,8 @@ private void configureComponents() { preProcessHooks.setUiEnabled(true); _entityServiceImpl = new EntityServiceImpl( - _aspectDao, - _mockProducer, - _testEntityRegistry, - true, - _mockUpdateIndicesService, - preProcessHooks, - true); + _aspectDao, _mockProducer, _testEntityRegistry, true, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); } /** diff --git a/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java index 5d7137a52eb21e..e9c79f06f37c6a 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java @@ -38,13 +38,8 @@ public void setupTest() { preProcessHooks.setUiEnabled(true); _entityServiceImpl = new EntityServiceImpl( - _aspectDao, - _mockProducer, - _testEntityRegistry, - true, - _mockUpdateIndicesService, - preProcessHooks, - true); + _aspectDao, _mockProducer, _testEntityRegistry, true, preProcessHooks, true); + _entityServiceImpl.setUpdateIndicesService(_mockUpdateIndicesService); } /** diff --git a/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java index 23ca4a4a4247e1..f6141a1d8803ff 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceTestBase.java @@ -3,6 +3,9 @@ import static com.linkedin.metadata.Constants.INGESTION_MAX_SERIALIZED_STRING_LENGTH; import static com.linkedin.metadata.Constants.MAX_JACKSON_STRING_SIZE; import static io.datahubproject.test.search.SearchTestUtils.syncAfterWrite; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; @@ -24,6 +27,7 @@ import com.linkedin.data.template.StringArrayArray; import com.linkedin.data.template.StringMap; import com.linkedin.data.template.StringMapArray; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.DataSchemaFactory; @@ -45,6 +49,7 @@ import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; +import com.linkedin.r2.RemoteInvocationException; import com.linkedin.timeseries.AggregationSpec; import com.linkedin.timeseries.AggregationType; import com.linkedin.timeseries.CalendarInterval; @@ -54,6 +59,7 @@ import com.linkedin.timeseries.GroupingBucketType; import com.linkedin.timeseries.TimeWindowSize; import com.linkedin.timeseries.TimeseriesIndexSizeResult; +import java.net.URISyntaxException; import java.util.Calendar; import java.util.List; import java.util.Map; @@ -99,43 +105,49 @@ public abstract class TimeseriesAspectServiceTestBase extends AbstractTestNGSpri @Nonnull protected abstract ESIndexBuilder getIndexBuilder(); - private EntityRegistry _entityRegistry; - private IndexConvention _indexConvention; - private ElasticSearchTimeseriesAspectService _elasticSearchTimeseriesAspectService; - private AspectSpec _aspectSpec; + private AspectRetriever aspectRetriever; + private IndexConvention indexConvention; + private ElasticSearchTimeseriesAspectService elasticSearchTimeseriesAspectService; + private AspectSpec aspectSpec; - private Map _testEntityProfiles; - private Long _startTime; + private Map testEntityProfiles; + private Long startTime; /* * Basic setup and teardown */ @BeforeClass - public void setup() { - _entityRegistry = + public void setup() throws RemoteInvocationException, URISyntaxException { + EntityRegistry entityRegistry = new ConfigEntityRegistry( new DataSchemaFactory("com.datahub.test"), List.of(), TestEntityProfile.class .getClassLoader() .getResourceAsStream("test-entity-registry.yml")); - _indexConvention = new IndexConventionImpl("es_timeseries_aspect_service_test"); - _elasticSearchTimeseriesAspectService = buildService(); - _elasticSearchTimeseriesAspectService.configure(); - EntitySpec entitySpec = _entityRegistry.getEntitySpec(ENTITY_NAME); - _aspectSpec = entitySpec.getAspectSpec(ASPECT_NAME); + aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()).thenReturn(entityRegistry); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + + indexConvention = new IndexConventionImpl("es_timeseries_aspect_service_test"); + elasticSearchTimeseriesAspectService = buildService(); + elasticSearchTimeseriesAspectService.configure(); + EntitySpec entitySpec = entityRegistry.getEntitySpec(ENTITY_NAME); + aspectSpec = entitySpec.getAspectSpec(ASPECT_NAME); } @Nonnull private ElasticSearchTimeseriesAspectService buildService() { return new ElasticSearchTimeseriesAspectService( - getSearchClient(), - _indexConvention, - new TimeseriesAspectIndexBuilders(getIndexBuilder(), _entityRegistry, _indexConvention), - _entityRegistry, - getBulkProcessor(), - 1); + getSearchClient(), + indexConvention, + new TimeseriesAspectIndexBuilders( + getIndexBuilder(), aspectRetriever.getEntityRegistry(), indexConvention), + aspectRetriever.getEntityRegistry(), + getBulkProcessor(), + 1) + .postConstruct(aspectRetriever); } /* @@ -144,11 +156,11 @@ private ElasticSearchTimeseriesAspectService buildService() { private void upsertDocument(TestEntityProfile dp, Urn urn) throws JsonProcessingException { Map documents = - TimeseriesAspectTransformer.transform(urn, dp, _aspectSpec, null); + TimeseriesAspectTransformer.transform(urn, dp, aspectSpec, null); assertEquals(documents.size(), 3); documents.forEach( (key, value) -> - _elasticSearchTimeseriesAspectService.upsertDocument( + elasticSearchTimeseriesAspectService.upsertDocument( ENTITY_NAME, ASPECT_NAME, key, value)); } @@ -190,10 +202,10 @@ private TestEntityProfile makeTestProfile(long eventTime, long stat, String mess @Test(groups = "upsert") public void testUpsertProfiles() throws Exception { // Create the testEntity profiles that we would like to use for testing. - _startTime = Calendar.getInstance().getTimeInMillis(); - _startTime = _startTime - _startTime % 86400000; + startTime = Calendar.getInstance().getTimeInMillis(); + startTime = startTime - startTime % 86400000; // Create the testEntity profiles that we would like to use for testing. - TestEntityProfile firstProfile = makeTestProfile(_startTime, 20, null); + TestEntityProfile firstProfile = makeTestProfile(startTime, 20, null); Stream testEntityProfileStream = Stream.iterate( firstProfile, @@ -201,17 +213,17 @@ public void testUpsertProfiles() throws Exception { makeTestProfile( prev.getTimestampMillis() + TIME_INCREMENT, prev.getStat() + 10, null)); - _testEntityProfiles = + testEntityProfiles = testEntityProfileStream .limit(NUM_PROFILES) .collect(Collectors.toMap(TestEntityProfile::getTimestampMillis, Function.identity())); - Long endTime = _startTime + (NUM_PROFILES - 1) * TIME_INCREMENT; + Long endTime = startTime + (NUM_PROFILES - 1) * TIME_INCREMENT; - assertNotNull(_testEntityProfiles.get(_startTime)); - assertNotNull(_testEntityProfiles.get(endTime)); + assertNotNull(testEntityProfiles.get(startTime)); + assertNotNull(testEntityProfiles.get(endTime)); // Upsert the documents into the index. - _testEntityProfiles + testEntityProfiles .values() .forEach( x -> { @@ -260,7 +272,7 @@ public void testUpsertProfilesWithUniqueMessageIds() throws Exception { syncAfterWrite(getBulkProcessor()); List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( urn, ENTITY_NAME, ASPECT_NAME, null, null, testEntityProfiles.size(), null); assertEquals(resultAspects.size(), testEntityProfiles.size()); } @@ -273,8 +285,8 @@ private void validateAspectValue(EnvelopedAspect envelopedAspectResult) { TestEntityProfile actualProfile = (TestEntityProfile) GenericRecordUtils.deserializeAspect( - envelopedAspectResult.getAspect().getValue(), CONTENT_TYPE, _aspectSpec); - TestEntityProfile expectedProfile = _testEntityProfiles.get(actualProfile.getTimestampMillis()); + envelopedAspectResult.getAspect().getValue(), CONTENT_TYPE, aspectSpec); + TestEntityProfile expectedProfile = testEntityProfiles.get(actualProfile.getTimestampMillis()); assertNotNull(expectedProfile); assertEquals(actualProfile.getStat(), expectedProfile.getStat()); assertEquals(actualProfile.getTimestampMillis(), expectedProfile.getTimestampMillis()); @@ -288,20 +300,20 @@ private void validateAspectValues(List aspects, long numResults @Test(groups = "getAspectValues", dependsOnGroups = "upsert") public void testGetAspectTimeseriesValuesAll() { List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, null, null, NUM_PROFILES, null); validateAspectValues(resultAspects, NUM_PROFILES); TestEntityProfile firstProfile = (TestEntityProfile) GenericRecordUtils.deserializeAspect( - resultAspects.get(0).getAspect().getValue(), CONTENT_TYPE, _aspectSpec); + resultAspects.get(0).getAspect().getValue(), CONTENT_TYPE, aspectSpec); TestEntityProfile lastProfile = (TestEntityProfile) GenericRecordUtils.deserializeAspect( resultAspects.get(resultAspects.size() - 1).getAspect().getValue(), CONTENT_TYPE, - _aspectSpec); + aspectSpec); // Now verify that the first index is the one with the highest stat value, and the last the one // with the lower. @@ -312,7 +324,7 @@ public void testGetAspectTimeseriesValuesAll() { @Test(groups = "getAspectValues", dependsOnGroups = "upsert") public void testGetAspectTimeseriesValuesAllSorted() { List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, @@ -326,13 +338,13 @@ public void testGetAspectTimeseriesValuesAllSorted() { TestEntityProfile firstProfile = (TestEntityProfile) GenericRecordUtils.deserializeAspect( - resultAspects.get(0).getAspect().getValue(), CONTENT_TYPE, _aspectSpec); + resultAspects.get(0).getAspect().getValue(), CONTENT_TYPE, aspectSpec); TestEntityProfile lastProfile = (TestEntityProfile) GenericRecordUtils.deserializeAspect( resultAspects.get(resultAspects.size() - 1).getAspect().getValue(), CONTENT_TYPE, - _aspectSpec); + aspectSpec); // Now verify that the first index is the one with the highest stat value, and the last the one // with the lower. @@ -347,7 +359,7 @@ public void testGetAspectTimeseriesValuesWithFilter() { new Criterion().setField("stat").setCondition(Condition.EQUAL).setValue("20"); filter.setCriteria(new CriterionArray(hasStatEqualsTwenty)); List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, null, null, NUM_PROFILES, filter); validateAspectValues(resultAspects, 1); } @@ -356,12 +368,12 @@ public void testGetAspectTimeseriesValuesWithFilter() { public void testGetAspectTimeseriesValuesSubRangeInclusiveOverlap() { int expectedNumRows = 10; List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, - _startTime, - _startTime + TIME_INCREMENT * (expectedNumRows - 1), + startTime, + startTime + TIME_INCREMENT * (expectedNumRows - 1), expectedNumRows, null); validateAspectValues(resultAspects, expectedNumRows); @@ -371,12 +383,12 @@ public void testGetAspectTimeseriesValuesSubRangeInclusiveOverlap() { public void testGetAspectTimeseriesValuesSubRangeExclusiveOverlap() { int expectedNumRows = 10; List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, - _startTime + TIME_INCREMENT / 2, - _startTime + TIME_INCREMENT * expectedNumRows + TIME_INCREMENT / 2, + startTime + TIME_INCREMENT / 2, + startTime + TIME_INCREMENT * expectedNumRows + TIME_INCREMENT / 2, expectedNumRows, null); validateAspectValues(resultAspects, expectedNumRows); @@ -386,12 +398,12 @@ public void testGetAspectTimeseriesValuesSubRangeExclusiveOverlap() { public void testGetAspectTimeseriesValuesSubRangeExclusiveOverlapLatestValueOnly() { int expectedNumRows = 1; List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, - _startTime + TIME_INCREMENT / 2, - _startTime + TIME_INCREMENT * expectedNumRows + TIME_INCREMENT / 2, + startTime + TIME_INCREMENT / 2, + startTime + TIME_INCREMENT * expectedNumRows + TIME_INCREMENT / 2, expectedNumRows, null); validateAspectValues(resultAspects, expectedNumRows); @@ -401,12 +413,12 @@ public void testGetAspectTimeseriesValuesSubRangeExclusiveOverlapLatestValueOnly public void testGetAspectTimeseriesValuesExactlyOneResponse() { int expectedNumRows = 1; List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( TEST_URN, ENTITY_NAME, ASPECT_NAME, - _startTime + TIME_INCREMENT / 2, - _startTime + TIME_INCREMENT * 3 / 2, + startTime + TIME_INCREMENT / 2, + startTime + TIME_INCREMENT * 3 / 2, expectedNumRows, null); validateAspectValues(resultAspects, expectedNumRows); @@ -418,7 +430,7 @@ public void testGetAspectTimeseriesValuesExactlyOneResponse() { public void testGetAspectTimeseriesValueMissingUrn() { Urn nonExistingUrn = new TestEntityUrn("missing", "missing", "missing"); List resultAspects = - _elasticSearchTimeseriesAspectService.getAspectValues( + elasticSearchTimeseriesAspectService.getAspectValues( nonExistingUrn, ENTITY_NAME, ASPECT_NAME, null, null, NUM_PROFILES, null); validateAspectValues(resultAspects, 0); } @@ -439,12 +451,12 @@ public void testGetAggregatedStatsLatestStatForDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -462,7 +474,7 @@ public void testGetAggregatedStatsLatestStatForDay1() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -481,8 +493,8 @@ public void testGetAggregatedStatsLatestStatForDay1() { resultTable.getRows(), new StringArrayArray( new StringArray( - _startTime.toString(), - _testEntityProfiles.get(_startTime + 23 * TIME_INCREMENT).getStat().toString()))); + startTime.toString(), + testEntityProfiles.get(startTime + 23 * TIME_INCREMENT).getStat().toString()))); } @Test( @@ -496,13 +508,13 @@ public void testGetAggregatedStatsLatestStatForDay1WithValues() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValues(new StringArray(_startTime.toString())) + .setValues(new StringArray(startTime.toString())) .setValue(""); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValues(new StringArray(String.valueOf(_startTime + 23 * TIME_INCREMENT))) + .setValues(new StringArray(String.valueOf(startTime + 23 * TIME_INCREMENT))) .setValue(""); Filter filter = @@ -521,7 +533,7 @@ public void testGetAggregatedStatsLatestStatForDay1WithValues() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -540,8 +552,8 @@ public void testGetAggregatedStatsLatestStatForDay1WithValues() { resultTable.getRows(), new StringArrayArray( new StringArray( - _startTime.toString(), - _testEntityProfiles.get(_startTime + 23 * TIME_INCREMENT).getStat().toString()))); + startTime.toString(), + testEntityProfiles.get(startTime + 23 * TIME_INCREMENT).getStat().toString()))); } @Test( @@ -555,12 +567,12 @@ public void testGetAggregatedStatsLatestAComplexNestedRecordForDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -580,7 +592,7 @@ public void testGetAggregatedStatsLatestAComplexNestedRecordForDay1() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -595,13 +607,13 @@ public void testGetAggregatedStatsLatestAComplexNestedRecordForDay1() { // Validate rows assertNotNull(resultTable.getRows()); assertEquals(resultTable.getRows().size(), 1); - assertEquals(resultTable.getRows().get(0).get(0), _startTime.toString()); + assertEquals(resultTable.getRows().get(0).get(0), startTime.toString()); try { ComplexNestedRecord latestAComplexNestedRecord = OBJECT_MAPPER.readValue(resultTable.getRows().get(0).get(1), ComplexNestedRecord.class); assertEquals( latestAComplexNestedRecord, - _testEntityProfiles.get(_startTime + 23 * TIME_INCREMENT).getAComplexNestedRecord()); + testEntityProfiles.get(startTime + 23 * TIME_INCREMENT).getAComplexNestedRecord()); } catch (JsonProcessingException e) { fail("Unexpected exception thrown" + e); } @@ -618,12 +630,12 @@ public void testGetAggregatedStatsLatestStrArrayDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -641,7 +653,7 @@ public void testGetAggregatedStatsLatestStrArrayDay1() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -656,7 +668,7 @@ public void testGetAggregatedStatsLatestStrArrayDay1() { assertNotNull(resultTable.getRows()); assertEquals(resultTable.getRows().size(), 1); StringArray expectedStrArray = - _testEntityProfiles.get(_startTime + 23 * TIME_INCREMENT).getStrArray(); + testEntityProfiles.get(startTime + 23 * TIME_INCREMENT).getStrArray(); // assertEquals(resultTable.getRows(), new StringArrayArray(new // StringArray(_startTime.toString(), // expectedStrArray.toString()))); @@ -681,12 +693,12 @@ public void testGetAggregatedStatsLatestStatForTwoDays() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 47 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 47 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -704,7 +716,7 @@ public void testGetAggregatedStatsLatestStatForTwoDays() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -719,16 +731,16 @@ public void testGetAggregatedStatsLatestStatForTwoDays() { // Validate rows assertNotNull(resultTable.getRows()); assertEquals(resultTable.getRows().size(), 2); - Long latestDay1Ts = _startTime + 23 * TIME_INCREMENT; - Long latestDay2Ts = _startTime + 47 * TIME_INCREMENT; + Long latestDay1Ts = startTime + 23 * TIME_INCREMENT; + Long latestDay2Ts = startTime + 47 * TIME_INCREMENT; assertEquals( resultTable.getRows(), new StringArrayArray( new StringArray( - _startTime.toString(), _testEntityProfiles.get(latestDay1Ts).getStat().toString()), + startTime.toString(), testEntityProfiles.get(latestDay1Ts).getStat().toString()), new StringArray( - String.valueOf(_startTime + 24 * TIME_INCREMENT), - _testEntityProfiles.get(latestDay2Ts).getStat().toString()))); + String.valueOf(startTime + 24 * TIME_INCREMENT), + testEntityProfiles.get(latestDay2Ts).getStat().toString()))); } @Test( @@ -741,12 +753,12 @@ public void testGetAggregatedStatsLatestStatForFirst10HoursOfDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 9 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 9 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -764,7 +776,7 @@ public void testGetAggregatedStatsLatestStatForFirst10HoursOfDay1() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -783,22 +795,22 @@ public void testGetAggregatedStatsLatestStatForFirst10HoursOfDay1() { resultTable.getRows(), new StringArrayArray( new StringArray( - _startTime.toString(), - _testEntityProfiles.get(_startTime + 9 * TIME_INCREMENT).getStat().toString()))); + startTime.toString(), + testEntityProfiles.get(startTime + 9 * TIME_INCREMENT).getStat().toString()))); } @Test( groups = {"getAggregatedStats"}, dependsOnGroups = {"upsert"}) public void testGetAggregatedStatsLatestStatForCol1Day1() { - Long lastEntryTimeStamp = _startTime + 23 * TIME_INCREMENT; + Long lastEntryTimeStamp = startTime + 23 * TIME_INCREMENT; Criterion hasUrnCriterion = new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(TEST_URN.toString()); Criterion startTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) @@ -833,7 +845,7 @@ public void testGetAggregatedStatsLatestStatForCol1Day1() { .setType(GroupingBucketType.STRING_GROUPING_BUCKET); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -853,9 +865,9 @@ public void testGetAggregatedStatsLatestStatForCol1Day1() { resultTable.getRows(), new StringArrayArray( new StringArray( - _startTime.toString(), + startTime.toString(), "col1", - _testEntityProfiles + testEntityProfiles .get(lastEntryTimeStamp) .getComponentProfiles() .get(0) @@ -867,14 +879,14 @@ public void testGetAggregatedStatsLatestStatForCol1Day1() { groups = {"getAggregatedStats"}, dependsOnGroups = {"upsert"}) public void testGetAggregatedStatsLatestStatForAllColumnsDay1() { - Long lastEntryTimeStamp = _startTime + 23 * TIME_INCREMENT; + Long lastEntryTimeStamp = startTime + 23 * TIME_INCREMENT; Criterion hasUrnCriterion = new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(TEST_URN.toString()); Criterion startTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) @@ -904,7 +916,7 @@ public void testGetAggregatedStatsLatestStatForAllColumnsDay1() { .setType(GroupingBucketType.STRING_GROUPING_BUCKET); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {latestStatAggregationSpec}, @@ -920,9 +932,9 @@ public void testGetAggregatedStatsLatestStatForAllColumnsDay1() { // Validate rows StringArray expectedRow1 = new StringArray( - _startTime.toString(), + startTime.toString(), "col1", - _testEntityProfiles + testEntityProfiles .get(lastEntryTimeStamp) .getComponentProfiles() .get(0) @@ -930,9 +942,9 @@ public void testGetAggregatedStatsLatestStatForAllColumnsDay1() { .toString()); StringArray expectedRow2 = new StringArray( - _startTime.toString(), + startTime.toString(), "col2", - _testEntityProfiles + testEntityProfiles .get(lastEntryTimeStamp) .getComponentProfiles() .get(1) @@ -955,12 +967,12 @@ public void testGetAggregatedStatsSumStatForFirst10HoursOfDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 9 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 9 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -978,7 +990,7 @@ public void testGetAggregatedStatsSumStatForFirst10HoursOfDay1() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {sumAggregationSpec}, @@ -996,21 +1008,21 @@ public void testGetAggregatedStatsSumStatForFirst10HoursOfDay1() { // TODO: Compute this caching the documents. assertEquals( resultTable.getRows(), - new StringArrayArray(new StringArray(_startTime.toString(), String.valueOf(650)))); + new StringArrayArray(new StringArray(startTime.toString(), String.valueOf(650)))); } @Test( groups = {"getAggregatedStats"}, dependsOnGroups = {"upsert"}) public void testGetAggregatedStatsSumStatForCol2Day1() { - Long lastEntryTimeStamp = _startTime + 23 * TIME_INCREMENT; + Long lastEntryTimeStamp = startTime + 23 * TIME_INCREMENT; Criterion hasUrnCriterion = new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(TEST_URN.toString()); Criterion startTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) @@ -1045,7 +1057,7 @@ public void testGetAggregatedStatsSumStatForCol2Day1() { .setType(GroupingBucketType.STRING_GROUPING_BUCKET); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {sumStatAggregationSpec}, @@ -1065,7 +1077,7 @@ public void testGetAggregatedStatsSumStatForCol2Day1() { // TODO: Compute this caching the documents. assertEquals( resultTable.getRows(), - new StringArrayArray(new StringArray(_startTime.toString(), "col2", String.valueOf(3288)))); + new StringArrayArray(new StringArray(startTime.toString(), "col2", String.valueOf(3288)))); } @Test( @@ -1079,12 +1091,12 @@ public void testGetAggregatedStatsCardinalityAggStrStatDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -1104,7 +1116,7 @@ public void testGetAggregatedStatsCardinalityAggStrStatDay1() { .setTimeWindowSize(new TimeWindowSize().setMultiple(1).setUnit(CalendarInterval.DAY)); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {cardinalityStatAggregationSpec}, @@ -1120,7 +1132,7 @@ public void testGetAggregatedStatsCardinalityAggStrStatDay1() { assertNotNull(resultTable.getRows()); assertEquals(resultTable.getRows().size(), 1); assertEquals( - resultTable.getRows(), new StringArrayArray(new StringArray(_startTime.toString(), "24"))); + resultTable.getRows(), new StringArrayArray(new StringArray(startTime.toString(), "24"))); } @Test( @@ -1134,12 +1146,12 @@ public void testGetAggregatedStatsSumStatsCollectionDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( @@ -1158,7 +1170,7 @@ public void testGetAggregatedStatsSumStatsCollectionDay1() { .setType(GroupingBucketType.STRING_GROUPING_BUCKET); GenericTable resultTable = - _elasticSearchTimeseriesAspectService.getAggregatedStats( + elasticSearchTimeseriesAspectService.getAggregatedStats( ENTITY_NAME, ASPECT_NAME, new AggregationSpec[] {cardinalityStatAggregationSpec}, @@ -1188,18 +1200,18 @@ public void testDeleteAspectValuesByUrnAndTimeRangeDay1() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter filter = QueryUtils.getFilterFromCriteria( ImmutableList.of(hasUrnCriterion, startTimeCriterion, endTimeCriterion)); DeleteAspectValuesResult result = - _elasticSearchTimeseriesAspectService.deleteAspectValues(ENTITY_NAME, ASPECT_NAME, filter); + elasticSearchTimeseriesAspectService.deleteAspectValues(ENTITY_NAME, ASPECT_NAME, filter); // For day1, we expect 24 (number of hours) * 3 (each testEntityProfile aspect expands 3 elastic // docs: // 1 original + 2 for componentProfiles) = 72 total. @@ -1214,7 +1226,7 @@ public void testDeleteAspectValuesByUrn() { new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(TEST_URN.toString()); Filter filter = QueryUtils.getFilterFromCriteria(ImmutableList.of(hasUrnCriterion)); DeleteAspectValuesResult result = - _elasticSearchTimeseriesAspectService.deleteAspectValues(ENTITY_NAME, ASPECT_NAME, filter); + elasticSearchTimeseriesAspectService.deleteAspectValues(ENTITY_NAME, ASPECT_NAME, filter); // Of the 300 elastic docs upserted for TEST_URN, 72 got deleted by deleteAspectValues1 test // group leaving 228. assertEquals(result.getNumDocsDeleted(), Long.valueOf(228L)); @@ -1229,7 +1241,7 @@ public void testCountByFilter() { new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(TEST_URN.toString()); Filter filter = QueryUtils.getFilterFromCriteria(ImmutableList.of(hasUrnCriterion)); long count = - _elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, filter); + elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, filter); assertEquals(count, 300L); // Test with filter with multiple criteria @@ -1237,24 +1249,24 @@ public void testCountByFilter() { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter urnAndTimeFilter = QueryUtils.getFilterFromCriteria( ImmutableList.of(hasUrnCriterion, startTimeCriterion, endTimeCriterion)); count = - _elasticSearchTimeseriesAspectService.countByFilter( + elasticSearchTimeseriesAspectService.countByFilter( ENTITY_NAME, ASPECT_NAME, urnAndTimeFilter); assertEquals(count, 72L); // test without filter count = - _elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, new Filter()); + elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, new Filter()); // There may be other entities in there from other tests assertTrue(count >= 300L); } @@ -1269,7 +1281,7 @@ public void testCountByFilterAfterDelete() throws InterruptedException { new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(TEST_URN.toString()); Filter filter = QueryUtils.getFilterFromCriteria(ImmutableList.of(hasUrnCriterion)); long count = - _elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, filter); + elasticSearchTimeseriesAspectService.countByFilter(ENTITY_NAME, ASPECT_NAME, filter); assertEquals(count, 228L); // Test with filter with multiple criteria @@ -1277,18 +1289,18 @@ public void testCountByFilterAfterDelete() throws InterruptedException { new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO) - .setValue(_startTime.toString()); + .setValue(startTime.toString()); Criterion endTimeCriterion = new Criterion() .setField(ES_FIELD_TIMESTAMP) .setCondition(Condition.LESS_THAN_OR_EQUAL_TO) - .setValue(String.valueOf(_startTime + 23 * TIME_INCREMENT)); + .setValue(String.valueOf(startTime + 23 * TIME_INCREMENT)); Filter urnAndTimeFilter = QueryUtils.getFilterFromCriteria( ImmutableList.of(hasUrnCriterion, startTimeCriterion, endTimeCriterion)); count = - _elasticSearchTimeseriesAspectService.countByFilter( + elasticSearchTimeseriesAspectService.countByFilter( ENTITY_NAME, ASPECT_NAME, urnAndTimeFilter); assertEquals(count, 0L); } @@ -1297,7 +1309,7 @@ public void testCountByFilterAfterDelete() throws InterruptedException { groups = {"getAggregatedStats"}, dependsOnGroups = {"upsert"}) public void testGetIndexSizes() { - List result = _elasticSearchTimeseriesAspectService.getIndexSizes(); + List result = elasticSearchTimeseriesAspectService.getIndexSizes(); // CHECKSTYLE:OFF /* Example result: diff --git a/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceUnitTest.java b/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceUnitTest.java index a23267dcf6f55e..d56ddb2cc808ea 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceUnitTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/timeseries/search/TimeseriesAspectServiceUnitTest.java @@ -5,7 +5,7 @@ import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.NumericNode; import com.fasterxml.jackson.databind.node.ObjectNode; -import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.search.elasticsearch.update.ESBulkProcessor; import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.metadata.timeseries.elastic.ElasticSearchTimeseriesAspectService; @@ -34,17 +34,18 @@ public class TimeseriesAspectServiceUnitTest { private final IndexConvention _indexConvention = mock(IndexConvention.class); private final TimeseriesAspectIndexBuilders _timeseriesAspectIndexBuilders = mock(TimeseriesAspectIndexBuilders.class); - private final EntityRegistry _entityRegistry = mock(EntityRegistry.class); + private final AspectRetriever aspectRetriever = mock(AspectRetriever.class); private final ESBulkProcessor _bulkProcessor = mock(ESBulkProcessor.class); private final RestClient _restClient = mock(RestClient.class); private final TimeseriesAspectService _timeseriesAspectService = new ElasticSearchTimeseriesAspectService( - _searchClient, - _indexConvention, - _timeseriesAspectIndexBuilders, - _entityRegistry, - _bulkProcessor, - 0); + _searchClient, + _indexConvention, + _timeseriesAspectIndexBuilders, + aspectRetriever.getEntityRegistry(), + _bulkProcessor, + 0) + .postConstruct(aspectRetriever); private static final String INDEX_PATTERN = "indexPattern"; diff --git a/metadata-io/src/test/java/io/datahubproject/test/DataGenerator.java b/metadata-io/src/test/java/io/datahubproject/test/DataGenerator.java index eb4c85209ce422..c27a1c337ed5c5 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/DataGenerator.java +++ b/metadata-io/src/test/java/io/datahubproject/test/DataGenerator.java @@ -16,7 +16,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.glossary.GlossaryTermInfo; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.aspect.batch.MCPBatchItem; +import com.linkedin.metadata.aspect.batch.MCPItem; import com.linkedin.metadata.aspect.utils.DefaultAspectsUtil; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.AspectDao; @@ -27,7 +27,6 @@ import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.service.UpdateIndicesService; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; @@ -75,7 +74,6 @@ public static DataGenerator build(EntityRegistry entityRegistry) { mock(EventProducer.class), entityRegistry, false, - mock(UpdateIndicesService.class), mock(PreProcessHooks.class), anyBoolean()); return new DataGenerator(mockEntityServiceImpl); @@ -172,7 +170,7 @@ public Stream> generateMCPs( entityService, true) .stream() - .map(MCPBatchItem::getMetadataChangeProposal)) + .map(MCPItem::getMetadataChangeProposal)) .collect(Collectors.toList()); } else { return List.of(mcp); diff --git a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java index b42cd89131f51f..24acb7bbcb4a70 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java +++ b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SampleDataFixtureConfiguration.java @@ -8,6 +8,7 @@ import com.fasterxml.jackson.dataformat.yaml.YAMLMapper; import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.config.cache.EntityDocCountCacheConfiguration; @@ -138,25 +139,29 @@ protected EntityIndexBuilders entityIndexBuildersHelper( protected ElasticSearchService entitySearchService( @Qualifier("entityRegistry") EntityRegistry entityRegistry, @Qualifier("sampleDataEntityIndexBuilders") EntityIndexBuilders indexBuilders, - @Qualifier("sampleDataIndexConvention") IndexConvention indexConvention) + @Qualifier("sampleDataIndexConvention") IndexConvention indexConvention, + @Qualifier("aspectRetriever") final AspectRetriever aspectRetriever) throws IOException { - return entitySearchServiceHelper(entityRegistry, indexBuilders, indexConvention); + return entitySearchServiceHelper( + entityRegistry, indexBuilders, indexConvention, aspectRetriever); } @Bean(name = "longTailEntitySearchService") protected ElasticSearchService longTailEntitySearchService( @Qualifier("entityRegistry") EntityRegistry longTailEntityRegistry, @Qualifier("longTailEntityIndexBuilders") EntityIndexBuilders longTailEndexBuilders, - @Qualifier("longTailIndexConvention") IndexConvention longTailIndexConvention) + @Qualifier("longTailIndexConvention") IndexConvention longTailIndexConvention, + @Qualifier("aspectRetriever") final AspectRetriever aspectRetriever) throws IOException { return entitySearchServiceHelper( - longTailEntityRegistry, longTailEndexBuilders, longTailIndexConvention); + longTailEntityRegistry, longTailEndexBuilders, longTailIndexConvention, aspectRetriever); } protected ElasticSearchService entitySearchServiceHelper( EntityRegistry entityRegistry, EntityIndexBuilders indexBuilders, - IndexConvention indexConvention) + IndexConvention indexConvention, + AspectRetriever aspectRetriever) throws IOException { CustomConfiguration customConfiguration = new CustomConfiguration(); customConfiguration.setEnabled(true); @@ -166,7 +171,6 @@ protected ElasticSearchService entitySearchServiceHelper( ESSearchDAO searchDAO = new ESSearchDAO( - entityRegistry, _searchClient, indexConvention, false, @@ -175,14 +179,11 @@ protected ElasticSearchService entitySearchServiceHelper( customSearchConfiguration); ESBrowseDAO browseDAO = new ESBrowseDAO( - entityRegistry, - _searchClient, - indexConvention, - _searchConfiguration, - _customSearchConfiguration); + _searchClient, indexConvention, _searchConfiguration, _customSearchConfiguration); ESWriteDAO writeDAO = new ESWriteDAO(entityRegistry, _searchClient, indexConvention, _bulkProcessor, 1); - return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); + return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO) + .postConstruct(aspectRetriever); } @Bean(name = "sampleDataSearchService") @@ -296,8 +297,7 @@ private EntityClient entityClientHelper( PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); return new JavaEntityClient( - new EntityServiceImpl( - mockAspectDao, null, entityRegistry, true, null, preProcessHooks, true), + new EntityServiceImpl(mockAspectDao, null, entityRegistry, true, preProcessHooks, true), null, entitySearchService, cachingEntitySearchService, diff --git a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java index 07d27245222b9e..1c43e623443c1e 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java +++ b/metadata-io/src/test/java/io/datahubproject/test/fixtures/search/SearchLineageFixtureConfiguration.java @@ -3,6 +3,7 @@ import static com.linkedin.metadata.Constants.*; import com.linkedin.entity.client.EntityClient; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.config.cache.EntityDocCountCacheConfiguration; @@ -53,13 +54,13 @@ @Import(SearchCommonTestConfiguration.class) public class SearchLineageFixtureConfiguration { - @Autowired private ESBulkProcessor _bulkProcessor; + @Autowired private ESBulkProcessor bulkProcessor; - @Autowired private RestHighLevelClient _searchClient; + @Autowired private RestHighLevelClient searchClient; - @Autowired private SearchConfiguration _searchConfiguration; + @Autowired private SearchConfiguration searchConfiguration; - @Autowired private CustomSearchConfiguration _customSearchConfiguration; + @Autowired private CustomSearchConfiguration customSearchConfiguration; @Bean(name = "searchLineagePrefix") protected String indexPrefix() { @@ -91,7 +92,7 @@ protected EntityIndexBuilders entityIndexBuilders( GitVersion gitVersion = new GitVersion("0.0.0-test", "123456", Optional.empty()); ESIndexBuilder indexBuilder = new ESIndexBuilder( - _searchClient, + searchClient, 1, 0, 1, @@ -107,28 +108,25 @@ protected EntityIndexBuilders entityIndexBuilders( @Bean(name = "searchLineageEntitySearchService") protected ElasticSearchService entitySearchService( - @Qualifier("entityRegistry") EntityRegistry entityRegistry, + @Qualifier("aspectRetriever") AspectRetriever aspectRetriever, @Qualifier("searchLineageEntityIndexBuilders") EntityIndexBuilders indexBuilders, @Qualifier("searchLineageIndexConvention") IndexConvention indexConvention) { ESSearchDAO searchDAO = new ESSearchDAO( - entityRegistry, - _searchClient, + searchClient, indexConvention, false, ELASTICSEARCH_IMPLEMENTATION_ELASTICSEARCH, - _searchConfiguration, + searchConfiguration, null); ESBrowseDAO browseDAO = new ESBrowseDAO( - entityRegistry, - _searchClient, - indexConvention, - _searchConfiguration, - _customSearchConfiguration); + searchClient, indexConvention, searchConfiguration, customSearchConfiguration); ESWriteDAO writeDAO = - new ESWriteDAO(entityRegistry, _searchClient, indexConvention, _bulkProcessor, 1); - return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO); + new ESWriteDAO( + aspectRetriever.getEntityRegistry(), searchClient, indexConvention, bulkProcessor, 1); + return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO) + .postConstruct(aspectRetriever); } @Bean(name = "searchLineageESIndexBuilder") @@ -136,7 +134,7 @@ protected ElasticSearchService entitySearchService( protected ESIndexBuilder esIndexBuilder() { GitVersion gitVersion = new GitVersion("0.0.0-test", "123456", Optional.empty()); return new ESIndexBuilder( - _searchClient, + searchClient, 1, 1, 1, @@ -158,11 +156,11 @@ protected ElasticSearchGraphService graphService( ElasticSearchGraphService graphService = new ElasticSearchGraphService( lineageRegistry, - _bulkProcessor, + bulkProcessor, indexConvention, - new ESGraphWriteDAO(indexConvention, _bulkProcessor, 1), + new ESGraphWriteDAO(indexConvention, bulkProcessor, 1), new ESGraphQueryDAO( - _searchClient, + searchClient, lineageRegistry, indexConvention, GraphQueryConfiguration.testDefaults), @@ -183,7 +181,7 @@ protected LineageSearchService lineageSearchService( // Load fixture data (after graphService mappings applied) FixtureReader.builder() - .bulkProcessor(_bulkProcessor) + .bulkProcessor(bulkProcessor) .fixtureName(fixtureName) .targetIndexPrefix(prefix) .refreshIntervalSeconds(SearchTestContainerConfiguration.REFRESH_INTERVAL_SECONDS) @@ -234,7 +232,7 @@ protected EntityClient entityClient( PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); return new JavaEntityClient( - new EntityServiceImpl(null, null, entityRegistry, true, null, preProcessHooks, true), + new EntityServiceImpl(null, null, entityRegistry, true, preProcessHooks, true), null, entitySearchService, cachingEntitySearchService, diff --git a/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java b/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java index 17747d9ba1cc9e..ae81eaf1ef3884 100644 --- a/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java +++ b/metadata-io/src/test/java/io/datahubproject/test/search/config/SearchCommonTestConfiguration.java @@ -1,6 +1,12 @@ package io.datahubproject.test.search.config; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.datahub.test.Snapshot; import com.fasterxml.jackson.dataformat.yaml.YAMLMapper; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.CustomConfiguration; import com.linkedin.metadata.config.search.ExactMatchConfiguration; import com.linkedin.metadata.config.search.PartialConfiguration; @@ -10,6 +16,10 @@ import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistryException; +import com.linkedin.metadata.models.registry.SnapshotEntityRegistry; +import com.linkedin.r2.RemoteInvocationException; +import java.net.URISyntaxException; +import java.util.Map; import org.springframework.boot.test.context.TestConfiguration; import org.springframework.context.annotation.Bean; @@ -59,4 +69,23 @@ public EntityRegistry entityRegistry() throws EntityRegistryException { .getClassLoader() .getResourceAsStream("entity-registry.yml")); } + + @Bean(name = "aspectRetriever") + protected AspectRetriever aspectRetriever(final EntityRegistry entityRegistry) + throws RemoteInvocationException, URISyntaxException { + AspectRetriever aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()).thenReturn(entityRegistry); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + return aspectRetriever; + } + + @Bean(name = "snapshotRegistryAspectRetriever") + protected AspectRetriever snapshotRegistryAspectRetriever() + throws RemoteInvocationException, URISyntaxException { + AspectRetriever aspectRetriever = mock(AspectRetriever.class); + when(aspectRetriever.getEntityRegistry()) + .thenReturn(new SnapshotEntityRegistry(new Snapshot())); + when(aspectRetriever.getLatestAspectObjects(any(), any())).thenReturn(Map.of()); + return aspectRetriever; + } } diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java index 278c52030b5fc0..91cba8f927a33d 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeLogProcessor.java @@ -54,6 +54,11 @@ public MetadataChangeLogProcessor(List metadataChangeLogH .filter(MetadataChangeLogHook::isEnabled) .sorted(Comparator.comparing(MetadataChangeLogHook::executionOrder)) .collect(Collectors.toList()); + log.info( + "Enabled hooks: {}", + this.hooks.stream() + .map(hook -> hook.getClass().getSimpleName()) + .collect(Collectors.toList())); this.hooks.forEach(MetadataChangeLogHook::init); } diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java index fc47679bebd395..a80017a0956b22 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java @@ -7,13 +7,16 @@ import com.datahub.metadata.ingestion.IngestionScheduler; import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.gms.factory.kafka.schemaregistry.SchemaRegistryConfig; +import com.linkedin.metadata.aspect.CachingAspectRetriever; import com.linkedin.metadata.boot.kafka.DataHubUpgradeKafkaListener; +import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.registry.SchemaRegistryService; import com.linkedin.metadata.search.elasticsearch.ElasticSearchService; import com.linkedin.metadata.search.elasticsearch.indexbuilder.EntityIndexBuilders; import com.linkedin.metadata.search.transformer.SearchDocumentTransformer; +import com.linkedin.metadata.service.FormService; import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; import org.apache.avro.generic.GenericRecord; @@ -45,7 +48,7 @@ public class MCLSpringTestConfiguration { @MockBean public IngestionScheduler ingestionScheduler; - @Bean + @Bean(name = "systemEntityClient") public SystemEntityClient systemEntityClient( @Qualifier("systemAuthentication") Authentication systemAuthentication) { SystemEntityClient systemEntityClient = mock(SystemEntityClient.class); @@ -55,6 +58,13 @@ public SystemEntityClient systemEntityClient( @MockBean public ElasticSearchService searchService; + @MockBean public EntityService entityService; + + @MockBean public FormService formService; + + @MockBean(name = "cachingAspectRetriever") + CachingAspectRetriever cachingAspectRetriever; + @MockBean(name = "systemAuthentication") public Authentication systemAuthentication; diff --git a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java index ba72a979088462..9ebcfb0ba0c6bb 100644 --- a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java @@ -1,12 +1,13 @@ package com.linkedin.metadata.aspect.plugins.hooks; -import com.linkedin.metadata.aspect.batch.MCLBatchItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.MCLItem; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.entity.ebean.batch.MCLBatchItemImpl; +import com.linkedin.metadata.entity.ebean.batch.MCLItemImpl; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeLog; import com.mycompany.dq.DataQualityRuleEvent; +import java.util.Collection; import java.util.Optional; import java.util.stream.Stream; import javax.annotation.Nonnull; @@ -18,30 +19,32 @@ public CustomDataQualityRulesMCLSideEffect(AspectPluginConfig config) { } @Override - protected Stream applyMCLSideEffect( - @Nonnull MCLBatchItem input, @Nonnull AspectRetriever aspectRetriever) { + protected Stream applyMCLSideEffect( + @Nonnull Collection mclItems, @Nonnull AspectRetriever aspectRetriever) { + return mclItems.stream() + .map( + item -> { + // Generate Timeseries event aspect based on non-Timeseries aspect + MetadataChangeLog originMCP = item.getMetadataChangeLog(); - // Generate Timeseries event aspect based on non-Timeseries aspect - MetadataChangeLog originMCP = input.getMetadataChangeLog(); - - Optional timeseriesOptional = - buildEvent(originMCP) - .map( - event -> { - try { - MetadataChangeLog eventMCP = originMCP.clone(); - eventMCP.setAspect(GenericRecordUtils.serializeAspect(event)); - eventMCP.setAspectName("customDataQualityRuleEvent"); - return eventMCP; - } catch (CloneNotSupportedException e) { - throw new RuntimeException(e); - } - }) - .map( - eventMCP -> - MCLBatchItemImpl.builder().metadataChangeLog(eventMCP).build(aspectRetriever)); - - return timeseriesOptional.stream(); + return buildEvent(originMCP) + .map( + event -> { + try { + MetadataChangeLog eventMCP = originMCP.clone(); + eventMCP.setAspect(GenericRecordUtils.serializeAspect(event)); + eventMCP.setAspectName("customDataQualityRuleEvent"); + return eventMCP; + } catch (CloneNotSupportedException e) { + throw new RuntimeException(e); + } + }) + .map( + eventMCP -> + MCLItemImpl.builder().metadataChangeLog(eventMCP).build(aspectRetriever)); + }) + .filter(Optional::isPresent) + .map(Optional::get); } private Optional buildEvent(MetadataChangeLog originMCP) { diff --git a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java index c21b64c8a4fc00..103584f7a01401 100644 --- a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java @@ -2,10 +2,11 @@ import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; +import java.util.Collection; import java.util.stream.Stream; import javax.annotation.Nonnull; @@ -16,17 +17,21 @@ public CustomDataQualityRulesMCPSideEffect(AspectPluginConfig aspectPluginConfig } @Override - protected Stream applyMCPSideEffect( - UpsertItem input, @Nonnull AspectRetriever aspectRetriever) { + protected Stream applyMCPSideEffect( + Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { // Mirror aspects to another URN in SQL & Search - Urn mirror = UrnUtils.getUrn(input.getUrn().toString().replace(",PROD)", ",DEV)")); - return Stream.of( - MCPUpsertBatchItem.builder() - .urn(mirror) - .aspectName(input.getAspectName()) - .recordTemplate(input.getRecordTemplate()) - .auditStamp(input.getAuditStamp()) - .systemMetadata(input.getSystemMetadata()) - .build(aspectRetriever)); + return changeMCPS.stream() + .map( + changeMCP -> { + Urn mirror = + UrnUtils.getUrn(changeMCP.getUrn().toString().replace(",PROD)", ",DEV)")); + return ChangeItemImpl.builder() + .urn(mirror) + .aspectName(changeMCP.getAspectName()) + .recordTemplate(changeMCP.getRecordTemplate()) + .auditStamp(changeMCP.getAuditStamp()) + .systemMetadata(changeMCP.getSystemMetadata()) + .build(aspectRetriever); + }); } } diff --git a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java index 576ba3bf305f53..35b99d6c02abd7 100644 --- a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java @@ -1,17 +1,15 @@ package com.linkedin.metadata.aspect.plugins.hooks; -import com.linkedin.common.AuditStamp; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; -import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; import com.mycompany.dq.DataQualityRule; import com.mycompany.dq.DataQualityRules; +import java.util.Collection; +import java.util.Objects; +import java.util.stream.Stream; import javax.annotation.Nonnull; -import javax.annotation.Nullable; public class CustomDataQualityRulesMutator extends MutationHook { @@ -20,26 +18,29 @@ public CustomDataQualityRulesMutator(AspectPluginConfig config) { } @Override - protected void mutate( - @Nonnull ChangeType changeType, - @Nonnull EntitySpec entitySpec, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate oldAspectValue, - @Nullable RecordTemplate newAspectValue, - @Nullable SystemMetadata oldSystemMetadata, - @Nullable SystemMetadata newSystemMetadata, - @Nonnull AuditStamp auditStamp, - @Nonnull AspectRetriever aspectRetriever) { + protected Stream> writeMutation( + @Nonnull Collection changeMCPS, @Nonnull AspectRetriever aspectRetriever) { + return changeMCPS.stream() + .map( + changeMCP -> { + boolean mutated = false; - if (newAspectValue != null) { - DataQualityRules newDataQualityRules = new DataQualityRules(newAspectValue.data()); + if (changeMCP.getRecordTemplate() != null) { + DataQualityRules newDataQualityRules = + new DataQualityRules(changeMCP.getRecordTemplate().data()); - for (DataQualityRule rule : newDataQualityRules.getRules()) { - // Ensure uniform lowercase - if (!rule.getType().toLowerCase().equals(rule.getType())) { - rule.setType(rule.getType().toLowerCase()); - } - } - } + for (DataQualityRule rule : newDataQualityRules.getRules()) { + // Ensure uniform lowercase + if (!rule.getType().toLowerCase().equals(rule.getType())) { + mutated = true; + rule.setType(rule.getType().toLowerCase()); + } + } + } + + return mutated ? changeMCP : null; + }) + .filter(Objects::nonNull) + .map(changeMCP -> Pair.of(changeMCP, true)); } } diff --git a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java index 667d7ad614a791..ca291c46971230 100644 --- a/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java @@ -1,16 +1,16 @@ package com.linkedin.metadata.aspect.plugins.validation; -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; -import com.linkedin.metadata.models.AspectSpec; -import com.mycompany.dq.DataQualityRule; import com.mycompany.dq.DataQualityRules; +import java.util.Collection; import java.util.Map; +import java.util.Objects; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; -import javax.annotation.Nullable; public class CustomDataQualityRulesValidator extends AspectPayloadValidator { @@ -19,52 +19,59 @@ public CustomDataQualityRulesValidator(AspectPluginConfig config) { } @Override - protected void validateProposedAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nonnull RecordTemplate aspectPayload, - @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { - DataQualityRules rules = new DataQualityRules(aspectPayload.data()); - - // Enforce at least 1 rule - if (rules.getRules().isEmpty()) { - throw new AspectValidationException("At least one rule is required."); - } + protected Stream validateProposedAspects( + @Nonnull Collection mcpItems, @Nonnull AspectRetriever aspectRetriever) { + return mcpItems.stream() + .map( + item -> { + DataQualityRules rules = new DataQualityRules(item.getRecordTemplate().data()); + // Enforce at least 1 rule + return rules.getRules().isEmpty() + ? new AspectValidationException( + item.getUrn(), item.getAspectName(), "At least one rule is required.") + : null; + }) + .filter(Objects::nonNull); } @Override - protected void validatePreCommitAspect( - @Nonnull ChangeType changeType, - @Nonnull Urn entityUrn, - @Nonnull AspectSpec aspectSpec, - @Nullable RecordTemplate previousAspect, - @Nonnull RecordTemplate proposedAspect, - @Nonnull AspectRetriever aspectRetriever) - throws AspectValidationException { + protected Stream validatePreCommitAspects( + @Nonnull Collection changeMCPs, AspectRetriever aspectRetriever) { + return changeMCPs.stream() + .flatMap( + changeMCP -> { + if (changeMCP.getPreviousSystemAspect() != null) { + DataQualityRules oldRules = changeMCP.getPreviousAspect(DataQualityRules.class); + DataQualityRules newRules = changeMCP.getAspect(DataQualityRules.class); - if (previousAspect != null) { - DataQualityRules oldRules = new DataQualityRules(previousAspect.data()); - DataQualityRules newRules = new DataQualityRules(proposedAspect.data()); + Map newFieldTypeMap = + newRules.getRules().stream() + .filter(rule -> rule.getField() != null) + .map(rule -> Map.entry(rule.getField(), rule.getType())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - Map newFieldTypeMap = - newRules.getRules().stream() - .filter(rule -> rule.getField() != null) - .map(rule -> Map.entry(rule.getField(), rule.getType())) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + // Ensure the old and new field type is the same + return oldRules.getRules().stream() + .map( + oldRule -> { + if (!newFieldTypeMap + .getOrDefault(oldRule.getField(), oldRule.getType()) + .equals(oldRule.getType())) { + return new AspectValidationException( + changeMCP.getUrn(), + changeMCP.getAspectName(), + String.format( + "Field type mismatch. Field: %s Old: %s New: %s", + oldRule.getField(), + oldRule.getType(), + newFieldTypeMap.get(oldRule.getField()))); + } + return null; + }) + .filter(Objects::nonNull); + } - // Ensure the old and new field type is the same - for (DataQualityRule oldRule : oldRules.getRules()) { - if (!newFieldTypeMap - .getOrDefault(oldRule.getField(), oldRule.getType()) - .equals(oldRule.getType())) { - throw new AspectValidationException( - String.format( - "Field type mismatch. Field: %s Old: %s New: %s", - oldRule.getField(), oldRule.getType(), newFieldTypeMap.get(oldRule.getField()))); - } - } - } + return Stream.empty(); + }); } } diff --git a/metadata-models/src/main/resources/entity-registry.yml b/metadata-models/src/main/resources/entity-registry.yml index 62eaa2af2d4a0a..bc7b84d04c2c00 100644 --- a/metadata-models/src/main/resources/entity-registry.yml +++ b/metadata-models/src/main/resources/entity-registry.yml @@ -525,13 +525,22 @@ plugins: enabled: true supportedOperations: - UPSERT + - DELETE supportedEntityAspectNames: - entityName: structuredProperty aspectName: propertyDefinition + - entityName: structuredProperty + aspectName: structuredPropertyKey - className: 'com.linkedin.metadata.aspect.validation.StructuredPropertiesValidator' enabled: true supportedOperations: - UPSERT supportedEntityAspectNames: - entityName: '*' - aspectName: structuredProperties \ No newline at end of file + aspectName: structuredProperties + mutationHooks: + - className: 'com.linkedin.metadata.aspect.hooks.StructuredPropertiesSoftDelete' + enabled: true + supportedEntityAspectNames: + - entityName: '*' + aspectName: structuredProperties diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/user/NativeUserServiceTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/user/NativeUserServiceTest.java index 2b584c3461452e..60b10e3c53ef47 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/user/NativeUserServiceTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/user/NativeUserServiceTest.java @@ -109,7 +109,7 @@ public void testCreateNativeUserUserSystemUser() throws Exception { @Test public void testCreateNativeUserPasses() throws Exception { - when(_entityService.exists(any(), any())).thenReturn(false); + when(_entityService.exists(any(Urn.class), anyBoolean())).thenReturn(false); when(_secretService.generateSalt(anyInt())).thenReturn(SALT); when(_secretService.encrypt(any())).thenReturn(ENCRYPTED_SALT); when(_secretService.getHashedPassword(any(), any())).thenReturn(HASHED_PASSWORD); diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml index b6382b422ad9e7..2b85e7bb355484 100644 --- a/metadata-service/configuration/src/main/resources/application.yml +++ b/metadata-service/configuration/src/main/resources/application.yml @@ -400,7 +400,8 @@ cache: corpUserCredentials: 20 corpUserSettings: 20 structuredProperty: - propertyDefinition: 86400 # 1 day + status: 300 # 5 min + propertyDefinition: 300 # 5 min structuredPropertyKey: 86400 # 1 day graphQL: diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java index 2ccdee5fb1dbf5..0c7808abe538b1 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityServiceFactory.java @@ -6,9 +6,8 @@ import com.linkedin.metadata.entity.AspectDao; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.EntityServiceImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.service.UpdateIndicesService; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; @@ -25,12 +24,11 @@ public class EntityServiceFactory { @Bean(name = "entityService") @DependsOn({"entityAspectDao", "kafkaEventProducer", "entityRegistry"}) @Nonnull - protected EntityService createInstance( + protected EntityService createInstance( @Qualifier("kafkaEventProducer") final KafkaEventProducer eventProducer, @Qualifier("entityAspectDao") AspectDao aspectDao, EntityRegistry entityRegistry, ConfigurationProvider configurationProvider, - UpdateIndicesService updateIndicesService, @Value("${featureFlags.showBrowseV2}") final boolean enableBrowsePathV2) { FeatureFlags featureFlags = configurationProvider.getFeatureFlags(); @@ -40,7 +38,6 @@ protected EntityService createInstance( eventProducer, entityRegistry, featureFlags.isAlwaysEmitChangeLog(), - updateIndicesService, featureFlags.getPreProcessHooks(), _ebeanMaxTransactionRetry, enableBrowsePathV2); diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java index 31ad933b9579d1..db9d9c8e657f84 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java @@ -5,7 +5,7 @@ import com.linkedin.metadata.entity.RetentionService; import com.linkedin.metadata.entity.cassandra.CassandraRetentionService; import com.linkedin.metadata.entity.ebean.EbeanRetentionService; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.spring.YamlPropertySourceFactory; import io.ebean.Database; import javax.annotation.Nonnull; @@ -24,7 +24,7 @@ public class RetentionServiceFactory { @Autowired @Qualifier("entityService") - private EntityService _entityService; + private EntityService _entityService; @Value("${RETENTION_APPLICATION_BATCH_SIZE:1000}") private Integer _batchSize; @@ -33,8 +33,8 @@ public class RetentionServiceFactory { @DependsOn({"cassandraSession", "entityService"}) @ConditionalOnProperty(name = "entityService.impl", havingValue = "cassandra") @Nonnull - protected RetentionService createCassandraInstance(CqlSession session) { - RetentionService retentionService = + protected RetentionService createCassandraInstance(CqlSession session) { + RetentionService retentionService = new CassandraRetentionService<>(_entityService, session, _batchSize); _entityService.setRetentionService(retentionService); return retentionService; @@ -44,8 +44,8 @@ protected RetentionService createCassandraInstance(CqlSessio @DependsOn({"ebeanServer", "entityService"}) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @Nonnull - protected RetentionService createEbeanInstance(Database server) { - RetentionService retentionService = + protected RetentionService createEbeanInstance(Database server) { + RetentionService retentionService = new EbeanRetentionService<>(_entityService, server, _batchSize); _entityService.setRetentionService(retentionService); return retentionService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/update/indices/UpdateIndicesServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/update/indices/UpdateIndicesServiceFactory.java index 34c1887d67c56f..d6b033e9268fc5 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/update/indices/UpdateIndicesServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/update/indices/UpdateIndicesServiceFactory.java @@ -1,19 +1,17 @@ package com.linkedin.gms.factory.entity.update.indices; -import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.gms.factory.search.EntityIndexBuildersFactory; -import com.linkedin.metadata.client.EntityClientAspectRetriever; +import com.linkedin.metadata.aspect.CachingAspectRetriever; +import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphService; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.elasticsearch.indexbuilder.EntityIndexBuilders; import com.linkedin.metadata.search.transformer.SearchDocumentTransformer; import com.linkedin.metadata.service.UpdateIndicesService; import com.linkedin.metadata.systemmetadata.SystemMetadataService; import com.linkedin.metadata.timeseries.TimeseriesAspectService; -import org.springframework.beans.factory.annotation.Autowired; -import org.springframework.beans.factory.annotation.Value; -import org.springframework.context.ApplicationContext; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.Import; @@ -21,20 +19,45 @@ @Configuration @Import(EntityIndexBuildersFactory.class) public class UpdateIndicesServiceFactory { - @Autowired private ApplicationContext context; - - @Value("${entityClient.impl:java}") - private String entityClientImpl; + /* + When restli mode the EntityService is not available. Wire in an AspectRetriever here instead + based on the entity client + */ @Bean - public UpdateIndicesService updateIndicesService( + @ConditionalOnProperty(name = "entityClient.impl", havingValue = "restli") + public UpdateIndicesService searchIndicesServiceNonGMS( GraphService graphService, EntitySearchService entitySearchService, TimeseriesAspectService timeseriesAspectService, SystemMetadataService systemMetadataService, - EntityRegistry entityRegistry, SearchDocumentTransformer searchDocumentTransformer, - EntityIndexBuilders entityIndexBuilders) { + EntityIndexBuilders entityIndexBuilders, + @Qualifier("cachingAspectRetriever") final CachingAspectRetriever aspectRetriever) { + + UpdateIndicesService updateIndicesService = + new UpdateIndicesService( + graphService, + entitySearchService, + timeseriesAspectService, + systemMetadataService, + searchDocumentTransformer, + entityIndexBuilders); + updateIndicesService.initializeAspectRetriever(aspectRetriever); + + return updateIndicesService; + } + + @Bean + @ConditionalOnProperty(name = "entityClient.impl", havingValue = "java", matchIfMissing = true) + public UpdateIndicesService searchIndicesServiceGMS( + final GraphService graphService, + final EntitySearchService entitySearchService, + final TimeseriesAspectService timeseriesAspectService, + final SystemMetadataService systemMetadataService, + final SearchDocumentTransformer searchDocumentTransformer, + final EntityIndexBuilders entityIndexBuilders, + final EntityService entityService) { UpdateIndicesService updateIndicesService = new UpdateIndicesService( @@ -45,18 +68,8 @@ public UpdateIndicesService updateIndicesService( searchDocumentTransformer, entityIndexBuilders); - if ("restli".equals(entityClientImpl)) { - /* - When restli mode the EntityService is not available. Wire in an AspectRetriever here instead - based on the entity client - */ - SystemEntityClient systemEntityClient = context.getBean(SystemEntityClient.class); - updateIndicesService.initializeAspectRetriever( - EntityClientAspectRetriever.builder() - .entityRegistry(entityRegistry) - .entityClient(systemEntityClient) - .build()); - } + updateIndicesService.initializeAspectRetriever(entityService); + entityService.setUpdateIndicesService(updateIndicesService); return updateIndicesService; } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/AspectRetrieverFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/AspectRetrieverFactory.java new file mode 100644 index 00000000000000..5892e1ec4e4204 --- /dev/null +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityclient/AspectRetrieverFactory.java @@ -0,0 +1,27 @@ +package com.linkedin.gms.factory.entityclient; + +import com.linkedin.entity.client.SystemEntityClient; +import com.linkedin.metadata.aspect.CachingAspectRetriever; +import com.linkedin.metadata.client.EntityClientAspectRetriever; +import com.linkedin.metadata.models.registry.EntityRegistry; +import javax.annotation.Nonnull; +import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +@Slf4j +@Configuration +public class AspectRetrieverFactory { + + @Bean(name = "cachingAspectRetriever") + @Nonnull + protected CachingAspectRetriever cachingAspectRetriever( + final EntityRegistry entityRegistry, + @Qualifier("systemEntityClient") final SystemEntityClient systemEntityClient) { + return EntityClientAspectRetriever.builder() + .entityRegistry(entityRegistry) + .entityClient(systemEntityClient) + .build(); + } +} diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java index bc82df9f8cdadc..15bf674581b6a8 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java @@ -172,7 +172,7 @@ public class GraphQLEngineFactory { @Bean(name = "graphQLEngine") @Nonnull - protected GraphQLEngine getInstance( + protected GraphQLEngine graphQLEngine( @Qualifier("entityClient") final EntityClient entityClient, @Qualifier("systemEntityClient") final SystemEntityClient systemEntityClient) { GmsGraphQLEngineArgs args = new GmsGraphQLEngineArgs(); diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java index 7b5f4e18d4d539..7649e3a1ada425 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java @@ -63,7 +63,7 @@ public class ElasticSearchServiceFactory { @Bean(name = "elasticSearchService") @Nonnull - protected ElasticSearchService getInstance(ConfigurationProvider configurationProvider) + protected ElasticSearchService getInstance(final ConfigurationProvider configurationProvider) throws IOException { log.info("Search configuration: {}", configurationProvider.getElasticSearch().getSearch()); @@ -77,7 +77,6 @@ protected ElasticSearchService getInstance(ConfigurationProvider configurationPr ESSearchDAO esSearchDAO = new ESSearchDAO( - entityRegistry, components.getSearchClient(), components.getIndexConvention(), configurationProvider.getFeatureFlags().isPointInTimeCreationEnabled(), @@ -88,7 +87,6 @@ protected ElasticSearchService getInstance(ConfigurationProvider configurationPr entityIndexBuilders, esSearchDAO, new ESBrowseDAO( - entityRegistry, components.getSearchClient(), components.getIndexConvention(), searchConfiguration, diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStep.java index 19efa5e9c4de20..9434928e1bfa03 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStep.java @@ -11,7 +11,7 @@ import com.linkedin.metadata.entity.AspectMigrationsDao; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.utils.DataPlatformInstanceUtils; import com.linkedin.metadata.utils.EntityKeyUtils; @@ -65,7 +65,7 @@ public void execute() throws Exception { start, start + BATCH_SIZE); - List items = new LinkedList<>(); + List items = new LinkedList<>(); final AuditStamp aspectAuditStamp = new AuditStamp() .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) @@ -76,7 +76,7 @@ public void execute() throws Exception { Optional dataPlatformInstance = getDataPlatformInstance(urn); if (dataPlatformInstance.isPresent()) { items.add( - MCPUpsertBatchItem.builder() + ChangeItemImpl.builder() .urn(urn) .aspectName(DATA_PLATFORM_INSTANCE_ASPECT_NAME) .recordTemplate(dataPlatformInstance.get()) @@ -85,7 +85,10 @@ public void execute() throws Exception { } } - _entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); + _entityService.ingestAspects( + AspectsBatchImpl.builder().aspectRetriever(_entityService).items(items).build(), + true, + true); log.info( "Finished ingesting DataPlatformInstance for urn {} to {}", start, start + BATCH_SIZE); diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformsStep.java index d2bb61ad7ade5d..11e86241e216a1 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformsStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformsStep.java @@ -13,7 +13,7 @@ import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import java.io.IOException; import java.net.URISyntaxException; import java.util.List; @@ -62,7 +62,7 @@ public void execute() throws IOException, URISyntaxException { } // 2. For each JSON object, cast into a DataPlatformSnapshot object. - List dataPlatformAspects = + List dataPlatformAspects = StreamSupport.stream( Spliterators.spliteratorUnknownSize(dataPlatforms.iterator(), Spliterator.ORDERED), false) @@ -83,7 +83,7 @@ public void execute() throws IOException, URISyntaxException { DataPlatformInfo.class, dataPlatform.get("aspect").toString()); try { - return MCPUpsertBatchItem.builder() + return ChangeItemImpl.builder() .urn(urn) .aspectName(PLATFORM_ASPECT_NAME) .recordTemplate(info) @@ -99,6 +99,11 @@ public void execute() throws IOException, URISyntaxException { .collect(Collectors.toList()); _entityService.ingestAspects( - AspectsBatchImpl.builder().items(dataPlatformAspects).build(), true, false); + AspectsBatchImpl.builder() + .aspectRetriever(_entityService) + .items(dataPlatformAspects) + .build(), + true, + false); } } diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java index 5617d7e9714b08..b9cbf2abe06730 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java @@ -8,7 +8,7 @@ import com.linkedin.common.urn.UrnUtils; import com.linkedin.metadata.entity.AspectMigrationsDao; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +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.ConfigEntityRegistry; @@ -122,7 +122,7 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw item.getUrn().getEntityType().equals("chart") && item.getAspectName() .equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - && ((MCPUpsertBatchItem) item).getRecordTemplate() + && ((ChangeItemImpl) item).getRecordTemplate() instanceof DataPlatformInstance)), anyBoolean(), anyBoolean()); @@ -136,7 +136,7 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw item.getUrn().getEntityType().equals("chart") && item.getAspectName() .equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - && ((MCPUpsertBatchItem) item).getRecordTemplate() + && ((ChangeItemImpl) item).getRecordTemplate() instanceof DataPlatformInstance)), anyBoolean(), anyBoolean()); diff --git a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java index 39a7e4722988e1..1e375f90fc38a9 100644 --- a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java +++ b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/v2/delegates/EntityApiDelegateImpl.java @@ -179,7 +179,7 @@ public ResponseEntity createAspect( public ResponseEntity headAspect(String urn, String aspect) { try { Urn entityUrn = Urn.createFromString(urn); - if (_entityService.exists(entityUrn, aspect)) { + if (_entityService.exists(entityUrn, aspect, true)) { return new ResponseEntity<>(HttpStatus.NO_CONTENT); } else { return new ResponseEntity<>(HttpStatus.NOT_FOUND); diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java index ff65db09c2682f..63e78c30383af3 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/entities/EntitiesController.java @@ -17,7 +17,7 @@ import com.linkedin.common.urn.UrnUtils; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.utils.metrics.MetricUtils; import com.linkedin.util.Pair; import io.datahubproject.openapi.dto.RollbackRunResultDto; @@ -65,7 +65,7 @@ description = "APIs for ingesting and accessing entities and their constituent aspects") public class EntitiesController { - private final EntityService _entityService; + private final EntityService _entityService; private final ObjectMapper _objectMapper; private final AuthorizerChain _authorizerChain; diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/platform/entities/PlatformEntitiesController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/platform/entities/PlatformEntitiesController.java index 3cc67e77ec27e1..7193da3bf85878 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/platform/entities/PlatformEntitiesController.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/platform/entities/PlatformEntitiesController.java @@ -9,7 +9,7 @@ import com.google.common.collect.ImmutableList; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.search.client.CachingEntitySearchService; import com.linkedin.util.Pair; import io.datahubproject.openapi.exception.UnauthorizedException; @@ -45,7 +45,7 @@ description = "Platform level APIs intended for lower level access to entities") public class PlatformEntitiesController { - private final EntityService _entityService; + private final EntityService _entityService; private final CachingEntitySearchService _cachingEntitySearchService; private final ObjectMapper _objectMapper; private final AuthorizerChain _authorizerChain; diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java index 13d2e501abf09f..6b31159a206652 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java +++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/util/MappingUtil.java @@ -30,7 +30,7 @@ import com.linkedin.metadata.entity.IngestResult; import com.linkedin.metadata.entity.RollbackRunResult; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; @@ -441,7 +441,7 @@ public static boolean authorizeProposals( public static Pair ingestProposal( com.linkedin.mxe.MetadataChangeProposal serviceProposal, String actorUrn, - EntityService entityService, + EntityService entityService, boolean async) { // TODO: Use the actor present in the IC. 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 44202c20ca6db7..656d6542483cf3 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 @@ -14,14 +14,14 @@ import com.linkedin.data.ByteString; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.EnvelopedAspect; -import com.linkedin.metadata.aspect.batch.UpsertItem; +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.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.UpdateAspectResult; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +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; @@ -306,12 +306,17 @@ public ResponseEntity createAspect( } AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); - UpsertItem upsert = + ChangeMCP upsert = toUpsertItem(UrnUtils.getUrn(entityUrn), aspectSpec, jsonAspect, authentication.getActor()); List results = entityService.ingestAspects( - AspectsBatchImpl.builder().items(List.of(upsert)).build(), true, true); + AspectsBatchImpl.builder() + .aspectRetriever(entityService) + .items(List.of(upsert)) + .build(), + true, + true); return ResponseEntity.of( results.stream() @@ -371,7 +376,7 @@ public ResponseEntity patchAspect( .templateDefault( aspectSpec.getDataTemplateClass().getDeclaredConstructor().newInstance()) .build(); - UpsertItem upsert = + ChangeMCP upsert = toUpsertItem( UrnUtils.getUrn(entityUrn), aspectSpec, @@ -381,7 +386,12 @@ public ResponseEntity patchAspect( List results = entityService.ingestAspects( - AspectsBatchImpl.builder().items(List.of(upsert)).build(), true, true); + AspectsBatchImpl.builder() + .aspectRetriever(entityService) + .items(List.of(upsert)) + .build(), + true, + true); return ResponseEntity.of( results.stream() @@ -409,7 +419,9 @@ private List toRecordTemplates( } private Boolean exists(Urn urn, @Nullable String aspect) { - return aspect == null ? entityService.exists(urn, true) : entityService.exists(urn, aspect); + return aspect == null + ? entityService.exists(urn, true) + : entityService.exists(urn, aspect, true); } private List toRecordTemplates( @@ -474,10 +486,10 @@ private RecordTemplate toRecordTemplate(AspectSpec aspectSpec, EnvelopedAspect e aspectSpec.getDataTemplateClass(), envelopedAspect.getValue().data()); } - private UpsertItem toUpsertItem( + private ChangeMCP toUpsertItem( Urn entityUrn, AspectSpec aspectSpec, String jsonAspect, Actor actor) throws URISyntaxException { - return MCPUpsertBatchItem.builder() + return ChangeItemImpl.builder() .urn(entityUrn) .aspectName(aspectSpec.getName()) .auditStamp(AuditStampUtils.createAuditStamp(actor.toUrnStr())) @@ -489,14 +501,14 @@ private UpsertItem toUpsertItem( .build(entityService); } - private UpsertItem toUpsertItem( + private ChangeMCP toUpsertItem( @Nonnull Urn urn, @Nonnull AspectSpec aspectSpec, @Nullable RecordTemplate currentValue, @Nonnull GenericPatchTemplate genericPatchTemplate, @Nonnull Actor actor) throws URISyntaxException { - return MCPUpsertBatchItem.fromPatch( + return ChangeItemImpl.fromPatch( urn, aspectSpec, currentValue, diff --git a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java index 20862bbc7f000d..4cecfe21281995 100644 --- a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java +++ b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java @@ -17,7 +17,6 @@ import com.linkedin.metadata.entity.UpdateAspectResult; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.service.UpdateIndicesService; import io.datahubproject.openapi.dto.UpsertAspectRequest; import io.datahubproject.openapi.entities.EntitiesController; import io.datahubproject.openapi.generated.AuditStamp; @@ -79,16 +78,11 @@ public void setup() .apply(Mockito.mock(Transaction.class)))); EventProducer mockEntityEventProducer = Mockito.mock(EventProducer.class); - UpdateIndicesService mockUpdateIndicesService = mock(UpdateIndicesService.class); PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); MockEntityService mockEntityService = new MockEntityService( - aspectDao, - mockEntityEventProducer, - mockEntityRegistry, - mockUpdateIndicesService, - preProcessHooks); + aspectDao, mockEntityEventProducer, mockEntityRegistry, preProcessHooks); AuthorizerChain authorizerChain = Mockito.mock(AuthorizerChain.class); _entitiesController = new EntitiesController(mockEntityService, new ObjectMapper(), authorizerChain); diff --git a/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java index be5f99bed8e630..8ed7c397c5ba4b 100644 --- a/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java +++ b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java @@ -1,6 +1,7 @@ package mock; import static entities.EntitiesControllerTest.*; +import static org.mockito.Mockito.mock; import com.linkedin.common.AuditStamp; import com.linkedin.common.GlobalTags; @@ -57,9 +58,9 @@ public MockEntityService( @Nonnull AspectDao aspectDao, @Nonnull EventProducer producer, @Nonnull EntityRegistry entityRegistry, - @Nonnull UpdateIndicesService updateIndicesService, PreProcessHooks preProcessHooks) { - super(aspectDao, producer, entityRegistry, true, updateIndicesService, preProcessHooks, true); + super(aspectDao, producer, entityRegistry, true, preProcessHooks, true); + setUpdateIndicesService(mock(UpdateIndicesService.class)); } @Override @@ -69,7 +70,7 @@ public Map> getLatestAspects( } @Override - public Map getLatestAspectsForUrn( + public @NotNull Map getLatestAspectsForUrn( @Nonnull Urn urn, @Nonnull Set aspectNames) { return Collections.emptyMap(); } diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java index 676b80c8bea32f..65169344776b7b 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java @@ -12,6 +12,7 @@ import com.linkedin.entity.Aspect; import com.linkedin.entity.Entity; import com.linkedin.entity.EntityResponse; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.browse.BrowseResult; @@ -44,6 +45,9 @@ // Consider renaming this to datahub client. public interface EntityClient { + /** Perform post construction asks if needed. Can be used to break circular dependencies */ + default void postConstruct(AspectRetriever aspectRetriever) {} + @Nullable public EntityResponse getV2( @Nonnull String entityName, diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java index 21a9f47a13f738..8658096b174378 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java @@ -18,14 +18,10 @@ import com.linkedin.metadata.aspect.EnvelopedAspectArray; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.authorization.PoliciesConfig; -import com.linkedin.metadata.entity.AspectUtils; -import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.IngestResult; import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.aspect.batch.AspectsBatch; -import com.linkedin.metadata.entity.ebean.batch.MCLBatchItemImpl; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; @@ -52,8 +48,6 @@ import java.time.Clock; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; -import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import javax.inject.Inject; @@ -83,10 +77,10 @@ public class AspectResource extends CollectionResourceTaskTemplate _entityService; + private EntityService _entityService; @VisibleForTesting - void setEntityService(EntityService entityService) { + void setEntityService(EntityService entityService) { _entityService = entityService; } diff --git a/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java b/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java index 17c51604947223..62edb9fdfa6281 100644 --- a/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java +++ b/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java @@ -20,7 +20,7 @@ import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.EntityServiceImpl; import com.linkedin.metadata.entity.UpdateAspectResult; -import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -35,28 +35,29 @@ import org.testng.annotations.Test; public class AspectResourceTest { - private AspectResource _aspectResource; - private EntityService _entityService; - private AspectDao _aspectDao; - private EventProducer _producer; - private EntityRegistry _entityRegistry; - private UpdateIndicesService _updateIndicesService; - private PreProcessHooks _preProcessHooks; - private Authorizer _authorizer; + private AspectResource aspectResource; + private EntityService entityService; + private AspectDao aspectDao; + private EventProducer producer; + private EntityRegistry entityRegistry; + private UpdateIndicesService updateIndicesService; + private PreProcessHooks preProcessHooks; + private Authorizer authorizer; @BeforeTest public void setup() { - _aspectResource = new AspectResource(); - _aspectDao = mock(AspectDao.class); - _producer = mock(EventProducer.class); - _entityRegistry = new MockEntityRegistry(); - _updateIndicesService = mock(UpdateIndicesService.class); - _preProcessHooks = mock(PreProcessHooks.class); - _entityService = new EntityServiceImpl(_aspectDao, _producer, _entityRegistry, false, - _updateIndicesService, _preProcessHooks, true); - _authorizer = mock(Authorizer.class); - _aspectResource.setAuthorizer(_authorizer); - _aspectResource.setEntityService(_entityService); + aspectResource = new AspectResource(); + aspectDao = mock(AspectDao.class); + producer = mock(EventProducer.class); + entityRegistry = new MockEntityRegistry(); + updateIndicesService = mock(UpdateIndicesService.class); + preProcessHooks = mock(PreProcessHooks.class); + entityService = new EntityServiceImpl(aspectDao, producer, entityRegistry, false, + preProcessHooks, true); + entityService.setUpdateIndicesService(updateIndicesService); + authorizer = mock(Authorizer.class); + aspectResource.setAuthorizer(authorizer); + aspectResource.setEntityService(entityService); } @Test @@ -74,21 +75,21 @@ public void testAsyncDefaultAspects() throws URISyntaxException { AuthenticationContext.setAuthentication(mockAuthentication); Actor actor = new Actor(ActorType.USER, "user"); when(mockAuthentication.getActor()).thenReturn(actor); - _aspectResource.ingestProposal(mcp, "true"); - verify(_producer, times(1)).produceMetadataChangeProposal(urn, mcp); - verifyNoMoreInteractions(_producer); - verifyNoMoreInteractions(_aspectDao); + aspectResource.ingestProposal(mcp, "true"); + verify(producer, times(1)).produceMetadataChangeProposal(urn, mcp); + verifyNoMoreInteractions(producer); + verifyNoMoreInteractions(aspectDao); - reset(_producer, _aspectDao); + reset(producer, aspectDao); - MCPUpsertBatchItem req = MCPUpsertBatchItem.builder() + ChangeItemImpl req = ChangeItemImpl.builder() .urn(urn) .aspectName(mcp.getAspectName()) .recordTemplate(mcp.getAspect()) .auditStamp(new AuditStamp()) .metadataChangeProposal(mcp) - .build(_entityService); - when(_aspectDao.runInTransactionWithRetry(any(), any(), anyInt())) + .build(entityService); + when(aspectDao.runInTransactionWithRetry(any(), any(), anyInt())) .thenReturn( List.of(List.of( UpdateAspectResult.builder() @@ -121,9 +122,9 @@ public void testAsyncDefaultAspects() throws URISyntaxException { .auditStamp(new AuditStamp()) .request(req) .build()))); - _aspectResource.ingestProposal(mcp, "false"); - verify(_producer, times(10)) + aspectResource.ingestProposal(mcp, "false"); + verify(producer, times(10)) .produceMetadataChangeLog(eq(urn), any(AspectSpec.class), any(MetadataChangeLog.class)); - verifyNoMoreInteractions(_producer); + verifyNoMoreInteractions(producer); } } diff --git a/metadata-service/restli-servlet-impl/src/test/java/mock/MockTimeseriesAspectService.java b/metadata-service/restli-servlet-impl/src/test/java/mock/MockTimeseriesAspectService.java index 5187cba0b91510..9d33551fa2af0b 100644 --- a/metadata-service/restli-servlet-impl/src/test/java/mock/MockTimeseriesAspectService.java +++ b/metadata-service/restli-servlet-impl/src/test/java/mock/MockTimeseriesAspectService.java @@ -2,6 +2,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; @@ -39,6 +40,11 @@ public MockTimeseriesAspectService(long count, long filteredCount, String taskId this._taskId = taskId; } + @Override + public TimeseriesAspectService postConstruct(AspectRetriever aspectRetriever) { + return this; + } + @Override public void configure() {} diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java index d9b0f4b73d5805..1b4b65baeecd62 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -8,10 +8,10 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.aspect.batch.AspectsBatch; -import com.linkedin.metadata.aspect.batch.UpsertItem; -import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult; import com.linkedin.metadata.models.AspectSpec; @@ -33,16 +33,60 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -public interface EntityService extends AspectRetriever { +public interface EntityService extends AspectRetriever { /** * Just whether the entity/aspect exists * + * @param urns urns for the entities + * @param aspectName aspect for the entity, if null, assumes key aspect + * @param includeSoftDelete including soft deleted entities + * @return set of urns with the specified aspect existing + */ + Set exists( + @Nonnull final Collection urns, @Nullable String aspectName, boolean includeSoftDelete); + + /** + * Just whether the entity/aspect exists, prefer batched method. + * * @param urn urn for the entity - * @param aspectName aspect for the entity - * @return exists or not + * @param aspectName aspect for the entity, if null use the key aspect + * @param includeSoftDelete including soft deleted entities + * @return boolean if the entity/aspect exists + */ + default boolean exists(@Nonnull Urn urn, @Nullable String aspectName, boolean includeSoftDelete) { + return exists(Set.of(urn), aspectName, includeSoftDelete).contains(urn); + } + + /** + * Returns a set of urns of entities that exist (has materialized aspects). + * + * @param urns the list of urns of the entities to check + * @return a set of urns of entities that exist. */ - Boolean exists(Urn urn, String aspectName); + default Set exists(@Nonnull final Collection urns, boolean includeSoftDelete) { + return exists(urns, null, includeSoftDelete); + } + + /** + * Returns a set of urns of entities that exist (has materialized aspects). + * + * @param urns the list of urns of the entities to check + * @return a set of urns of entities that exist. + */ + default Set exists(@Nonnull final Collection urns) { + return exists(urns, true); + } + + /** + * Returns whether the urn of the entity exists (has materialized aspects). + * + * @param urn the urn of the entity to check + * @return entities exists. + */ + default boolean exists(@Nonnull Urn urn, boolean includeSoftDelete) { + return exists(List.of(urn), includeSoftDelete).contains(urn); + } /** * Retrieves the latest aspects corresponding to a batch of {@link Urn}s based on a provided set @@ -285,29 +329,11 @@ RollbackRunResult rollbackWithConditions( IngestResult ingestProposal( MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async); - /** - * Returns a set of urns of entities that exist (has materialized aspects). - * - * @param urns the list of urns of the entities to check - * @return a set of urns of entities that exist. - */ - Set exists(@Nonnull final Collection urns, boolean includeSoftDelete); - - /** - * Returns a set of urns of entities that exist (has materialized aspects). - * - * @param urns the list of urns of the entities to check - * @return a set of urns of entities that exist. - */ - default Set exists(@Nonnull final Collection urns) { - return exists(urns, true); - } - - default boolean exists(@Nonnull Urn urn, boolean includeSoftDelete) { - return exists(List.of(urn), includeSoftDelete).contains(urn); - } - void setWritable(boolean canWrite); RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName); + + SearchIndicesService getUpdateIndicesService(); + + void setUpdateIndicesService(@Nullable SearchIndicesService updateIndicesService); } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/RetentionService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/RetentionService.java index ae33b72010ce2a..ef30e4c82046eb 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/RetentionService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/RetentionService.java @@ -8,7 +8,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.batch.AspectsBatch; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; import com.linkedin.metadata.key.DataHubRetentionKey; @@ -37,7 +37,7 @@ * storage and retention concerns apart, let AspectDaos deal with storage, and merge all retention * concerns into a single class. */ -public abstract class RetentionService { +public abstract class RetentionService { protected static final String ALL = "*"; protected abstract EntityService getEntityService(); diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/SearchIndicesService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/SearchIndicesService.java new file mode 100644 index 00000000000000..def5bb2730ba84 --- /dev/null +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/SearchIndicesService.java @@ -0,0 +1,11 @@ +package com.linkedin.metadata.entity; + +import com.linkedin.metadata.aspect.AspectRetriever; +import com.linkedin.mxe.MetadataChangeLog; +import javax.annotation.Nonnull; + +public interface SearchIndicesService { + void handleChangeEvent(@Nonnull MetadataChangeLog metadataChangeLog); + + void initializeAspectRetriever(@Nonnull AspectRetriever aspectRetriever); +} diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/UpdateAspectResult.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/UpdateAspectResult.java index 515e08646f9ed3..e85e0567f963ba 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/UpdateAspectResult.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/UpdateAspectResult.java @@ -3,7 +3,7 @@ import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.aspect.batch.UpsertItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; import com.linkedin.mxe.MetadataAuditOperation; import com.linkedin.mxe.SystemMetadata; import java.util.concurrent.Future; @@ -14,7 +14,7 @@ @Value public class UpdateAspectResult { Urn urn; - UpsertItem request; + ChangeMCP request; RecordTemplate oldValue; RecordTemplate newValue; SystemMetadata oldSystemMetadata; diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java index 0d1c031db136e4..b0b36ce868eb63 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.search; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.browse.BrowseResultV2; import com.linkedin.metadata.query.AutoCompleteResult; @@ -14,6 +15,13 @@ public interface EntitySearchService { + /** + * Set aspect retriever after construction to prevent circular dependencies + * + * @param aspectRetriever + */ + EntitySearchService postConstruct(AspectRetriever aspectRetriever); + void configure(); /** Clear all data within the service */ diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/timeseries/TimeseriesAspectService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/timeseries/TimeseriesAspectService.java index 529e8e00ecf570..77fa2720a68be1 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/timeseries/TimeseriesAspectService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/timeseries/TimeseriesAspectService.java @@ -2,6 +2,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.query.filter.Filter; import com.linkedin.metadata.query.filter.SortCriterion; @@ -16,6 +17,13 @@ public interface TimeseriesAspectService { + /** + * Set aspect retriever after construction to prevent circular dependencies + * + * @param aspectRetriever + */ + TimeseriesAspectService postConstruct(AspectRetriever aspectRetriever); + /** Configure the Time-Series aspect service one time at boot-up. */ void configure(); diff --git a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java index 27aa9ee04cc756..afaeb9c81039bf 100644 --- a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java +++ b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java @@ -6,6 +6,7 @@ import com.datahub.gms.util.CSVWriter; import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper; import com.linkedin.gms.factory.config.ConfigurationProvider; +import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.config.search.SearchConfiguration; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -38,13 +39,14 @@ private ConfigurationProvider getConfigProvider(WebApplicationContext ctx) { return (ConfigurationProvider) ctx.getBean("configurationProvider"); } - private EntityRegistry getEntityRegistry(WebApplicationContext ctx) { - return (EntityRegistry) ctx.getBean("entityRegistry"); + private AspectRetriever getAspectRetriever(WebApplicationContext ctx) { + return (AspectRetriever) ctx.getBean("aspectRetriever"); } private void writeSearchCsv(WebApplicationContext ctx, PrintWriter pw) { SearchConfiguration searchConfiguration = getConfigProvider(ctx).getElasticSearch().getSearch(); - EntityRegistry entityRegistry = getEntityRegistry(ctx); + AspectRetriever aspectRetriever = getAspectRetriever(ctx); + EntityRegistry entityRegistry = aspectRetriever.getEntityRegistry(); CSVWriter writer = CSVWriter.builder().printWriter(pw).build(); @@ -79,7 +81,8 @@ private void writeSearchCsv(WebApplicationContext ctx, PrintWriter pw) { entitySpecOpt -> { EntitySpec entitySpec = entitySpecOpt.get(); SearchRequest searchRequest = - SearchRequestHandler.getBuilder(entitySpec, searchConfiguration, null) + SearchRequestHandler.getBuilder( + entitySpec, searchConfiguration, null, aspectRetriever) .getSearchRequest( "*", null, diff --git a/smoke-test/build.gradle b/smoke-test/build.gradle index a6f3cd793ddd63..3cba93c452a101 100644 --- a/smoke-test/build.gradle +++ b/smoke-test/build.gradle @@ -72,3 +72,29 @@ task lintFix(type: Exec, dependsOn: installDev) { "ruff --fix tests/ && " + "mypy tests/" } + +/** + * The following tasks assume an already running quickstart. + * ./gradlew quickstart (or another variation) + */ +task quickstartNoCypressSuite0(type: Exec, dependsOn: [installDev, ':metadata-ingestion:installDev']) { + environment 'RUN_QUICKSTART', 'false' + environment 'DATAHUB_KAFKA_SCHEMA_REGISTRY_URL', 'http://localhost:8080/schema-registry/api/' + environment 'TEST_STRATEGY', 'no_cypress_suite0' + + workingDir = project.projectDir + commandLine 'bash', '-c', + "source ${venv_name}/bin/activate && set -x && " + + "./smoke.sh" +} + +task quickstartNoCypressSuite1(type: Exec, dependsOn: [installDev, ':metadata-ingestion:installDev']) { + environment 'RUN_QUICKSTART', 'false' + environment 'DATAHUB_KAFKA_SCHEMA_REGISTRY_URL', 'http://localhost:8080/schema-registry/api/' + environment 'TEST_STRATEGY', 'no_cypress_suite1' + + workingDir = project.projectDir + commandLine 'bash', '-c', + "source ${venv_name}/bin/activate && set -x && " + + "./smoke.sh" +} diff --git a/smoke-test/tests/structured_properties/test_structured_properties.py b/smoke-test/tests/structured_properties/test_structured_properties.py index de85d2af95e034..49472056a280a8 100644 --- a/smoke-test/tests/structured_properties/test_structured_properties.py +++ b/smoke-test/tests/structured_properties/test_structured_properties.py @@ -182,12 +182,9 @@ def get_property_from_entity( # ) @pytest.mark.dependency(depends=["test_healthchecks"]) def test_structured_property_string(ingest_cleanup_data, graph): - property_name = "retentionPolicy" + property_name = f"retention{randint(10, 10000)}Policy" create_property_definition(property_name, graph) - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.retentionPolicy" - ) attach_property_to_entity(dataset_urns[0], property_name, ["30d"], graph=graph) @@ -209,10 +206,8 @@ def test_structured_property_string(ingest_cleanup_data, graph): # ) @pytest.mark.dependency(depends=["test_healthchecks"]) def test_structured_property_double(ingest_cleanup_data, graph): - property_name = "expiryTime" - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.{property_name}" - ) + property_name = f"expiryTime{randint(10, 10000)}" + create_property_definition(property_name, graph, value_type="number") attach_property_to_entity(dataset_urns[0], property_name, 2000034, graph=graph) @@ -248,10 +243,7 @@ def test_structured_property_double(ingest_cleanup_data, graph): # ) @pytest.mark.dependency(depends=["test_healthchecks"]) def test_structured_property_double_multiple(ingest_cleanup_data, graph): - property_name = "versions" - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.{property_name}" - ) + property_name = f"versions{randint(10, 10000)}" create_property_definition( property_name, graph, value_type="number", cardinality="MULTIPLE" @@ -266,10 +258,7 @@ def test_structured_property_double_multiple(ingest_cleanup_data, graph): # ) @pytest.mark.dependency(depends=["test_healthchecks"]) def test_structured_property_string_allowed_values(ingest_cleanup_data, graph): - property_name = "enumProperty" - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.{property_name}" - ) + property_name = f"enumProperty{randint(10, 10000)}" create_property_definition( property_name, @@ -302,7 +291,7 @@ def test_structured_property_string_allowed_values(ingest_cleanup_data, graph): @pytest.mark.dependency(depends=["test_healthchecks"]) def test_structured_property_definition_evolution(ingest_cleanup_data, graph): - property_name = "enumProperty1234" + property_name = f"enumProperty{randint(10, 10000)}" create_property_definition( property_name, @@ -314,9 +303,6 @@ def test_structured_property_definition_evolution(ingest_cleanup_data, graph): PropertyValueClass(value="bar"), ], ) - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.{property_name}" - ) try: create_property_definition( @@ -354,9 +340,6 @@ def test_structured_property_schema_field(ingest_cleanup_data, graph): value_type="date", entity_types=["schemaField"], ) - generated_urns.append( - f"urn:li:structuredProperty:io.datahubproject.test.{property_name}" - ) attach_property_to_entity( schema_field_urns[0], @@ -425,16 +408,13 @@ def test_dataset_yaml_loader(ingest_cleanup_data, graph): def test_dataset_structured_property_validation(ingest_cleanup_data, graph, caplog): from datahub.api.entities.dataset.dataset import Dataset - property_name = "replicationSLA" + property_name = f"replicationSLA{randint(10, 10000)}" property_value = 30 value_type = "number" create_property_definition( property_name=property_name, graph=graph, value_type=value_type ) - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.replicationSLA" - ) attach_property_to_entity( dataset_urns[0], property_name, [property_value], graph=graph @@ -470,9 +450,6 @@ def to_es_name(property_name, namespace=default_namespace): value_type="date", entity_types=["schemaField"], ) - generated_urns.append( - f"urn:li:structuredProperty:io.datahubproject.test.{field_property_name}" - ) attach_property_to_entity( schema_field_urns[0], @@ -481,16 +458,13 @@ def to_es_name(property_name, namespace=default_namespace): graph=graph, namespace="io.datahubproject.test", ) - dataset_property_name = "replicationSLA" + dataset_property_name = f"replicationSLA{randint(10, 10000)}" property_value = 30 value_type = "number" create_property_definition( property_name=dataset_property_name, graph=graph, value_type=value_type ) - generated_urns.append( - f"urn:li:structuredProperty:{default_namespace}.{dataset_property_name}" - ) attach_property_to_entity( dataset_urns[0], dataset_property_name, [property_value], graph=graph @@ -558,37 +532,245 @@ def to_es_name(property_name, namespace=default_namespace): assert dataset_urns[0] in field_urns -@pytest.mark.skip(reason="Functionality and test needs to be validated for correctness") def test_dataset_structured_property_patch(ingest_cleanup_data, graph, caplog): - property_name = "replicationSLA" - property_value = 30 + # Create 1st Property + property_name = f"replicationSLA{randint(10, 10000)}" + property_value1 = 30.0 + property_value2 = 100.0 value_type = "number" + cardinality = "MULTIPLE" + + create_property_definition( + property_name=property_name, + graph=graph, + value_type=value_type, + cardinality=cardinality, + ) + + # Create 2nd Property + property_name_other = f"replicationSLAOther{randint(10, 10000)}" + property_value_other = 200.0 + create_property_definition( + property_name=property_name_other, + graph=graph, + value_type=value_type, + cardinality=cardinality, + ) + + def patch_one(prop_name, prop_value): + dataset_patcher: DatasetPatchBuilder = DatasetPatchBuilder(urn=dataset_urns[0]) + dataset_patcher.set_structured_property( + StructuredPropertyUrn.make_structured_property_urn( + f"{default_namespace}.{prop_name}" + ), + prop_value, + ) + + for mcp in dataset_patcher.build(): + graph.emit(mcp) + wait_for_writes_to_sync() + + # Add 1 value for property 1 + patch_one(property_name, property_value1) + + actual_property_values = get_property_from_entity( + dataset_urns[0], f"{default_namespace}.{property_name}", graph=graph + ) + assert actual_property_values == [property_value1] + + # Add 1 value for property 2 + patch_one(property_name_other, property_value_other) + + actual_property_values = get_property_from_entity( + dataset_urns[0], f"{default_namespace}.{property_name_other}", graph=graph + ) + assert actual_property_values == [property_value_other] + + # Add 2 values to property 1 + patch_one(property_name, [property_value1, property_value2]) + + actual_property_values = set( + get_property_from_entity( + dataset_urns[0], f"{default_namespace}.{property_name}", graph=graph + ) + ) + assert actual_property_values == {property_value1, property_value2} + + # Validate property 2 is the same + actual_property_values = get_property_from_entity( + dataset_urns[0], f"{default_namespace}.{property_name_other}", graph=graph + ) + assert actual_property_values == [property_value_other] + + +def test_dataset_structured_property_hard_delete(ingest_cleanup_data, graph, caplog): + property_name = f"hardDeleteTest{randint(10, 10000)}Property" + value_type = "string" + property_urn = f"urn:li:structuredProperty:{default_namespace}.{property_name}" create_property_definition( property_name=property_name, graph=graph, value_type=value_type ) - dataset_patcher: DatasetPatchBuilder = DatasetPatchBuilder(urn=dataset_urns[0]) + test_property = StructuredProperties.from_datahub(graph=graph, urn=property_urn) + assert test_property is not None - dataset_patcher.set_structured_property( - StructuredPropertyUrn.make_structured_property_urn( - f"{default_namespace}.{property_name}" - ), - property_value, + try: + graph.hard_delete_entity(urn=property_urn) + raise AssertionError("Should not be able to HARD delete structured property") + except Exception as e: + if "Hard delete of Structured Property Definitions is not supported" in str(e): + pass + else: + raise e + + +def test_dataset_structured_property_soft_delete_validation( + ingest_cleanup_data, graph, caplog +): + property_name = f"softDeleteTest{randint(10, 10000)}Property" + value_type = "string" + property_urn = f"urn:li:structuredProperty:{default_namespace}.{property_name}" + + create_property_definition( + property_name=property_name, + graph=graph, + value_type=value_type, + cardinality="SINGLE", ) - for mcp in dataset_patcher.build(): - graph.emit(mcp) + test_property = StructuredProperties.from_datahub(graph=graph, urn=property_urn) + assert test_property is not None + + graph.soft_delete_entity(urn=property_urn) + + # Attempt to modify soft deleted definition + try: + create_property_definition( + property_name=property_name, + graph=graph, + value_type=value_type, + cardinality="SINGLE", + ) + raise AssertionError( + "Should not be able to modify soft deleted structured property" + ) + except Exception as e: + if "Cannot mutate a soft deleted Structured Property Definition" in str(e): + pass + else: + raise e + + # Attempt to add soft deleted structured property to entity + try: + attach_property_to_entity( + dataset_urns[0], property_name, "test string", graph=graph + ) + raise AssertionError( + "Should not be able to apply a soft deleted structured property to another entity" + ) + except Exception as e: + if "Cannot apply a soft deleted Structured Property value" in str(e): + pass + else: + raise e + + +def test_dataset_structured_property_soft_delete_read_mutation( + ingest_cleanup_data, graph, caplog +): + property_name = f"softDeleteReadTest{randint(10, 10000)}Property" + value_type = "string" + property_urn = f"urn:li:structuredProperty:{default_namespace}.{property_name}" + property_value = "test string" + + # Create property on a dataset + create_property_definition( + property_name=property_name, + graph=graph, + value_type=value_type, + cardinality="SINGLE", + ) + attach_property_to_entity( + dataset_urns[0], property_name, property_value, graph=graph + ) + + # Make sure it exists on the dataset + actual_property_values = get_property_from_entity( + dataset_urns[0], f"{default_namespace}.{property_name}", graph=graph + ) + assert actual_property_values == [property_value] + + # Soft delete the structured property + graph.soft_delete_entity(urn=property_urn) wait_for_writes_to_sync() - dataset = Dataset.from_datahub(graph=graph, urn=dataset_urns[0]) - assert dataset.structured_properties is not None - assert isinstance(dataset.structured_properties, list) - assert [ - int(float(k)) - for k in dataset.structured_properties[ - StructuredPropertyUrn.make_structured_property_urn( - f"{default_namespace}.{property_name}" + # Make sure it is no longer returned on the dataset + actual_property_values = get_property_from_entity( + dataset_urns[0], f"{default_namespace}.{property_name}", graph=graph + ) + assert actual_property_values is None + + +def test_dataset_structured_property_soft_delete_search_filter_validation( + ingest_cleanup_data, graph, caplog +): + def to_es_name(property_name, namespace=default_namespace): + namespace_field = namespace.replace(".", "_") + return f"structuredProperties.{namespace_field}_{property_name}" + + # Create a test structured property + dataset_property_name = f"softDeleteSearchFilter{randint(10, 10000)}" + property_value = 30 + value_type = "number" + property_urn = ( + f"urn:li:structuredProperty:{default_namespace}.{dataset_property_name}" + ) + + create_property_definition( + property_name=dataset_property_name, graph=graph, value_type=value_type + ) + attach_property_to_entity( + dataset_urns[0], dataset_property_name, [property_value], graph=graph + ) + + # Perform search, make sure it works + entity_urns = list( + graph.get_urns_by_filter( + extraFilters=[ + { + "field": to_es_name(dataset_property_name), + "negated": "false", + "condition": "EXISTS", + } + ] + ) + ) + assert len(entity_urns) == 1 + assert entity_urns[0] == dataset_urns[0] + + # Soft delete the structured property + graph.soft_delete_entity(urn=property_urn) + wait_for_writes_to_sync() + + # Perform search, make sure it validates filter and rejects as invalid request + try: + list( + graph.get_urns_by_filter( + extraFilters=[ + { + "field": to_es_name(dataset_property_name), + "negated": "false", + "condition": "EXISTS", + } + ] ) - ] - ] == [property_value] + ) + raise AssertionError( + "Should not be able to filter by soft deleted structured property" + ) + except Exception as e: + if "Cannot filter on deleted Structured Property" in str(e): + pass + else: + raise e diff --git a/smoke-test/tests/tokens/revokable_access_token_test.py b/smoke-test/tests/tokens/revokable_access_token_test.py index 6e8deb41f177ea..e2fabce8ac4e8d 100644 --- a/smoke-test/tests/tokens/revokable_access_token_test.py +++ b/smoke-test/tests/tokens/revokable_access_token_test.py @@ -10,6 +10,8 @@ wait_for_writes_to_sync, ) +pytestmark = pytest.mark.no_cypress_suite1 + # Disable telemetry os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false"