From 8415fc214b1b9d0f99eafbffc01f3e46043f54dc Mon Sep 17 00:00:00 2001 From: david-leifker <114954101+david-leifker@users.noreply.github.com> Date: Mon, 8 Jan 2024 14:20:03 -0600 Subject: [PATCH] feat(entity-registry): entity registry plugins (#9538) --- build.gradle | 3 +- .../linkedin/metadata/aspect/plugins/config | 1 + .../linkedin/datahub/graphql/TestUtils.java | 71 ++-- .../BatchUpdateSoftDeletedResolverTest.java | 8 +- .../BatchUpdateDeprecationResolverTest.java | 8 +- .../domain/BatchSetDomainResolverTest.java | 8 +- .../embed/UpdateEmbedResolverTest.java | 9 +- .../owner/AddOwnersResolverTest.java | 8 +- .../owner/BatchAddOwnersResolverTest.java | 8 +- .../owner/BatchRemoveOwnersResolverTest.java | 8 +- .../resolvers/tag/AddTagsResolverTest.java | 9 +- .../tag/BatchAddTagsResolverTest.java | 23 +- .../tag/BatchRemoveTagsResolverTest.java | 18 +- .../resolvers/term/AddTermsResolverTest.java | 29 +- .../term/BatchAddTermsResolverTest.java | 8 +- .../term/BatchRemoveTermsResolverTest.java | 8 +- .../restorebackup/RestoreStorageStep.java | 2 +- .../upgrade/restoreindices/SendMAEStep.java | 2 +- entity-registry/build.gradle | 3 + .../metadata/aspect/batch/AspectsBatch.java | 100 +++++ .../metadata/aspect/batch/BatchItem.java | 66 +++ .../metadata/aspect/batch/MCLBatchItem.java | 58 +++ .../metadata/aspect/batch/MCPBatchItem.java | 46 ++ .../metadata/aspect/batch/PatchItem.java | 26 ++ .../metadata/aspect/batch/SystemAspect.java | 25 ++ .../metadata/aspect/batch/UpsertItem.java | 24 ++ .../aspect/plugins/PluginFactory.java | 269 ++++++++++++ .../metadata/aspect/plugins/PluginSpec.java | 56 +++ .../plugins/config/AspectPluginConfig.java | 50 +++ .../plugins/config/PluginConfiguration.java | 33 ++ .../aspect/plugins/hooks/MCLSideEffect.java | 38 ++ .../aspect/plugins/hooks/MCPSideEffect.java | 36 ++ .../aspect/plugins/hooks/MutationHook.java | 68 +++ .../validation/AspectPayloadValidator.java | 83 ++++ .../plugins/validation/AspectRetriever.java | 13 + .../validation/AspectValidationException.java | 12 + .../metadata/models/DataSchemaFactory.java | 68 +-- .../models/registry/ConfigEntityRegistry.java | 21 +- .../models/registry/EntityRegistry.java | 107 +++++ .../models/registry/MergedEntityRegistry.java | 20 + .../models/registry/PatchEntityRegistry.java | 8 + .../registry/PluginEntityRegistryLoader.java | 5 + .../models/registry/config/Entities.java | 2 + .../config/EntityRegistryLoadResult.java | 17 + .../metadata/aspect/plugins/PluginsTest.java | 211 ++++++++++ .../plugins/hooks/MCLSideEffectTest.java | 69 +++ .../plugins/hooks/MCPSideEffectTest.java | 67 +++ .../plugins/hooks/MutationPluginTest.java | 76 ++++ .../validation/ValidatorPluginTest.java | 97 +++++ .../registry/PatchEntityRegistryTest.java | 19 +- .../test-entity-registry-plugins-1.yml | 67 +++ .../test-entity-registry-plugins-2.yml | 45 ++ .../test-entity-registry-plugins-3.yml | 38 ++ .../metadata/client/JavaEntityClient.java | 16 +- .../client/SystemJavaEntityClient.java | 3 +- .../linkedin/metadata/entity/AspectDao.java | 2 +- .../metadata/entity/EntityAspect.java | 77 ++++ .../metadata/entity/EntityServiceImpl.java | 392 ++++++++++-------- .../linkedin/metadata/entity/EntityUtils.java | 11 +- .../entity/cassandra/CassandraAspectDao.java | 26 +- .../cassandra/CassandraRetentionService.java | 23 +- .../metadata/entity/ebean/EbeanAspectDao.java | 2 +- .../entity/ebean/EbeanRetentionService.java | 23 +- .../entity/ebean/batch/AspectsBatchImpl.java | 143 +++++++ .../entity/ebean/batch/MCLBatchItemImpl.java | 157 +++++++ .../MCPPatchBatchItem.java} | 77 ++-- .../MCPUpsertBatchItem.java} | 137 ++++-- .../ebean/transactions/AspectsBatchImpl.java | 71 ---- .../EntityRegistryUrnValidator.java | 12 +- .../entity/validation/ValidationUtils.java | 34 +- .../service/UpdateIndicesService.java | 89 ++-- .../metadata/AspectIngestionUtils.java | 45 +- .../entity/CassandraEntityServiceTest.java | 2 +- .../entity/EbeanEntityServiceTest.java | 56 ++- .../metadata/entity/EntityServiceTest.java | 190 +++++---- .../TimeseriesAspectServiceTestBase.java | 1 + .../kafka/hook/UpdateIndicesHookTest.java | 12 + .../test/resources/test-entity-registry.yml | 9 + metadata-models-custom/README.md | 246 +++++++++++ metadata-models-custom/build.gradle | 21 +- .../registry/entity-registry.yaml | 36 +- .../CustomDataQualityRulesMCLSideEffect.java | 72 ++++ .../CustomDataQualityRulesMCPSideEffect.java | 33 ++ .../hooks/CustomDataQualityRulesMutator.java | 45 ++ .../CustomDataQualityRulesValidator.java | 70 ++++ .../com/mycompany/dq/DataQualityRuleEvent.pdl | 44 ++ .../token/StatefulTokenService.java | 14 +- .../factory/entity/EntityServiceFactory.java | 3 +- .../entity/JavaEntityClientFactory.java | 3 +- .../entity/RetentionServiceFactory.java | 3 +- .../linkedin/metadata/boot/UpgradeStep.java | 3 +- .../IngestDataPlatformInstancesStep.java | 25 +- .../boot/steps/IngestDataPlatformsStep.java | 35 +- .../boot/steps/IngestOwnershipTypesStep.java | 9 +- .../boot/steps/IngestPoliciesStep.java | 13 +- .../metadata/boot/steps/IngestRolesStep.java | 13 +- .../steps/RestoreColumnLineageIndices.java | 4 +- .../boot/steps/RestoreDbtSiblingsIndices.java | 3 +- .../IngestDataPlatformInstancesStepTest.java | 10 +- .../delegates/EntityApiDelegateImpl.java | 8 +- .../openapi/entities/EntitiesController.java | 22 +- .../entities/PlatformEntitiesController.java | 17 +- .../openapi/util/MappingUtil.java | 17 +- .../java/entities/EntitiesControllerTest.java | 2 +- .../src/test/java/mock/MockEntityService.java | 3 +- .../linkedin/entity/client/EntityClient.java | 12 +- .../entity/client/SystemEntityClient.java | 12 +- .../resources/entity/AspectResource.java | 54 ++- .../resources/entity/EntityResource.java | 2 +- .../metadata/resources/operations/Utils.java | 2 +- .../resources/entity/AspectResourceTest.java | 13 +- .../linkedin/metadata/entity/AspectUtils.java | 42 +- .../metadata/entity/DeleteEntityService.java | 2 +- .../metadata/entity/EntityService.java | 26 +- .../metadata/entity/IngestResult.java | 4 +- .../metadata/entity/RetentionService.java | 20 +- .../metadata/entity/UpdateAspectResult.java | 4 +- .../transactions/AbstractBatchItem.java | 94 ----- .../entity/transactions/AspectsBatch.java | 26 -- .../metadata/utils/SystemMetadataUtils.java | 6 + 120 files changed, 3877 insertions(+), 1000 deletions(-) create mode 120000 buildSrc/src/main/java/com/linkedin/metadata/aspect/plugins/config create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPBatchItem.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchItem.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/SystemAspect.java create 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/plugins/PluginFactory.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginSpec.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/AspectPluginConfig.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffect.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffect.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MutationHook.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectPayloadValidator.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectRetriever.java create mode 100644 entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectValidationException.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffectTest.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffectTest.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MutationPluginTest.java create mode 100644 entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/validation/ValidatorPluginTest.java create mode 100644 entity-registry/src/test/resources/test-entity-registry-plugins-1.yml create mode 100644 entity-registry/src/test/resources/test-entity-registry-plugins-2.yml create mode 100644 entity-registry/src/test/resources/test-entity-registry-plugins-3.yml create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLBatchItemImpl.java rename metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/{transactions/PatchBatchItem.java => batch/MCPPatchBatchItem.java} (71%) rename metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/{transactions/UpsertBatchItem.java => batch/MCPUpsertBatchItem.java} (52%) delete mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java create mode 100644 metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java create mode 100644 metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java create mode 100644 metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java create mode 100644 metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java create mode 100644 metadata-models-custom/src/main/pegasus/com/mycompany/dq/DataQualityRuleEvent.pdl delete mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java delete mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java diff --git a/build.gradle b/build.gradle index 4680598165d285..ee9e5307532710 100644 --- a/build.gradle +++ b/build.gradle @@ -231,7 +231,8 @@ project.ext.externalDependency = [ 'common': 'commons-io:commons-io:2.7', 'jline':'jline:jline:1.4.1', 'jetbrains':' org.jetbrains.kotlin:kotlin-stdlib:1.6.0', - 'annotationApi': 'javax.annotation:javax.annotation-api:1.3.2' + 'annotationApi': 'javax.annotation:javax.annotation-api:1.3.2', + 'classGraph': 'io.github.classgraph:classgraph:4.8.165', ] allprojects { diff --git a/buildSrc/src/main/java/com/linkedin/metadata/aspect/plugins/config b/buildSrc/src/main/java/com/linkedin/metadata/aspect/plugins/config new file mode 120000 index 00000000000000..087629f8ac1df2 --- /dev/null +++ b/buildSrc/src/main/java/com/linkedin/metadata/aspect/plugins/config @@ -0,0 +1 @@ +../../../../../../../../../entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config \ No newline at end of file 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 69cd73ecd7d68d..de507eda8cdef7 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 @@ -1,5 +1,7 @@ package com.linkedin.datahub.graphql; +import static org.mockito.Mockito.mock; + import com.datahub.authentication.Actor; import com.datahub.authentication.ActorType; import com.datahub.authentication.Authentication; @@ -10,7 +12,8 @@ import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.mxe.MetadataChangeProposal; @@ -19,13 +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 = Mockito.mock(EntityService.class); + EntityService mockEntityService = + (EntityService) Mockito.mock(EntityService.class); Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry); return mockEntityService; } @@ -35,11 +39,11 @@ public static QueryContext getMockAllowContext() { } public static QueryContext getMockAllowContext(String actorUrn) { - QueryContext mockContext = Mockito.mock(QueryContext.class); + QueryContext mockContext = mock(QueryContext.class); Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn); - Authorizer mockAuthorizer = Mockito.mock(Authorizer.class); - AuthorizationResult result = Mockito.mock(AuthorizationResult.class); + Authorizer mockAuthorizer = mock(Authorizer.class); + AuthorizationResult result = mock(AuthorizationResult.class); Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.ALLOW); Mockito.when(mockAuthorizer.authorize(Mockito.any())).thenReturn(result); @@ -52,11 +56,11 @@ public static QueryContext getMockAllowContext(String actorUrn) { } public static QueryContext getMockAllowContext(String actorUrn, AuthorizationRequest request) { - QueryContext mockContext = Mockito.mock(QueryContext.class); + QueryContext mockContext = mock(QueryContext.class); Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn); - Authorizer mockAuthorizer = Mockito.mock(Authorizer.class); - AuthorizationResult result = Mockito.mock(AuthorizationResult.class); + Authorizer mockAuthorizer = mock(Authorizer.class); + AuthorizationResult result = mock(AuthorizationResult.class); Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.ALLOW); Mockito.when(mockAuthorizer.authorize(Mockito.eq(request))).thenReturn(result); @@ -73,11 +77,11 @@ public static QueryContext getMockDenyContext() { } public static QueryContext getMockDenyContext(String actorUrn) { - QueryContext mockContext = Mockito.mock(QueryContext.class); + QueryContext mockContext = mock(QueryContext.class); Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn); - Authorizer mockAuthorizer = Mockito.mock(Authorizer.class); - AuthorizationResult result = Mockito.mock(AuthorizationResult.class); + Authorizer mockAuthorizer = mock(Authorizer.class); + AuthorizationResult result = mock(AuthorizationResult.class); Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.DENY); Mockito.when(mockAuthorizer.authorize(Mockito.any())).thenReturn(result); @@ -90,11 +94,11 @@ public static QueryContext getMockDenyContext(String actorUrn) { } public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequest request) { - QueryContext mockContext = Mockito.mock(QueryContext.class); + QueryContext mockContext = mock(QueryContext.class); Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn); - Authorizer mockAuthorizer = Mockito.mock(Authorizer.class); - AuthorizationResult result = Mockito.mock(AuthorizationResult.class); + Authorizer mockAuthorizer = mock(Authorizer.class); + AuthorizationResult result = mock(AuthorizationResult.class); Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.DENY); Mockito.when(mockAuthorizer.authorize(Mockito.eq(request))).thenReturn(result); @@ -107,32 +111,44 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ } public static void verifyIngestProposal( - EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { + EntityService mockService, + int numberOfInvocations, + MetadataChangeProposal proposal) { verifyIngestProposal(mockService, numberOfInvocations, List.of(proposal)); } public static void verifyIngestProposal( - EntityService mockService, int numberOfInvocations, List proposals) { + EntityService mockService, + int numberOfInvocations, + List proposals) { AspectsBatchImpl batch = - AspectsBatchImpl.builder().mcps(proposals, mockService.getEntityRegistry()).build(); + AspectsBatchImpl.builder() + .mcps( + proposals, + mock(AuditStamp.class), + mockService.getEntityRegistry(), + mockService.getSystemEntityClient()) + .build(); Mockito.verify(mockService, Mockito.times(numberOfInvocations)) - .ingestProposal(Mockito.eq(batch), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .ingestProposal(Mockito.eq(batch), Mockito.eq(false)); } public static void verifySingleIngestProposal( - EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { + EntityService mockService, + int numberOfInvocations, + MetadataChangeProposal proposal) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)) .ingestProposal(Mockito.eq(proposal), Mockito.any(AuditStamp.class), Mockito.eq(false)); } - public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations) { + public static void verifyIngestProposal( + EntityService mockService, int numberOfInvocations) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .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), @@ -140,12 +156,9 @@ 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.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } private TestUtils() {} diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java index 49ccc751d35f63..56b01be29e1633 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java @@ -5,7 +5,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.Status; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -15,7 +14,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.List; @@ -184,10 +183,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchUpdateSoftDeletedResolver resolver = new BatchUpdateSoftDeletedResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java index 8c3620fa978a98..be7f200a6b9d72 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java @@ -5,7 +5,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.Deprecation; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -16,7 +15,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.List; @@ -217,10 +216,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java index d5ba88066e8461..32f0d30e7751a0 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java @@ -5,7 +5,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.UrnArray; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -18,7 +17,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -311,10 +310,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchSetDomainResolver resolver = new BatchSetDomainResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java index 45a17744a26971..241951319c75ed 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java @@ -7,7 +7,6 @@ import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import com.linkedin.common.AuditStamp; import com.linkedin.common.Embed; import com.linkedin.common.urn.CorpuserUrn; import com.linkedin.common.urn.Urn; @@ -19,7 +18,7 @@ import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.r2.RemoteInvocationException; import graphql.schema.DataFetchingEnvironment; @@ -142,8 +141,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false)); ; } @@ -161,8 +159,7 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false)); } @Test diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java index 74f88f95fc171e..5e199f2c6b2c71 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java @@ -4,7 +4,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.Owner; import com.linkedin.common.OwnerArray; import com.linkedin.common.Ownership; @@ -21,7 +20,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -399,10 +398,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); AddOwnersResolver resolver = new AddOwnersResolver(Mockito.mock(EntityService.class)); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchAddOwnersResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchAddOwnersResolverTest.java index 92a789530d6e4f..92960f45232b5a 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchAddOwnersResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchAddOwnersResolverTest.java @@ -4,7 +4,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.Owner; import com.linkedin.common.OwnerArray; import com.linkedin.common.Ownership; @@ -20,7 +19,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -337,10 +336,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchAddOwnersResolver resolver = new BatchAddOwnersResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchRemoveOwnersResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchRemoveOwnersResolverTest.java index 7cef90ffee5121..10c95c1bac648e 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchRemoveOwnersResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchRemoveOwnersResolverTest.java @@ -4,7 +4,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.Owner; import com.linkedin.common.OwnerArray; import com.linkedin.common.Ownership; @@ -17,7 +16,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.BatchRemoveOwnersResolver; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -204,10 +203,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchRemoveOwnersResolver resolver = new BatchRemoveOwnersResolver(mockService); 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 340802cde467b8..2468cef0e1216f 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 @@ -5,7 +5,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.GlobalTags; import com.linkedin.common.TagAssociation; import com.linkedin.common.TagAssociationArray; @@ -17,7 +16,8 @@ import com.linkedin.datahub.graphql.resolvers.mutate.AddTagsResolver; import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; @@ -210,12 +210,11 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = getMockEntityService(); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false)); AddTagsResolver resolver = new AddTagsResolver(Mockito.mock(EntityService.class)); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchAddTagsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchAddTagsResolverTest.java index 71354627b11452..c174d917748eba 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchAddTagsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchAddTagsResolverTest.java @@ -5,7 +5,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.GlobalTags; import com.linkedin.common.TagAssociation; import com.linkedin.common.TagAssociationArray; @@ -19,7 +18,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.List; @@ -197,10 +196,7 @@ public void testGetFailureTagDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -240,10 +236,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -266,10 +259,7 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -278,10 +268,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchAddTagsResolver resolver = new BatchAddTagsResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchRemoveTagsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchRemoveTagsResolverTest.java index 8cd10afee293ea..ba75b41388587c 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchRemoveTagsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/BatchRemoveTagsResolverTest.java @@ -5,7 +5,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.GlobalTags; import com.linkedin.common.TagAssociation; import com.linkedin.common.TagAssociationArray; @@ -20,7 +19,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -199,10 +198,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -225,10 +221,7 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -237,10 +230,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchRemoveTagsResolver resolver = new BatchRemoveTagsResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/AddTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/AddTermsResolverTest.java index cb827a42333b23..397bb533ff871b 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/AddTermsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/AddTermsResolverTest.java @@ -4,7 +4,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.GlossaryTermAssociation; import com.linkedin.common.GlossaryTermAssociationArray; import com.linkedin.common.GlossaryTerms; @@ -16,7 +15,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.AddTermsResolver; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -58,8 +57,7 @@ public void testGetSuccessNoExistingTerms() throws Exception { // Unable to easily validate exact payload due to the injected timestamp Mockito.verify(mockService, Mockito.times(1)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false)); Mockito.verify(mockService, Mockito.times(1)) .exists(Mockito.eq(Urn.createFromString(TEST_TERM_1_URN))); @@ -105,8 +103,7 @@ public void testGetSuccessExistingTerms() throws Exception { // Unable to easily validate exact payload due to the injected timestamp Mockito.verify(mockService, Mockito.times(1)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false)); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false)); Mockito.verify(mockService, Mockito.times(1)) .exists(Mockito.eq(Urn.createFromString(TEST_TERM_1_URN))); @@ -141,10 +138,7 @@ public void testGetFailureTermDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -173,10 +167,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -195,10 +186,7 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); } @Test @@ -207,10 +195,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); AddTermsResolver resolver = new AddTermsResolver(Mockito.mock(EntityService.class)); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchAddTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchAddTermsResolverTest.java index 7df19fad52689f..2c85e870dd6acb 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchAddTermsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchAddTermsResolverTest.java @@ -4,7 +4,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.GlossaryTermAssociation; import com.linkedin.common.GlossaryTermAssociationArray; import com.linkedin.common.GlossaryTerms; @@ -17,7 +16,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.BatchAddTermsResolver; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -239,10 +238,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchAddTermsResolver resolver = new BatchAddTermsResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchRemoveTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchRemoveTermsResolverTest.java index 659ce40542a9cf..c2520f4dfb7121 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchRemoveTermsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/BatchRemoveTermsResolverTest.java @@ -4,7 +4,6 @@ import static org.testng.Assert.*; import com.google.common.collect.ImmutableList; -import com.linkedin.common.AuditStamp; import com.linkedin.common.GlossaryTermAssociation; import com.linkedin.common.GlossaryTermAssociationArray; import com.linkedin.common.GlossaryTerms; @@ -17,7 +16,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.BatchRemoveTermsResolver; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -200,10 +199,7 @@ public void testGetEntityClientException() throws Exception { Mockito.doThrow(RuntimeException.class) .when(mockService) - .ingestProposal( - Mockito.any(AspectsBatchImpl.class), - Mockito.any(AuditStamp.class), - Mockito.anyBoolean()); + .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean()); BatchRemoveTermsResolver resolver = new BatchRemoveTermsResolver(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 5c4567c856d0ed..5c4e8cdc47e345 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 @@ -39,7 +39,7 @@ public class RestoreStorageStep implements UpgradeStep { private static final int REPORT_BATCH_SIZE = 1000; private static final int DEFAULT_THREAD_POOL = 4; - private final EntityService _entityService; + private final EntityService _entityService; private final EntityRegistry _entityRegistry; private final Map>>> _backupReaders; diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java index 574b1f08b5f543..bedf200a1c0553 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java @@ -34,7 +34,7 @@ public class SendMAEStep implements UpgradeStep { private static final boolean DEFAULT_URN_BASED_PAGINATION = false; private final Database _server; - private final EntityService _entityService; + private final EntityService _entityService; public class KafkaJob implements Callable { UpgradeContext context; diff --git a/entity-registry/build.gradle b/entity-registry/build.gradle index 77cca24c0e7234..315a29e305b77c 100644 --- a/entity-registry/build.gradle +++ b/entity-registry/build.gradle @@ -8,6 +8,7 @@ dependencies { implementation spec.product.pegasus.generator api project(path: ':metadata-models') api project(path: ':metadata-models', configuration: "dataTemplate") + api externalDependency.classGraph implementation externalDependency.slf4jApi compileOnly externalDependency.lombok implementation externalDependency.guava @@ -30,6 +31,8 @@ dependencies { testImplementation externalDependency.testng testImplementation externalDependency.mockito testImplementation externalDependency.mockitoInline + testCompileOnly externalDependency.lombok + testImplementation externalDependency.classGraph } compileTestJava.dependsOn tasks.getByPath(':entity-registry:custom-test-model:modelDeploy') 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 new file mode 100644 index 00000000000000..83e40b22a5e447 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java @@ -0,0 +1,100 @@ +package com.linkedin.metadata.aspect.batch; + +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.annotation.Nonnull; + +/** + * A batch of aspects in the context of either an MCP or MCL write path to a data store. The item is + * a record that encapsulates the change type, raw aspect and ancillary information like {@link + * SystemMetadata} and record/message created time + */ +public interface AspectsBatch { + List getItems(); + + /** + * Returns MCP items. Can be patch, upsert, etc. + * + * @return batch items + */ + default List getMCPItems() { + return getItems().stream() + .filter(item -> item instanceof MCPBatchItem) + .map(item -> (MCPBatchItem) item) + .collect(Collectors.toList()); + } + + Pair>, List> toUpsertBatchItems( + Map> latestAspects, + EntityRegistry entityRegistry, + AspectRetriever aspectRetriever); + + default Stream applyMCPSideEffects( + List items, EntityRegistry entityRegistry, AspectRetriever aspectRetriever) { + return entityRegistry.getAllMCPSideEffects().stream() + .flatMap(mcpSideEffect -> mcpSideEffect.apply(items, entityRegistry, aspectRetriever)); + } + + default boolean containsDuplicateAspects() { + return getItems().stream() + .map(i -> String.format("%s_%s", i.getClass().getName(), i.hashCode())) + .distinct() + .count() + != getItems().size(); + } + + default Map> getUrnAspectsMap() { + return getItems().stream() + .map(aspect -> Map.entry(aspect.getUrn().toString(), aspect.getAspectName())) + .collect( + Collectors.groupingBy( + Map.Entry::getKey, Collectors.mapping(Map.Entry::getValue, Collectors.toSet()))); + } + + default Map> getNewUrnAspectsMap( + Map> existingMap, List items) { + Map> newItemsMap = + items.stream() + .map(aspect -> Map.entry(aspect.getUrn().toString(), aspect.getAspectName())) + .collect( + Collectors.groupingBy( + Map.Entry::getKey, + Collectors.mapping( + Map.Entry::getValue, Collectors.toCollection(HashSet::new)))); + + return newItemsMap.entrySet().stream() + .filter( + entry -> + !existingMap.containsKey(entry.getKey()) + || !existingMap.get(entry.getKey()).containsAll(entry.getValue())) + .peek( + entry -> { + if (existingMap.containsKey(entry.getKey())) { + entry.getValue().removeAll(existingMap.get(entry.getKey())); + } + }) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + default Map> merge( + @Nonnull Map> a, @Nonnull Map> b) { + return Stream.concat(a.entrySet().stream(), b.entrySet().stream()) + .flatMap( + entry -> + entry.getValue().entrySet().stream() + .map(innerEntry -> Pair.of(entry.getKey(), innerEntry))) + .collect( + Collectors.groupingBy( + Pair::getKey, + Collectors.mapping( + Pair::getValue, Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))); + } +} 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 new file mode 100644 index 00000000000000..a4c0624150532c --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java @@ -0,0 +1,66 @@ +package com.linkedin.metadata.aspect.batch; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.mxe.SystemMetadata; +import javax.annotation.Nonnull; + +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(); + + /** + * Timestamp and actor + * + * @return the audit information + */ + AuditStamp getAuditStamp(); + + /** + * The type of change + * + * @return change type + */ + @Nonnull + ChangeType getChangeType(); + + /** + * 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/batch/MCLBatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java new file mode 100644 index 00000000000000..30e882705da453 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java @@ -0,0 +1,58 @@ +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.mxe.MetadataChangeLog; +import com.linkedin.mxe.SystemMetadata; +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 { + + @Nonnull + MetadataChangeLog getMetadataChangeLog(); + + @Override + default Urn getUrn() { + return getMetadataChangeLog().getEntityUrn(); + } + + @Nonnull + @Override + default String getAspectName() { + if (getMetadataChangeLog().getAspectName() != null) { + return getMetadataChangeLog().getAspectName(); + } else { + return getAspect().schema().getName(); + } + } + + @Override + default SystemMetadata getSystemMetadata() { + return getMetadataChangeLog().getSystemMetadata(); + } + + default SystemMetadata getPreviousSystemMetadata() { + return getMetadataChangeLog().getPreviousSystemMetadata(); + } + + @Nullable + RecordTemplate getPreviousAspect(); + + @Nonnull + RecordTemplate getAspect(); + + @Override + @Nonnull + default ChangeType getChangeType() { + return getMetadataChangeLog().getChangeType(); + } + + @Override + default AuditStamp getAuditStamp() { + return getMetadataChangeLog().getCreated(); + } +} 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/MCPBatchItem.java new file mode 100644 index 00000000000000..bb5e0ac53934af --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCPBatchItem.java @@ -0,0 +1,46 @@ +package com.linkedin.metadata.aspect.batch; + +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.registry.template.AspectTemplateEngine; +import com.linkedin.mxe.MetadataChangeProposal; +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 { + + @Nullable + public abstract MetadataChangeProposal getMetadataChangeProposal(); + + /** + * Validates that a change type is valid for the given aspect + * + * @param changeType + * @param aspectSpec + * @return + */ + protected static boolean isValidChangeType(ChangeType changeType, AspectSpec aspectSpec) { + if (aspectSpec.isTimeseries()) { + // Timeseries aspects only support UPSERT + return ChangeType.UPSERT.equals(changeType); + } else { + if (ChangeType.PATCH.equals(changeType)) { + return supportsPatch(aspectSpec); + } else { + return ChangeType.UPSERT.equals(changeType); + } + } + } + + protected 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, + // specifically having array based fields that require merging without specifying merge + // behavior can get into bad states + throw new UnsupportedOperationException( + "Aspect: " + aspectSpec.getName() + " does not currently support patch " + "operations."); + } + return true; + } +} 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/PatchItem.java new file mode 100644 index 00000000000000..f790c12ee53354 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/PatchItem.java @@ -0,0 +1,26 @@ +package com.linkedin.metadata.aspect.batch; + +import com.github.fge.jsonpatch.Patch; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.models.registry.EntityRegistry; + +/** + * A change proposal represented as a patch to an exiting stored object in the primary data store. + */ +public abstract class PatchItem extends MCPBatchItem { + + /** + * Convert a Patch to an Upsert + * + * @param entityRegistry the entity registry + * @param recordTemplate the current value record template + * @return the upsert + */ + public abstract UpsertItem applyPatch( + EntityRegistry entityRegistry, + RecordTemplate recordTemplate, + AspectRetriever aspectRetriever); + + public abstract 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 new file mode 100644 index 00000000000000..88ac902ae52fed --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/SystemAspect.java @@ -0,0 +1,25 @@ +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 new file mode 100644 index 00000000000000..4e4d2a38799dcc --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java @@ -0,0 +1,24 @@ +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 com.linkedin.metadata.models.registry.EntityRegistry; +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 RecordTemplate getAspect(); + + public abstract SystemAspect toLatestEntityAspect(); + + public abstract void validatePreCommit( + @Nullable RecordTemplate previous, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) + throws AspectValidationException; +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java new file mode 100644 index 00000000000000..dd9bbcda8f4af7 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java @@ -0,0 +1,269 @@ +package com.linkedin.metadata.aspect.plugins; + +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.aspect.plugins.config.PluginConfiguration; +import com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffect; +import com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffect; +import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; +import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; +import com.linkedin.metadata.models.registry.config.EntityRegistryLoadResult; +import io.github.classgraph.ClassGraph; +import io.github.classgraph.ClassInfo; +import io.github.classgraph.MethodInfo; +import io.github.classgraph.ScanResult; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class PluginFactory { + + public static PluginFactory withCustomClasspath( + @Nullable PluginConfiguration pluginConfiguration, @Nonnull List classLoaders) { + return new PluginFactory(pluginConfiguration, classLoaders); + } + + public static PluginFactory withConfig(@Nullable PluginConfiguration pluginConfiguration) { + return PluginFactory.withCustomClasspath(pluginConfiguration, List.of()); + } + + public static PluginFactory empty() { + return PluginFactory.withConfig(PluginConfiguration.EMPTY); + } + + public static PluginFactory merge(PluginFactory a, PluginFactory b) { + return PluginFactory.withCustomClasspath( + PluginConfiguration.merge(a.getPluginConfiguration(), b.getPluginConfiguration()), + Stream.concat(a.getClassLoaders().stream(), b.getClassLoaders().stream()) + .collect(Collectors.toList())); + } + + @Getter private final PluginConfiguration pluginConfiguration; + @Nonnull @Getter private final List classLoaders; + @Getter private final List aspectPayloadValidators; + @Getter private final List mutationHooks; + @Getter private final List mclSideEffects; + @Getter private final List mcpSideEffects; + + private final ClassGraph classGraph; + + public PluginFactory( + @Nullable PluginConfiguration pluginConfiguration, @Nonnull List classLoaders) { + this.classGraph = + new ClassGraph() + .enableRemoteJarScanning() + .enableExternalClasses() + .enableClassInfo() + .enableMethodInfo(); + + this.classLoaders = classLoaders; + + if (!this.classLoaders.isEmpty()) { + classLoaders.forEach(this.classGraph::addClassLoader); + } + + this.pluginConfiguration = + pluginConfiguration == null ? PluginConfiguration.EMPTY : pluginConfiguration; + this.aspectPayloadValidators = buildAspectPayloadValidators(this.pluginConfiguration); + this.mutationHooks = buildMutationHooks(this.pluginConfiguration); + this.mclSideEffects = buildMCLSideEffects(this.pluginConfiguration); + this.mcpSideEffects = buildMCPSideEffects(this.pluginConfiguration); + } + + /** + * Returns applicable {@link AspectPayloadValidator} implementations given the change type and + * entity/aspect information. + * + * @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 + public List getAspectPayloadValidators( + @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { + return aspectPayloadValidators.stream() + .filter(plugin -> plugin.shouldApply(changeType, entityName, aspectName)) + .collect(Collectors.toList()); + } + + /** + * Return mutation hooks for {@link com.linkedin.data.template.RecordTemplate} + * + * @param changeType The type of change + * @param entityName The entity name + * @param aspectName The aspect name + * @return Mutation hooks + */ + @Nonnull + public List getMutationHooks( + @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { + return mutationHooks.stream() + .filter(plugin -> plugin.shouldApply(changeType, entityName, aspectName)) + .collect(Collectors.toList()); + } + + /** + * 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 + * @return MCP side effects + */ + @Nonnull + public List getMCPSideEffects( + @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { + return mcpSideEffects.stream() + .filter(plugin -> plugin.shouldApply(changeType, entityName, aspectName)) + .collect(Collectors.toList()); + } + + /** + * 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 + * @return MCL side effects + */ + @Nonnull + public List getMCLSideEffects( + @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { + return mclSideEffects.stream() + .filter(plugin -> plugin.shouldApply(changeType, entityName, aspectName)) + .collect(Collectors.toList()); + } + + @Nonnull + public EntityRegistryLoadResult.PluginLoadResult getPluginLoadResult() { + return EntityRegistryLoadResult.PluginLoadResult.builder() + .validatorCount(aspectPayloadValidators.size()) + .mutationHookCount(mutationHooks.size()) + .mcpSideEffectCount(mcpSideEffects.size()) + .mclSideEffectCount(mclSideEffects.size()) + .validatorClasses( + aspectPayloadValidators.stream() + .map(cls -> cls.getClass().getName()) + .collect(Collectors.toSet())) + .mutationHookClasses( + mutationHooks.stream().map(cls -> cls.getClass().getName()).collect(Collectors.toSet())) + .mcpSideEffectClasses( + mcpSideEffects.stream() + .map(cls -> cls.getClass().getName()) + .collect(Collectors.toSet())) + .mclSideEffectClasses( + mclSideEffects.stream() + .map(cls -> cls.getClass().getName()) + .collect(Collectors.toSet())) + .build(); + } + + private List buildAspectPayloadValidators( + @Nullable PluginConfiguration pluginConfiguration) { + return pluginConfiguration == null + ? List.of() + : applyDisable( + build( + AspectPayloadValidator.class, + pluginConfiguration.getAspectPayloadValidators(), + "com.linkedin.metadata.aspect.plugins.validation")); + } + + private List buildMutationHooks(@Nullable PluginConfiguration pluginConfiguration) { + return pluginConfiguration == null + ? List.of() + : applyDisable( + build( + MutationHook.class, + pluginConfiguration.getMutationHooks(), + "com.linkedin.metadata.aspect.plugins.hooks")); + } + + private List buildMCLSideEffects( + @Nullable PluginConfiguration pluginConfiguration) { + return pluginConfiguration == null + ? List.of() + : applyDisable( + build( + MCLSideEffect.class, + pluginConfiguration.getMclSideEffects(), + "com.linkedin.metadata.aspect.plugins.hooks")); + } + + private List buildMCPSideEffects( + @Nullable PluginConfiguration pluginConfiguration) { + return pluginConfiguration == null + ? List.of() + : applyDisable( + build( + MCPSideEffect.class, + pluginConfiguration.getMcpSideEffects(), + "com.linkedin.metadata.aspect.plugins.hooks")); + } + + private List build( + Class baseClazz, List configs, String... packageNames) { + try (ScanResult scanResult = classGraph.acceptPackages(packageNames).scan()) { + + Map classMap = + scanResult.getSubclasses(baseClazz).stream() + .collect(Collectors.toMap(ClassInfo::getName, Function.identity())); + + return configs.stream() + .flatMap( + config -> { + try { + ClassInfo classInfo = classMap.get(config.getClassName()); + MethodInfo constructorMethod = classInfo.getConstructorInfo().get(0); + return Stream.of( + (T) constructorMethod.loadClassAndGetConstructor().newInstance(config)); + } catch (Exception e) { + log.error( + "Error constructing entity registry plugin class: {}", + config.getClassName(), + e); + return Stream.empty(); + } + }) + .collect(Collectors.toList()); + + } catch (Exception e) { + throw new IllegalArgumentException( + String.format("Failed to load entity registry plugins: %s.", baseClazz.getName()), e); + } + } + + @Nonnull + private static List applyDisable(@Nonnull List plugins) { + return IntStream.range(0, plugins.size()) + .mapToObj( + idx -> { + List subsequentPlugins = plugins.subList(idx + 1, plugins.size()); + T thisPlugin = plugins.get(idx); + AspectPluginConfig thisPluginConfig = thisPlugin.getConfig(); + + if (subsequentPlugins.stream() + .anyMatch( + otherPlugin -> thisPluginConfig.isDisabledBy(otherPlugin.getConfig()))) { + return null; + } + + return thisPlugin; + }) + .filter(Objects::nonNull) + .filter(p -> p.getConfig().isEnabled()) + .collect(Collectors.toList()); + } +} 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 new file mode 100644 index 00000000000000..03a0473677fb81 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginSpec.java @@ -0,0 +1,56 @@ +package com.linkedin.metadata.aspect.plugins; + +import com.linkedin.common.urn.Urn; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.models.AspectSpec; +import javax.annotation.Nonnull; +import lombok.AllArgsConstructor; +import lombok.EqualsAndHashCode; + +@AllArgsConstructor +@EqualsAndHashCode +public abstract class PluginSpec { + protected static String ENTITY_WILDCARD = "*"; + + private final AspectPluginConfig aspectPluginConfig; + + protected AspectPluginConfig getConfig() { + return this.aspectPluginConfig; + } + + public boolean shouldApply( + @Nonnull 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) { + return shouldApply(changeType, entityName, aspectSpec.getName()); + } + + public boolean shouldApply( + @Nonnull ChangeType changeType, @Nonnull String entityName, @Nonnull String aspectName) { + return getConfig().isEnabled() + && isChangeTypeSupported(changeType) + && isEntityAspectSupported(entityName, aspectName); + } + + protected boolean isEntityAspectSupported( + @Nonnull String entityName, @Nonnull String aspectName) { + return (ENTITY_WILDCARD.equals(entityName) + || getConfig().getSupportedEntityAspectNames().stream() + .anyMatch(supported -> supported.getEntityName().equals(entityName))) + && isAspectSupported(aspectName); + } + + protected boolean isAspectSupported(@Nonnull String aspectName) { + return getConfig().getSupportedEntityAspectNames().stream() + .anyMatch(supported -> supported.getAspectName().equals(aspectName)); + } + + protected boolean isChangeTypeSupported(@Nonnull ChangeType changeType) { + return getConfig().getSupportedOperations().stream() + .anyMatch(supported -> changeType.toString().equals(supported)); + } +} 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 new file mode 100644 index 00000000000000..059f133ad27760 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/AspectPluginConfig.java @@ -0,0 +1,50 @@ +package com.linkedin.metadata.aspect.plugins.config; + +import java.util.List; +import javax.annotation.Nonnull; +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Data; +import lombok.NoArgsConstructor; + +@Data +@NoArgsConstructor +@AllArgsConstructor +@Builder +public class AspectPluginConfig { + @Nonnull private String className; + private boolean enabled; + + @Nonnull private List supportedOperations; + @Nonnull private List supportedEntityAspectNames; + + @Data + @NoArgsConstructor + @AllArgsConstructor + @Builder + public static class EntityAspectName { + @Nonnull private String entityName; + @Nonnull private String aspectName; + } + + /** + * Used to determine is an earlier plugin is disabled by a subsequent plugin + * + * @param o the other plugin + * @return whether this plugin should be disabled based on another plugin + */ + public boolean isDisabledBy(AspectPluginConfig o) { + return enabled && this.isEqualExcludingEnabled(o) && !o.enabled; + } + + private boolean isEqualExcludingEnabled(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + AspectPluginConfig that = (AspectPluginConfig) o; + + if (!className.equals(that.className)) return false; + if (!supportedOperations.equals(that.supportedOperations)) return false; + return supportedEntityAspectNames.equals(that.supportedEntityAspectNames); + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java new file mode 100644 index 00000000000000..a4d0678c130f3b --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java @@ -0,0 +1,33 @@ +package com.linkedin.metadata.aspect.plugins.config; + +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; + +@Data +@AllArgsConstructor +@NoArgsConstructor +public class PluginConfiguration { + private List aspectPayloadValidators = List.of(); + private List mutationHooks = List.of(); + private List mclSideEffects = List.of(); + private List mcpSideEffects = List.of(); + + public static PluginConfiguration EMPTY = new PluginConfiguration(); + + public static PluginConfiguration merge(PluginConfiguration a, PluginConfiguration b) { + return new PluginConfiguration( + Stream.concat( + a.getAspectPayloadValidators().stream(), b.getAspectPayloadValidators().stream()) + .collect(Collectors.toList()), + Stream.concat(a.getMutationHooks().stream(), b.getMutationHooks().stream()) + .collect(Collectors.toList()), + Stream.concat(a.getMclSideEffects().stream(), b.getMclSideEffects().stream()) + .collect(Collectors.toList()), + Stream.concat(a.getMcpSideEffects().stream(), b.getMcpSideEffects().stream()) + .collect(Collectors.toList())); + } +} 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 new file mode 100644 index 00000000000000..ef9786f8d711ed --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffect.java @@ -0,0 +1,38 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +import com.linkedin.metadata.aspect.batch.MCLBatchItem; +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.registry.EntityRegistry; +import java.util.List; +import java.util.stream.Stream; +import javax.annotation.Nonnull; + +/** Given an MCL produce additional MCLs for writing */ +public abstract class MCLSideEffect extends PluginSpec { + + public MCLSideEffect(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + /** + * Given a list of MCLs, output additional MCLs + * + * @param input list + * @return additional upserts + */ + public final Stream apply( + @Nonnull List input, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) { + return input.stream() + .filter(item -> shouldApply(item.getChangeType(), item.getUrn(), item.getAspectSpec())) + .flatMap(i -> applyMCLSideEffect(i, entityRegistry, aspectRetriever)); + } + + protected abstract Stream applyMCLSideEffect( + @Nonnull MCLBatchItem input, + @Nonnull EntityRegistry entityRegistry, + @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 new file mode 100644 index 00000000000000..fc1d1587d10fb0 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffect.java @@ -0,0 +1,36 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +import com.linkedin.metadata.aspect.batch.UpsertItem; +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.registry.EntityRegistry; +import java.util.List; +import java.util.stream.Stream; +import javax.annotation.Nonnull; + +/** Given an MCP produce additional MCPs to write */ +public abstract class MCPSideEffect extends PluginSpec { + + public MCPSideEffect(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + /** + * Given the list of MCP upserts, output additional upserts + * + * @param input list + * @return additional upserts + */ + public final Stream apply( + List input, + EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) { + return input.stream() + .filter(item -> shouldApply(item.getChangeType(), item.getUrn(), item.getAspectSpec())) + .flatMap(i -> applyMCPSideEffect(i, entityRegistry, aspectRetriever)); + } + + protected abstract Stream applyMCPSideEffect( + UpsertItem input, EntityRegistry entityRegistry, @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 new file mode 100644 index 00000000000000..730a494c03d7b9 --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/MutationHook.java @@ -0,0 +1,68 @@ +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.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 javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** Applies changes to the RecordTemplate prior to write */ +public abstract class MutationHook extends PluginSpec { + + public MutationHook(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + /** + * Mutating hook + * + * @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 + */ + 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); + } + } + + 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); +} 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 new file mode 100644 index 00000000000000..656d017724571e --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectPayloadValidator.java @@ -0,0 +1,83 @@ +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.plugins.PluginSpec; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.models.AspectSpec; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +public abstract class AspectPayloadValidator extends PluginSpec { + + public AspectPayloadValidator(AspectPluginConfig aspectPluginConfig) { + super(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); + } + } + + /** + * 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); + } + } + + protected abstract void validateProposedAspect( + @Nonnull ChangeType changeType, + @Nonnull Urn entityUrn, + @Nonnull AspectSpec aspectSpec, + @Nonnull RecordTemplate aspectPayload, + @Nonnull AspectRetriever aspectRetriever) + throws AspectValidationException; + + protected abstract void validatePreCommitAspect( + @Nonnull ChangeType changeType, + @Nonnull Urn entityUrn, + @Nonnull AspectSpec aspectSpec, + @Nullable RecordTemplate previousAspect, + @Nonnull RecordTemplate proposedAspect, + AspectRetriever aspectRetriever) + throws AspectValidationException; +} 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/plugins/validation/AspectRetriever.java new file mode 100644 index 00000000000000..78aa4689472f5d --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectRetriever.java @@ -0,0 +1,13 @@ +package com.linkedin.metadata.aspect.plugins.validation; + +import com.linkedin.common.urn.Urn; +import com.linkedin.entity.Aspect; +import com.linkedin.r2.RemoteInvocationException; +import java.net.URISyntaxException; +import javax.annotation.Nonnull; + +public interface AspectRetriever { + + Aspect getLatestAspectObject(@Nonnull final Urn urn, @Nonnull final String aspectName) + throws RemoteInvocationException, URISyntaxException; +} 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 new file mode 100644 index 00000000000000..f858bdcf141aeb --- /dev/null +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/validation/AspectValidationException.java @@ -0,0 +1,12 @@ +package com.linkedin.metadata.aspect.plugins.validation; + +public class AspectValidationException extends Exception { + + public AspectValidationException(String msg) { + super(msg); + } + + public AspectValidationException(String msg, Exception e) { + super(msg, e); + } +} diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/DataSchemaFactory.java b/entity-registry/src/main/java/com/linkedin/metadata/models/DataSchemaFactory.java index b9766d0ca8640b..e41e8159f64f23 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/DataSchemaFactory.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/DataSchemaFactory.java @@ -20,6 +20,7 @@ import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; +import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; import org.reflections.Reflections; @@ -64,37 +65,48 @@ public static DataSchemaFactory withCustomClasspath(Path pluginLocation) throws // no custom classpath, just return the default factory return INSTANCE; } - // first we load up classes from the classpath - File pluginDir = pluginLocation.toFile(); - if (!pluginDir.exists()) { - throw new RuntimeException( - "Failed to find plugin directory " - + pluginDir.getAbsolutePath() - + ". Current directory is " - + new File(".").getAbsolutePath()); - } - List urls = new ArrayList(); - if (pluginDir.isDirectory()) { - List jarFiles = - Files.walk(pluginLocation) - .filter(Files::isRegularFile) - .filter(p -> p.toString().endsWith(".jar")) - .collect(Collectors.toList()); - for (Path f : jarFiles) { - URL url = f.toUri().toURL(); - if (url != null) { - urls.add(url); + + return new DataSchemaFactory( + DEFAULT_TOP_LEVEL_NAMESPACES, getClassLoader(pluginLocation).get()); + } + + public static Optional getClassLoader(@Nullable Path pluginLocation) + throws IOException { + if (pluginLocation == null) { + return Optional.empty(); + } else { + // first we load up classes from the classpath + File pluginDir = pluginLocation.toFile(); + if (!pluginDir.exists()) { + throw new RuntimeException( + "Failed to find plugin directory " + + pluginDir.getAbsolutePath() + + ". Current directory is " + + new File(".").getAbsolutePath()); + } + List urls = new ArrayList(); + if (pluginDir.isDirectory()) { + List jarFiles = + Files.walk(pluginLocation) + .filter(Files::isRegularFile) + .filter(p -> p.toString().endsWith(".jar")) + .collect(Collectors.toList()); + for (Path f : jarFiles) { + URL url = f.toUri().toURL(); + if (url != null) { + urls.add(url); + } } + } else { + URL url = (pluginLocation.toUri().toURL()); + urls.add(url); } - } else { - URL url = (pluginLocation.toUri().toURL()); - urls.add(url); + URL[] urlsArray = new URL[urls.size()]; + urls.toArray(urlsArray); + URLClassLoader classLoader = + new URLClassLoader(urlsArray, Thread.currentThread().getContextClassLoader()); + return Optional.of(classLoader); } - URL[] urlsArray = new URL[urls.size()]; - urls.toArray(urlsArray); - URLClassLoader classLoader = - new URLClassLoader(urlsArray, Thread.currentThread().getContextClassLoader()); - return new DataSchemaFactory(DEFAULT_TOP_LEVEL_NAMESPACES, classLoader); } /** diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/ConfigEntityRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/ConfigEntityRegistry.java index fba916abd24306..ce8718c536fbe1 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/ConfigEntityRegistry.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/ConfigEntityRegistry.java @@ -7,6 +7,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import com.linkedin.data.schema.DataSchema; +import com.linkedin.metadata.aspect.plugins.PluginFactory; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.DataSchemaFactory; import com.linkedin.metadata.models.DefaultEntitySpec; @@ -33,6 +34,7 @@ import java.util.Optional; import java.util.stream.Collectors; import javax.annotation.Nonnull; +import lombok.Getter; import lombok.extern.slf4j.Slf4j; /** @@ -43,6 +45,7 @@ public class ConfigEntityRegistry implements EntityRegistry { private final DataSchemaFactory dataSchemaFactory; + @Getter private final PluginFactory pluginFactory; private final Map entityNameToSpec; private final Map eventNameToSpec; private final List entitySpecs; @@ -64,6 +67,7 @@ public class ConfigEntityRegistry implements EntityRegistry { public ConfigEntityRegistry(Pair configFileClassPathPair) throws IOException { this( DataSchemaFactory.withCustomClasspath(configFileClassPathPair.getSecond()), + DataSchemaFactory.getClassLoader(configFileClassPathPair.getSecond()).stream().toList(), configFileClassPathPair.getFirst()); } @@ -108,24 +112,29 @@ private static Pair getFileAndClassPath(String entityRegistryRoot) } public ConfigEntityRegistry(InputStream configFileInputStream) { - this(DataSchemaFactory.getInstance(), configFileInputStream); + this(DataSchemaFactory.getInstance(), List.of(), configFileInputStream); } - public ConfigEntityRegistry(DataSchemaFactory dataSchemaFactory, Path configFilePath) + public ConfigEntityRegistry( + DataSchemaFactory dataSchemaFactory, List classLoaders, Path configFilePath) throws FileNotFoundException { - this(dataSchemaFactory, new FileInputStream(configFilePath.toString())); + this(dataSchemaFactory, classLoaders, new FileInputStream(configFilePath.toString())); } - public ConfigEntityRegistry(DataSchemaFactory dataSchemaFactory, InputStream configFileStream) { + public ConfigEntityRegistry( + DataSchemaFactory dataSchemaFactory, + List classLoaders, + InputStream configFileStream) { this.dataSchemaFactory = dataSchemaFactory; Entities entities; try { entities = OBJECT_MAPPER.readValue(configFileStream, Entities.class); + this.pluginFactory = PluginFactory.withCustomClasspath(entities.getPlugins(), classLoaders); } catch (IOException e) { - e.printStackTrace(); throw new IllegalArgumentException( String.format( - "Error while reading config file in path %s: %s", configFileStream, e.getMessage())); + "Error while reading config file in path %s: %s", configFileStream, e.getMessage()), + e); } if (entities.getId() != null) { identifier = entities.getId(); 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 8c415d56f0d5f4..fbc3285579cc08 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,11 +1,19 @@ package com.linkedin.metadata.models.registry; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.plugins.PluginFactory; +import com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffect; +import com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffect; +import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; +import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.DefaultEntitySpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.EventSpec; import com.linkedin.metadata.models.registry.template.AspectTemplateEngine; +import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -68,4 +76,103 @@ default String getIdentifier() { */ @Nonnull AspectTemplateEngine getAspectTemplateEngine(); + + /** + * Returns applicable {@link AspectPayloadValidator} implementations given the change type and + * entity/aspect information. + * + * @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} + * + * @param changeType The type of change + * @param entityName The entity name + * @param aspectName The aspect name + * @return 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(); + } + + /** + * 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 + * @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(); + } + + /** + * 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 + * @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(); + } + + /** + * Returns underlying plugin factory + * + * @return the plugin factory + */ + @Nonnull + default PluginFactory getPluginFactory() { + return PluginFactory.empty(); + } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java index 06aeefc2e5aa09..285b96b93d1d60 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java @@ -3,6 +3,7 @@ import com.linkedin.data.schema.compatibility.CompatibilityChecker; import com.linkedin.data.schema.compatibility.CompatibilityOptions; import com.linkedin.data.schema.compatibility.CompatibilityResult; +import com.linkedin.metadata.aspect.plugins.PluginFactory; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.ConfigEntitySpec; import com.linkedin.metadata.models.DefaultEntitySpec; @@ -27,6 +28,7 @@ public class MergedEntityRegistry implements EntityRegistry { private final Map eventNameToSpec; private final AspectTemplateEngine _aspectTemplateEngine; private final Map _aspectNameToSpec; + @Nonnull private PluginFactory pluginFactory; public MergedEntityRegistry(EntityRegistry baseEntityRegistry) { // baseEntityRegistry.get*Specs() can return immutable Collections.emptyMap() which fails @@ -42,6 +44,13 @@ public MergedEntityRegistry(EntityRegistry baseEntityRegistry) { baseEntityRegistry.getAspectTemplateEngine(); _aspectTemplateEngine = baseEntityRegistry.getAspectTemplateEngine(); _aspectNameToSpec = baseEntityRegistry.getAspectSpecs(); + if (baseEntityRegistry instanceof ConfigEntityRegistry) { + this.pluginFactory = ((ConfigEntityRegistry) baseEntityRegistry).getPluginFactory(); + } else if (baseEntityRegistry instanceof PatchEntityRegistry) { + this.pluginFactory = ((PatchEntityRegistry) baseEntityRegistry).getPluginFactory(); + } else { + this.pluginFactory = PluginFactory.empty(); + } } private void validateEntitySpec(EntitySpec entitySpec, final ValidationResult validationResult) { @@ -81,6 +90,11 @@ public MergedEntityRegistry apply(EntityRegistry patchEntityRegistry) eventNameToSpec.putAll(patchEntityRegistry.getEventSpecs()); } // TODO: Validate that the entity registries don't have conflicts among each other + + // Merge Plugins + this.pluginFactory = + PluginFactory.merge(this.pluginFactory, patchEntityRegistry.getPluginFactory()); + return this; } @@ -200,6 +214,12 @@ public AspectTemplateEngine getAspectTemplateEngine() { return _aspectTemplateEngine; } + @Nonnull + @Override + public PluginFactory getPluginFactory() { + return this.pluginFactory; + } + @Setter @Getter private class ValidationResult { diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java index 9eafbe05a4fc6d..c605cfa188fc83 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java @@ -7,6 +7,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import com.linkedin.data.schema.DataSchema; +import com.linkedin.metadata.aspect.plugins.PluginFactory; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.DataSchemaFactory; import com.linkedin.metadata.models.EntitySpec; @@ -32,6 +33,7 @@ import java.util.Optional; import java.util.stream.Collectors; import javax.annotation.Nonnull; +import lombok.Getter; import lombok.extern.slf4j.Slf4j; import org.apache.maven.artifact.versioning.ComparableVersion; @@ -44,6 +46,7 @@ public class PatchEntityRegistry implements EntityRegistry { private final DataSchemaFactory dataSchemaFactory; + @Getter private final PluginFactory pluginFactory; private final Map entityNameToSpec; private final Map eventNameToSpec; private final Map _aspectNameToSpec; @@ -90,6 +93,7 @@ public PatchEntityRegistry( throws IOException, EntityRegistryException { this( DataSchemaFactory.withCustomClasspath(configFileClassPathPair.getSecond()), + DataSchemaFactory.getClassLoader(configFileClassPathPair.getSecond()).stream().toList(), configFileClassPathPair.getFirst(), registryName, registryVersion); @@ -138,12 +142,14 @@ private static Pair getFileAndClassPath(String entityRegistryRoot) public PatchEntityRegistry( DataSchemaFactory dataSchemaFactory, + List classLoaders, Path configFilePath, String registryName, ComparableVersion registryVersion) throws FileNotFoundException, EntityRegistryException { this( dataSchemaFactory, + classLoaders, new FileInputStream(configFilePath.toString()), registryName, registryVersion); @@ -151,6 +157,7 @@ public PatchEntityRegistry( private PatchEntityRegistry( DataSchemaFactory dataSchemaFactory, + List classLoaders, InputStream configFileStream, String registryName, ComparableVersion registryVersion) @@ -162,6 +169,7 @@ private PatchEntityRegistry( Entities entities; try { entities = OBJECT_MAPPER.readValue(configFileStream, Entities.class); + this.pluginFactory = PluginFactory.withCustomClasspath(entities.getPlugins(), classLoaders); } catch (IOException e) { e.printStackTrace(); throw new IllegalArgumentException( diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoader.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoader.java index 05c752a5c15753..b90e5eb72400bb 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoader.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PluginEntityRegistryLoader.java @@ -181,6 +181,10 @@ private void loadOneRegistry( entityRegistry = new PatchEntityRegistry(patchDirectory, registryName, registryVersion); parentRegistry.apply(entityRegistry); loadResultBuilder.loadResult(LoadStatus.SUCCESS); + + // Load plugin information + loadResultBuilder.plugins(entityRegistry.getPluginFactory().getPluginLoadResult()); + log.info("Loaded registry {} successfully", entityRegistry); } catch (RuntimeException | EntityRegistryException | IOException e) { log.debug("{}: Failed to load registry {} with {}", this, registryName, e.getMessage()); @@ -189,6 +193,7 @@ private void loadOneRegistry( e.printStackTrace(pw); loadResultBuilder.loadResult(LoadStatus.FAILURE).failureReason(sw.toString()).failureCount(1); } + addLoadResult(registryName, registryVersion, loadResultBuilder.build(), entityRegistry); } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/Entities.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/Entities.java index e55bfb69d6848e..94f705d4f1193e 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/Entities.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/Entities.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.models.registry.config; +import com.linkedin.metadata.aspect.plugins.config.PluginConfiguration; import java.util.List; import lombok.AccessLevel; import lombok.AllArgsConstructor; @@ -13,4 +14,5 @@ public class Entities { String id; List entities; List events; + PluginConfiguration plugins; } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/EntityRegistryLoadResult.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/EntityRegistryLoadResult.java index f08fa5ba0a4772..076387909326bd 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/EntityRegistryLoadResult.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/config/EntityRegistryLoadResult.java @@ -1,6 +1,8 @@ package com.linkedin.metadata.models.registry.config; +import java.util.Set; import lombok.Builder; +import lombok.Data; import lombok.Getter; import lombok.Setter; @@ -11,4 +13,19 @@ public class EntityRegistryLoadResult { private String registryLocation; private String failureReason; @Setter private int failureCount; + private PluginLoadResult plugins; + + @Builder + @Data + public static class PluginLoadResult { + private int validatorCount; + private int mutationHookCount; + private int mcpSideEffectCount; + private int mclSideEffectCount; + + @Builder.Default private Set validatorClasses = Set.of(); + @Builder.Default private Set mutationHookClasses = Set.of(); + @Builder.Default private Set mcpSideEffectClasses = Set.of(); + @Builder.Default private Set mclSideEffectClasses = Set.of(); + } } 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 new file mode 100644 index 00000000000000..8c3f71fcc8019b --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java @@ -0,0 +1,211 @@ +package com.linkedin.metadata.aspect.plugins; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + +import com.datahub.test.TestEntityProfile; +import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.EventSpec; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import com.linkedin.metadata.models.registry.EntityRegistryException; +import com.linkedin.metadata.models.registry.MergedEntityRegistry; +import java.io.FileNotFoundException; +import java.util.Map; +import org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +public class PluginsTest { + public static String REGISTRY_FILE_1 = "test-entity-registry-plugins-1.yml"; + public static String REGISTRY_FILE_2 = "test-entity-registry-plugins-2.yml"; + public static String REGISTRY_FILE_3 = "test-entity-registry-plugins-3.yml"; + + @BeforeTest + public void disableAssert() { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + @Test + public void testConfigEntityRegistry() throws FileNotFoundException { + ConfigEntityRegistry configEntityRegistry = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_1)); + + Map entitySpecs = configEntityRegistry.getEntitySpecs(); + Map eventSpecs = configEntityRegistry.getEventSpecs(); + assertEquals(entitySpecs.values().size(), 2); + assertEquals(eventSpecs.values().size(), 1); + + EntitySpec entitySpec = configEntityRegistry.getEntitySpec("dataset"); + assertEquals(entitySpec.getName(), "dataset"); + assertEquals(entitySpec.getKeyAspectSpec().getName(), "datasetKey"); + assertEquals(entitySpec.getAspectSpecs().size(), 4); + assertNotNull(entitySpec.getAspectSpec("datasetKey")); + assertNotNull(entitySpec.getAspectSpec("datasetProperties")); + assertNotNull(entitySpec.getAspectSpec("schemaMetadata")); + assertNotNull(entitySpec.getAspectSpec("status")); + + entitySpec = configEntityRegistry.getEntitySpec("chart"); + assertEquals(entitySpec.getName(), "chart"); + assertEquals(entitySpec.getKeyAspectSpec().getName(), "chartKey"); + assertEquals(entitySpec.getAspectSpecs().size(), 3); + assertNotNull(entitySpec.getAspectSpec("chartKey")); + assertNotNull(entitySpec.getAspectSpec("chartInfo")); + assertNotNull(entitySpec.getAspectSpec("status")); + + EventSpec eventSpec = configEntityRegistry.getEventSpec("testEvent"); + assertEquals(eventSpec.getName(), "testEvent"); + assertNotNull(eventSpec.getPegasusSchema()); + + assertEquals( + configEntityRegistry.getAspectPayloadValidators(ChangeType.UPSERT, "*", "status").size(), + 2); + assertEquals( + configEntityRegistry.getAspectPayloadValidators(ChangeType.DELETE, "*", "status").size(), + 0); + + assertEquals( + configEntityRegistry.getMCLSideEffects(ChangeType.UPSERT, "chart", "chartInfo").size(), 1); + assertEquals( + configEntityRegistry.getMCLSideEffects(ChangeType.DELETE, "chart", "chartInfo").size(), 0); + + assertEquals( + configEntityRegistry.getMCPSideEffects(ChangeType.UPSERT, "dataset", "datasetKey").size(), + 1); + assertEquals( + configEntityRegistry.getMCPSideEffects(ChangeType.DELETE, "dataset", "datasetKey").size(), + 0); + + assertEquals( + configEntityRegistry.getMutationHooks(ChangeType.UPSERT, "*", "schemaMetadata").size(), 1); + assertEquals( + configEntityRegistry.getMutationHooks(ChangeType.DELETE, "*", "schemaMetadata").size(), 0); + } + + @Test + public void testMergedEntityRegistry() throws EntityRegistryException { + ConfigEntityRegistry configEntityRegistry1 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_1)); + ConfigEntityRegistry configEntityRegistry2 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_2)); + + MergedEntityRegistry mergedEntityRegistry = new MergedEntityRegistry(configEntityRegistry1); + mergedEntityRegistry.apply(configEntityRegistry2); + + Map entitySpecs = mergedEntityRegistry.getEntitySpecs(); + Map eventSpecs = mergedEntityRegistry.getEventSpecs(); + assertEquals(entitySpecs.values().size(), 2); + assertEquals(eventSpecs.values().size(), 1); + + EntitySpec entitySpec = mergedEntityRegistry.getEntitySpec("dataset"); + assertEquals(entitySpec.getName(), "dataset"); + assertEquals(entitySpec.getKeyAspectSpec().getName(), "datasetKey"); + assertEquals(entitySpec.getAspectSpecs().size(), 4); + assertNotNull(entitySpec.getAspectSpec("datasetKey")); + assertNotNull(entitySpec.getAspectSpec("datasetProperties")); + assertNotNull(entitySpec.getAspectSpec("schemaMetadata")); + assertNotNull(entitySpec.getAspectSpec("status")); + + entitySpec = mergedEntityRegistry.getEntitySpec("chart"); + assertEquals(entitySpec.getName(), "chart"); + assertEquals(entitySpec.getKeyAspectSpec().getName(), "chartKey"); + assertEquals(entitySpec.getAspectSpecs().size(), 3); + assertNotNull(entitySpec.getAspectSpec("chartKey")); + assertNotNull(entitySpec.getAspectSpec("chartInfo")); + assertNotNull(entitySpec.getAspectSpec("status")); + + EventSpec eventSpec = mergedEntityRegistry.getEventSpec("testEvent"); + assertEquals(eventSpec.getName(), "testEvent"); + assertNotNull(eventSpec.getPegasusSchema()); + + assertEquals( + mergedEntityRegistry.getAspectPayloadValidators(ChangeType.UPSERT, "*", "status").size(), + 3); + assertEquals( + mergedEntityRegistry.getAspectPayloadValidators(ChangeType.DELETE, "*", "status").size(), + 1); + + assertEquals( + mergedEntityRegistry.getMCLSideEffects(ChangeType.UPSERT, "chart", "chartInfo").size(), 2); + assertEquals( + mergedEntityRegistry.getMCLSideEffects(ChangeType.DELETE, "chart", "chartInfo").size(), 1); + + assertEquals( + mergedEntityRegistry.getMCPSideEffects(ChangeType.UPSERT, "dataset", "datasetKey").size(), + 2); + assertEquals( + mergedEntityRegistry.getMCPSideEffects(ChangeType.DELETE, "dataset", "datasetKey").size(), + 1); + + assertEquals( + mergedEntityRegistry.getMutationHooks(ChangeType.UPSERT, "*", "schemaMetadata").size(), 2); + assertEquals( + mergedEntityRegistry.getMutationHooks(ChangeType.DELETE, "*", "schemaMetadata").size(), 1); + } + + @Test + public void tripleMergeWithDisabled() throws EntityRegistryException { + ConfigEntityRegistry configEntityRegistry1 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_1)); + ConfigEntityRegistry configEntityRegistry2 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_2)); + ConfigEntityRegistry configEntityRegistry3 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_3)); + + MergedEntityRegistry mergedEntityRegistry = new MergedEntityRegistry(configEntityRegistry1); + mergedEntityRegistry.apply(configEntityRegistry2); + + assertEquals( + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + assertEquals( + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + assertEquals( + mergedEntityRegistry.getAllMCLSideEffects().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + assertEquals( + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + + // This one disables earlier plugins that are delete + mergedEntityRegistry.apply(configEntityRegistry3); + + assertEquals( + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 0); + assertEquals( + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 0); + assertEquals( + mergedEntityRegistry.getAllMCLSideEffects().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 0); + assertEquals( + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 0); + } +} 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 new file mode 100644 index 00000000000000..ce904142fecfeb --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCLSideEffectTest.java @@ -0,0 +1,69 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +import static org.testng.Assert.assertEquals; + +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.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import com.linkedin.metadata.models.registry.EntityRegistry; +import java.util.List; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +public class MCLSideEffectTest { + public static String REGISTRY_FILE = "test-entity-registry-plugins-1.yml"; + + @BeforeTest + public void disableAssert() { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + @Test + public void testCustomMCLSideEffect() { + ConfigEntityRegistry configEntityRegistry = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); + + List mclSideEffects = + configEntityRegistry.getMCLSideEffects(ChangeType.UPSERT, "chart", "chartInfo"); + assertEquals( + mclSideEffects, + List.of( + new TestMCLSideEffect( + AspectPluginConfig.builder() + .className( + "com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffectTest$TestMCLSideEffect") + .supportedOperations(List.of("UPSERT")) + .enabled(true) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("chart") + .aspectName("chartInfo") + .build())) + .build()))); + } + + public static class TestMCLSideEffect extends MCLSideEffect { + + public TestMCLSideEffect(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + @Override + protected Stream applyMCLSideEffect( + @Nonnull MCLBatchItem input, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) { + return Stream.of(input); + } + } +} 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 new file mode 100644 index 00000000000000..ee8f947e0e994a --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MCPSideEffectTest.java @@ -0,0 +1,67 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +import static org.testng.Assert.assertEquals; + +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.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import com.linkedin.metadata.models.registry.EntityRegistry; +import java.util.List; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +public class MCPSideEffectTest { + public static String REGISTRY_FILE = "test-entity-registry-plugins-1.yml"; + + @BeforeTest + public void disableAssert() { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + @Test + public void testCustomMCPSideEffect() { + ConfigEntityRegistry configEntityRegistry = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); + + List mcpSideEffects = + configEntityRegistry.getMCPSideEffects(ChangeType.UPSERT, "dataset", "datasetKey"); + assertEquals( + mcpSideEffects, + List.of( + new MCPSideEffectTest.TestMCPSideEffect( + AspectPluginConfig.builder() + .className( + "com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffectTest$TestMCPSideEffect") + .supportedOperations(List.of("UPSERT")) + .enabled(true) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("dataset") + .aspectName("datasetKey") + .build())) + .build()))); + } + + public static class TestMCPSideEffect extends MCPSideEffect { + + public TestMCPSideEffect(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + @Override + protected Stream applyMCPSideEffect( + UpsertItem input, EntityRegistry entityRegistry, @Nonnull AspectRetriever aspectRetriever) { + return Stream.of(input); + } + } +} 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 new file mode 100644 index 00000000000000..5094fd7fdd443d --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/hooks/MutationPluginTest.java @@ -0,0 +1,76 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +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 org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +public class MutationPluginTest { + public static String REGISTRY_FILE = "test-entity-registry-plugins-1.yml"; + + @BeforeTest + public void disableAssert() { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + @Test + public void testCustomMutator() { + ConfigEntityRegistry configEntityRegistry = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); + + List mutators = + configEntityRegistry.getMutationHooks(ChangeType.UPSERT, "*", "schemaMetadata"); + assertEquals( + mutators, + List.of( + new TestMutator( + AspectPluginConfig.builder() + .className( + "com.linkedin.metadata.aspect.plugins.hooks.MutationPluginTest$TestMutator") + .supportedOperations(List.of("UPSERT")) + .enabled(true) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("*") + .aspectName("schemaMetadata") + .build())) + .build()))); + } + + 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 new file mode 100644 index 00000000000000..07c99ee8546be0 --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/validation/ValidatorPluginTest.java @@ -0,0 +1,97 @@ +package com.linkedin.metadata.aspect.plugins.validation; + +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.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import java.util.List; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +public class ValidatorPluginTest { + public static String REGISTRY_FILE = "test-entity-registry-plugins-1.yml"; + + @BeforeTest + public void disableAssert() { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + @Test + public void testCustomValidator() { + ConfigEntityRegistry configEntityRegistry = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE)); + + List validators = + configEntityRegistry.getAspectPayloadValidators(ChangeType.UPSERT, "*", "status"); + assertEquals( + validators, + List.of( + new TestValidator( + AspectPluginConfig.builder() + .className( + "com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator") + .supportedOperations(List.of("UPSERT")) + .enabled(true) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("*") + .aspectName("status") + .build())) + .build()), + new TestValidator( + AspectPluginConfig.builder() + .className( + "com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator") + .supportedOperations(List.of("UPSERT")) + .enabled(true) + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("chart") + .aspectName("status") + .build())) + .build()))); + } + + public static class TestValidator extends AspectPayloadValidator { + + public TestValidator(AspectPluginConfig config) { + super(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"); + } + } + + @Override + protected void validatePreCommitAspect( + @Nonnull ChangeType changeType, + @Nonnull Urn entityUrn, + @Nonnull AspectSpec aspectSpec, + @Nullable RecordTemplate previousAspect, + @Nonnull RecordTemplate proposedAspect, + AspectRetriever aspectRetriever) + throws AspectValidationException {} + } +} diff --git a/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java b/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java index 1652a512905978..27227f133ab55e 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java @@ -5,6 +5,7 @@ import com.linkedin.metadata.models.DataSchemaFactory; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.EventSpec; +import java.nio.file.Path; import java.nio.file.Paths; import java.util.Map; import org.testng.annotations.Test; @@ -47,18 +48,20 @@ public void testEntityRegistryLoad() throws Exception, EntityRegistryException { */ @Test public void testEntityRegistryWithKeyLoad() throws Exception, EntityRegistryException { - DataSchemaFactory dataSchemaFactory = - DataSchemaFactory.withCustomClasspath( - Paths.get( - TestConstants.BASE_DIRECTORY - + "/" - + TestConstants.TEST_REGISTRY - + "/" - + TestConstants.TEST_VERSION.toString())); + Path pluginLocation = + Paths.get( + TestConstants.BASE_DIRECTORY + + "/" + + TestConstants.TEST_REGISTRY + + "/" + + TestConstants.TEST_VERSION.toString()); + + DataSchemaFactory dataSchemaFactory = DataSchemaFactory.withCustomClasspath(pluginLocation); PatchEntityRegistry patchEntityRegistry = new PatchEntityRegistry( dataSchemaFactory, + DataSchemaFactory.getClassLoader(pluginLocation).stream().toList(), Paths.get("src/test_plugins/mycompany-full-model/0.0.1/entity-registry.yaml"), TestConstants.TEST_REGISTRY, TestConstants.TEST_VERSION); diff --git a/entity-registry/src/test/resources/test-entity-registry-plugins-1.yml b/entity-registry/src/test/resources/test-entity-registry-plugins-1.yml new file mode 100644 index 00000000000000..7ef21bce3144ce --- /dev/null +++ b/entity-registry/src/test/resources/test-entity-registry-plugins-1.yml @@ -0,0 +1,67 @@ +id: test-registry-1 +entities: + - name: dataset + keyAspect: datasetKey + category: core + aspects: + - datasetProperties + - schemaMetadata + - status + - name: chart + keyAspect: chartKey + aspects: + - chartInfo + - status +events: + - name: testEvent + +plugins: + aspectPayloadValidators: + # All status aspects on any entity + - className: 'com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + # Chart status only + - className: 'com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: chart + aspectName: status + # Disabled + - className: 'com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator' + enabled: false + supportedOperations: + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + mutationHooks: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MutationPluginTest$TestMutator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: schemaMetadata + mclSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffectTest$TestMCLSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: chart + aspectName: chartInfo + mcpSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffectTest$TestMCPSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: dataset + aspectName: datasetKey diff --git a/entity-registry/src/test/resources/test-entity-registry-plugins-2.yml b/entity-registry/src/test/resources/test-entity-registry-plugins-2.yml new file mode 100644 index 00000000000000..b35b17d3bd7db1 --- /dev/null +++ b/entity-registry/src/test/resources/test-entity-registry-plugins-2.yml @@ -0,0 +1,45 @@ +id: test-registry-2 +entities: [] +plugins: + aspectPayloadValidators: + - className: 'com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: dataset + aspectName: status + - className: 'com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator' + enabled: true + supportedOperations: + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + mutationHooks: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MutationPluginTest$TestMutator' + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: schemaMetadata + mclSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffectTest$TestMCLSideEffect' + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: chart + aspectName: chartInfo + mcpSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffectTest$TestMCPSideEffect' + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: dataset + aspectName: datasetKey diff --git a/entity-registry/src/test/resources/test-entity-registry-plugins-3.yml b/entity-registry/src/test/resources/test-entity-registry-plugins-3.yml new file mode 100644 index 00000000000000..8ce21e27d0a1cc --- /dev/null +++ b/entity-registry/src/test/resources/test-entity-registry-plugins-3.yml @@ -0,0 +1,38 @@ +id: test-registry-3 +entities: [] +plugins: + aspectPayloadValidators: + - className: 'com.linkedin.metadata.aspect.plugins.validation.ValidatorPluginTest$TestValidator' + enabled: false + supportedOperations: + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + mutationHooks: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MutationPluginTest$TestMutator' + enabled: false + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: schemaMetadata + mclSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffectTest$TestMCLSideEffect' + enabled: false + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: chart + aspectName: chartInfo + mcpSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffectTest$TestMCPSideEffect' + enabled: false + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: dataset + aspectName: datasetKey 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 e7ec4d313b5f58..34921e4182b10d 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 @@ -23,14 +23,15 @@ import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.aspect.EnvelopedAspectArray; import com.linkedin.metadata.aspect.VersionedAspect; +import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.browse.BrowseResultV2; import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.DeleteEntityService; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.IngestResult; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.graph.LineageDirection; import com.linkedin.metadata.query.AutoCompleteResult; @@ -84,7 +85,7 @@ public class JavaEntityClient implements EntityClient { private final Clock _clock = Clock.systemUTC(); - private final EntityService _entityService; + private final EntityService _entityService; private final DeleteEntityService _deleteEntityService; private final EntitySearchService _entitySearchService; private final CachingEntitySearchService _cachingEntitySearchService; @@ -712,11 +713,14 @@ public String ingestProposal( Stream.concat(Stream.of(metadataChangeProposal), additionalChanges.stream()); AspectsBatch batch = AspectsBatchImpl.builder() - .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .mcps( + proposalStream.collect(Collectors.toList()), + auditStamp, + _entityService.getEntityRegistry(), + this) .build(); - IngestResult one = - _entityService.ingestProposal(batch, auditStamp, async).stream().findFirst().get(); + IngestResult one = _entityService.ingestProposal(batch, async).stream().findFirst().get(); Urn urn = one.getUrn(); tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata()); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java index 0ac18b4aacc04e..31c2846a9c9f3b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java @@ -7,6 +7,7 @@ import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; import com.linkedin.metadata.entity.DeleteEntityService; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.LineageSearchService; @@ -23,7 +24,7 @@ public class SystemJavaEntityClient extends JavaEntityClient implements SystemEn private final Authentication systemAuthentication; public SystemJavaEntityClient( - EntityService entityService, + EntityService entityService, DeleteEntityService deleteEntityService, EntitySearchService entitySearchService, CachingEntitySearchService cachingEntitySearchService, diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java index ae27f9f7e6f1a3..e00a696a095a15 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java @@ -1,9 +1,9 @@ package com.linkedin.metadata.entity; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.metadata.utils.metrics.MetricUtils; import io.ebean.PagedList; import io.ebean.Transaction; 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 eaf9b1a2cc415a..d72586e289ea78 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,7 +1,14 @@ package com.linkedin.metadata.entity; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.metadata.aspect.batch.SystemAspect; +import com.linkedin.metadata.models.registry.EntityRegistry; +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.EqualsAndHashCode; import lombok.Getter; @@ -40,4 +47,74 @@ public class EntityAspect { public EntityAspectIdentifier toAspectIdentifier() { 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. + */ + @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 + public Urn getUrn() { + try { + return Urn.createFromString(entityAspect.getUrn()); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + + @Nonnull + public String getUrnRaw() { + return entityAspect.getUrn(); + } + + @Override + public SystemMetadata getSystemMetadata() { + return EntityUtils.parseSystemMetadata(entityAspect.getSystemMetadata()); + } + + @Nullable + public String getSystemMetadataRaw() { + return entityAspect.getSystemMetadata(); + } + + @Override + public Timestamp getCreatedOn() { + return entityAspect.getCreatedOn(); + } + + @Override + public String getAspectName() { + return entityAspect.aspect; + } + + @Override + public long getVersion() { + return entityAspect.getVersion(); + } + + @Override + public RecordTemplate getRecordTemplate(EntityRegistry entityRegistry) { + return EntityUtils.toAspectRecord( + getUrn().getEntityType(), getAspectName(), entityAspect.getMetadata(), entityRegistry); + } + + public EntityAspect asRaw() { + return entityAspect; + } + } } 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 7bd8e763cdc27a..2e19916ee3c8f2 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 @@ -1,8 +1,20 @@ package com.linkedin.metadata.entity; -import static com.linkedin.metadata.Constants.*; -import static com.linkedin.metadata.search.utils.BrowsePathUtils.*; -import static com.linkedin.metadata.utils.PegasusUtils.*; +import static com.linkedin.metadata.Constants.APP_SOURCE; +import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; +import static com.linkedin.metadata.Constants.BROWSE_PATHS_ASPECT_NAME; +import static com.linkedin.metadata.Constants.BROWSE_PATHS_V2_ASPECT_NAME; +import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ASPECT_NAME; +import static com.linkedin.metadata.Constants.DEFAULT_RUN_ID; +import static com.linkedin.metadata.Constants.FORCE_INDEXING_KEY; +import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME; +import static com.linkedin.metadata.Constants.SYSTEM_ACTOR; +import static com.linkedin.metadata.Constants.UI_SOURCE; +import static com.linkedin.metadata.search.utils.BrowsePathUtils.buildDataPlatformUrn; +import static com.linkedin.metadata.search.utils.BrowsePathUtils.getDefaultBrowsePath; +import static com.linkedin.metadata.utils.PegasusUtils.constructMCL; +import static com.linkedin.metadata.utils.PegasusUtils.getDataTemplateClassFromSchema; +import static com.linkedin.metadata.utils.PegasusUtils.urnToEntityName; import com.codahale.metrics.Timer; import com.datahub.util.RecordUtils; @@ -39,17 +51,20 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.Aspect; 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.config.PreProcessHooks; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.transactions.PatchBatchItem; -import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; 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.transactions.AbstractBatchItem; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -72,6 +87,7 @@ import com.linkedin.util.Pair; import io.ebean.PagedList; import io.ebean.Transaction; +import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.nio.charset.StandardCharsets; import java.sql.Timestamp; @@ -129,7 +145,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 @@ -141,7 +157,7 @@ public class EntityServiceImpl implements EntityService { private final EventProducer _producer; private final EntityRegistry _entityRegistry; private final Map> _entityToValidAspects; - private RetentionService _retentionService; + private RetentionService _retentionService; private final Boolean _alwaysEmitChangeLog; @Getter private final UpdateIndicesService _updateIndicesService; private final PreProcessHooks _preProcessHooks; @@ -149,6 +165,8 @@ public class EntityServiceImpl implements EntityService { private final Integer ebeanMaxTransactionRetry; + private SystemEntityClient systemEntityClient; + public EntityServiceImpl( @Nonnull final AspectDao aspectDao, @Nonnull final EventProducer producer, @@ -187,9 +205,21 @@ public EntityServiceImpl( @Override public void setSystemEntityClient(SystemEntityClient systemEntityClient) { + this.systemEntityClient = systemEntityClient; this._updateIndicesService.setSystemEntityClient(systemEntityClient); } + @Override + public SystemEntityClient getSystemEntityClient() { + return this.systemEntityClient; + } + + @Override + public RecordTemplate getLatestAspect(@Nonnull Urn urn, @Nonnull String aspectName) { + log.debug("Invoked getLatestAspect with urn {}, aspect {}", urn, aspectName); + return getAspect(urn, aspectName, ASPECT_LATEST_VERSION); + } + /** * Retrieves the latest aspects corresponding to a batch of {@link Urn}s based on a provided set * of aspect names. @@ -231,8 +261,7 @@ public Map> getLatestAspects( } final RecordTemplate aspectRecord = - EntityUtils.toAspectRecord( - urn, aspectName, aspectEntry.getMetadata(), getEntityRegistry()); + aspectEntry.asSystemAspect().getRecordTemplate(getEntityRegistry()); urnToAspects.putIfAbsent(urn, new ArrayList<>()); urnToAspects.get(urn).add(aspectRecord); }); @@ -252,8 +281,7 @@ public Map getLatestAspectsForUrn( (key, aspectEntry) -> { final String aspectName = key.getAspect(); final RecordTemplate aspectRecord = - EntityUtils.toAspectRecord( - urn, aspectName, aspectEntry.getMetadata(), getEntityRegistry()); + aspectEntry.asSystemAspect().getRecordTemplate(getEntityRegistry()); result.put(aspectName, aspectRecord); }); return result; @@ -320,13 +348,14 @@ public EntityResponse getEntityV2( * @param aspectNames set of aspects to fetch * @return a map of {@link Urn} to {@link Entity} object */ + @WithSpan @Override public Map getEntitiesV2( @Nonnull final String entityName, @Nonnull final Set urns, @Nonnull final Set aspectNames) throws URISyntaxException { - return getLatestEnvelopedAspects(entityName, urns, aspectNames).entrySet().stream() + return getLatestEnvelopedAspects(urns, aspectNames).entrySet().stream() .collect( Collectors.toMap( Map.Entry::getKey, entry -> toEntityResponse(entry.getKey(), entry.getValue()))); @@ -354,16 +383,13 @@ public Map getEntitiesVersionedV2( /** * Retrieves the latest aspects for the given set of urns as a list of enveloped aspects * - * @param entityName name of the entity to fetch * @param urns set of urns to fetch * @param aspectNames set of aspects to fetch - * @return a map of {@link Urn} to {@link EnvelopedAspect} object + * @return a map of {@link Urn} to {@link EntityAspect.EntitySystemAspect} object */ @Override public Map> getLatestEnvelopedAspects( - // TODO: entityName is unused, can we remove this as a param? - @Nonnull String entityName, @Nonnull Set urns, @Nonnull Set aspectNames) - throws URISyntaxException { + @Nonnull Set urns, @Nonnull Set aspectNames) throws URISyntaxException { final Set dbKeys = urns.stream() @@ -483,7 +509,7 @@ private Map> getCorrespondingAspects( public EnvelopedAspect getLatestEnvelopedAspect( @Nonnull final String entityName, @Nonnull final Urn urn, @Nonnull final String aspectName) throws Exception { - return getLatestEnvelopedAspects(entityName, ImmutableSet.of(urn), ImmutableSet.of(aspectName)) + return getLatestEnvelopedAspects(ImmutableSet.of(urn), ImmutableSet.of(aspectName)) .getOrDefault(urn, Collections.emptyList()) .stream() .filter(envelopedAspect -> envelopedAspect.getName().equals(aspectName)) @@ -597,18 +623,19 @@ public List ingestAspects( List> pairList, @Nonnull final AuditStamp auditStamp, SystemMetadata systemMetadata) { - List items = + List items = pairList.stream() .map( pair -> - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(pair.getKey()) .aspect(pair.getValue()) .systemMetadata(systemMetadata) - .build(_entityRegistry)) + .auditStamp(auditStamp) + .build(_entityRegistry, systemEntityClient)) .collect(Collectors.toList()); - return ingestAspects(AspectsBatchImpl.builder().items(items).build(), auditStamp, true, true); + return ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); } /** @@ -616,22 +643,17 @@ public List ingestAspects( * com.linkedin.mxe.MetadataChangeLog}. * * @param aspectsBatch aspects to write - * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time * @param emitMCL whether a {@link com.linkedin.mxe.MetadataChangeLog} should be emitted in * correspondence upon successful update * @return the {@link RecordTemplate} representation of the written aspect object */ @Override public List ingestAspects( - @Nonnull final AspectsBatch aspectsBatch, - @Nonnull final AuditStamp auditStamp, - boolean emitMCL, - boolean overwrite) { + @Nonnull final AspectsBatch aspectsBatch, boolean emitMCL, boolean overwrite) { Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectsToLocalDB").time(); - List ingestResults = - ingestAspectsToLocalDB(aspectsBatch, auditStamp, overwrite); + List ingestResults = ingestAspectsToLocalDB(aspectsBatch, overwrite); List mclResults = emitMCL(ingestResults, emitMCL); ingestToLocalDBTimer.stop(); @@ -646,14 +668,11 @@ public List ingestAspects( * @param aspectsBatch Collection of the following: an urn associated with the new aspect, name of * the aspect being inserted, and a function to apply to the latest version of the aspect to * get the updated version - * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time * @return Details about the new and old version of the aspect */ @Nonnull private List ingestAspectsToLocalDB( - @Nonnull final AspectsBatch aspectsBatch, - @Nonnull final AuditStamp auditStamp, - boolean overwrite) { + @Nonnull final AspectsBatch aspectsBatch, boolean overwrite) { if (aspectsBatch.containsDuplicateAspects()) { log.warn(String.format("Batch contains duplicates: %s", aspectsBatch)); @@ -662,50 +681,68 @@ private List ingestAspectsToLocalDB( return _aspectDao.runInTransactionWithRetry( (tx) -> { // Read before write is unfortunate, however batch it - Map> urnAspects = aspectsBatch.getUrnAspectsMap(); + final Map> urnAspects = aspectsBatch.getUrnAspectsMap(); // read #1 - Map> latestAspects = - _aspectDao.getLatestAspects(urnAspects); + final Map> latestAspects = + toSystemEntityAspects(_aspectDao.getLatestAspects(urnAspects)); // read #2 - Map> nextVersions = _aspectDao.getNextVersions(urnAspects); + final Map> nextVersions = + _aspectDao.getNextVersions(urnAspects); + + // 1. Convert patches to full upserts + // 2. Run any entity/aspect level hooks + Pair>, List> updatedItems = + aspectsBatch.toUpsertBatchItems(latestAspects, _entityRegistry, systemEntityClient); + + // Fetch additional information if needed + final Map> updatedLatestAspects; + final Map> updatedNextVersions; + if (!updatedItems.getFirst().isEmpty()) { + Map> newLatestAspects = + toSystemEntityAspects(_aspectDao.getLatestAspects(updatedItems.getFirst())); + Map> newNextVersions = + _aspectDao.getNextVersions(updatedItems.getFirst()); + // merge + updatedLatestAspects = aspectsBatch.merge(latestAspects, newLatestAspects); + updatedNextVersions = aspectsBatch.merge(nextVersions, newNextVersions); + } else { + updatedLatestAspects = latestAspects; + updatedNextVersions = nextVersions; + } - List items = - aspectsBatch.getItems().stream() - .map( - item -> { - if (item instanceof UpsertBatchItem) { - return (UpsertBatchItem) item; - } else { - // patch to upsert - PatchBatchItem patchBatchItem = (PatchBatchItem) item; - final String urnStr = patchBatchItem.getUrn().toString(); - final EntityAspect latest = - latestAspects - .getOrDefault(urnStr, Map.of()) - .get(patchBatchItem.getAspectName()); - final RecordTemplate currentValue = - latest != null - ? EntityUtils.toAspectRecord( - patchBatchItem.getUrn(), - patchBatchItem.getAspectName(), - latest.getMetadata(), - _entityRegistry) - : null; - return patchBatchItem.applyPatch(_entityRegistry, currentValue); - } - }) - .collect(Collectors.toList()); + // do final pre-commit checks with previous aspect value + updatedItems + .getSecond() + .forEach( + item -> { + SystemAspect previousAspect = + updatedLatestAspects + .getOrDefault(item.getUrn().toString(), Map.of()) + .get(item.getAspectSpec().getName()); + try { + item.validatePreCommit( + previousAspect == null + ? null + : previousAspect.getRecordTemplate(_entityRegistry), + _entityRegistry, + systemEntityClient); + } catch (AspectValidationException e) { + throw new RuntimeException(e); + } + }); // Database Upsert results List upsertResults = - items.stream() + updatedItems.getSecond().stream() .map( item -> { final String urnStr = item.getUrn().toString(); - final EntityAspect latest = - latestAspects.getOrDefault(urnStr, Map.of()).get(item.getAspectName()); + final SystemAspect latest = + updatedLatestAspects + .getOrDefault(urnStr, Map.of()) + .get(item.getAspectName()); final long nextVersion = - nextVersions + updatedNextVersions .getOrDefault(urnStr, Map.of()) .getOrDefault(item.getAspectName(), 0L); @@ -717,9 +754,11 @@ private List ingestAspectsToLocalDB( item.getUrn(), item.getAspectName(), item.getAspect(), - auditStamp, + item.getAuditStamp(), item.getSystemMetadata(), - latest, + latest == null + ? null + : ((EntityAspect.EntitySystemAspect) latest).asRaw(), nextVersion) .toBuilder() .request(item) @@ -728,21 +767,15 @@ private List ingestAspectsToLocalDB( // support inner-batch upserts latestAspects .computeIfAbsent(urnStr, key -> new HashMap<>()) - .put(item.getAspectName(), item.toLatestEntityAspect(auditStamp)); + .put(item.getAspectName(), item.toLatestEntityAspect()); nextVersions .computeIfAbsent(urnStr, key -> new HashMap<>()) .put(item.getAspectName(), nextVersion + 1); } else { - RecordTemplate oldValue = - EntityUtils.toAspectRecord( - item.getUrn().getEntityType(), - item.getAspectName(), - latest.getMetadata(), - getEntityRegistry()); - SystemMetadata oldMetadata = - EntityUtils.parseSystemMetadata(latest.getSystemMetadata()); + RecordTemplate oldValue = latest.getRecordTemplate(_entityRegistry); + SystemMetadata oldMetadata = latest.getSystemMetadata(); result = - UpdateAspectResult.builder() + UpdateAspectResult.builder() .urn(item.getUrn()) .request(item) .oldValue(oldValue) @@ -750,7 +783,7 @@ private List ingestAspectsToLocalDB( .oldSystemMetadata(oldMetadata) .newSystemMetadata(oldMetadata) .operation(MetadataAuditOperation.UPDATE) - .auditStamp(auditStamp) + .auditStamp(item.getAuditStamp()) .maxVersion(latest.getVersion()) .build(); } @@ -804,6 +837,25 @@ private List ingestAspectsToLocalDB( DEFAULT_MAX_TRANSACTION_RETRY); } + /** + * 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 = @@ -875,14 +927,15 @@ public RecordTemplate ingestAspectIfNotPresent( AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() .one( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) .aspect(newValue) .systemMetadata(systemMetadata) - .build(_entityRegistry)) + .auditStamp(auditStamp) + .build(_entityRegistry, systemEntityClient)) .build(); - List ingested = ingestAspects(aspectsBatch, auditStamp, true, false); + List ingested = ingestAspects(aspectsBatch, true, false); return ingested.stream().findFirst().get().getNewValue(); } @@ -900,8 +953,9 @@ public RecordTemplate ingestAspectIfNotPresent( public IngestResult ingestProposal( MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async) { return ingestProposal( - AspectsBatchImpl.builder().mcps(List.of(proposal), getEntityRegistry()).build(), - auditStamp, + AspectsBatchImpl.builder() + .mcps(List.of(proposal), auditStamp, getEntityRegistry(), systemEntityClient) + .build(), async) .stream() .findFirst() @@ -917,19 +971,16 @@ public IngestResult ingestProposal( * Key aspect in the DB. Instead, use an Entity Client. * * @param aspectsBatch the proposals to ingest - * @param auditStamp an audit stamp representing the time and actor proposing the change * @param async a flag to control whether we commit to primary store or just write to proposal log * before returning * @return an {@link IngestResult} containing the results */ @Override - public Set ingestProposal( - AspectsBatch aspectsBatch, AuditStamp auditStamp, final boolean async) { + public Set ingestProposal(AspectsBatch aspectsBatch, final boolean async) { - Stream timeseriesIngestResults = - ingestTimeseriesProposal(aspectsBatch, auditStamp); + Stream timeseriesIngestResults = ingestTimeseriesProposal(aspectsBatch); Stream nonTimeseriesIngestResults = - async ? ingestProposalAsync(aspectsBatch) : ingestProposalSync(aspectsBatch, auditStamp); + async ? ingestProposalAsync(aspectsBatch) : ingestProposalSync(aspectsBatch); return Stream.concat(timeseriesIngestResults, nonTimeseriesIngestResults) .collect(Collectors.toSet()); @@ -939,12 +990,10 @@ public Set ingestProposal( * Timeseries is pass through to MCL, no MCP * * @param aspectsBatch timeseries upserts batch - * @param auditStamp provided audit information * @return returns ingest proposal result, however was never in the MCP topic */ - private Stream ingestTimeseriesProposal( - AspectsBatch aspectsBatch, AuditStamp auditStamp) { - List unsupported = + private Stream ingestTimeseriesProposal(AspectsBatch aspectsBatch) { + List unsupported = aspectsBatch.getItems().stream() .filter( item -> @@ -954,15 +1003,13 @@ private Stream ingestTimeseriesProposal( if (!unsupported.isEmpty()) { throw new UnsupportedOperationException( "ChangeType not supported: " - + unsupported.stream() - .map(AbstractBatchItem::getChangeType) - .collect(Collectors.toSet())); + + unsupported.stream().map(BatchItem::getChangeType).collect(Collectors.toSet())); } - List, Boolean>>>> timeseriesResults = + List, Boolean>>>> timeseriesResults = aspectsBatch.getItems().stream() .filter(item -> item.getAspectSpec().isTimeseries()) - .map(item -> (UpsertBatchItem) item) + .map(item -> (MCPUpsertBatchItem) item) .map( item -> Pair.of( @@ -974,7 +1021,7 @@ private Stream ingestTimeseriesProposal( item.getSystemMetadata(), item.getMetadataChangeProposal(), item.getUrn(), - auditStamp, + item.getAuditStamp(), item.getAspectSpec()))) .collect(Collectors.toList()); @@ -992,7 +1039,7 @@ private Stream ingestTimeseriesProposal( } }); - UpsertBatchItem request = result.getFirst(); + MCPUpsertBatchItem request = result.getFirst(); return IngestResult.builder() .urn(request.getUrn()) .request(request) @@ -1010,8 +1057,8 @@ private Stream ingestTimeseriesProposal( * @return produced items to the MCP topic */ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { - List nonTimeseries = - aspectsBatch.getItems().stream() + List nonTimeseries = + aspectsBatch.getMCPItems().stream() .filter(item -> !item.getAspectSpec().isTimeseries()) .collect(Collectors.toList()); @@ -1029,7 +1076,7 @@ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { return nonTimeseries.stream() .map( item -> - IngestResult.builder() + IngestResult.builder() .urn(item.getUrn()) .request(item) .publishedMCP(true) @@ -1046,8 +1093,7 @@ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { } } - private Stream ingestProposalSync( - AspectsBatch aspectsBatch, AuditStamp auditStamp) { + private Stream ingestProposalSync(AspectsBatch aspectsBatch) { AspectsBatchImpl nonTimeseries = AspectsBatchImpl.builder() .items( @@ -1056,8 +1102,8 @@ private Stream ingestProposalSync( .collect(Collectors.toList())) .build(); - List unsupported = - nonTimeseries.getItems().stream() + List unsupported = + nonTimeseries.getMCPItems().stream() .filter( item -> item.getMetadataChangeProposal().getChangeType() != ChangeType.PATCH @@ -1071,12 +1117,12 @@ private Stream ingestProposalSync( .collect(Collectors.toSet())); } - List upsertResults = ingestAspects(nonTimeseries, auditStamp, true, true); + List upsertResults = ingestAspects(nonTimeseries, true, true); return upsertResults.stream() .map( result -> { - AbstractBatchItem item = result.getRequest(); + UpsertItem item = result.getRequest(); return IngestResult.builder() .urn(item.getUrn()) @@ -1421,7 +1467,7 @@ public Optional, Boolean>> conditionallyProduceMCLAsync( } private UpdateAspectResult conditionallyProduceMCLAsync(UpdateAspectResult result) { - AbstractBatchItem request = result.getRequest(); + UpsertItem request = result.getRequest(); Optional, Boolean>> emissionStatus = conditionallyProduceMCLAsync( result.getOldValue(), @@ -1443,12 +1489,6 @@ private UpdateAspectResult conditionallyProduceMCLAsync(UpdateAspectResult resul .orElse(result); } - @Override - public RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName) { - log.debug("Invoked getLatestAspect with urn {}, aspect {}", urn, aspectName); - return getAspect(urn, aspectName, ASPECT_LATEST_VERSION); - } - @Override public void ingestEntities( @Nonnull final List entities, @@ -1647,16 +1687,17 @@ private void ingestSnapshotUnion( aspectRecordsToIngest.stream() .map( pair -> - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(urn) .aspectName(pair.getKey()) .aspect(pair.getValue()) + .auditStamp(auditStamp) .systemMetadata(systemMetadata) - .build(_entityRegistry)) + .build(_entityRegistry, systemEntityClient)) .collect(Collectors.toList())) .build(); - ingestAspects(aspectsBatch, auditStamp, true, true); + ingestAspects(aspectsBatch, true, true); } @Override @@ -1758,7 +1799,7 @@ public EntityRegistry getEntityRegistry() { } @Override - public void setRetentionService(RetentionService retentionService) { + public void setRetentionService(RetentionService retentionService) { _retentionService = retentionService; } @@ -1863,8 +1904,7 @@ public RollbackRunResult deleteUrn(Urn urn) { return new RollbackRunResult(removedAspects, rowsDeletedFromEntityDeletion); } - SystemMetadata latestKeySystemMetadata = - EntityUtils.parseSystemMetadata(latestKey.getSystemMetadata()); + SystemMetadata latestKeySystemMetadata = latestKey.asSystemAspect().getSystemMetadata(); RollbackResult result = deleteAspect( urn.toString(), @@ -1980,20 +2020,14 @@ public RollbackResult deleteAspect( } // 2. Compare the match conditions, if they don't match, ignore. - SystemMetadata latestSystemMetadata = - EntityUtils.parseSystemMetadata(latest.getSystemMetadata()); + SystemMetadata latestSystemMetadata = latest.asSystemAspect().getSystemMetadata(); if (!filterMatch(latestSystemMetadata, conditions)) { return null; } String latestMetadata = latest.getMetadata(); // 3. Check if this is a key aspect - Boolean isKeyAspect = false; - try { - isKeyAspect = getKeyAspectName(Urn.createFromString(urn)).equals(aspectName); - } catch (URISyntaxException e) { - log.error("Error occurred while parsing urn: {}", urn, e); - } + Boolean isKeyAspect = getKeyAspectName(entityUrn).equals(aspectName); // 4. Fetch all preceding aspects, that match List aspectsToDelete = new ArrayList<>(); @@ -2004,8 +2038,11 @@ public RollbackResult deleteAspect( while (maxVersion > 0 && filterMatch) { EntityAspect candidateAspect = _aspectDao.getAspect(urn, aspectName, maxVersion); SystemMetadata previousSysMetadata = - EntityUtils.parseSystemMetadata(candidateAspect.getSystemMetadata()); - filterMatch = filterMatch(previousSysMetadata, conditions); + candidateAspect != null + ? candidateAspect.asSystemAspect().getSystemMetadata() + : null; + filterMatch = + previousSysMetadata != null && filterMatch(previousSysMetadata, conditions); if (filterMatch) { aspectsToDelete.add(candidateAspect); maxVersion = maxVersion - 1; @@ -2069,7 +2106,7 @@ public RollbackResult deleteAspect( latest == null ? null : EntityUtils.toAspectRecord( - Urn.createFromString(latest.getUrn()), + entitySpec.getName(), latest.getAspect(), latestMetadata, getEntityRegistry()); @@ -2078,7 +2115,7 @@ public RollbackResult deleteAspect( survivingAspect == null ? null : EntityUtils.toAspectRecord( - Urn.createFromString(survivingAspect.getUrn()), + entitySpec.getName(), survivingAspect.getAspect(), previousMetadata, getEntityRegistry()); @@ -2098,7 +2135,7 @@ public RollbackResult deleteAspect( latestSystemMetadata, previousValue == null ? null - : EntityUtils.parseSystemMetadata(survivingAspect.getSystemMetadata()), + : survivingAspect.asSystemAspect().getSystemMetadata(), survivingAspect == null ? ChangeType.DELETE : ChangeType.UPSERT, isKeyAspect, additionalRowsDeleted); @@ -2117,7 +2154,8 @@ public RollbackResult deleteAspect( return result; } - protected boolean filterMatch(SystemMetadata systemMetadata, Map conditions) { + protected boolean filterMatch( + @Nonnull SystemMetadata systemMetadata, Map conditions) { String runIdCondition = conditions.getOrDefault("runId", null); if (runIdCondition != null) { if (!runIdCondition.equals(systemMetadata.getRunId())) { @@ -2202,40 +2240,42 @@ private Map getEnvelopedAspects( continue; } - // Aspect found. Now turn it into an EnvelopedAspect - final com.linkedin.entity.Aspect aspect = - RecordUtils.toRecordTemplate( - com.linkedin.entity.Aspect.class, currAspectEntry.getMetadata()); - final EnvelopedAspect envelopedAspect = new EnvelopedAspect(); - envelopedAspect.setName(currAspectEntry.getAspect()); - envelopedAspect.setVersion(currAspectEntry.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); + result.put(currKey, toEnvelopedAspect(currAspectEntry)); + } + return result; + } - try { - if (currAspectEntry.getSystemMetadata() != null) { - final SystemMetadata systemMetadata = - RecordUtils.toRecordTemplate( - SystemMetadata.class, currAspectEntry.getSystemMetadata()); - envelopedAspect.setSystemMetadata(systemMetadata); - } - } catch (Exception e) { - log.warn( - "Exception encountered when setting system metadata on enveloped aspect {}. Error: {}", - envelopedAspect.getName(), - e); - } + 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); - envelopedAspect.setCreated( - new AuditStamp() - .setActor(UrnUtils.getUrn(currAspectEntry.getCreatedBy())) - .setTime(currAspectEntry.getCreatedOn().getTime())); - result.put(currKey, envelopedAspect); + 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()); } - return result; + + envelopedAspect.setCreated( + new AuditStamp() + .setActor(UrnUtils.getUrn(entityAspect.getCreatedBy())) + .setTime(entityAspect.getCreatedOn().getTime())); + + return envelopedAspect; } private EnvelopedAspect getKeyEnvelopedAspect(final Urn urn) { @@ -2287,8 +2327,7 @@ private UpdateAspectResult ingestAspectToLocalDB( // 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 = - EntityUtils.parseSystemMetadata(latest.getSystemMetadata()); + SystemMetadata latestSystemMetadata = latest.asSystemAspect().getSystemMetadata(); latestSystemMetadata.setLastObserved(providedSystemMetadata.getLastObserved()); latestSystemMetadata.setLastRunId( providedSystemMetadata.getLastRunId(GetMode.NULL), SetMode.IGNORE_NULL); @@ -2306,7 +2345,7 @@ private UpdateAspectResult ingestAspectToLocalDB( .urn(urn) .oldValue(oldValue) .newValue(oldValue) - .oldSystemMetadata(EntityUtils.parseSystemMetadata(latest.getSystemMetadata())) + .oldSystemMetadata(latest.asSystemAspect().getSystemMetadata()) .newSystemMetadata(latestSystemMetadata) .operation(MetadataAuditOperation.UPDATE) .auditStamp(auditStamp) @@ -2342,8 +2381,7 @@ private UpdateAspectResult ingestAspectToLocalDB( .urn(urn) .oldValue(oldValue) .newValue(newValue) - .oldSystemMetadata( - latest == null ? null : EntityUtils.parseSystemMetadata(latest.getSystemMetadata())) + .oldSystemMetadata(latest == null ? null : latest.asSystemAspect().getSystemMetadata()) .newSystemMetadata(providedSystemMetadata) .operation(MetadataAuditOperation.UPDATE) .auditStamp(auditStamp) 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 c2a0a211f9e765..459b2d183d7ac2 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 @@ -11,7 +11,7 @@ import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.EnvelopedAspect; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +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; @@ -64,8 +64,13 @@ public static void ingestChangeProposals( @Nonnull Urn actor, @Nonnull Boolean async) { entityService.ingestProposal( - AspectsBatchImpl.builder().mcps(changes, entityService.getEntityRegistry()).build(), - getAuditStamp(actor), + AspectsBatchImpl.builder() + .mcps( + changes, + getAuditStamp(actor), + entityService.getEntityRegistry(), + entityService.getSystemEntityClient()) + .build(), async); } 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 3293bc6178e430..f37f63913abe45 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 @@ -500,8 +500,30 @@ public PagedList getPagedAspects(final RestoreIndicesArgs args) { @Nonnull @Override public Stream streamAspects(String entityName, String aspectName) { - // Not implemented - return null; + SimpleStatement ss = + selectFrom(CassandraAspect.TABLE_NAME) + .all() + // assumes alpha characters after the entityType prefix + .whereColumn(CassandraAspect.URN_COLUMN) + .isGreaterThan(literal(String.join(":", List.of("urn", "li", entityName, "")))) + .whereColumn(CassandraAspect.URN_COLUMN) + .isLessThan( + literal( + String.join( + ":", + List.of( + "urn", + "li", + entityName, + "|")))) // this is used for slicing prefixes with alpha characters + .whereColumn(CassandraAspect.ASPECT_COLUMN) + .isEqualTo(literal(aspectName)) + .allowFiltering() // performance impact, however # of properties expected to be + // relatively small + .build(); + + ResultSet rs = _cqlSession.execute(ss); + return rs.all().stream().map(CassandraAspect::rowToEntityAspect); } @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 6a1ba72c376768..f1b7d761087b47 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 @@ -13,16 +13,18 @@ import com.datastax.oss.driver.api.querybuilder.select.Select; import com.datastax.oss.driver.api.querybuilder.select.Selector; import com.google.common.collect.ImmutableList; +import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.EntityAspectIdentifier; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RetentionService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.retention.DataHubRetentionConfig; import com.linkedin.retention.Retention; @@ -43,21 +45,28 @@ @Slf4j @RequiredArgsConstructor -public class CassandraRetentionService extends RetentionService { - private final EntityService _entityService; +public class CassandraRetentionService extends RetentionService { + private final EntityService _entityService; private final CqlSession _cqlSession; private final int _batchSize; private final Clock _clock = Clock.systemUTC(); @Override - public EntityService getEntityService() { + public EntityService getEntityService() { return _entityService; } @Override - protected AspectsBatch buildAspectsBatch(List mcps) { - return AspectsBatchImpl.builder().mcps(mcps, _entityService.getEntityRegistry()).build(); + protected AspectsBatch buildAspectsBatch( + List mcps, @Nonnull AuditStamp auditStamp) { + return AspectsBatchImpl.builder() + .mcps( + mcps, + auditStamp, + _entityService.getEntityRegistry(), + _entityService.getSystemEntityClient()) + .build(); } @Override 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 26946890daa3b7..176a99d8d3a498 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 @@ -7,13 +7,13 @@ import com.datahub.util.exception.RetryLimitReached; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.entity.AspectDao; import com.linkedin.metadata.entity.AspectMigrationsDao; import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.EntityAspectIdentifier; import com.linkedin.metadata.entity.ListResult; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.query.ExtraInfo; 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 e12f0f8f1b5d96..d1f54f8a7e6e52 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 @@ -1,14 +1,16 @@ package com.linkedin.metadata.entity.ebean; import com.datahub.util.RecordUtils; +import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RetentionService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.retention.DataHubRetentionConfig; import com.linkedin.retention.Retention; @@ -38,21 +40,28 @@ @Slf4j @RequiredArgsConstructor -public class EbeanRetentionService extends RetentionService { - private final EntityService _entityService; +public class EbeanRetentionService extends RetentionService { + private final EntityService _entityService; private final Database _server; private final int _batchSize; private final Clock _clock = Clock.systemUTC(); @Override - public EntityService getEntityService() { + public EntityService getEntityService() { return _entityService; } @Override - protected AspectsBatch buildAspectsBatch(List mcps) { - return AspectsBatchImpl.builder().mcps(mcps, _entityService.getEntityRegistry()).build(); + protected AspectsBatch buildAspectsBatch( + List mcps, @Nonnull AuditStamp auditStamp) { + return AspectsBatchImpl.builder() + .mcps( + mcps, + auditStamp, + _entityService.getEntityRegistry(), + _entityService.getSystemEntityClient()) + .build(); } @Override 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 new file mode 100644 index 00000000000000..4b75fe73a12e5d --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImpl.java @@ -0,0 +1,143 @@ +package com.linkedin.metadata.entity.ebean.batch; + +import com.linkedin.common.AuditStamp; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.events.metadata.ChangeType; +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.models.registry.EntityRegistry; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; +import lombok.Builder; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class AspectsBatchImpl implements AspectsBatch { + + private final List items; + + /** + * Convert patches to upserts, apply hooks at the aspect and batch level. + * + * @param latestAspects latest version in the database + * @param entityRegistry entity registry + * @return The new urn/aspectnames and the uniform upserts, possibly expanded/mutated by the + * various hooks + */ + @Override + public Pair>, List> toUpsertBatchItems( + final Map> latestAspects, + EntityRegistry entityRegistry, + AspectRetriever aspectRetriever) { + + LinkedList upsertBatchItems = + items.stream() + .map( + item -> { + final String urnStr = item.getUrn().toString(); + // latest is also the old aspect + final SystemAspect latest = + latestAspects.getOrDefault(urnStr, Map.of()).get(item.getAspectName()); + + final MCPUpsertBatchItem upsertItem; + if (item instanceof MCPUpsertBatchItem) { + upsertItem = (MCPUpsertBatchItem) item; + } else { + // patch to upsert + MCPPatchBatchItem patchBatchItem = (MCPPatchBatchItem) item; + final RecordTemplate currentValue = + latest != null ? latest.getRecordTemplate(entityRegistry) : null; + upsertItem = + patchBatchItem.applyPatch(entityRegistry, currentValue, aspectRetriever); + } + + // Apply hooks + final SystemMetadata oldSystemMetadata = + latest != null ? latest.getSystemMetadata() : null; + final RecordTemplate oldAspectValue = + latest != null ? latest.getRecordTemplate(entityRegistry) : null; + upsertItem.applyMutationHooks( + oldAspectValue, oldSystemMetadata, entityRegistry, aspectRetriever); + + return upsertItem; + }) + .collect(Collectors.toCollection(LinkedList::new)); + + LinkedList newItems = + applyMCPSideEffects(upsertBatchItems, entityRegistry, aspectRetriever) + .collect(Collectors.toCollection(LinkedList::new)); + Map> newUrnAspectNames = getNewUrnAspectsMap(getUrnAspectsMap(), newItems); + upsertBatchItems.addAll(newItems); + + return Pair.of(newUrnAspectNames, upsertBatchItems); + } + + public static class AspectsBatchImplBuilder { + /** + * Just one aspect record template + * + * @param data aspect data + * @return builder + */ + public AspectsBatchImplBuilder one(BatchItem data) { + this.items = List.of(data); + return this; + } + + public AspectsBatchImplBuilder mcps( + List mcps, + AuditStamp auditStamp, + EntityRegistry entityRegistry, + AspectRetriever aspectRetriever) { + this.items = + mcps.stream() + .map( + mcp -> { + if (mcp.getChangeType().equals(ChangeType.PATCH)) { + return MCPPatchBatchItem.MCPPatchBatchItemBuilder.build( + mcp, auditStamp, entityRegistry); + } else { + return MCPUpsertBatchItem.MCPUpsertBatchItemBuilder.build( + mcp, auditStamp, entityRegistry, aspectRetriever); + } + }) + .collect(Collectors.toList()); + return this; + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AspectsBatchImpl that = (AspectsBatchImpl) o; + return Objects.equals(items, that.items); + } + + @Override + public int hashCode() { + return Objects.hash(items); + } + + @Override + public String toString() { + return "AspectsBatchImpl{" + "items=" + items + '}'; + } +} 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/MCLBatchItemImpl.java new file mode 100644 index 00000000000000..f61280bac4b223 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCLBatchItemImpl.java @@ -0,0 +1,157 @@ +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.entity.validation.ValidationUtils; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.utils.EntityKeyUtils; +import com.linkedin.metadata.utils.GenericRecordUtils; +import com.linkedin.mxe.MetadataChangeLog; +import com.linkedin.util.Pair; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import lombok.Builder; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class MCLBatchItemImpl implements MCLBatchItem { + + @Nonnull private final MetadataChangeLog metadataChangeLog; + + @Nullable private final RecordTemplate aspect; + + @Nullable private final RecordTemplate previousAspect; + + // derived + private final EntitySpec entitySpec; + private final AspectSpec aspectSpec; + + public static class MCLBatchItemImplBuilder { + + public MCLBatchItemImpl build( + MetadataChangeLog metadataChangeLog, + EntityRegistry entityRegistry, + AspectRetriever aspectRetriever) { + return MCLBatchItemImpl.builder() + .metadataChangeLog(metadataChangeLog) + .build(entityRegistry, aspectRetriever); + } + + public MCLBatchItemImpl build(EntityRegistry entityRegistry, AspectRetriever aspectRetriever) { + log.debug("entity type = {}", this.metadataChangeLog.getEntityType()); + entitySpec(entityRegistry.getEntitySpec(this.metadataChangeLog.getEntityType())); + aspectSpec(validateAspect(this.metadataChangeLog, this.entitySpec)); + + Urn urn = this.metadataChangeLog.getEntityUrn(); + if (urn == null) { + urn = + EntityKeyUtils.getUrnFromLog( + this.metadataChangeLog, this.entitySpec.getKeyAspectSpec()); + } + EntityUtils.validateUrn(entityRegistry, urn); + log.debug("entity type = {}", urn.getEntityType()); + + entitySpec(entityRegistry.getEntitySpec(urn.getEntityType())); + log.debug("entity spec = {}", this.entitySpec); + + aspectSpec(ValidationUtils.validate(this.entitySpec, this.metadataChangeLog.getAspectName())); + log.debug("aspect spec = {}", this.aspectSpec); + + Pair aspects = + convertToRecordTemplate(this.metadataChangeLog, aspectSpec); + + // validate new + ValidationUtils.validateRecordTemplate( + this.metadataChangeLog.getChangeType(), + entityRegistry, + this.entitySpec, + this.aspectSpec, + urn, + aspects.getFirst(), + aspectRetriever); + + return new MCLBatchItemImpl( + this.metadataChangeLog, + aspects.getFirst(), + aspects.getSecond(), + this.entitySpec, + this.aspectSpec); + } + + private MCLBatchItemImplBuilder entitySpec(EntitySpec entitySpec) { + this.entitySpec = entitySpec; + return this; + } + + private MCLBatchItemImplBuilder aspectSpec(AspectSpec aspectSpec) { + this.aspectSpec = aspectSpec; + return this; + } + + private static Pair convertToRecordTemplate( + MetadataChangeLog mcl, AspectSpec aspectSpec) { + final RecordTemplate aspect; + final RecordTemplate prevAspect; + try { + + if (!ChangeType.DELETE.equals(mcl.getChangeType())) { + aspect = + GenericRecordUtils.deserializeAspect( + mcl.getAspect().getValue(), mcl.getAspect().getContentType(), aspectSpec); + ValidationUtils.validateOrThrow(aspect); + } else { + aspect = null; + } + + if (mcl.getPreviousAspectValue() != null) { + prevAspect = + GenericRecordUtils.deserializeAspect( + mcl.getPreviousAspectValue().getValue(), + mcl.getPreviousAspectValue().getContentType(), + aspectSpec); + ValidationUtils.validateOrThrow(prevAspect); + } else { + prevAspect = null; + } + } catch (ModelConversionException e) { + throw new RuntimeException( + String.format( + "Could not deserialize %s for aspect %s", + mcl.getAspect().getValue(), mcl.getAspectName())); + } + + return Pair.of(aspect, prevAspect); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + MCLBatchItemImpl that = (MCLBatchItemImpl) o; + + return metadataChangeLog.equals(that.metadataChangeLog); + } + + @Override + public int hashCode() { + return metadataChangeLog.hashCode(); + } +} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPPatchBatchItem.java similarity index 71% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java rename to metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPPatchBatchItem.java index f9b1e340d5541b..3adf384f3b0ed8 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPPatchBatchItem.java @@ -1,6 +1,8 @@ -package com.linkedin.metadata.entity.ebean.transactions; +package com.linkedin.metadata.entity.ebean.batch; -import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.Constants.INGESTION_MAX_SERIALIZED_STRING_LENGTH; +import static com.linkedin.metadata.Constants.MAX_JACKSON_STRING_SIZE; +import static com.linkedin.metadata.entity.AspectUtils.validateAspect; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.StreamReadConstraints; @@ -9,22 +11,26 @@ import com.github.fge.jsonpatch.JsonPatch; import com.github.fge.jsonpatch.JsonPatchException; import com.github.fge.jsonpatch.Patch; +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.batch.PatchItem; +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.entity.EntityUtils; -import com.linkedin.metadata.entity.transactions.AbstractBatchItem; import com.linkedin.metadata.entity.validation.ValidationUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.models.registry.template.AspectTemplateEngine; import com.linkedin.metadata.utils.EntityKeyUtils; +import com.linkedin.metadata.utils.SystemMetadataUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.Objects; +import javax.annotation.Nonnull; import lombok.Builder; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -32,7 +38,7 @@ @Slf4j @Getter @Builder(toBuilder = true) -public class PatchBatchItem extends AbstractBatchItem { +public class MCPPatchBatchItem extends PatchItem { private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); static { @@ -50,6 +56,7 @@ public class PatchBatchItem extends AbstractBatchItem { // aspectName name of the aspect being inserted private final String aspectName; private final SystemMetadata systemMetadata; + private final AuditStamp auditStamp; private final Patch patch; @@ -59,22 +66,22 @@ public class PatchBatchItem extends AbstractBatchItem { private final EntitySpec entitySpec; private final AspectSpec aspectSpec; + @Nonnull @Override public ChangeType getChangeType() { return ChangeType.PATCH; } - @Override - public void validateUrn(EntityRegistry entityRegistry, Urn urn) { - EntityUtils.validateUrn(entityRegistry, urn); - } - - public UpsertBatchItem applyPatch(EntityRegistry entityRegistry, RecordTemplate recordTemplate) { - UpsertBatchItem.UpsertBatchItemBuilder builder = - UpsertBatchItem.builder() + public MCPUpsertBatchItem applyPatch( + EntityRegistry entityRegistry, + RecordTemplate recordTemplate, + AspectRetriever aspectRetriever) { + MCPUpsertBatchItem.MCPUpsertBatchItemBuilder builder = + MCPUpsertBatchItem.builder() .urn(getUrn()) .aspectName(getAspectName()) .metadataChangeProposal(getMetadataChangeProposal()) + .auditStamp(auditStamp) .systemMetadata(getSystemMetadata()); AspectTemplateEngine aspectTemplateEngine = entityRegistry.getAspectTemplateEngine(); @@ -99,12 +106,18 @@ public UpsertBatchItem applyPatch(EntityRegistry entityRegistry, RecordTemplate throw new RuntimeException(e); } - return builder.build(entityRegistry); + return builder.build(entityRegistry, aspectRetriever); } - public static class PatchBatchItemBuilder { + public static class MCPPatchBatchItemBuilder { - public PatchBatchItem build(EntityRegistry entityRegistry) { + public MCPPatchBatchItem.MCPPatchBatchItemBuilder systemMetadata( + SystemMetadata systemMetadata) { + this.systemMetadata = SystemMetadataUtils.generateSystemMetadataIfEmpty(systemMetadata); + return this; + } + + public MCPPatchBatchItem build(EntityRegistry entityRegistry) { EntityUtils.validateUrn(entityRegistry, this.urn); log.debug("entity type = {}", this.urn.getEntityType()); @@ -119,22 +132,24 @@ public PatchBatchItem build(EntityRegistry entityRegistry) { String.format("Missing patch to apply. Aspect: %s", this.aspectSpec.getName())); } - return new PatchBatchItem( + return new MCPPatchBatchItem( this.urn, this.aspectName, - generateSystemMetadataIfEmpty(this.systemMetadata), + SystemMetadataUtils.generateSystemMetadataIfEmpty(this.systemMetadata), + this.auditStamp, this.patch, this.metadataChangeProposal, this.entitySpec, this.aspectSpec); } - public static PatchBatchItem build(MetadataChangeProposal mcp, EntityRegistry entityRegistry) { + public static MCPPatchBatchItem 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 (!isValidChangeType(ChangeType.PATCH, aspectSpec)) { + if (!PatchItem.isValidChangeType(ChangeType.PATCH, aspectSpec)) { throw new UnsupportedOperationException( "ChangeType not supported: " + mcp.getChangeType() @@ -147,23 +162,23 @@ public static PatchBatchItem build(MetadataChangeProposal mcp, EntityRegistry en urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } - PatchBatchItemBuilder builder = - PatchBatchItem.builder() - .urn(urn) - .aspectName(mcp.getAspectName()) - .systemMetadata(mcp.getSystemMetadata()) - .metadataChangeProposal(mcp) - .patch(convertToJsonPatch(mcp)); - - return builder.build(entityRegistry); + return MCPPatchBatchItem.builder() + .urn(urn) + .aspectName(mcp.getAspectName()) + .systemMetadata( + SystemMetadataUtils.generateSystemMetadataIfEmpty(mcp.getSystemMetadata())) + .metadataChangeProposal(mcp) + .auditStamp(auditStamp) + .patch(convertToJsonPatch(mcp)) + .build(entityRegistry); } - private PatchBatchItemBuilder entitySpec(EntitySpec entitySpec) { + private MCPPatchBatchItemBuilder entitySpec(EntitySpec entitySpec) { this.entitySpec = entitySpec; return this; } - private PatchBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { + private MCPPatchBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { this.aspectSpec = aspectSpec; return this; } @@ -187,7 +202,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { return false; } - PatchBatchItem that = (PatchBatchItem) o; + MCPPatchBatchItem that = (MCPPatchBatchItem) 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/transactions/UpsertBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPUpsertBatchItem.java similarity index 52% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java rename to metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPUpsertBatchItem.java index c232e4846f7d19..9d41b141dcd608 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/batch/MCPUpsertBatchItem.java @@ -1,59 +1,92 @@ -package com.linkedin.metadata.entity.ebean.transactions; +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; 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.batch.SystemAspect; +import com.linkedin.metadata.aspect.batch.UpsertItem; +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.transactions.AbstractBatchItem; import com.linkedin.metadata.entity.validation.ValidationUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; +import com.linkedin.metadata.utils.SystemMetadataUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; import java.sql.Timestamp; import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.Builder; import lombok.Getter; +import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; @Slf4j @Getter @Builder(toBuilder = true) -public class UpsertBatchItem extends AbstractBatchItem { +public class MCPUpsertBatchItem extends UpsertItem { // urn an urn associated with the new aspect - private final Urn urn; + @Nonnull private final Urn urn; + // aspectName name of the aspect being inserted - private final String aspectName; - private final SystemMetadata systemMetadata; + @Nonnull private final String aspectName; + + @Nonnull private final RecordTemplate aspect; + + @Nonnull private final SystemMetadata systemMetadata; - private final RecordTemplate aspect; + @Nonnull private final AuditStamp auditStamp; - private final MetadataChangeProposal metadataChangeProposal; + @Nullable private final MetadataChangeProposal metadataChangeProposal; // derived - private final EntitySpec entitySpec; - private final AspectSpec aspectSpec; + @Nonnull private final EntitySpec entitySpec; + @Nonnull private final AspectSpec aspectSpec; + @Nonnull @Override public ChangeType getChangeType() { return ChangeType.UPSERT; } - @Override - public void validateUrn(EntityRegistry entityRegistry, Urn urn) { - EntityUtils.validateUrn(entityRegistry, urn); + public void applyMutationHooks( + @Nullable RecordTemplate oldAspectValue, + @Nullable SystemMetadata oldSystemMetadata, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) { + // add audit stamp/system meta if needed + for (MutationHook mutationHook : + entityRegistry.getMutationHooks( + getChangeType(), entitySpec.getName(), aspectSpec.getName())) { + mutationHook.applyMutation( + getChangeType(), + entitySpec, + aspectSpec, + oldAspectValue, + aspect, + oldSystemMetadata, + systemMetadata, + auditStamp, + aspectRetriever); + } } - public EntityAspect toLatestEntityAspect(AuditStamp auditStamp) { + @Override + public SystemAspect toLatestEntityAspect() { EntityAspect latest = new EntityAspect(); latest.setAspect(getAspectName()); latest.setMetadata(EntityUtils.toJsonAspect(getAspect())); @@ -61,12 +94,39 @@ public EntityAspect toLatestEntityAspect(AuditStamp auditStamp) { latest.setVersion(ASPECT_LATEST_VERSION); latest.setCreatedOn(new Timestamp(auditStamp.getTime())); latest.setCreatedBy(auditStamp.getActor().toString()); - return latest; + return latest.asSystemAspect(); + } + + @Override + public void validatePreCommit( + @Nullable RecordTemplate previous, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) + throws AspectValidationException { + + for (AspectPayloadValidator validator : + entityRegistry.getAspectPayloadValidators( + getChangeType(), entitySpec.getName(), aspectSpec.getName())) { + validator.validatePreCommit( + getChangeType(), urn, getAspectSpec(), previous, this.aspect, aspectRetriever); + } } - public static class UpsertBatchItemBuilder { + public static class MCPUpsertBatchItemBuilder { + + // Ensure use of other builders + private MCPUpsertBatchItem build() { + return null; + } + + public MCPUpsertBatchItemBuilder systemMetadata(SystemMetadata systemMetadata) { + this.systemMetadata = SystemMetadataUtils.generateSystemMetadataIfEmpty(systemMetadata); + return this; + } - public UpsertBatchItem build(EntityRegistry entityRegistry) { + @SneakyThrows + public MCPUpsertBatchItem build( + EntityRegistry entityRegistry, AspectRetriever aspectRetriever) { EntityUtils.validateUrn(entityRegistry, this.urn); log.debug("entity type = {}", this.urn.getEntityType()); @@ -77,19 +137,30 @@ public UpsertBatchItem build(EntityRegistry entityRegistry) { log.debug("aspect spec = {}", this.aspectSpec); ValidationUtils.validateRecordTemplate( - entityRegistry, this.entitySpec, this.urn, this.aspect); + ChangeType.UPSERT, + entityRegistry, + this.entitySpec, + this.aspectSpec, + this.urn, + this.aspect, + aspectRetriever); - return new UpsertBatchItem( + return new MCPUpsertBatchItem( this.urn, this.aspectName, - AbstractBatchItem.generateSystemMetadataIfEmpty(this.systemMetadata), this.aspect, + SystemMetadataUtils.generateSystemMetadataIfEmpty(this.systemMetadata), + this.auditStamp, this.metadataChangeProposal, this.entitySpec, this.aspectSpec); } - public static UpsertBatchItem build(MetadataChangeProposal mcp, EntityRegistry entityRegistry) { + public static MCPUpsertBatchItem build( + MetadataChangeProposal mcp, + AuditStamp auditStamp, + EntityRegistry entityRegistry, + AspectRetriever aspectRetriever) { if (!mcp.getChangeType().equals(ChangeType.UPSERT)) { throw new IllegalArgumentException( "Invalid MCP, this class only supports change type of UPSERT."); @@ -112,23 +183,23 @@ public static UpsertBatchItem build(MetadataChangeProposal mcp, EntityRegistry e urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } - UpsertBatchItemBuilder builder = - UpsertBatchItem.builder() - .urn(urn) - .aspectName(mcp.getAspectName()) - .systemMetadata(mcp.getSystemMetadata()) - .metadataChangeProposal(mcp) - .aspect(convertToRecordTemplate(mcp, aspectSpec)); - - return builder.build(entityRegistry); + return MCPUpsertBatchItem.builder() + .urn(urn) + .aspectName(mcp.getAspectName()) + .systemMetadata( + SystemMetadataUtils.generateSystemMetadataIfEmpty(mcp.getSystemMetadata())) + .metadataChangeProposal(mcp) + .auditStamp(auditStamp) + .aspect(convertToRecordTemplate(mcp, aspectSpec)) + .build(entityRegistry, aspectRetriever); } - private UpsertBatchItemBuilder entitySpec(EntitySpec entitySpec) { + private MCPUpsertBatchItemBuilder entitySpec(EntitySpec entitySpec) { this.entitySpec = entitySpec; return this; } - private UpsertBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { + private MCPUpsertBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { this.aspectSpec = aspectSpec; return this; } @@ -160,7 +231,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { return false; } - UpsertBatchItem that = (UpsertBatchItem) o; + MCPUpsertBatchItem that = (MCPUpsertBatchItem) 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/transactions/AspectsBatchImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java deleted file mode 100644 index 11261afdaa0b27..00000000000000 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java +++ /dev/null @@ -1,71 +0,0 @@ -package com.linkedin.metadata.entity.ebean.transactions; - -import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.entity.transactions.AbstractBatchItem; -import com.linkedin.metadata.entity.transactions.AspectsBatch; -import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.mxe.MetadataChangeProposal; -import java.util.List; -import java.util.Objects; -import java.util.stream.Collectors; -import lombok.Builder; -import lombok.Getter; -import lombok.extern.slf4j.Slf4j; - -@Slf4j -@Getter -@Builder(toBuilder = true) -public class AspectsBatchImpl implements AspectsBatch { - private final List items; - - public static class AspectsBatchImplBuilder { - /** - * Just one aspect record template - * - * @param data aspect data - * @return builder - */ - public AspectsBatchImplBuilder one(AbstractBatchItem data) { - this.items = List.of(data); - return this; - } - - public AspectsBatchImplBuilder mcps( - List mcps, EntityRegistry entityRegistry) { - this.items = - mcps.stream() - .map( - mcp -> { - if (mcp.getChangeType().equals(ChangeType.PATCH)) { - return PatchBatchItem.PatchBatchItemBuilder.build(mcp, entityRegistry); - } else { - return UpsertBatchItem.UpsertBatchItemBuilder.build(mcp, entityRegistry); - } - }) - .collect(Collectors.toList()); - return this; - } - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - AspectsBatchImpl that = (AspectsBatchImpl) o; - return Objects.equals(items, that.items); - } - - @Override - public int hashCode() { - return Objects.hash(items); - } - - @Override - public String toString() { - return "AspectsBatchImpl{" + "items=" + items + '}'; - } -} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/EntityRegistryUrnValidator.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/EntityRegistryUrnValidator.java index ad8fbfdf2eddd5..3d7abee556290f 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/EntityRegistryUrnValidator.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/EntityRegistryUrnValidator.java @@ -12,7 +12,6 @@ import com.linkedin.data.schema.PathSpec; import com.linkedin.data.schema.validator.Validator; import com.linkedin.data.schema.validator.ValidatorContext; -import com.linkedin.data.template.RecordTemplate; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.RelationshipFieldSpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -48,8 +47,15 @@ protected void validateUrnField(ValidatorContext context) { String urnStr = (String) context.dataElement().getValue(); Urn urn = Urn.createFromString(urnStr); EntitySpec entitySpec = _entityRegistry.getEntitySpec(urn.getEntityType()); - RecordTemplate entityKey = - EntityKeyUtils.convertUrnToEntityKey(urn, entitySpec.getKeyAspectSpec()); + // This is not always false + if (entitySpec == null) { + throw new IllegalArgumentException( + String.format("Entity type %s is missing from entity registry", urn.getEntityType())); + } + + // Ensure urn conversion is successful + EntityKeyUtils.convertUrnToEntityKey(urn, entitySpec.getKeyAspectSpec()); + NamedDataSchema namedDataSchema = ((NamedDataSchema) context.dataElement().getSchema()); Class urnClass; try { 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 7f23bacdc47582..97f7aa06340d2d 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,11 +3,17 @@ 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.entity.EntityUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; import java.util.function.Consumer; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; @Slf4j @@ -60,7 +66,13 @@ public static AspectSpec validate(EntitySpec entitySpec, String aspectName) { } public static void validateRecordTemplate( - EntityRegistry entityRegistry, EntitySpec entitySpec, Urn urn, RecordTemplate aspect) { + ChangeType changeType, + EntityRegistry entityRegistry, + EntitySpec entitySpec, + AspectSpec aspectSpec, + Urn urn, + @Nullable RecordTemplate aspect, + @Nonnull AspectRetriever aspectRetriever) { EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(entityRegistry); validator.setCurrentEntitySpec(entitySpec); Consumer resultFunction = @@ -73,13 +85,21 @@ public static void validateRecordTemplate( }; RecordTemplateValidator.validate( EntityUtils.buildKeyAspect(entityRegistry, urn), resultFunction, validator); - RecordTemplateValidator.validate(aspect, resultFunction, validator); - } - public static void validateRecordTemplate( - EntityRegistry entityRegistry, Urn urn, RecordTemplate aspect) { - EntitySpec entitySpec = entityRegistry.getEntitySpec(urn.getEntityType()); - validateRecordTemplate(entityRegistry, entitySpec, urn, aspect); + 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) { + throw new IllegalArgumentException( + "Failed to validate aspect due to: " + e.getMessage(), e); + } + } + } } private ValidationUtils() {} 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 b2c615c1f47f57..247d542604da70 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 @@ -18,6 +18,8 @@ import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.aspect.batch.MCLBatchItem; +import com.linkedin.metadata.entity.ebean.batch.MCLBatchItemImpl; import com.linkedin.metadata.graph.Edge; import com.linkedin.metadata.graph.GraphIndexUtils; import com.linkedin.metadata.graph.GraphService; @@ -39,8 +41,6 @@ import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.metadata.timeseries.transformer.TimeseriesAspectTransformer; import com.linkedin.metadata.utils.EntityKeyUtils; -import com.linkedin.metadata.utils.GenericRecordUtils; -import com.linkedin.mxe.GenericAspect; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; @@ -56,6 +56,7 @@ 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; @@ -73,6 +74,8 @@ public class UpdateIndicesService { private final SearchDocumentTransformer _searchDocumentTransformer; private final EntityIndexBuilders _entityIndexBuilders; + private SystemEntityClient systemEntityClient; + @Value("${featureFlags.graphServiceDiffModeEnabled:true}") private boolean _graphDiffMode; @@ -111,10 +114,25 @@ public UpdateIndicesService( public void handleChangeEvent(@Nonnull final MetadataChangeLog event) { try { - if (UPDATE_CHANGE_TYPES.contains(event.getChangeType())) { - handleUpdateChangeEvent(event); - } else if (event.getChangeType() == ChangeType.DELETE) { - handleDeleteChangeEvent(event); + MCLBatchItemImpl batch = + MCLBatchItemImpl.builder().build(event, _entityRegistry, systemEntityClient); + + Stream sideEffects = + _entityRegistry + .getMCLSideEffects( + event.getChangeType(), event.getEntityType(), event.getAspectName()) + .stream() + .flatMap( + mclSideEffect -> + mclSideEffect.apply(List.of(batch), _entityRegistry, systemEntityClient)); + + for (MCLBatchItem mclBatchItem : Stream.concat(Stream.of(batch), sideEffects).toList()) { + MetadataChangeLog hookEvent = mclBatchItem.getMetadataChangeLog(); + if (UPDATE_CHANGE_TYPES.contains(hookEvent.getChangeType())) { + handleUpdateChangeEvent(mclBatchItem); + } else if (hookEvent.getChangeType() == ChangeType.DELETE) { + handleDeleteChangeEvent(mclBatchItem); + } } } catch (IOException e) { throw new RuntimeException(e); @@ -130,38 +148,19 @@ public void handleChangeEvent(@Nonnull final MetadataChangeLog event) { * * @param event the change event to be processed. */ - public void handleUpdateChangeEvent(@Nonnull final MetadataChangeLog event) throws IOException { + private void handleUpdateChangeEvent(@Nonnull final MCLBatchItem event) throws IOException { - final EntitySpec entitySpec = getEventEntitySpec(event); - final Urn urn = EntityKeyUtils.getUrnFromLog(event, entitySpec.getKeyAspectSpec()); + final EntitySpec entitySpec = event.getEntitySpec(); + final AspectSpec aspectSpec = event.getAspectSpec(); + final Urn urn = event.getUrn(); - if (!event.hasAspectName() || !event.hasAspect()) { - log.error("Aspect or aspect name is missing. Skipping aspect processing..."); - return; - } - - AspectSpec aspectSpec = entitySpec.getAspectSpec(event.getAspectName()); - if (aspectSpec == null) { - throw new RuntimeException( - String.format( - "Failed to retrieve Aspect Spec for entity with name %s, aspect with name %s. Cannot update indices for MCL.", - event.getEntityType(), event.getAspectName())); - } - - RecordTemplate aspect = - GenericRecordUtils.deserializeAspect( - event.getAspect().getValue(), event.getAspect().getContentType(), aspectSpec); - GenericAspect previousAspectValue = event.getPreviousAspectValue(); - RecordTemplate previousAspect = - previousAspectValue != null - ? GenericRecordUtils.deserializeAspect( - previousAspectValue.getValue(), previousAspectValue.getContentType(), aspectSpec) - : null; + RecordTemplate aspect = event.getAspect(); + RecordTemplate previousAspect = event.getPreviousAspect(); // Step 0. If the aspect is timeseries, add to its timeseries index. if (aspectSpec.isTimeseries()) { updateTimeseriesFields( - event.getEntityType(), + urn.getEntityType(), event.getAspectName(), urn, aspect, @@ -185,9 +184,9 @@ public void handleUpdateChangeEvent(@Nonnull final MetadataChangeLog event) thro && (systemMetadata == null || systemMetadata.getProperties() == null || !Boolean.parseBoolean(systemMetadata.getProperties().get(FORCE_INDEXING_KEY)))) { - updateGraphServiceDiff(urn, aspectSpec, previousAspect, aspect, event); + updateGraphServiceDiff(urn, aspectSpec, previousAspect, aspect, event.getMetadataChangeLog()); } else { - updateGraphService(urn, aspectSpec, aspect, event); + updateGraphService(urn, aspectSpec, aspect, event.getMetadataChangeLog()); } } @@ -203,34 +202,25 @@ public void handleUpdateChangeEvent(@Nonnull final MetadataChangeLog event) thro * * @param event the change event to be processed. */ - public void handleDeleteChangeEvent(@Nonnull final MetadataChangeLog event) { + private void handleDeleteChangeEvent(@Nonnull final MCLBatchItem event) { - final EntitySpec entitySpec = getEventEntitySpec(event); - final Urn urn = EntityKeyUtils.getUrnFromLog(event, entitySpec.getKeyAspectSpec()); - - if (!event.hasAspectName() || !event.hasPreviousAspectValue()) { - log.error("Previous aspect or aspect name is missing. Skipping aspect processing..."); - return; - } + final EntitySpec entitySpec = event.getEntitySpec(); + final Urn urn = event.getUrn(); AspectSpec aspectSpec = entitySpec.getAspectSpec(event.getAspectName()); if (aspectSpec == null) { throw new RuntimeException( String.format( "Failed to retrieve Aspect Spec for entity with name %s, aspect with name %s. Cannot update indices for MCL.", - event.getEntityType(), event.getAspectName())); + urn.getEntityType(), event.getAspectName())); } - RecordTemplate aspect = - GenericRecordUtils.deserializeAspect( - event.getPreviousAspectValue().getValue(), - event.getPreviousAspectValue().getContentType(), - aspectSpec); + RecordTemplate aspect = event.getAspect(); Boolean isDeletingKey = event.getAspectName().equals(entitySpec.getKeyAspectName()); if (!aspectSpec.isTimeseries()) { deleteSystemMetadata(urn, aspectSpec, isDeletingKey); - deleteGraphData(urn, aspectSpec, aspect, isDeletingKey, event); + deleteGraphData(urn, aspectSpec, aspect, isDeletingKey, event.getMetadataChangeLog()); deleteSearchData( _entitySearchService, urn, entitySpec.getName(), aspectSpec, aspect, isDeletingKey); } @@ -633,6 +623,7 @@ private EntitySpec getEventEntitySpec(@Nonnull final MetadataChangeLog event) { * @param systemEntityClient system entity client */ public void setSystemEntityClient(SystemEntityClient systemEntityClient) { + this.systemEntityClient = systemEntityClient; _searchDocumentTransformer.setEntityClient(systemEntityClient); } } 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 2113e5a04f3a2c..252ac2d633b98e 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java @@ -5,8 +5,8 @@ import com.linkedin.common.urn.UrnUtils; import com.linkedin.identity.CorpUserInfo; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.key.CorpUserKey; import java.util.HashMap; import java.util.LinkedList; @@ -26,27 +26,24 @@ 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( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) .aspect(aspect) + .auditStamp(AspectGenerationUtils.createAuditStamp()) .systemMetadata(AspectGenerationUtils.createSystemMetadata()) - .build(entityService.getEntityRegistry())); + .build(entityService.getEntityRegistry(), entityService.getSystemEntityClient())); } - entityService.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), - AspectGenerationUtils.createAuditStamp(), - true, - true); + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); return aspects; } @@ -61,25 +58,22 @@ 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( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) .aspect(aspect) + .auditStamp(AspectGenerationUtils.createAuditStamp()) .systemMetadata(AspectGenerationUtils.createSystemMetadata()) - .build(entityService.getEntityRegistry())); + .build(entityService.getEntityRegistry(), entityService.getSystemEntityClient())); } - entityService.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), - AspectGenerationUtils.createAuditStamp(), - true, - true); + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); return aspects; } @@ -94,7 +88,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); @@ -102,18 +96,15 @@ public static Map ingestChartInfoAspects( ChartInfo aspect = AspectGenerationUtils.createChartInfo(title, description); aspects.put(urn, aspect); items.add( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) .aspect(aspect) + .auditStamp(AspectGenerationUtils.createAuditStamp()) .systemMetadata(AspectGenerationUtils.createSystemMetadata()) - .build(entityService.getEntityRegistry())); + .build(entityService.getEntityRegistry(), entityService.getSystemEntityClient())); } - entityService.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), - AspectGenerationUtils.createAuditStamp(), - true, - true); + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); return aspects; } } 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 74c81ff2e86025..bad47f9acf507c 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 @@ -75,7 +75,7 @@ private void configureComponents() { _aspectDao, _mockProducer, _testEntityRegistry, - true, + false, _mockUpdateIndicesService, preProcessHooks); _retentionService = new CassandraRetentionService(_entityServiceImpl, session, 1000); 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 eeb014f7afdc2f..45e992576676d6 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 @@ -16,8 +16,8 @@ import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.ebean.EbeanAspectDao; import com.linkedin.metadata.entity.ebean.EbeanRetentionService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; import com.linkedin.metadata.models.registry.EntityRegistryException; @@ -73,7 +73,7 @@ public void setupTest() { _aspectDao, _mockProducer, _testEntityRegistry, - true, + false, _mockUpdateIndicesService, preProcessHooks); _retentionService = new EbeanRetentionService(_entityServiceImpl, server, 1000); @@ -116,28 +116,30 @@ public void testIngestListLatestAspects() throws AssertionError { // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) .aspect(writeAspect2) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) .aspect(writeAspect3) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // List aspects ListResult batch1 = @@ -183,28 +185,30 @@ public void testIngestListUrns() throws AssertionError { // Ingest CorpUserInfo Aspect #3 RecordTemplate writeAspect3 = AspectGenerationUtils.createCorpUserKey(entityUrn3); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) .aspect(writeAspect2) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) .aspect(writeAspect3) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // List aspects urns ListUrnsResult batch1 = _entityServiceImpl.listUrns(entityUrn1.getEntityType(), 0, 2); @@ -447,8 +451,14 @@ public void run() { auditStamp.setActor(Urn.createFromString(Constants.DATAHUB_ACTOR)); auditStamp.setTime(System.currentTimeMillis()); AspectsBatchImpl batch = - AspectsBatchImpl.builder().mcps(mcps, entityService.getEntityRegistry()).build(); - entityService.ingestProposal(batch, auditStamp, false); + AspectsBatchImpl.builder() + .mcps( + mcps, + auditStamp, + entityService.getEntityRegistry(), + entityService.getSystemEntityClient()) + .build(); + entityService.ingestProposal(batch, false); } } catch (InterruptedException | URISyntaxException ie) { throw new RuntimeException(ie); 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 f03811da35ea81..e9e67f4b2114ed 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 @@ -36,8 +36,8 @@ import com.linkedin.metadata.aspect.CorpUserAspect; import com.linkedin.metadata.aspect.CorpUserAspectArray; import com.linkedin.metadata.aspect.VersionedAspect; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; @@ -525,9 +525,8 @@ public void testReingestAspectsGetLatestAspects() throws Exception { _entityServiceImpl.ingestAspects(entityUrn, pairToIngest, TEST_AUDIT_STAMP, metadata1); - verify(_mockProducer, times(1)) - .produceMetadataChangeLog( - Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(restateChangeLog)); + verify(_mockProducer, times(0)) + .produceMetadataChangeLog(Mockito.any(), Mockito.any(), Mockito.any()); verifyNoMoreInteractions(_mockProducer); } @@ -840,34 +839,37 @@ public void testRollbackAspect() throws AssertionError { CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) .aspect(writeAspect2) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) .aspect(writeAspect3) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1Overwrite) .systemMetadata(metadata2) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // this should no-op since this run has been overwritten AspectRowSummary rollbackOverwrittenAspect = new AspectRowSummary(); @@ -916,28 +918,30 @@ public void testRollbackKey() throws AssertionError { CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(keyAspectName) .aspect(writeKey1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1Overwrite) .systemMetadata(metadata2) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // this should no-op since the key should have been written in the furst run AspectRowSummary rollbackKeyWithWrongRunId = new AspectRowSummary(); @@ -994,40 +998,44 @@ public void testRollbackUrn() throws AssertionError { CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(keyAspectName) .aspect(writeKey1) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) .aspect(writeAspect2) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) .aspect(writeAspect3) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) .aspect(writeAspect1Overwrite) .systemMetadata(metadata2) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // this should no-op since the key should have been written in the furst run AspectRowSummary rollbackKeyWithWrongRunId = new AspectRowSummary(); @@ -1057,16 +1065,16 @@ 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( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // Validate retrieval of CorpUserInfo Aspect #1 RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1090,14 +1098,14 @@ public void testIngestGetLatestAspect() throws AssertionError { items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect2) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata2) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // Validate retrieval of CorpUserInfo Aspect #2 RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1134,16 +1142,16 @@ 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( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1) + .auditStamp(TEST_AUDIT_STAMP) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // Validate retrieval of CorpUserInfo Aspect #1 EnvelopedAspect readAspect1 = @@ -1156,14 +1164,14 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect2) .systemMetadata(metadata2) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // Validate retrieval of CorpUserInfo Aspect #2 EnvelopedAspect readAspect2 = @@ -1199,16 +1207,16 @@ public void testIngestSameAspect() throws AssertionError { SystemMetadata metadata3 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-123", "run-456"); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // Validate retrieval of CorpUserInfo Aspect #1 RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1232,14 +1240,14 @@ public void testIngestSameAspect() throws AssertionError { items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect2) .systemMetadata(metadata2) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); // Validate retrieval of CorpUserInfo Aspect #2 RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); @@ -1258,8 +1266,8 @@ public void testIngestSameAspect() throws AssertionError { DataTemplateUtil.areEqual( EntityUtils.parseSystemMetadata(readAspectDao2.getSystemMetadata()), metadata3)); - verify(_mockProducer, times(1)) - .produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), mclCaptor.capture()); + verify(_mockProducer, times(0)) + .produceMetadataChangeLog(Mockito.any(), Mockito.any(), Mockito.any()); verifyNoMoreInteractions(_mockProducer); } @@ -1283,46 +1291,51 @@ public void testRetention() throws AssertionError { Status writeAspect2a = new Status().setRemoved(false); Status writeAspect2b = new Status().setRemoved(true); - List items = + List items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1a) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1b) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) .aspect(writeAspect2) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) .aspect(writeAspect2a) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) .aspect(writeAspect2b) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName, 1), writeAspect1); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1347,20 +1360,21 @@ public void testRetention() throws AssertionError { items = List.of( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) .aspect(writeAspect1c) .systemMetadata(metadata1) - .build(_testEntityRegistry), - UpsertBatchItem.builder() + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient()), + MCPUpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) .aspect(writeAspect2c) .systemMetadata(metadata1) - .build(_testEntityRegistry)); - _entityServiceImpl.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + .auditStamp(TEST_AUDIT_STAMP) + .build(_testEntityRegistry, _entityServiceImpl.getSystemEntityClient())); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), true, true); assertNull(_entityServiceImpl.getAspect(entityUrn, aspectName, 1)); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName2, 1), writeAspect2); 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 13236e302c2594..8d7701f6d174f8 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 @@ -116,6 +116,7 @@ public void setup() { _entityRegistry = new ConfigEntityRegistry( new DataSchemaFactory("com.datahub.test"), + List.of(), TestEntityProfile.class .getClassLoader() .getResourceAsStream("test-entity-registry.yml")); diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHookTest.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHookTest.java index 12c8ad7d0c69ba..a227668e22e9b4 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHookTest.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/UpdateIndicesHookTest.java @@ -19,6 +19,8 @@ import com.linkedin.dataset.DatasetLineageType; import com.linkedin.dataset.FineGrainedLineage; import com.linkedin.dataset.FineGrainedLineageArray; +import com.linkedin.dataset.FineGrainedLineageDownstreamType; +import com.linkedin.dataset.FineGrainedLineageUpstreamType; import com.linkedin.dataset.Upstream; import com.linkedin.dataset.UpstreamArray; import com.linkedin.dataset.UpstreamLineage; @@ -47,7 +49,9 @@ import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.SystemMetadata; +import com.linkedin.schema.NumberType; import com.linkedin.schema.SchemaField; +import com.linkedin.schema.SchemaFieldDataType; import java.net.URISyntaxException; import java.net.URLEncoder; import java.nio.charset.StandardCharsets; @@ -425,6 +429,9 @@ private EntityRegistry createMockEntityRegistry() { .thenReturn(entitySpec); Mockito.when(mockEntityRegistry.getEntitySpec(Constants.DATASET_ENTITY_NAME)) .thenReturn(entitySpec); + Mockito.when(mockEntityRegistry.getEntitySpec(SCHEMA_FIELD_ENTITY_NAME)).thenReturn(entitySpec); + Mockito.when(mockEntityRegistry.getEntitySpec(DATA_PLATFORM_ENTITY_NAME)) + .thenReturn(entitySpec); Mockito.when(entitySpec.getAspectSpec(Constants.INPUT_FIELDS_ASPECT_NAME)) .thenReturn(aspectSpec); Mockito.when(entitySpec.getAspectSpec(Constants.UPSTREAM_LINEAGE_ASPECT_NAME)) @@ -462,6 +469,8 @@ private MetadataChangeLog createUpstreamLineageMCL( UpstreamLineage upstreamLineage = new UpstreamLineage(); FineGrainedLineageArray fineGrainedLineages = new FineGrainedLineageArray(); FineGrainedLineage fineGrainedLineage = new FineGrainedLineage(); + fineGrainedLineage.setDownstreamType(FineGrainedLineageDownstreamType.FIELD); + fineGrainedLineage.setUpstreamType(FineGrainedLineageUpstreamType.DATASET); UrnArray upstreamUrns = new UrnArray(); upstreamUrns.add(upstreamUrn); fineGrainedLineage.setUpstreams(upstreamUrns); @@ -509,6 +518,9 @@ private MetadataChangeLog createInputFieldsMCL(Urn upstreamUrn, String downstrea inputField.setSchemaFieldUrn(upstreamUrn); SchemaField schemaField = new SchemaField(); schemaField.setFieldPath(downstreamFieldPath); + schemaField.setNativeDataType("int"); + schemaField.setType( + new SchemaFieldDataType().setType(SchemaFieldDataType.Type.create(new NumberType()))); inputField.setSchemaField(schemaField); inputFieldsArray.add(inputField); inputFields.setFields(inputFieldsArray); diff --git a/metadata-jobs/mae-consumer/src/test/resources/test-entity-registry.yml b/metadata-jobs/mae-consumer/src/test/resources/test-entity-registry.yml index 0a1afcb235c29c..c0af9e705d712b 100644 --- a/metadata-jobs/mae-consumer/src/test/resources/test-entity-registry.yml +++ b/metadata-jobs/mae-consumer/src/test/resources/test-entity-registry.yml @@ -1,4 +1,9 @@ entities: + - name: dataPlatform + category: core + keyAspect: dataPlatformKey + aspects: + - dataPlatformInfo - name: dataHubIngestionSource keyAspect: dataHubIngestionSourceKey aspects: @@ -21,5 +26,9 @@ entities: keyAspect: chartKey aspects: - domains + - name: schemaField + category: core + keyAspect: schemaFieldKey + aspects: [] events: - name: entityChangeEvent \ No newline at end of file diff --git a/metadata-models-custom/README.md b/metadata-models-custom/README.md index 7223451f31b58b..94399a67806a65 100644 --- a/metadata-models-custom/README.md +++ b/metadata-models-custom/README.md @@ -165,6 +165,252 @@ e.g. `datahub delete by-registry --registry-id=mycompany-dq-model:0.0.1 --hard` As you evolve the metadata model, you can publish new versions of the repository and deploy it into DataHub as well using the same steps outlined above. DataHub will check whether your new models are backwards compatible with the previous versioned model and decline loading models that are backwards incompatible. +### Custom Plugins + +Adding custom aspects to DataHub's existing data model is a powerful way to extend DataHub without forking the entire repo. Often however extending +just the data model is not enough and additional custom code might be required. For a few of these use cases a plugin framework was developed +to control how instances of custom aspects can be validated, mutated, and generate side effects (additional aspects). + +It should be noted that validation, mutation, and generation of the *core* DataHub aspects can lead to system corruption and should be used +by advanced users only. + +The `/config` endpoint documented above has been extended to return information on the instances of the various plugins as well as the classes +that were loaded for debugging purposes. + +```json +{ + "mycompany-dq-model": { + "0.0.0-dev": { + "plugins": { + "validatorCount": 1, + "mutationHookCount": 1, + "mcpSideEffectCount": 1, + "mclSideEffectCount": 1, + "validatorClasses": [ + "com.linkedin.metadata.aspect.plugins.validation.CustomDataQualityRulesValidator" + ], + "mutationHookClasses": [ + "com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMutator" + ], + "mcpSideEffectClasses": [ + "com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMCPSideEffect" + ], + "mclSideEffectClasses": [ + "com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMCLSideEffect" + ] + } + } + } +} +``` + +#### Custom Validators + +Custom aspects might require that instances of those aspects adhere to specific conditions or rules. These conditions could vary wildly depending on the use case however they could be as simple +as a null or range check for one or more fields within the custom aspect. Additionally, a lookup can be done on other aspects in order to validate the current aspect using the `AspectRetriever`. + +There are two integration points for validation. The first integration point is `on request` via the `validateProposedAspect` method where the aspect is validated independent of the previous value. This validation is performed +outside of any kind of database transaction and can perform more intensive checks without introducing added latency within a transaction. + +The second integration point for validation occurs within the database transaction using the `validatePreCommitAspect` and has access to the new aspect as well as the old aspect. See the included +example in [`CustomDataQualityRulesValidator.java`](src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java). + +Shown below is the interface to be implemented for a custom validator. + +```java +public class CustomDataQualityRulesValidator extends AspectPayloadValidator { + @Override + protected void validateProposedAspect( + @Nonnull ChangeType changeType, + @Nonnull Urn entityUrn, + @Nonnull AspectSpec aspectSpec, + @Nonnull RecordTemplate aspectPayload, + @Nonnull AspectRetriever aspectRetriever) + throws AspectValidationException { + + } + + @Override + protected void validatePreCommitAspect( + @Nonnull ChangeType changeType, + @Nonnull Urn entityUrn, + @Nonnull AspectSpec aspectSpec, + @Nullable RecordTemplate previousAspect, + @Nonnull RecordTemplate proposedAspect, + @Nonnull AspectRetriever aspectRetriever) + throws AspectValidationException { + + } +} +``` + +In order to register this custom validator add the following to your `entity-registry.yml` file. This will activate +the validator to run on upsert operations for any entity with the custom aspect `customDataQualityRules`. Alternatively separate +validators could be written within the context of specific entities, in this case simply specify the entity name instead of `*`. + +```yaml + +plugins: + aspectPayloadValidators: + - className: 'com.linkedin.metadata.aspect.plugins.validation.CustomDataQualityRulesValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: customDataQualityRules +``` + +#### Custom Mutator + +**Warning: This hook is for advanced users only. It is possible to corrupt data and render your system inoperable.** + +In this example, we want to make sure that the field type is always lowercase regardless of the string being provided +by ingestion. The full example can be found in [`CustomDataQualityMutator.java`](src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java). + +```java +public class CustomDataQualityRulesMutator extends MutationHook { + @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) { + + if (newAspectValue != null) { + DataQualityRules newDataQualityRules = new DataQualityRules(newAspectValue.data()); + + for (DataQualityRule rule : newDataQualityRules.getRules()) { + // Ensure uniform lowercase + if (!rule.getType().toLowerCase().equals(rule.getType())) { + rule.setType(rule.getType().toLowerCase()); + } + } + } + } +} +``` + +```yaml +plugins: + mutationHooks: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMutator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: customDataQualityRules +``` + +#### MetadataChangeProposal (MCP) Side Effects + +**Warning: This hook is for advanced users only. It is possible to corrupt data and render your system inoperable.** + +MCP Side Effects allow for the creation of new aspects based on an input aspect. + +Notes: +* MCPs will write aspects to the primary data store (SQL for example) as well as the search indices. +* Side effects in general must include a dependency on the `metadata-io` module since it deals with lower level storage primitives. + +The full example can be found in [`CustomDataQualityRulesMCPSideEffect.java`](src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java). + +```java +public class CustomDataQualityRulesMCPSideEffect extends MCPSideEffect { + @Override + protected Stream applyMCPSideEffect( + UpsertItem input, EntityRegistry entityRegistry, @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()) + .aspect(input.getAspect()) + .auditStamp(input.getAuditStamp()) + .systemMetadata(input.getSystemMetadata()) + .build(entityRegistry, aspectRetriever)); + } +} +``` + +```yaml +plugins: + mcpSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMCPSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: customDataQualityRules +``` + +#### MetadataChangeLog (MCL) Side Effects + +**Warning: This hook is for advanced users only. It is possible to corrupt data and render your system inoperable.** + +MCL Side Effects allow for the creation of new aspects based on an input aspect. In this example, we are generating a timeseries aspect to represent an event. When a DataQualityRule is created +or modified we'll record the actor, event type, and timestamp in a timeseries aspect index. + +Notes: +* MCLs are only persisted to the search indices which allows for adding to the search documents only. +* Dependency on the `metadata-io` module since it deals with lower level storage primitives. + +The full example can be found in [`CustomDataQualityRulesMCLSideEffect.java`](src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java). + +```java +public class CustomDataQualityRulesMCLSideEffect extends MCLSideEffect { + @Override + protected Stream applyMCLSideEffect( + @Nonnull MCLBatchItem input, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) { + + // 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(entityRegistry, aspectRetriever)); + + return timeseriesOptional.stream(); + } +} +``` + +```yaml +plugins: + mclSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMCLSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: 'dataset' + aspectName: customDataQualityRules +``` + ## The Future Hopefully this repository shows you how easily you can extend and customize DataHub's metadata model! diff --git a/metadata-models-custom/build.gradle b/metadata-models-custom/build.gradle index 3ac08dca7c0dbe..8bf9d3b2f491ed 100644 --- a/metadata-models-custom/build.gradle +++ b/metadata-models-custom/build.gradle @@ -14,7 +14,7 @@ buildscript { } plugins { - id 'base' + id 'java-library' id 'maven-publish' id 'pegasus' } @@ -25,15 +25,19 @@ if (project.hasProperty('projVersion')) { project.version = '0.0.0-dev' } - dependencies { implementation spec.product.pegasus.data // Uncomment these if you want to depend on models defined in core datahub - //implementation project(':li-utils') - //dataModel project(':li-utils') - //implementation project(':metadata-models') - //dataModel project(':metadata-models') - + // DataQualityRuleEvent in this example uses Urn and TimeseriesAspectBase + implementation project(':li-utils') + dataModel project(':li-utils') + implementation project(':metadata-models') + dataModel project(':metadata-models') + + // Required for custom code plugins + implementation project(':entity-registry') + // Required for MCL/MCP hooks + implementation project (':metadata-io') } def deployBaseDir = findProperty('pluginModelsDir') ?: file(project.gradle.gradleUserHomeDir.parent + "/.datahub/plugins/models") @@ -43,9 +47,10 @@ pegasus.main.generationModes = [PegasusGenerationMode.PEGASUS, PegasusGeneration task modelArtifact(type: Zip) { + dependsOn jar from(layout.buildDirectory.dir("libs")) { - include "*-data-template-*.jar" + include "*.jar" exclude "*-test-data-template-*.jar" into "libs" } diff --git a/metadata-models-custom/registry/entity-registry.yaml b/metadata-models-custom/registry/entity-registry.yaml index 2b501946ca858c..e6180172837e04 100644 --- a/metadata-models-custom/registry/entity-registry.yaml +++ b/metadata-models-custom/registry/entity-registry.yaml @@ -3,6 +3,40 @@ entities: - name: dataset aspects: - customDataQualityRules + - customDataQualityRuleEvent - name: container aspects: - - customDataQualityRules \ No newline at end of file + - customDataQualityRules +plugins: + aspectPayloadValidators: + - className: 'com.linkedin.metadata.aspect.plugins.validation.CustomDataQualityRulesValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: 'dataset' + aspectName: customDataQualityRules + mutationHooks: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMutator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: 'dataset' + aspectName: customDataQualityRules + mclSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMCLSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: 'dataset' + aspectName: customDataQualityRules + mcpSideEffects: + - className: 'com.linkedin.metadata.aspect.plugins.hooks.CustomDataQualityRulesMCPSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: 'dataset' + aspectName: customDataQualityRules \ No newline at end of file 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 new file mode 100644 index 00000000000000..a8735bae1521a9 --- /dev/null +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCLSideEffect.java @@ -0,0 +1,72 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +import com.linkedin.metadata.aspect.batch.MCLBatchItem; +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.models.registry.EntityRegistry; +import com.linkedin.metadata.utils.GenericRecordUtils; +import com.linkedin.mxe.MetadataChangeLog; +import com.mycompany.dq.DataQualityRuleEvent; +import java.util.Optional; +import java.util.stream.Stream; +import javax.annotation.Nonnull; + +public class CustomDataQualityRulesMCLSideEffect extends MCLSideEffect { + + public CustomDataQualityRulesMCLSideEffect(AspectPluginConfig config) { + super(config); + } + + @Override + protected Stream applyMCLSideEffect( + @Nonnull MCLBatchItem input, + @Nonnull EntityRegistry entityRegistry, + @Nonnull AspectRetriever aspectRetriever) { + + // 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(entityRegistry, aspectRetriever)); + + return timeseriesOptional.stream(); + } + + private Optional buildEvent(MetadataChangeLog originMCP) { + if (originMCP.getAspect() != null) { + DataQualityRuleEvent event = new DataQualityRuleEvent(); + if (event.getActor() != null) { + event.setActor(event.getActor()); + } + event.setEventTimestamp(originMCP.getSystemMetadata().getLastObserved()); + event.setTimestampMillis(originMCP.getSystemMetadata().getLastObserved()); + if (originMCP.getPreviousAspectValue() == null) { + event.setEventType("RuleCreated"); + } else { + event.setEventType("RuleUpdated"); + } + event.setAffectedDataset(originMCP.getEntityUrn()); + + return Optional.of(event); + } + + return Optional.empty(); + } +} 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 new file mode 100644 index 00000000000000..2c989725f4f9de --- /dev/null +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMCPSideEffect.java @@ -0,0 +1,33 @@ +package com.linkedin.metadata.aspect.plugins.hooks; + +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.metadata.aspect.batch.UpsertItem; +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.models.registry.EntityRegistry; +import java.util.stream.Stream; +import javax.annotation.Nonnull; + +public class CustomDataQualityRulesMCPSideEffect extends MCPSideEffect { + + public CustomDataQualityRulesMCPSideEffect(AspectPluginConfig aspectPluginConfig) { + super(aspectPluginConfig); + } + + @Override + protected Stream applyMCPSideEffect( + UpsertItem input, EntityRegistry entityRegistry, @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()) + .aspect(input.getAspect()) + .auditStamp(input.getAuditStamp()) + .systemMetadata(input.getSystemMetadata()) + .build(entityRegistry, 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 new file mode 100644 index 00000000000000..576ba3bf305f53 --- /dev/null +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/hooks/CustomDataQualityRulesMutator.java @@ -0,0 +1,45 @@ +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.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.mycompany.dq.DataQualityRule; +import com.mycompany.dq.DataQualityRules; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +public class CustomDataQualityRulesMutator extends MutationHook { + + public CustomDataQualityRulesMutator(AspectPluginConfig config) { + super(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) { + + if (newAspectValue != null) { + DataQualityRules newDataQualityRules = new DataQualityRules(newAspectValue.data()); + + for (DataQualityRule rule : newDataQualityRules.getRules()) { + // Ensure uniform lowercase + if (!rule.getType().toLowerCase().equals(rule.getType())) { + rule.setType(rule.getType().toLowerCase()); + } + } + } + } +} 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 new file mode 100644 index 00000000000000..667d7ad614a791 --- /dev/null +++ b/metadata-models-custom/src/main/java/com/linkedin/metadata/aspect/plugins/validation/CustomDataQualityRulesValidator.java @@ -0,0 +1,70 @@ +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.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.models.AspectSpec; +import com.mycompany.dq.DataQualityRule; +import com.mycompany.dq.DataQualityRules; +import java.util.Map; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +public class CustomDataQualityRulesValidator extends AspectPayloadValidator { + + public CustomDataQualityRulesValidator(AspectPluginConfig config) { + super(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."); + } + } + + @Override + protected void validatePreCommitAspect( + @Nonnull ChangeType changeType, + @Nonnull Urn entityUrn, + @Nonnull AspectSpec aspectSpec, + @Nullable RecordTemplate previousAspect, + @Nonnull RecordTemplate proposedAspect, + @Nonnull AspectRetriever aspectRetriever) + throws AspectValidationException { + + 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)); + + // 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()))); + } + } + } + } +} diff --git a/metadata-models-custom/src/main/pegasus/com/mycompany/dq/DataQualityRuleEvent.pdl b/metadata-models-custom/src/main/pegasus/com/mycompany/dq/DataQualityRuleEvent.pdl new file mode 100644 index 00000000000000..075c90898c43fa --- /dev/null +++ b/metadata-models-custom/src/main/pegasus/com/mycompany/dq/DataQualityRuleEvent.pdl @@ -0,0 +1,44 @@ +namespace com.mycompany.dq + +import com.linkedin.common.Urn +import com.linkedin.timeseries.TimeseriesAspectBase + +/** + * Operational info for an entity. + */ + @Aspect = { + "name": "customDataQualityRuleEvent", + "type": "timeseries" + } +record DataQualityRuleEvent includes TimeseriesAspectBase { + + /** + * Actor who issued this operation. + */ + @TimeseriesField = {} + actor: optional Urn + + /** + * Event type. + */ + @TimeseriesField = {} + eventType: string + + /** + * Which dataset was affected by this event. + */ + @TimeseriesFieldCollection = {"key":"datasetUrn"} + affectedDataset: optional Urn + + /** + * Custom properties + */ + customProperties: optional map[string, string] + + /** + * The time at which the event occurred. + */ + @TimeseriesField = {} + @Searchable = { "fieldType": "DATETIME", "fieldName": "eventTimestamp" } + eventTimestamp: long +} diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java index 2879f157843708..c631bede453642 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java @@ -12,7 +12,8 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.key.DataHubAccessTokenKey; import com.linkedin.metadata.utils.AuditStampUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -40,7 +41,7 @@ @Slf4j public class StatefulTokenService extends StatelessTokenService { - private final EntityService _entityService; + private final EntityService _entityService; private final LoadingCache _revokedTokenCache; private final String salt; @@ -48,7 +49,7 @@ public StatefulTokenService( @Nonnull final String signingKey, @Nonnull final String signingAlgorithm, @Nullable final String iss, - @Nonnull final EntityService entityService, + @Nonnull final EntityService entityService, @Nonnull final String salt) { super(signingKey, signingAlgorithm, iss); this._entityService = entityService; @@ -153,9 +154,12 @@ public String generateAccessToken( _entityService.ingestProposal( AspectsBatchImpl.builder() - .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .mcps( + proposalStream.collect(Collectors.toList()), + auditStamp, + _entityService.getEntityRegistry(), + _entityService.getSystemEntityClient()) .build(), - auditStamp, false); return accessToken; 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 e75ec0c0dc44a5..88a3f5749343b1 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 @@ -8,6 +8,7 @@ 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.models.registry.EntityRegistry; import com.linkedin.metadata.service.UpdateIndicesService; import com.linkedin.mxe.TopicConvention; @@ -35,7 +36,7 @@ public class EntityServiceFactory { "entityRegistry" }) @Nonnull - protected EntityService createInstance( + protected EntityService createInstance( Producer producer, TopicConvention convention, KafkaHealthChecker kafkaHealthChecker, diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java index 080845147766f8..c550fc161b6062 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java @@ -8,6 +8,7 @@ import com.linkedin.metadata.client.SystemJavaEntityClient; import com.linkedin.metadata.entity.DeleteEntityService; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.LineageSearchService; @@ -28,7 +29,7 @@ public class JavaEntityClientFactory { @Autowired @Qualifier("entityService") - private EntityService _entityService; + private EntityService _entityService; @Autowired @Qualifier("deleteEntityService") 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 b02541586de494..dae5f903d7d803 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,6 +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.spring.YamlPropertySourceFactory; import io.ebean.Database; import javax.annotation.Nonnull; @@ -23,7 +24,7 @@ public class RetentionServiceFactory { @Autowired @Qualifier("entityService") - private EntityService _entityService; + private EntityService _entityService; @Value("${RETENTION_APPLICATION_BATCH_SIZE:1000}") private Integer _batchSize; diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java index 9ccb2c3f650bd7..ff5d3f215d86bb 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java @@ -7,6 +7,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.key.DataHubUpgradeKey; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -20,7 +21,7 @@ @Slf4j public abstract class UpgradeStep implements BootstrapStep { - protected final EntityService _entityService; + protected final EntityService _entityService; private final String _version; private final String _upgradeId; private final Urn _upgradeUrn; 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 ae4baee37c8224..e2f0b70526af52 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 @@ -10,8 +10,8 @@ import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.AspectMigrationsDao; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.utils.DataPlatformInstanceUtils; import com.linkedin.metadata.utils.EntityKeyUtils; @@ -28,7 +28,7 @@ public class IngestDataPlatformInstancesStep implements BootstrapStep { private static final int BATCH_SIZE = 1000; - private final EntityService _entityService; + private final EntityService _entityService; private final AspectMigrationsDao _migrationsDao; @Override @@ -65,27 +65,28 @@ 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)) + .setTime(System.currentTimeMillis()); for (String urnStr : _migrationsDao.listAllUrns(start, start + BATCH_SIZE)) { Urn urn = Urn.createFromString(urnStr); Optional dataPlatformInstance = getDataPlatformInstance(urn); if (dataPlatformInstance.isPresent()) { items.add( - UpsertBatchItem.builder() + MCPUpsertBatchItem.builder() .urn(urn) .aspectName(DATA_PLATFORM_INSTANCE_ASPECT_NAME) .aspect(dataPlatformInstance.get()) - .build(_entityService.getEntityRegistry())); + .auditStamp(aspectAuditStamp) + .build( + _entityService.getEntityRegistry(), _entityService.getSystemEntityClient())); } } - final AuditStamp aspectAuditStamp = - new AuditStamp() - .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) - .setTime(System.currentTimeMillis()); - _entityService.ingestAspects( - AspectsBatchImpl.builder().items(items).build(), aspectAuditStamp, true, true); + _entityService.ingestAspects(AspectsBatchImpl.builder().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 db8cad65caa8a7..37eac6d5ec4708 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 @@ -12,8 +12,8 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import java.io.IOException; import java.net.URISyntaxException; import java.util.List; @@ -31,7 +31,7 @@ public class IngestDataPlatformsStep implements BootstrapStep { private static final String PLATFORM_ASPECT_NAME = "dataPlatformInfo"; - private final EntityService _entityService; + private final EntityService _entityService; @Override public String name() { @@ -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) @@ -82,20 +82,25 @@ public void execute() throws IOException, URISyntaxException { RecordUtils.toRecordTemplate( DataPlatformInfo.class, dataPlatform.get("aspect").toString()); - return UpsertBatchItem.builder() - .urn(urn) - .aspectName(PLATFORM_ASPECT_NAME) - .aspect(info) - .build(_entityService.getEntityRegistry()); + try { + return MCPUpsertBatchItem.builder() + .urn(urn) + .aspectName(PLATFORM_ASPECT_NAME) + .aspect(info) + .auditStamp( + new AuditStamp() + .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) + .setTime(System.currentTimeMillis())) + .build( + _entityService.getEntityRegistry(), + _entityService.getSystemEntityClient()); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } }) .collect(Collectors.toList()); _entityService.ingestAspects( - AspectsBatchImpl.builder().items(dataPlatformAspects).build(), - new AuditStamp() - .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) - .setTime(System.currentTimeMillis()), - true, - false); + AspectsBatchImpl.builder().items(dataPlatformAspects).build(), true, false); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java index f5a76b5f757783..fc1c82fc6d631d 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java @@ -11,7 +11,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -100,9 +100,12 @@ private void ingestOwnershipType( _entityService.ingestProposal( AspectsBatchImpl.builder() - .mcps(List.of(keyAspectProposal, proposal), _entityService.getEntityRegistry()) + .mcps( + List.of(keyAspectProposal, proposal), + auditStamp, + _entityService.getEntityRegistry(), + _entityService.getSystemEntityClient()) .build(), - auditStamp, false); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java index 2aa5fe4f46b65c..9b9feb8e146389 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java @@ -15,7 +15,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; @@ -205,11 +205,14 @@ private void ingestPolicy(final Urn urn, final DataHubPolicyInfo info) throws UR _entityService.ingestProposal( AspectsBatchImpl.builder() - .mcps(List.of(keyAspectProposal, proposal), _entityRegistry) + .mcps( + List.of(keyAspectProposal, proposal), + new AuditStamp() + .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) + .setTime(System.currentTimeMillis()), + _entityRegistry, + _entityService.getSystemEntityClient()) .build(), - new AuditStamp() - .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) - .setTime(System.currentTimeMillis()), false); } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRolesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRolesStep.java index f3c395abdfc3a5..67c3cca3384e39 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRolesStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRolesStep.java @@ -12,7 +12,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.utils.EntityKeyUtils; @@ -125,11 +125,14 @@ private void ingestRole( _entityService.ingestProposal( AspectsBatchImpl.builder() - .mcps(List.of(keyAspectProposal, proposal), _entityRegistry) + .mcps( + List.of(keyAspectProposal, proposal), + new AuditStamp() + .setActor(Urn.createFromString(SYSTEM_ACTOR)) + .setTime(System.currentTimeMillis()), + _entityRegistry, + _entityService.getSystemEntityClient()) .build(), - new AuditStamp() - .setActor(Urn.createFromString(SYSTEM_ACTOR)) - .setTime(System.currentTimeMillis()), false); _entityService.alwaysProduceMCLAsync( diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndices.java index 333928999f4539..919ba93c9213e7 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndices.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndices.java @@ -10,6 +10,7 @@ import com.linkedin.metadata.boot.UpgradeStep; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ListResult; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ExtraInfo; @@ -30,7 +31,8 @@ public class RestoreColumnLineageIndices extends UpgradeStep { private final EntityRegistry _entityRegistry; public RestoreColumnLineageIndices( - @Nonnull final EntityService entityService, @Nonnull final EntityRegistry entityRegistry) { + @Nonnull final EntityService entityService, + @Nonnull final EntityRegistry entityRegistry) { super(entityService, VERSION, UPGRADE_ID); _entityRegistry = Objects.requireNonNull(entityRegistry, "entityRegistry must not be null"); } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreDbtSiblingsIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreDbtSiblingsIndices.java index bb7ad80ef73d29..e2d367a034491b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreDbtSiblingsIndices.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreDbtSiblingsIndices.java @@ -13,6 +13,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.BootstrapStep; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.key.DataHubUpgradeKey; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -46,7 +47,7 @@ public class RestoreDbtSiblingsIndices implements BootstrapStep { private static final Integer BATCH_SIZE = 1000; private static final Integer SLEEP_SECONDS = 120; - private final EntityService _entityService; + private final EntityService _entityService; private final EntityRegistry _entityRegistry; @Override 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 976698f3032d22..41672a07a23898 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.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; @@ -96,7 +96,7 @@ public void testExecuteChecksKeySpecForAllUrns() throws Exception { @Test public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throws Exception { final EntityRegistry entityRegistry = getTestEntityRegistry(); - final EntityService entityService = mock(EntityService.class); + final EntityService entityService = mock(EntityService.class); final AspectMigrationsDao migrationsDao = mock(AspectMigrationsDao.class); final int countOfCorpUserEntities = 5; final int countOfChartEntities = 7; @@ -122,9 +122,8 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw item.getUrn().getEntityType().equals("chart") && item.getAspectName() .equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - && ((UpsertBatchItem) item).getAspect() + && ((MCPUpsertBatchItem) item).getAspect() instanceof DataPlatformInstance)), - any(), anyBoolean(), anyBoolean()); verify(entityService, times(0)) @@ -137,9 +136,8 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw item.getUrn().getEntityType().equals("chart") && item.getAspectName() .equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - && ((UpsertBatchItem) item).getAspect() + && ((MCPUpsertBatchItem) item).getAspect() instanceof DataPlatformInstance)), - any(), anyBoolean(), anyBoolean()); } diff --git a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java index 31cd3e6c69e503..fc935514f4138f 100644 --- a/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java +++ b/metadata-service/openapi-entity-servlet/src/main/java/io/datahubproject/openapi/delegates/EntityApiDelegateImpl.java @@ -120,7 +120,7 @@ public ResponseEntity> create(List body) { OpenApiEntitiesUtil.convertEntityToUpsert(b, _reqClazz, _entityRegistry) .stream()) .collect(Collectors.toList()); - _v1Controller.postEntities(aspects); + _v1Controller.postEntities(aspects, false); List responses = body.stream() .map(req -> OpenApiEntitiesUtil.convertToResponse(req, _respClazz, _entityRegistry)) @@ -129,7 +129,7 @@ public ResponseEntity> create(List body) { } public ResponseEntity delete(String urn) { - _v1Controller.deleteEntities(new String[] {urn}, false); + _v1Controller.deleteEntities(new String[] {urn}, false, false); return new ResponseEntity<>(HttpStatus.OK); } @@ -165,7 +165,7 @@ public ResponseEntity createAspect( UpsertAspectRequest aspectUpsert = OpenApiEntitiesUtil.convertAspectToUpsert(urn, body, reqClazz); _v1Controller.postEntities( - Stream.of(aspectUpsert).filter(Objects::nonNull).collect(Collectors.toList())); + Stream.of(aspectUpsert).filter(Objects::nonNull).collect(Collectors.toList()), false); AR response = OpenApiEntitiesUtil.convertToResponseAspect(body, respClazz); return ResponseEntity.ok(response); } @@ -185,7 +185,7 @@ public ResponseEntity headAspect(String urn, String aspect) { public ResponseEntity deleteAspect(String urn, String aspect) { _entityService.deleteAspect(urn, aspect, Map.of(), false); - _v1Controller.deleteEntities(new String[] {urn}, false); + _v1Controller.deleteEntities(new String[] {urn}, false, false); return new ResponseEntity<>(HttpStatus.OK); } 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 6e0fc5deb0b3ce..ff65db09c2682f 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,6 +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.utils.metrics.MetricUtils; import com.linkedin.util.Pair; import io.datahubproject.openapi.dto.RollbackRunResultDto; @@ -32,6 +33,7 @@ import java.util.Collections; import java.util.HashSet; import java.util.List; +import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; @@ -63,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; @@ -152,7 +154,8 @@ public ResponseEntity getEntities( @PostMapping(value = "/", produces = MediaType.APPLICATION_JSON_VALUE) public ResponseEntity> postEntities( - @RequestBody @Nonnull List aspectRequests) { + @RequestBody @Nonnull List aspectRequests, + @RequestParam(required = false, name = "async") Boolean async) { log.info("INGEST PROPOSAL proposal: {}", aspectRequests); Authentication authentication = AuthenticationContext.getAuthentication(); @@ -174,9 +177,14 @@ public ResponseEntity> postEntities( throw new UnauthorizedException(actorUrnStr + " is unauthorized to edit entities."); } + boolean asyncBool = + Objects.requireNonNullElseGet( + async, () -> Boolean.parseBoolean(System.getenv("ASYNC_INGEST_DEFAULT"))); List> responses = proposals.stream() - .map(proposal -> MappingUtil.ingestProposal(proposal, actorUrnStr, _entityService)) + .map( + proposal -> + MappingUtil.ingestProposal(proposal, actorUrnStr, _entityService, asyncBool)) .collect(Collectors.toList()); if (responses.stream().anyMatch(Pair::getSecond)) { return ResponseEntity.status(HttpStatus.CREATED) @@ -205,7 +213,8 @@ public ResponseEntity> deleteEntities( description = "Determines whether the delete will be soft or hard, defaults to true for soft delete") @RequestParam(value = "soft", defaultValue = "true") - boolean soft) { + boolean soft, + @RequestParam(required = false, name = "async") Boolean async) { Throwable exceptionally = null; try (Timer.Context context = MetricUtils.timer("deleteEntities").time()) { Authentication authentication = AuthenticationContext.getAuthentication(); @@ -250,6 +259,9 @@ public ResponseEntity> deleteEntities( .map(entityUrn -> MappingUtil.createStatusRemoval(entityUrn, _entityService)) .collect(Collectors.toList()); + boolean asyncBool = + Objects.requireNonNullElseGet( + async, () -> Boolean.parseBoolean(System.getenv("ASYNC_INGEST_DEFAULT"))); return ResponseEntity.ok( Collections.singletonList( RollbackRunResultDto.builder() @@ -262,7 +274,7 @@ public ResponseEntity> deleteEntities( .map( proposal -> MappingUtil.ingestProposal( - proposal, actorUrnStr, _entityService)) + proposal, actorUrnStr, _entityService, asyncBool)) .filter(Pair::getSecond) .map(Pair::getFirst) .map(urnString -> AspectRowSummary.builder().urn(urnString).build()) 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 370f2019a42ddc..3cc67e77ec27e1 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,6 +9,8 @@ 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.search.client.CachingEntitySearchService; import com.linkedin.util.Pair; import io.datahubproject.openapi.exception.UnauthorizedException; import io.datahubproject.openapi.generated.MetadataChangeProposal; @@ -16,6 +18,7 @@ import io.swagger.v3.oas.annotations.tags.Tag; import java.util.Collections; import java.util.List; +import java.util.Objects; import java.util.stream.Collectors; import javax.annotation.Nonnull; import lombok.RequiredArgsConstructor; @@ -30,6 +33,7 @@ import org.springframework.web.bind.annotation.PostMapping; import org.springframework.web.bind.annotation.RequestBody; import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RequestParam; import org.springframework.web.bind.annotation.RestController; @RestController @@ -41,7 +45,8 @@ 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; @@ -55,7 +60,8 @@ public void initBinder(WebDataBinder binder) { @PostMapping(value = "/", produces = MediaType.APPLICATION_JSON_VALUE) public ResponseEntity> postEntities( - @RequestBody @Nonnull List metadataChangeProposals) { + @RequestBody @Nonnull List metadataChangeProposals, + @RequestParam(required = false, name = "async") Boolean async) { log.info("INGEST PROPOSAL proposal: {}", metadataChangeProposals); Authentication authentication = AuthenticationContext.getAuthentication(); @@ -77,9 +83,14 @@ public ResponseEntity> postEntities( throw new UnauthorizedException(actorUrnStr + " is unauthorized to edit entities."); } + boolean asyncBool = + Objects.requireNonNullElseGet( + async, () -> Boolean.parseBoolean(System.getenv("ASYNC_INGEST_DEFAULT"))); List> responses = proposals.stream() - .map(proposal -> MappingUtil.ingestProposal(proposal, actorUrnStr, _entityService)) + .map( + proposal -> + MappingUtil.ingestProposal(proposal, actorUrnStr, _entityService, asyncBool)) .collect(Collectors.toList()); if (responses.stream().anyMatch(Pair::getSecond)) { return ResponseEntity.status(HttpStatus.CREATED) 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 0eb3e2d6b8c6ee..c87820465dc889 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 @@ -25,12 +25,13 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.Aspect; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.batch.AspectsBatch; import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.IngestResult; import com.linkedin.metadata.entity.RollbackRunResult; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; @@ -441,7 +442,9 @@ 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. Timer.Context context = MetricUtils.timer("postEntity").time(); final com.linkedin.common.AuditStamp auditStamp = @@ -462,10 +465,14 @@ public static Pair ingestProposal( AspectsBatch batch = AspectsBatchImpl.builder() - .mcps(proposalStream.collect(Collectors.toList()), entityService.getEntityRegistry()) + .mcps( + proposalStream.collect(Collectors.toList()), + auditStamp, + entityService.getEntityRegistry(), + entityService.getSystemEntityClient()) .build(); - Set proposalResult = entityService.ingestProposal(batch, auditStamp, false); + Set proposalResult = entityService.ingestProposal(batch, async); Urn urn = proposalResult.stream().findFirst().get().getUrn(); return new Pair<>( 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 06640ba13fb8b7..17be5a60816d30 100644 --- a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java +++ b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java @@ -217,7 +217,7 @@ public void testIngestDataset() { .build(); datasetAspects.add(glossaryTerms); - _entitiesController.postEntities(datasetAspects); + _entitiesController.postEntities(datasetAspects, false); } // @Test 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 91e9e4fd4671e1..fdf99cdc303c19 100644 --- a/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java +++ b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java @@ -80,8 +80,7 @@ public RecordTemplate getAspect(@Nonnull Urn urn, @Nonnull String aspectName, lo @Override public Map> getLatestEnvelopedAspects( - @Nonnull String entityName, @Nonnull Set urns, @Nonnull Set aspectNames) - throws URISyntaxException { + @Nonnull Set urns, @Nonnull Set aspectNames) throws URISyntaxException { Urn urn = UrnUtils.getUrn(DATASET_URN); Map> envelopedAspectMap = new HashMap<>(); List aspects = new ArrayList<>(); 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 598c252b4f7664..64ae3632c353a2 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 @@ -6,10 +6,12 @@ import com.linkedin.data.DataMap; import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.template.StringArray; +import com.linkedin.entity.Aspect; import com.linkedin.entity.Entity; import com.linkedin.entity.EntityResponse; import com.linkedin.metadata.aspect.EnvelopedAspect; import com.linkedin.metadata.aspect.VersionedAspect; +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.browse.BrowseResult; import com.linkedin.metadata.browse.BrowseResultV2; import com.linkedin.metadata.graph.LineageDirection; @@ -38,7 +40,7 @@ import javax.annotation.Nullable; // Consider renaming this to datahub client. -public interface EntityClient { +public interface EntityClient extends AspectRetriever { @Nullable public EntityResponse getV2( @@ -623,4 +625,12 @@ public void producePlatformEvent( public void rollbackIngestion(@Nonnull String runId, @Nonnull Authentication authentication) throws Exception; + + default Aspect getLatestAspectObject(@Nonnull Urn urn, @Nonnull String aspectName) + throws RemoteInvocationException, URISyntaxException { + return getV2(urn.getEntityType(), urn, Set.of(aspectName), null) + .getAspects() + .get(aspectName) + .getValue(); + } } diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java index babb290655d3d4..dfad20b5f52b29 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java @@ -2,7 +2,9 @@ import com.datahub.authentication.Authentication; import com.linkedin.common.urn.Urn; +import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; +import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever; import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.PlatformEvent; @@ -14,7 +16,7 @@ import javax.annotation.Nullable; /** Adds entity/aspect cache and assumes system authentication */ -public interface SystemEntityClient extends EntityClient { +public interface SystemEntityClient extends EntityClient, AspectRetriever { EntityClientCache getEntityClientCache(); @@ -98,4 +100,12 @@ default String ingestProposal( default void setWritable(boolean canWrite) throws RemoteInvocationException { setWritable(canWrite, getSystemAuthentication()); } + + default Aspect getLatestAspectObject(@Nonnull Urn urn, @Nonnull String aspectName) + throws RemoteInvocationException, URISyntaxException { + return getV2(urn.getEntityType(), urn, Set.of(aspectName), getSystemAuthentication()) + .getAspects() + .get(aspectName) + .getValue(); + } } 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 f14dc2e8b29184..c5b019e85e0c9d 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 @@ -19,10 +19,13 @@ 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.transactions.AspectsBatchImpl; -import com.linkedin.metadata.entity.transactions.AspectsBatch; +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; @@ -80,10 +83,10 @@ public class AspectResource extends CollectionResourceTaskTemplate _entityService; @VisibleForTesting - void setEntityService(EntityService entityService) { + void setEntityService(EntityService entityService) { _entityService = entityService; } @@ -242,33 +245,26 @@ public Task ingestProposal( final AuditStamp auditStamp = new AuditStamp().setTime(_clock.millis()).setActor(Urn.createFromString(actorUrnStr)); - return RestliUtil.toTask( - () -> { - log.debug("Proposal: {}", metadataChangeProposal); - try { - final AspectsBatch batch; - if (asyncBool) { - // if async we'll expand the getAdditionalChanges later, no need to do this early - batch = - AspectsBatchImpl.builder() - .mcps(List.of(metadataChangeProposal), _entityService.getEntityRegistry()) - .build(); - } else { - Stream proposalStream = - Stream.concat( - Stream.of(metadataChangeProposal), - AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService) - .stream()); + return RestliUtil.toTask(() -> { + log.debug("Proposal: {}", metadataChangeProposal); + try { + final AspectsBatch batch; + if (asyncBool) { + // if async we'll expand the getAdditionalChanges later, no need to do this early + batch = AspectsBatchImpl.builder() + .mcps(List.of(metadataChangeProposal), auditStamp, _entityService.getEntityRegistry(), _entityService.getSystemEntityClient()) + .build(); + } else { + Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), + AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService).stream()); - batch = - AspectsBatchImpl.builder() - .mcps( - proposalStream.collect(Collectors.toList()), - _entityService.getEntityRegistry()) - .build(); - } + batch = AspectsBatchImpl.builder() + .mcps(proposalStream.collect(Collectors.toList()), auditStamp, _entityService.getEntityRegistry(), _entityService.getSystemEntityClient()) + .build(); + } - Set results = _entityService.ingestProposal(batch, auditStamp, asyncBool); + Set results = + _entityService.ingestProposal(batch, asyncBool); IngestResult one = results.stream().findFirst().get(); diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java index ddf5efa5027ca9..dfd986c2ebea08 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java @@ -121,7 +121,7 @@ public class EntityResource extends CollectionResourceTaskTemplate _entityService; @Inject @Named("searchService") diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java index bf07d0eb9dd5b2..7c7c25ad3492c0 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/operations/Utils.java @@ -34,7 +34,7 @@ public static String restoreIndices( @Nullable Integer start, @Nullable Integer batchSize, @Nonnull Authorizer authorizer, - @Nonnull EntityService entityService) { + @Nonnull EntityService entityService) { Authentication authentication = AuthenticationContext.getAuthentication(); EntitySpec resourceSpec = null; if (StringUtils.isNotBlank(urn)) { 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 d6eeb1a01ac153..e3534875c6cd26 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.transactions.UpsertBatchItem; +import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -52,9 +52,8 @@ public void setup() { _entityRegistry = new MockEntityRegistry(); _updateIndicesService = mock(UpdateIndicesService.class); _preProcessHooks = mock(PreProcessHooks.class); - _entityService = - new EntityServiceImpl( - _aspectDao, _producer, _entityRegistry, false, _updateIndicesService, _preProcessHooks); + _entityService = new EntityServiceImpl(_aspectDao, _producer, _entityRegistry, false, + _updateIndicesService, _preProcessHooks); _authorizer = mock(Authorizer.class); _aspectResource.setAuthorizer(_authorizer); _aspectResource.setEntityService(_entityService); @@ -82,13 +81,13 @@ public void testAsyncDefaultAspects() throws URISyntaxException { reset(_producer, _aspectDao); - UpsertBatchItem req = - UpsertBatchItem.builder() + MCPUpsertBatchItem req = MCPUpsertBatchItem.builder() .urn(urn) .aspectName(mcp.getAspectName()) .aspect(mcp.getAspect()) + .auditStamp(new AuditStamp()) .metadataChangeProposal(mcp) - .build(_entityRegistry); + .build(_entityRegistry, _entityService.getSystemEntityClient()); when(_aspectDao.runInTransactionWithRetry(any(), any(), anyInt())) .thenReturn( List.of( diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java index eab482c7bab275..c4216962c134cd 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/AspectUtils.java @@ -10,9 +10,12 @@ import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.GenericAspect; +import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import java.util.Collections; import java.util.HashMap; @@ -35,7 +38,7 @@ private AspectUtils() {} public static List getAdditionalChanges( @Nonnull MetadataChangeProposal metadataChangeProposal, - @Nonnull EntityService entityService, + @Nonnull EntityService entityService, boolean onPrimaryKeyInsertOnly) { // No additional changes for unsupported operations @@ -174,4 +177,41 @@ public static AuditStamp getAuditStamp(Urn actor) { auditStamp.setActor(actor); return auditStamp; } + + public static AspectSpec validateAspect(MetadataChangeLog mcl, EntitySpec entitySpec) { + if (!mcl.hasAspectName() + || (!ChangeType.DELETE.equals(mcl.getChangeType()) && !mcl.hasAspect())) { + throw new UnsupportedOperationException( + String.format( + "Aspect and aspect name is required for create and update operations. changeType: %s entityName: %s hasAspectName: %s hasAspect: %s", + mcl.getChangeType(), entitySpec.getName(), mcl.hasAspectName(), mcl.hasAspect())); + } + + AspectSpec aspectSpec = entitySpec.getAspectSpec(mcl.getAspectName()); + + if (aspectSpec == null) { + throw new RuntimeException( + String.format( + "Unknown aspect %s for entity %s", mcl.getAspectName(), mcl.getEntityType())); + } + + return aspectSpec; + } + + public static AspectSpec validateAspect(MetadataChangeProposal mcp, EntitySpec entitySpec) { + if (!mcp.hasAspectName() || !mcp.hasAspect()) { + throw new UnsupportedOperationException( + "Aspect and aspect name is required for create and update operations"); + } + + AspectSpec aspectSpec = entitySpec.getAspectSpec(mcp.getAspectName()); + + if (aspectSpec == null) { + throw new RuntimeException( + String.format( + "Unknown aspect %s for entity %s", mcp.getAspectName(), mcp.getEntityType())); + } + + return aspectSpec; + } } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java index 3b71c698e0c9f9..2cd1aadf7665d6 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java @@ -46,7 +46,7 @@ @RequiredArgsConstructor public class DeleteEntityService { - private final EntityService _entityService; + private final EntityService _entityService; private final GraphService _graphService; private static final Integer ELASTIC_BATCH_DELETE_SLEEP_SEC = 5; 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 8654df4435cd60..89b0e5ba9a5587 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 @@ -12,9 +12,10 @@ import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.aspect.VersionedAspect; +import com.linkedin.metadata.aspect.batch.AspectsBatch; +import com.linkedin.metadata.aspect.batch.UpsertItem; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; @@ -33,7 +34,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -public interface EntityService { +public interface EntityService { /** * Just whether the entity/aspect exists @@ -119,15 +120,12 @@ Map getEntitiesVersionedV2( /** * Retrieves the latest aspects for the given set of urns as a list of enveloped aspects * - * @param entityName name of the entity to fetch * @param urns set of urns to fetch * @param aspectNames set of aspects to fetch * @return a map of {@link Urn} to {@link EnvelopedAspect} object */ Map> getLatestEnvelopedAspects( - // TODO: entityName is unused, can we remove this as a param? - @Nonnull String entityName, @Nonnull Set urns, @Nonnull Set aspectNames) - throws URISyntaxException; + @Nonnull Set urns, @Nonnull Set aspectNames) throws URISyntaxException; /** * Retrieves the latest aspects for the given set of urns as a list of enveloped aspects @@ -169,10 +167,7 @@ List ingestAspects( @Nullable SystemMetadata systemMetadata); List ingestAspects( - @Nonnull final AspectsBatch aspectsBatch, - @Nonnull final AuditStamp auditStamp, - boolean emitMCL, - boolean overwrite); + @Nonnull final AspectsBatch aspectsBatch, boolean emitMCL, boolean overwrite); /** * Ingests (inserts) a new version of an entity aspect & emits a {@link @@ -233,7 +228,7 @@ Pair, Boolean> alwaysProduceMCLAsync( @Nonnull AuditStamp auditStamp, @Nonnull final ChangeType changeType); - RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName); + // RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName); @Deprecated void ingestEntities( @@ -250,7 +245,7 @@ void ingestEntity( @Nonnull AuditStamp auditStamp, @Nonnull SystemMetadata systemMetadata); - void setRetentionService(RetentionService retentionService); + void setRetentionService(RetentionService retentionService); AspectSpec getKeyAspectSpec(@Nonnull final Urn urn); @@ -304,8 +299,7 @@ RollbackRunResult rollbackRun( RollbackRunResult rollbackWithConditions( List aspectRows, Map conditions, boolean hardDelete); - Set ingestProposal( - AspectsBatch aspectsBatch, AuditStamp auditStamp, final boolean async); + Set ingestProposal(AspectsBatch aspectsBatch, final boolean async); /** * If you have more than 1 proposal use the {AspectsBatch} method @@ -343,4 +337,8 @@ BrowsePathsV2 buildDefaultBrowsePathV2(final @Nonnull Urn urn, boolean useContai * @param systemEntityClient system entity client */ void setSystemEntityClient(SystemEntityClient systemEntityClient); + + SystemEntityClient getSystemEntityClient(); + + RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName); } diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java index 3e72a763fb17cf..d3f8b507bb14ac 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java @@ -1,7 +1,7 @@ package com.linkedin.metadata.entity; import com.linkedin.common.urn.Urn; -import com.linkedin.metadata.entity.transactions.AbstractBatchItem; +import com.linkedin.metadata.aspect.batch.BatchItem; import lombok.Builder; import lombok.Value; @@ -9,7 +9,7 @@ @Value public class IngestResult { Urn urn; - AbstractBatchItem request; + BatchItem request; boolean publishedMCL; boolean processedMCL; boolean publishedMCP; 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 51519f48bd975e..ae33b72010ce2a 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 @@ -7,9 +7,10 @@ import com.linkedin.data.template.RecordTemplate; 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.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; -import com.linkedin.metadata.entity.transactions.AspectsBatch; import com.linkedin.metadata.key.DataHubRetentionKey; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -30,16 +31,16 @@ import lombok.Value; /** - * Service coupled with an {@link EntityServiceImpl} to handle aspect record retention. + * Service coupled with an {@link EntityService} to handle aspect record retention. * *

TODO: This class is abstract with storage-specific implementations. It'd be nice to pull - * storage and retention concerns apart, let (into {@link AspectDao}) deal with storage, and merge - * all retention concerns into a single class. + * 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(); + protected abstract EntityService getEntityService(); /** * Fetch retention policies given the entityName and aspectName Uses the entity service to fetch @@ -120,13 +121,14 @@ public boolean setRetention( new AuditStamp() .setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)) .setTime(System.currentTimeMillis()); - AspectsBatch batch = buildAspectsBatch(List.of(keyProposal, aspectProposal)); + AspectsBatch batch = buildAspectsBatch(List.of(keyProposal, aspectProposal), auditStamp); - return getEntityService().ingestProposal(batch, auditStamp, false).stream() + return getEntityService().ingestProposal(batch, false).stream() .anyMatch(IngestResult::isSqlCommitted); } - protected abstract AspectsBatch buildAspectsBatch(List mcps); + protected abstract AspectsBatch buildAspectsBatch( + List mcps, @Nonnull AuditStamp auditStamp); /** * Delete the retention policy set for given entity and aspect. 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 a10c90bc453204..515e08646f9ed3 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.entity.transactions.AbstractBatchItem; +import com.linkedin.metadata.aspect.batch.UpsertItem; 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; - AbstractBatchItem request; + UpsertItem request; RecordTemplate oldValue; RecordTemplate newValue; SystemMetadata oldSystemMetadata; diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java deleted file mode 100644 index 155385c62ecef5..00000000000000 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java +++ /dev/null @@ -1,94 +0,0 @@ -package com.linkedin.metadata.entity.transactions; - -import static com.linkedin.metadata.Constants.*; - -import com.linkedin.common.urn.Urn; -import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; -import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.models.registry.template.AspectTemplateEngine; -import com.linkedin.mxe.MetadataChangeProposal; -import com.linkedin.mxe.SystemMetadata; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; - -public abstract class AbstractBatchItem { - // urn an urn associated with the new aspect - public abstract Urn getUrn(); - - // aspectName name of the aspect being inserted - public abstract String getAspectName(); - - public abstract SystemMetadata getSystemMetadata(); - - public abstract ChangeType getChangeType(); - - public abstract EntitySpec getEntitySpec(); - - public abstract AspectSpec getAspectSpec(); - - public abstract MetadataChangeProposal getMetadataChangeProposal(); - - public abstract void validateUrn(EntityRegistry entityRegistry, Urn urn); - - @Nonnull - protected static SystemMetadata generateSystemMetadataIfEmpty( - @Nullable SystemMetadata systemMetadata) { - if (systemMetadata == null) { - systemMetadata = new SystemMetadata(); - systemMetadata.setRunId(DEFAULT_RUN_ID); - systemMetadata.setLastObserved(System.currentTimeMillis()); - } - return systemMetadata; - } - - protected static AspectSpec validateAspect(MetadataChangeProposal mcp, EntitySpec entitySpec) { - if (!mcp.hasAspectName() || !mcp.hasAspect()) { - throw new UnsupportedOperationException( - "Aspect and aspect name is required for create and update operations"); - } - - AspectSpec aspectSpec = entitySpec.getAspectSpec(mcp.getAspectName()); - - if (aspectSpec == null) { - throw new RuntimeException( - String.format( - "Unknown aspect %s for entity %s", mcp.getAspectName(), mcp.getEntityType())); - } - - return aspectSpec; - } - - /** - * Validates that a change type is valid for the given aspect - * - * @param changeType - * @param aspectSpec - * @return - */ - protected static boolean isValidChangeType(ChangeType changeType, AspectSpec aspectSpec) { - if (aspectSpec.isTimeseries()) { - // Timeseries aspects only support UPSERT - return ChangeType.UPSERT.equals(changeType); - } else { - if (ChangeType.PATCH.equals(changeType)) { - return supportsPatch(aspectSpec); - } else { - return ChangeType.UPSERT.equals(changeType); - } - } - } - - protected 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, - // specifically having array based fields that require merging without specifying merge - // behavior can get into bad states - throw new UnsupportedOperationException( - "Aspect: " + aspectSpec.getName() + " does not currently support patch " + "operations."); - } - return true; - } -} diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java deleted file mode 100644 index 4f2cf6073bdacb..00000000000000 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java +++ /dev/null @@ -1,26 +0,0 @@ -package com.linkedin.metadata.entity.transactions; - -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; - -public interface AspectsBatch { - List getItems(); - - default boolean containsDuplicateAspects() { - return getItems().stream() - .map(i -> String.format("%s_%s", i.getClass().getName(), i.hashCode())) - .distinct() - .count() - != getItems().size(); - } - - default Map> getUrnAspectsMap() { - return getItems().stream() - .map(aspect -> Map.entry(aspect.getUrn().toString(), aspect.getAspectName())) - .collect( - Collectors.groupingBy( - Map.Entry::getKey, Collectors.mapping(Map.Entry::getValue, Collectors.toSet()))); - } -} diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java index b0f42231b27f37..81bfcaab74ddb8 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/SystemMetadataUtils.java @@ -2,6 +2,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.mxe.SystemMetadata; +import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; @Slf4j @@ -14,4 +15,9 @@ public static SystemMetadata createDefaultSystemMetadata() { .setRunId(Constants.DEFAULT_RUN_ID) .setLastObserved(System.currentTimeMillis()); } + + public static SystemMetadata generateSystemMetadataIfEmpty( + @Nullable SystemMetadata systemMetadata) { + return systemMetadata == null ? createDefaultSystemMetadata() : systemMetadata; + } }