From 1b79142d9ea6563fa1723d1beb4e341d1fc898a3 Mon Sep 17 00:00:00 2001 From: david-leifker <114954101+david-leifker@users.noreply.github.com> Date: Sat, 2 Sep 2023 19:25:44 -0500 Subject: [PATCH 1/2] feat(EntityService): batched transactions and ebean updates (#8456) --- build.gradle | 6 +- datahub-frontend/build.gradle | 2 +- .../linkedin/datahub/graphql/TestUtils.java | 46 +- .../DeleteAssertionResolverTest.java | 10 +- .../BatchUpdateSoftDeletedResolverTest.java | 25 +- .../BatchUpdateDeprecationResolverTest.java | 26 +- .../UpdateDeprecationResolverTest.java | 2 +- .../domain/BatchSetDomainResolverTest.java | 36 +- .../domain/CreateDomainResolverTest.java | 6 +- .../domain/SetDomainResolverTest.java | 10 +- .../domain/UnsetDomainResolverTest.java | 8 +- .../embed/UpdateEmbedResolverTest.java | 27 +- .../glossary/AddRelatedTermsResolverTest.java | 6 +- .../CreateGlossaryNodeResolverTest.java | 7 +- .../CreateGlossaryTermResolverTest.java | 9 +- .../DeleteGlossaryEntityResolverTest.java | 5 +- .../RemoveRelatedTermsResolverTest.java | 12 +- .../glossary/UpdateNameResolverTest.java | 14 +- .../UpdateParentNodeResolverTest.java | 14 +- .../mutate/UpdateUserSettingResolverTest.java | 4 +- .../owner/AddOwnersResolverTest.java | 15 +- .../owner/BatchAddOwnersResolverTest.java | 19 +- .../owner/BatchRemoveOwnersResolverTest.java | 17 +- .../resolvers/tag/AddTagsResolverTest.java | 15 +- .../tag/BatchAddTagsResolverTest.java | 33 +- .../tag/BatchRemoveTagsResolverTest.java | 34 +- .../resolvers/tag/CreateTagResolverTest.java | 6 +- .../tag/SetTagColorResolverTest.java | 8 +- .../resolvers/term/AddTermsResolverTest.java | 26 +- .../term/BatchAddTermsResolverTest.java | 19 +- .../term/BatchRemoveTermsResolverTest.java | 17 +- .../test/resources/test-entity-registry.yaml | 295 ++++ datahub-upgrade/build.gradle | 2 +- .../upgrade/config/NoCodeCleanupConfig.java | 4 +- .../upgrade/config/NoCodeUpgradeConfig.java | 4 +- .../upgrade/config/RestoreBackupConfig.java | 4 +- .../upgrade/config/RestoreIndicesConfig.java | 4 +- .../upgrade/nocode/CreateAspectTableStep.java | 6 +- .../upgrade/nocode/DataMigrationStep.java | 20 +- .../datahub/upgrade/nocode/NoCodeUpgrade.java | 8 +- .../nocode/RemoveAspectV2TableStep.java | 8 +- .../nocode/UpgradeQualificationStep.java | 8 +- .../nocodecleanup/DeleteAspectTableStep.java | 8 +- .../nocodecleanup/NoCodeCleanupUpgrade.java | 8 +- .../NoCodeUpgradeQualificationStep.java | 6 +- .../restorebackup/ClearAspectV2TableStep.java | 6 +- .../upgrade/restorebackup/RestoreBackup.java | 6 +- .../restorebackup/RestoreStorageStep.java | 5 +- .../restoreindices/RestoreIndices.java | 6 +- .../upgrade/restoreindices/SendMAEStep.java | 6 +- ...pgradeCliApplicationTestConfiguration.java | 4 +- docs/deploy/environment-vars.md | 2 +- .../dao/producer/KafkaEventProducer.java | 53 +- metadata-io/build.gradle | 2 + .../metadata/client/JavaEntityClient.java | 16 +- .../linkedin/metadata/entity/AspectDao.java | 83 +- .../metadata/entity/EntityServiceImpl.java | 1275 +++++++---------- .../linkedin/metadata/entity/EntityUtils.java | 44 +- .../entity/cassandra/CassandraAspectDao.java | 101 +- .../cassandra/CassandraRetentionService.java | 45 +- .../ebean/AspectStorageValidationUtil.java | 18 +- .../metadata/entity/ebean/EbeanAspectDao.java | 254 ++-- .../entity/ebean/EbeanRetentionService.java | 147 +- .../ebean/transactions/AspectsBatchImpl.java | 67 + .../ebean/transactions/PatchBatchItem.java | 188 +++ .../ebean/transactions/UpsertBatchItem.java | 173 +++ .../entity/validation/ValidationUtils.java | 38 + .../metadata/event/EventProducer.java | 25 - .../metadata/AspectIngestionUtils.java | 31 +- .../linkedin/metadata/AspectUtilsTest.java | 4 +- .../com/linkedin/metadata/ESTestUtils.java | 5 +- .../com/linkedin/metadata/EbeanTestUtils.java | 10 +- .../entity/EbeanAspectMigrationsDaoTest.java | 4 +- .../entity/EbeanEntityServiceTest.java | 249 +++- .../metadata/entity/EntityServiceTest.java | 363 +++-- .../timeline/EbeanTimelineServiceTest.java | 4 +- .../java/io/datahub/test/DataGenerator.java | 359 +++++ metadata-jobs/mae-consumer-job/build.gradle | 2 +- ...eConsumerApplicationTestConfiguration.java | 4 +- metadata-jobs/mce-consumer-job/build.gradle | 2 +- ...eConsumerApplicationTestConfiguration.java | 4 +- .../token/StatefulTokenService.java | 13 +- .../DataHubTokenAuthenticatorTest.java | 1 + .../token/StatefulTokenServiceTest.java | 6 + .../test/resources/test-entity-registry.yaml | 10 + .../factory/entity/EbeanServerFactory.java | 6 +- .../entity/EntityAspectDaoFactory.java | 4 +- .../EntityAspectMigrationsDaoFactory.java | 4 +- .../entity/RetentionServiceFactory.java | 4 +- .../boot/steps/IndexDataPlatformsStep.java | 18 +- .../IngestDataPlatformInstancesStep.java | 26 +- .../boot/steps/IngestDataPlatformsStep.java | 62 +- .../boot/steps/IngestOwnershipTypesStep.java | 9 +- .../boot/steps/IngestPoliciesStep.java | 13 +- .../metadata/boot/steps/IngestRolesStep.java | 13 +- .../boot/steps/IngestRootUserStep.java | 9 +- .../steps/RestoreColumnLineageIndices.java | 30 +- .../boot/steps/RestoreDbtSiblingsIndices.java | 17 +- .../boot/steps/RestoreGlossaryIndices.java | 31 +- .../IngestDataPlatformInstancesStepTest.java | 25 +- .../RestoreColumnLineageIndicesTest.java | 26 +- .../steps/RestoreGlossaryIndicesTest.java | 25 +- .../test/resources/test-entity-registry.yaml | 6 + .../openapi/util/MappingUtil.java | 20 +- .../java/entities/EntitiesControllerTest.java | 10 + .../src/test/java/mock/MockEntityService.java | 22 +- .../resources/entity/AspectResource.java | 42 +- .../resources/entity/AspectResourceTest.java | 36 +- .../linkedin/metadata/entity/AspectUtils.java | 51 +- .../metadata/entity/DeleteEntityService.java | 4 +- .../metadata/entity/EntityService.java | 113 +- .../metadata/entity/IngestResult.java | 18 + .../metadata/entity/RetentionService.java | 88 +- .../metadata/entity/UpdateAspectResult.java | 8 + .../transactions/AbstractBatchItem.java | 92 ++ .../entity/transactions/AspectsBatch.java | 22 + metadata-service/war/build.gradle | 2 +- .../src/main/java/mock/MockEntitySpec.java | 2 +- settings.gradle | 1 - smoke-test/run-quickstart.sh | 2 +- smoke-test/tests/tags-and-terms/data.json | 39 + 121 files changed, 3582 insertions(+), 1822 deletions(-) create mode 100644 datahub-graphql-core/src/test/resources/test-entity-registry.yaml create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java create mode 100644 metadata-io/src/test/java/io/datahub/test/DataGenerator.java create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java diff --git a/build.gradle b/build.gradle index f73fe42d45956..e12d520e12de6 100644 --- a/build.gradle +++ b/build.gradle @@ -18,6 +18,7 @@ buildscript { ext.hadoop3Version = '3.3.5' ext.kafkaVersion = '2.3.0' ext.hazelcastVersion = '5.3.1' + ext.ebeanVersion = '12.16.1' ext.docker_registry = 'linkedin' @@ -86,8 +87,9 @@ project.ext.externalDependency = [ 'dgraph4j' : 'io.dgraph:dgraph4j:21.03.1', 'dropwizardMetricsCore': 'io.dropwizard.metrics:metrics-core:4.2.3', 'dropwizardMetricsJmx': 'io.dropwizard.metrics:metrics-jmx:4.2.3', - 'ebean': 'io.ebean:ebean:11.33.3', - 'ebeanAgent': 'io.ebean:ebean-agent:11.27.1', + 'ebean': 'io.ebean:ebean:' + ebeanVersion, + 'ebeanAgent': 'io.ebean:ebean-agent:' + ebeanVersion, + 'ebeanDdl': 'io.ebean:ebean-ddl-generator:' + ebeanVersion, 'elasticSearchRest': 'org.elasticsearch.client:elasticsearch-rest-high-level-client:' + elasticsearchVersion, 'elasticSearchTransport': 'org.elasticsearch.client:transport:' + elasticsearchVersion, 'findbugsAnnotations': 'com.google.code.findbugs:annotations:3.0.1', diff --git a/datahub-frontend/build.gradle b/datahub-frontend/build.gradle index cf1f8ca3cdd84..fdf13bac0accc 100644 --- a/datahub-frontend/build.gradle +++ b/datahub-frontend/build.gradle @@ -96,7 +96,7 @@ task unversionZip(type: Copy, dependsOn: [':datahub-web-react:build', dist]) { into "${buildDir}/docker/" rename "datahub-frontend-${version}.zip", "datahub-frontend.zip" } -tasks.getByName("docker").dependsOn(unversionZip) +tasks.getByPath(":datahub-frontend:docker").dependsOn(unversionZip) task cleanLocalDockerImages { doLast { 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 0718cc5b01d7e..272a93fa1989c 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 @@ -9,12 +9,24 @@ import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.UrnUtils; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.mxe.MetadataChangeProposal; import org.mockito.Mockito; +import java.util.List; + public class TestUtils { + public static EntityService getMockEntityService() { + EntityRegistry registry = new ConfigEntityRegistry(TestUtils.class.getResourceAsStream("/test-entity-registry.yaml")); + EntityService mockEntityService = Mockito.mock(EntityService.class); + Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry); + return mockEntityService; + } + public static QueryContext getMockAllowContext() { return getMockAllowContext("urn:li:corpuser:test"); } @@ -88,25 +100,47 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ } public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { + verifyIngestProposal(mockService, numberOfInvocations, List.of(proposal)); + } + + public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, List proposals) { + AspectsBatchImpl batch = AspectsBatchImpl.builder() + .mcps(proposals, mockService.getEntityRegistry()) + .build(); + Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( + Mockito.eq(batch), + Mockito.any(AuditStamp.class), + Mockito.eq(false) + ); + } + + public static void verifySingleIngestProposal(EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( - Mockito.eq(proposal), - Mockito.any(AuditStamp.class), - Mockito.eq(false) + Mockito.eq(proposal), + Mockito.any(AuditStamp.class), + Mockito.eq(false) ); } public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( - Mockito.any(MetadataChangeProposal.class), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); } + public static void verifySingleIngestProposal(EntityService mockService, int numberOfInvocations) { + Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( + Mockito.any(MetadataChangeProposal.class), + Mockito.any(AuditStamp.class), + Mockito.eq(false) + ); + } + public static void verifyNoIngestProposal(EntityService mockService) { Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), - Mockito.any(AuditStamp.class), Mockito.anyBoolean()); + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } private TestUtils() { } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolverTest.java index 42d2b864309ed..8afec0a889577 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolverTest.java @@ -31,7 +31,7 @@ public class DeleteAssertionResolverTest { public void testGetSuccess() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ASSERTION_URN))).thenReturn(true); Mockito.when(mockService.getAspect( Urn.createFromString(TEST_ASSERTION_URN), @@ -78,7 +78,7 @@ public void testGetSuccess() throws Exception { public void testGetSuccessNoAssertionInfoFound() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ASSERTION_URN))).thenReturn(true); Mockito.when(mockService.getAspect( Urn.createFromString(TEST_ASSERTION_URN), @@ -117,7 +117,7 @@ public void testGetSuccessAssertionAlreadyRemoved() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ASSERTION_URN))).thenReturn(false); DeleteAssertionResolver resolver = new DeleteAssertionResolver(mockClient, mockService); @@ -151,7 +151,7 @@ public void testGetSuccessAssertionAlreadyRemoved() throws Exception { public void testGetUnauthorized() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ASSERTION_URN))).thenReturn(true); Mockito.when(mockService.getAspect( Urn.createFromString(TEST_ASSERTION_URN), @@ -189,7 +189,7 @@ public void testGetEntityClientException() throws Exception { Mockito.any(), Mockito.any(Authentication.class)); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ASSERTION_URN))).thenReturn(true); DeleteAssertionResolver resolver = new DeleteAssertionResolver(mockClient, mockService); 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 7e549f201c2c7..bae6f27a854bc 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 @@ -11,8 +11,11 @@ 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.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; + +import java.util.List; import java.util.concurrent.CompletionException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -29,7 +32,7 @@ public class BatchUpdateSoftDeletedResolverTest { @Test public void testGetSuccessNoExistingStatus() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -61,20 +64,17 @@ public void testGetSuccessNoExistingStatus() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), STATUS_ASPECT_NAME, newStatus); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), STATUS_ASPECT_NAME, newStatus); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test public void testGetSuccessExistingStatus() throws Exception { final Status originalStatus = new Status().setRemoved(true); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -105,18 +105,15 @@ public void testGetSuccessExistingStatus() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), STATUS_ASPECT_NAME, newStatus); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), STATUS_ASPECT_NAME, newStatus); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -148,7 +145,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchUpdateSoftDeletedResolver resolver = new BatchUpdateSoftDeletedResolver(mockService); @@ -166,10 +163,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 634fd59a857a6..ce5a02bb573e1 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 @@ -12,8 +12,11 @@ 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.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; + +import java.util.List; import java.util.concurrent.CompletionException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -30,7 +33,7 @@ public class BatchUpdateDeprecationResolverTest { @Test public void testGetSuccessNoExistingDeprecation() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -68,12 +71,10 @@ public void testGetSuccessNoExistingDeprecation() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), DEPRECATION_ASPECT_NAME, newDeprecation); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), DEPRECATION_ASPECT_NAME, newDeprecation); - verifyIngestProposal(mockService, 1, proposal2); + + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test @@ -83,7 +84,7 @@ public void testGetSuccessExistingDeprecation() throws Exception { .setNote("") .setActor(UrnUtils.getUrn("urn:li:corpuser:test")); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -120,18 +121,15 @@ public void testGetSuccessExistingDeprecation() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), DEPRECATION_ASPECT_NAME, newDeprecation); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), DEPRECATION_ASPECT_NAME, newDeprecation); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -164,7 +162,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService); @@ -183,10 +181,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolverTest.java index cf11ecf865085..5d30ae08d6dea 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/UpdateDeprecationResolverTest.java @@ -58,7 +58,7 @@ public void testGetSuccessNoExistingDeprecation() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); UpdateDeprecationResolver resolver = new UpdateDeprecationResolver(mockClient, 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 9241661ccfc87..8cd3c71a21555 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 @@ -14,9 +14,12 @@ 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.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; + +import java.util.List; import java.util.concurrent.CompletionException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -35,7 +38,7 @@ public class BatchSetDomainResolverTest { @Test public void testGetSuccessNoExistingDomains() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -74,13 +77,10 @@ public void testGetSuccessNoExistingDomains() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), DOMAINS_ASPECT_NAME, newDomains); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), DOMAINS_ASPECT_NAME, newDomains); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_DOMAIN_2_URN)) @@ -92,7 +92,7 @@ public void testGetSuccessExistingDomains() throws Exception { final Domains originalDomain = new Domains().setDomains(new UrnArray(ImmutableList.of( Urn.createFromString(TEST_DOMAIN_1_URN)))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -135,13 +135,10 @@ public void testGetSuccessExistingDomains() throws Exception { proposal1.setAspectName(Constants.DOMAINS_ASPECT_NAME); proposal1.setAspect(GenericRecordUtils.serializeAspect(newDomains)); proposal1.setChangeType(ChangeType.UPSERT); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), DOMAINS_ASPECT_NAME, newDomains); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_DOMAIN_2_URN)) @@ -153,7 +150,7 @@ public void testGetSuccessUnsetDomains() throws Exception { final Domains originalDomain = new Domains().setDomains(new UrnArray(ImmutableList.of( Urn.createFromString(TEST_DOMAIN_1_URN)))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -189,18 +186,15 @@ public void testGetSuccessUnsetDomains() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), DOMAINS_ASPECT_NAME, newDomains); - - verifyIngestProposal(mockService, 1, proposal1); - - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), + final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), DOMAINS_ASPECT_NAME, newDomains); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test public void testGetFailureDomainDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -228,7 +222,7 @@ public void testGetFailureDomainDoesNotExist() throws Exception { @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -262,7 +256,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchSetDomainResolver resolver = new BatchSetDomainResolver(mockService); @@ -281,10 +275,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); BatchSetDomainResolver resolver = new BatchSetDomainResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java index 9343e5d772826..8c19f1dc3eb34 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/CreateDomainResolverTest.java @@ -40,7 +40,7 @@ public class CreateDomainResolverTest { public void testGetSuccess() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); CreateDomainResolver resolver = new CreateDomainResolver(mockClient, mockService); // Execute resolver @@ -76,7 +76,7 @@ public void testGetSuccess() throws Exception { public void testGetUnauthorized() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); CreateDomainResolver resolver = new CreateDomainResolver(mockClient, mockService); // Execute resolver @@ -95,7 +95,7 @@ public void testGetUnauthorized() throws Exception { public void testGetEntityClientException() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RemoteInvocationException.class).when(mockClient).ingestProposal( Mockito.any(), Mockito.any(Authentication.class), Mockito.eq(false)); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/SetDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/SetDomainResolverTest.java index 73d1f699dfd80..92fb26288aa1d 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/SetDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/SetDomainResolverTest.java @@ -54,7 +54,7 @@ public void testGetSuccessNoExistingDomains() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); Mockito.when(mockService.exists(Urn.createFromString(TEST_NEW_DOMAIN_URN))).thenReturn(true); @@ -110,7 +110,7 @@ public void testGetSuccessExistingDomains() throws Exception { new EnvelopedAspect().setValue(new Aspect(originalDomains.data())) ))))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); Mockito.when(mockService.exists(Urn.createFromString(TEST_NEW_DOMAIN_URN))).thenReturn(true); @@ -160,7 +160,7 @@ public void testGetFailureDomainDoesNotExist() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); Mockito.when(mockService.exists(Urn.createFromString(TEST_NEW_DOMAIN_URN))).thenReturn(false); @@ -196,7 +196,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(false); Mockito.when(mockService.exists(Urn.createFromString(TEST_NEW_DOMAIN_URN))).thenReturn(true); @@ -219,7 +219,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { public void testGetUnauthorized() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); SetDomainResolver resolver = new SetDomainResolver(mockClient, mockService); // Execute resolver diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/UnsetDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/UnsetDomainResolverTest.java index 18b2b9a2747e2..decda39943dde 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/UnsetDomainResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/UnsetDomainResolverTest.java @@ -53,7 +53,7 @@ public void testGetSuccessNoExistingDomains() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); UnsetDomainResolver resolver = new UnsetDomainResolver(mockClient, mockService); @@ -104,7 +104,7 @@ public void testGetSuccessExistingDomains() throws Exception { new EnvelopedAspect().setValue(new Aspect(originalDomains.data())) ))))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(true); UnsetDomainResolver resolver = new UnsetDomainResolver(mockClient, mockService); @@ -148,7 +148,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(false); UnsetDomainResolver resolver = new UnsetDomainResolver(mockClient, mockService); @@ -169,7 +169,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { public void testGetUnauthorized() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); UnsetDomainResolver resolver = new UnsetDomainResolver(mockClient, mockService); // Execute resolver 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 17f7a1968fdf6..f1d44fcb47255 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 @@ -15,6 +15,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.mxe.MetadataChangeProposal; import com.linkedin.r2.RemoteInvocationException; import graphql.schema.DataFetchingEnvironment; @@ -41,7 +42,7 @@ public class UpdateEmbedResolverTest { @Test public void testGetSuccessNoExistingEmbed() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)), @@ -64,11 +65,7 @@ public void testGetSuccessNoExistingEmbed() throws Exception { final MetadataChangeProposal proposal = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN), EMBED_ASPECT_NAME, newEmbed); - Mockito.verify(mockService, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal), - Mockito.any(AuditStamp.class), - Mockito.eq(false) - ); + verifySingleIngestProposal(mockService, 1, proposal);; Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)) @@ -80,7 +77,7 @@ public void testGetSuccessExistingEmbed() throws Exception { Embed originalEmbed = new Embed().setRenderUrl("https://otherurl.com"); // Create resolver - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)), @@ -103,11 +100,7 @@ public void testGetSuccessExistingEmbed() throws Exception { final MetadataChangeProposal proposal = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN), EMBED_ASPECT_NAME, newEmbed); - Mockito.verify(mockService, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal), - Mockito.any(AuditStamp.class), - Mockito.eq(false) - ); + verifySingleIngestProposal(mockService, 1, proposal); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)) @@ -130,7 +123,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { .setUrn(Urn.createFromString(TEST_ENTITY_URN)) .setAspects(new EnvelopedAspectMap(Collections.emptyMap())))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(false); UpdateEmbedResolver resolver = new UpdateEmbedResolver(mockService); @@ -144,7 +137,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) );; @@ -153,7 +146,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { // Create resolver - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); UpdateEmbedResolver resolver = new UpdateEmbedResolver(mockService); // Execute resolver @@ -164,7 +157,7 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -173,7 +166,7 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RemoteInvocationException.class).when(mockClient).ingestProposal( Mockito.any(), Mockito.any(Authentication.class)); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/AddRelatedTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/AddRelatedTermsResolverTest.java index 6bbf4f4797560..26c13186c4a81 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/AddRelatedTermsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/AddRelatedTermsResolverTest.java @@ -27,7 +27,7 @@ public class AddRelatedTermsResolverTest { private static final String DATASET_URN = "urn:li:dataset:(test,test,test)"; private EntityService setUpService() { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), Mockito.eq(Constants.GLOSSARY_RELATED_TERM_ASPECT_NAME), @@ -56,7 +56,7 @@ public void testGetSuccessIsRelatedNonExistent() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1); + verifySingleIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)) ); @@ -88,7 +88,7 @@ public void testGetSuccessHasRelatedNonExistent() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1); + verifySingleIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)) ); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryNodeResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryNodeResolverTest.java index 392ddf6ac4c74..3b47514d87181 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryNodeResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryNodeResolverTest.java @@ -15,6 +15,7 @@ import org.testng.annotations.Test; import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static com.linkedin.datahub.graphql.TestUtils.getMockEntityService; import static com.linkedin.metadata.Constants.*; @@ -69,7 +70,7 @@ private MetadataChangeProposal setupTest( @Test public void testGetSuccess() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); final MetadataChangeProposal proposal = setupTest(mockEnv, TEST_INPUT, "test-description", parentNodeUrn); @@ -86,7 +87,7 @@ public void testGetSuccess() throws Exception { @Test public void testGetSuccessNoDescription() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); final MetadataChangeProposal proposal = setupTest(mockEnv, TEST_INPUT_NO_DESCRIPTION, "", parentNodeUrn); @@ -103,7 +104,7 @@ public void testGetSuccessNoDescription() throws Exception { @Test public void testGetSuccessNoParentNode() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); final MetadataChangeProposal proposal = setupTest(mockEnv, TEST_INPUT_NO_PARENT_NODE, "test-description", null); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java index e4f32133b4b51..2dbe637d16057 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/CreateGlossaryTermResolverTest.java @@ -29,6 +29,7 @@ import java.util.concurrent.CompletionException; import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static com.linkedin.datahub.graphql.TestUtils.getMockEntityService; import static org.testng.Assert.assertThrows; import static com.linkedin.metadata.Constants.*; @@ -86,7 +87,7 @@ private MetadataChangeProposal setupTest( @Test public void testGetSuccess() throws Exception { EntityClient mockClient = initMockClient(); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); final MetadataChangeProposal proposal = setupTest(mockEnv, TEST_INPUT, "test-description", parentNodeUrn); @@ -103,7 +104,7 @@ public void testGetSuccess() throws Exception { @Test public void testGetSuccessNoDescription() throws Exception { EntityClient mockClient = initMockClient(); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); final MetadataChangeProposal proposal = setupTest(mockEnv, TEST_INPUT_NO_DESCRIPTION, "", parentNodeUrn); @@ -120,7 +121,7 @@ public void testGetSuccessNoDescription() throws Exception { @Test public void testGetSuccessNoParentNode() throws Exception { EntityClient mockClient = initMockClient(); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); final MetadataChangeProposal proposal = setupTest(mockEnv, TEST_INPUT_NO_PARENT_NODE, "test-description", null); @@ -166,7 +167,7 @@ public void testGetFailureExistingTermSameName() throws Exception { ) ).thenReturn(result); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); CreateGlossaryEntityInput input = new CreateGlossaryEntityInput( diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/DeleteGlossaryEntityResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/DeleteGlossaryEntityResolverTest.java index df24c23e89ae6..94f0d0b7a1143 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/DeleteGlossaryEntityResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/DeleteGlossaryEntityResolverTest.java @@ -13,6 +13,7 @@ import java.util.concurrent.CompletionException; import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static com.linkedin.datahub.graphql.TestUtils.getMockEntityService; import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; @@ -23,7 +24,7 @@ public class DeleteGlossaryEntityResolverTest { @Test public void testGetSuccess() throws Exception { EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_TERM_URN))).thenReturn(true); @@ -48,7 +49,7 @@ public void testGetEntityClientException() throws Exception { Mockito.any(), Mockito.any(Authentication.class)); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_TERM_URN))).thenReturn(true); DeleteGlossaryEntityResolver resolver = new DeleteGlossaryEntityResolver(mockClient, mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/RemoveRelatedTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/RemoveRelatedTermsResolverTest.java index dd54d7f9835c1..3906d1188cb17 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/RemoveRelatedTermsResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/RemoveRelatedTermsResolverTest.java @@ -34,7 +34,7 @@ public void testGetSuccessIsA() throws Exception { GlossaryTermUrn term2Urn = GlossaryTermUrn.createFromString(TEST_TERM_2_URN); final GlossaryRelatedTerms relatedTerms = new GlossaryRelatedTerms(); relatedTerms.setIsRelatedTerms(new GlossaryTermUrnArray(Arrays.asList(term1Urn, term2Urn))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), Mockito.eq(Constants.GLOSSARY_RELATED_TERM_ASPECT_NAME), @@ -54,7 +54,7 @@ public void testGetSuccessIsA() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1); + verifySingleIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)) ); @@ -66,7 +66,7 @@ public void testGetSuccessHasA() throws Exception { GlossaryTermUrn term2Urn = GlossaryTermUrn.createFromString(TEST_TERM_2_URN); final GlossaryRelatedTerms relatedTerms = new GlossaryRelatedTerms(); relatedTerms.setHasRelatedTerms(new GlossaryTermUrnArray(Arrays.asList(term1Urn, term2Urn))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), Mockito.eq(Constants.GLOSSARY_RELATED_TERM_ASPECT_NAME), @@ -86,7 +86,7 @@ public void testGetSuccessHasA() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1); + verifySingleIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_ENTITY_URN)) ); @@ -94,7 +94,7 @@ public void testGetSuccessHasA() throws Exception { @Test public void testFailAspectDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), Mockito.eq(Constants.GLOSSARY_RELATED_TERM_ASPECT_NAME), @@ -123,7 +123,7 @@ public void testFailNoPermissions() throws Exception { GlossaryTermUrn term2Urn = GlossaryTermUrn.createFromString(TEST_TERM_2_URN); final GlossaryRelatedTerms relatedTerms = new GlossaryRelatedTerms(); relatedTerms.setIsRelatedTerms(new GlossaryTermUrnArray(Arrays.asList(term1Urn, term2Urn))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), Mockito.eq(Constants.GLOSSARY_RELATED_TERM_ASPECT_NAME), diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java index 36f909bd7ebe5..064e2dd3bd59b 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateNameResolverTest.java @@ -58,7 +58,7 @@ private MetadataChangeProposal setupTests(DataFetchingEnvironment mockEnv, Entit @Test public void testGetSuccess() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(TERM_URN))).thenReturn(true); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); @@ -68,12 +68,12 @@ public void testGetSuccess() throws Exception { final MetadataChangeProposal proposal = setupTests(mockEnv, mockService); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1, proposal); + verifySingleIngestProposal(mockService, 1, proposal); } @Test public void testGetSuccessForNode() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(NODE_URN))).thenReturn(true); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); @@ -98,12 +98,12 @@ public void testGetSuccessForNode() throws Exception { UpdateNameResolver resolver = new UpdateNameResolver(mockService, mockClient); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1, proposal); + verifySingleIngestProposal(mockService, 1, proposal); } @Test public void testGetSuccessForDomain() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(DOMAIN_URN))).thenReturn(true); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); @@ -129,12 +129,12 @@ public void testGetSuccessForDomain() throws Exception { UpdateNameResolver resolver = new UpdateNameResolver(mockService, mockClient); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1, proposal); + verifySingleIngestProposal(mockService, 1, proposal); } @Test public void testGetFailureEntityDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(TERM_URN))).thenReturn(false); DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateParentNodeResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateParentNodeResolverTest.java index 43c6113d194a5..a78c28890fecf 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateParentNodeResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/glossary/UpdateParentNodeResolverTest.java @@ -58,7 +58,7 @@ private MetadataChangeProposal setupTests(DataFetchingEnvironment mockEnv, Entit @Test public void testGetSuccess() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(TERM_URN))).thenReturn(true); Mockito.when(mockService.exists(GlossaryNodeUrn.createFromString(PARENT_NODE_URN))).thenReturn(true); @@ -69,12 +69,12 @@ public void testGetSuccess() throws Exception { final MetadataChangeProposal proposal = setupTests(mockEnv, mockService); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1, proposal); + verifySingleIngestProposal(mockService, 1, proposal); } @Test public void testGetSuccessForNode() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(NODE_URN))).thenReturn(true); Mockito.when(mockService.exists(GlossaryNodeUrn.createFromString(PARENT_NODE_URN))).thenReturn(true); @@ -102,12 +102,12 @@ public void testGetSuccessForNode() throws Exception { UpdateParentNodeResolver resolver = new UpdateParentNodeResolver(mockService, mockClient); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 1, proposal); + verifySingleIngestProposal(mockService, 1, proposal); } @Test public void testGetFailureEntityDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(TERM_URN))).thenReturn(false); Mockito.when(mockService.exists(GlossaryNodeUrn.createFromString(PARENT_NODE_URN))).thenReturn(true); @@ -123,7 +123,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception { @Test public void testGetFailureNodeDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(TERM_URN))).thenReturn(true); Mockito.when(mockService.exists(GlossaryNodeUrn.createFromString(PARENT_NODE_URN))).thenReturn(false); @@ -139,7 +139,7 @@ public void testGetFailureNodeDoesNotExist() throws Exception { @Test public void testGetFailureParentIsNotNode() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockService.exists(Urn.createFromString(TERM_URN))).thenReturn(true); Mockito.when(mockService.exists(GlossaryNodeUrn.createFromString(PARENT_NODE_URN))).thenReturn(true); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolverTest.java index c7f1e16a0ea61..9bd44e9ab0906 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolverTest.java @@ -21,7 +21,7 @@ public class UpdateUserSettingResolverTest { private static final String TEST_USER_URN = "urn:li:corpuser:test"; @Test public void testWriteCorpUserSettings() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_USER_URN))).thenReturn(true); UpdateUserSettingResolver resolver = new UpdateUserSettingResolver(mockService); @@ -40,6 +40,6 @@ public void testWriteCorpUserSettings() throws Exception { final MetadataChangeProposal proposal = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_USER_URN), CORP_USER_SETTINGS_ASPECT_NAME, newSettings); - verifyIngestProposal(mockService, 1, proposal); + verifySingleIngestProposal(mockService, 1, proposal); } } 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 9bdb3c1db2b0a..efc0c5dfcf36d 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 @@ -13,6 +13,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 graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -30,7 +31,7 @@ public class AddOwnersResolverTest { @Test public void testGetSuccessNoExistingOwners() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -75,7 +76,7 @@ public void testGetSuccessNoExistingOwners() throws Exception { @Test public void testGetSuccessExistingOwners() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -120,7 +121,7 @@ public void testGetSuccessExistingOwners() throws Exception { @Test public void testGetFailureOwnerDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -148,7 +149,7 @@ public void testGetFailureOwnerDoesNotExist() throws Exception { @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -176,7 +177,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); AddOwnersResolver resolver = new AddOwnersResolver(mockService); @@ -195,10 +196,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 e38eb9eb677c2..79fc62742f444 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 @@ -17,6 +17,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 graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -35,7 +36,7 @@ public class BatchAddOwnersResolverTest { @Test public void testGetSuccessNoExistingOwners() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -82,7 +83,7 @@ public void testGetSuccessNoExistingOwners() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_OWNER_URN_1)) @@ -98,7 +99,7 @@ public void testGetSuccessExistingOwners() throws Exception { final Ownership originalOwnership = new Ownership().setOwners(new OwnerArray(ImmutableList.of( new Owner().setOwner(Urn.createFromString(TEST_OWNER_URN_1)).setType(OwnershipType.TECHNICAL_OWNER) ))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -150,7 +151,7 @@ public void testGetSuccessExistingOwners() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_OWNER_URN_1)) @@ -163,7 +164,7 @@ public void testGetSuccessExistingOwners() throws Exception { @Test public void testGetFailureOwnerDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -202,7 +203,7 @@ public void testGetFailureOwnerDoesNotExist() throws Exception { @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -247,7 +248,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchAddOwnersResolver resolver = new BatchAddOwnersResolver(mockService); @@ -277,10 +278,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 0884d442ea531..9dc2ec8127806 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 @@ -14,6 +14,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 graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -32,7 +33,7 @@ public class BatchRemoveOwnersResolverTest { @Test public void testGetSuccessNoExistingOwners() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -66,12 +67,12 @@ public void testGetSuccessNoExistingOwners() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); } @Test public void testGetSuccessExistingOwners() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); final Ownership oldOwners1 = new Ownership().setOwners(new OwnerArray(ImmutableList.of( new Owner().setOwner(Urn.createFromString(TEST_OWNER_URN_1)).setType(OwnershipType.TECHNICAL_OWNER) @@ -112,12 +113,12 @@ public void testGetSuccessExistingOwners() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -152,7 +153,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchRemoveOwnersResolver resolver = new BatchRemoveOwnersResolver(mockService); @@ -172,10 +173,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 06d9df3278847..268d6a6bc4268 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 @@ -13,6 +13,7 @@ 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.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; @@ -32,7 +33,7 @@ public class AddTagsResolverTest { @Test public void testGetSuccessNoExistingTags() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -82,7 +83,7 @@ public void testGetSuccessExistingTags() throws Exception { new TagAssociation().setTag(TagUrn.createFromString(TEST_TAG_1_URN)))) ); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -128,7 +129,7 @@ public void testGetSuccessExistingTags() throws Exception { @Test public void testGetFailureTagDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -156,7 +157,7 @@ public void testGetFailureTagDoesNotExist() throws Exception { @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -184,7 +185,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); AddTagsResolver resolver = new AddTagsResolver(mockService); @@ -203,10 +204,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 d47e6164fe221..651b89359c83f 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 @@ -15,8 +15,11 @@ 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.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; + +import java.util.List; import java.util.concurrent.CompletionException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -35,7 +38,7 @@ public class BatchAddTagsResolverTest { @Test public void testGetSuccessNoExistingTags() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -78,13 +81,10 @@ public void testGetSuccessNoExistingTags() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), GLOBAL_TAGS_ASPECT_NAME, newTags); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), GLOBAL_TAGS_ASPECT_NAME, newTags); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_TAG_1_URN)) @@ -101,7 +101,7 @@ public void testGetSuccessExistingTags() throws Exception { new TagAssociation().setTag(TagUrn.createFromString(TEST_TAG_1_URN)))) ); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -143,13 +143,10 @@ public void testGetSuccessExistingTags() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), GLOBAL_TAGS_ASPECT_NAME, newTags); - - verifyIngestProposal(mockService, 1, proposal1); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), GLOBAL_TAGS_ASPECT_NAME, newTags); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_TAG_1_URN)) @@ -162,7 +159,7 @@ public void testGetSuccessExistingTags() throws Exception { @Test public void testGetFailureTagDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -188,13 +185,13 @@ public void testGetFailureTagDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -227,13 +224,13 @@ public void testGetFailureResourceDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchAddTagsResolver resolver = new BatchAddTagsResolver(mockService); @@ -251,16 +248,16 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 44160cfbe1273..f302540eba904 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 @@ -16,10 +16,12 @@ 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.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.Collections; +import java.util.List; import java.util.concurrent.CompletionException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -38,7 +40,7 @@ public class BatchRemoveTagsResolverTest { @Test public void testGetSuccessNoExistingTags() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -76,12 +78,6 @@ public void testGetSuccessNoExistingTags() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), GLOBAL_TAGS_ASPECT_NAME, emptyTags); - - Mockito.verify(mockService, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal1), - Mockito.any(AuditStamp.class), Mockito.eq(false) - ); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), GLOBAL_TAGS_ASPECT_NAME, emptyTags); proposal2.setEntityUrn(Urn.createFromString(TEST_ENTITY_URN_2)); @@ -90,12 +86,12 @@ public void testGetSuccessNoExistingTags() throws Exception { proposal2.setAspect(GenericRecordUtils.serializeAspect(emptyTags)); proposal2.setChangeType(ChangeType.UPSERT); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test public void testGetSuccessExistingTags() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); final GlobalTags oldTags1 = new GlobalTags().setTags(new TagAssociationArray(ImmutableList.of( new TagAssociation().setTag(TagUrn.createFromString(TEST_TAG_1_URN)), @@ -143,21 +139,15 @@ public void testGetSuccessExistingTags() throws Exception { final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1), GLOBAL_TAGS_ASPECT_NAME, emptyTags); - - Mockito.verify(mockService, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal1), - Mockito.any(AuditStamp.class), Mockito.eq(false) - ); - final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2), GLOBAL_TAGS_ASPECT_NAME, emptyTags); - verifyIngestProposal(mockService, 1, proposal2); + verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2)); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -190,13 +180,13 @@ public void testGetFailureResourceDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchRemoveTagsResolver resolver = new BatchRemoveTagsResolver(mockService); @@ -214,16 +204,16 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); BatchRemoveTagsResolver resolver = new BatchRemoveTagsResolver(mockService); diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/CreateTagResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/CreateTagResolverTest.java index d294f806d1af7..f801daf4f2a3f 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/CreateTagResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/CreateTagResolverTest.java @@ -31,7 +31,7 @@ public class CreateTagResolverTest { @Test public void testGetSuccess() throws Exception { // Create resolver - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.when(mockClient.ingestProposal(Mockito.any(MetadataChangeProposal.class), Mockito.any(Authentication.class))) .thenReturn(String.format("urn:li:tag:%s", TEST_INPUT.getId())); @@ -64,7 +64,7 @@ public void testGetSuccess() throws Exception { @Test public void testGetUnauthorized() throws Exception { // Create resolver - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); CreateTagResolver resolver = new CreateTagResolver(mockClient, mockService); @@ -83,7 +83,7 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { // Create resolver - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); EntityClient mockClient = Mockito.mock(EntityClient.class); Mockito.doThrow(RuntimeException.class).when(mockClient).ingestProposal( Mockito.any(), diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/SetTagColorResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/SetTagColorResolverTest.java index da474ca3e0e56..b5bbf0775a8ba 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/SetTagColorResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/tag/SetTagColorResolverTest.java @@ -37,7 +37,7 @@ public class SetTagColorResolverTest { public void testGetSuccessExistingProperties() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); // Test setting the domain final TagProperties oldTagProperties = new TagProperties().setName("Test Tag"); @@ -78,7 +78,7 @@ public void testGetSuccessExistingProperties() throws Exception { public void testGetFailureNoExistingProperties() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); // Test setting the domain Mockito.when(mockService.getAspect( @@ -127,7 +127,7 @@ public void testGetFailureTagDoesNotExist() throws Exception { Constants.TAG_PROPERTIES_ASPECT_NAME, oldTagPropertiesAspect))))); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.exists(Urn.createFromString(TEST_ENTITY_URN))).thenReturn(false); SetTagColorResolver resolver = new SetTagColorResolver(mockClient, mockService); @@ -148,7 +148,7 @@ public void testGetFailureTagDoesNotExist() throws Exception { public void testGetUnauthorized() throws Exception { // Create resolver EntityClient mockClient = Mockito.mock(EntityClient.class); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); SetTagColorResolver resolver = new SetTagColorResolver(mockClient, mockService); // Execute resolver 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 c9ec92001f89b..213d21fd35dc1 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 @@ -13,7 +13,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.AddTermsResolver; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -31,7 +31,7 @@ public class AddTermsResolverTest { @Test public void testGetSuccessNoExistingTerms() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -58,7 +58,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(MetadataChangeProposal.class), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -77,7 +77,7 @@ public void testGetSuccessExistingTerms() throws Exception { new GlossaryTermAssociation().setUrn(GlossaryTermUrn.createFromString(TEST_TERM_1_URN)))) ); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -104,7 +104,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(MetadataChangeProposal.class), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -119,7 +119,7 @@ public void testGetSuccessExistingTerms() throws Exception { @Test public void testGetFailureTermDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -143,13 +143,13 @@ public void testGetFailureTermDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN)), @@ -173,13 +173,13 @@ public void testGetFailureResourceDoesNotExist() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); AddTermsResolver resolver = new AddTermsResolver(mockService); @@ -194,16 +194,16 @@ public void testGetUnauthorized() throws Exception { assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join()); Mockito.verify(mockService, Mockito.times(0)).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); } @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 dfe1394635c4e..8887bb452b478 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 @@ -14,6 +14,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 graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -32,7 +33,7 @@ public class BatchAddTermsResolverTest { @Test public void testGetSuccessNoExistingTerms() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -66,7 +67,7 @@ public void testGetSuccessNoExistingTerms() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_GLOSSARY_TERM_1_URN)) @@ -83,7 +84,7 @@ public void testGetSuccessExistingTerms() throws Exception { new GlossaryTermAssociation().setUrn(GlossaryTermUrn.createFromString(TEST_GLOSSARY_TERM_1_URN)))) ); - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -118,7 +119,7 @@ public void testGetSuccessExistingTerms() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); Mockito.verify(mockService, Mockito.times(1)).exists( Mockito.eq(Urn.createFromString(TEST_GLOSSARY_TERM_1_URN)) @@ -131,7 +132,7 @@ public void testGetSuccessExistingTerms() throws Exception { @Test public void testGetFailureTagDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -160,7 +161,7 @@ public void testGetFailureTagDoesNotExist() throws Exception { @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -196,7 +197,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchAddTermsResolver resolver = new BatchAddTermsResolver(mockService); @@ -217,10 +218,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.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 dcc8659c1baf3..995a4acb8a467 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 @@ -14,6 +14,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 graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -32,7 +33,7 @@ public class BatchRemoveTermsResolverTest { @Test public void testGetSuccessNoExistingTerms() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -66,12 +67,12 @@ public void testGetSuccessNoExistingTerms() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); } @Test public void testGetSuccessExistingTerms() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); final GlossaryTerms oldTerms1 = new GlossaryTerms().setTerms(new GlossaryTermAssociationArray(ImmutableList.of( new GlossaryTermAssociation().setUrn(GlossaryTermUrn.createFromString(TEST_TERM_1_URN)), @@ -115,12 +116,12 @@ public void testGetSuccessExistingTerms() throws Exception { Mockito.when(mockEnv.getContext()).thenReturn(mockContext); assertTrue(resolver.get(mockEnv).get()); - verifyIngestProposal(mockService, 2); + verifyIngestProposal(mockService, 1); } @Test public void testGetFailureResourceDoesNotExist() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.when(mockService.getAspect( Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)), @@ -157,7 +158,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchRemoveTermsResolver resolver = new BatchRemoveTermsResolver(mockService); @@ -179,10 +180,10 @@ public void testGetUnauthorized() throws Exception { @Test public void testGetEntityClientException() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal( - Mockito.any(), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean()); BatchRemoveTermsResolver resolver = new BatchRemoveTermsResolver(mockService); diff --git a/datahub-graphql-core/src/test/resources/test-entity-registry.yaml b/datahub-graphql-core/src/test/resources/test-entity-registry.yaml new file mode 100644 index 0000000000000..d694ae53ac42f --- /dev/null +++ b/datahub-graphql-core/src/test/resources/test-entity-registry.yaml @@ -0,0 +1,295 @@ +entities: +- name: dataPlatform + category: core + keyAspect: dataPlatformKey + aspects: + - dataPlatformInfo +- name: dataset + doc: Datasets represent logical or physical data assets stored or represented in various data platforms. Tables, Views, Streams are all instances of datasets. + category: core + keyAspect: datasetKey + aspects: + - viewProperties + - subTypes + - datasetProfile + - datasetUsageStatistics + - operation + - domains + - status + - container + - deprecation + - testResults + - siblings + - embed + - ownership + - glossaryTerms + - globalTags +- name: dataHubPolicy + doc: DataHub Policies represent access policies granted to users or groups on metadata operations like edit, view etc. + category: internal + keyAspect: dataHubPolicyKey + aspects: + - dataHubPolicyInfo +- name: dataJob + keyAspect: dataJobKey + aspects: + - datahubIngestionRunSummary + - datahubIngestionCheckpoint + - domains + - deprecation + - versionInfo +- name: dataFlow + category: core + keyAspect: dataFlowKey + aspects: + - domains + - deprecation + - versionInfo +- name: dataProcessInstance + doc: DataProcessInstance represents an instance of a datajob/jobflow run + keyAspect: dataProcessInstanceKey + aspects: + - dataProcessInstanceInput + - dataProcessInstanceOutput + - dataProcessInstanceProperties + - dataProcessInstanceRelationships + - dataProcessInstanceRunEvent +- name: chart + keyAspect: chartKey + aspects: + - domains + - container + - deprecation + - chartUsageStatistics + - embed +- name: dashboard + keyAspect: dashboardKey + aspects: + - domains + - container + - deprecation + - dashboardUsageStatistics + - subTypes + - embed +- name: notebook + doc: Notebook represents a combination of query, text, chart and etc. This is in BETA version + keyAspect: notebookKey + aspects: + - notebookInfo + - notebookContent + - editableNotebookProperties + - ownership + - status + - globalTags + - glossaryTerms + - browsePaths + - institutionalMemory + - domains + - subTypes + - dataPlatformInstance +- name: corpuser + doc: CorpUser represents an identity of a person (or an account) in the enterprise. + keyAspect: corpUserKey + aspects: + - corpUserInfo + - corpUserEditableInfo + - corpUserStatus + - groupMembership + - globalTags + - status + - corpUserCredentials + - nativeGroupMembership + - corpUserSettings + - origin + - roleMembership +- name: corpGroup + doc: CorpGroup represents an identity of a group of users in the enterprise. + keyAspect: corpGroupKey + aspects: + - corpGroupInfo + - corpGroupEditableInfo + - globalTags + - ownership + - status + - origin +- name: domain + doc: A data domain within an organization. + keyAspect: domainKey + aspects: + - domainProperties + - institutionalMemory + - ownership +- name: container + doc: A container of related data assets. + keyAspect: containerKey + aspects: + - containerProperties + - editableContainerProperties + - dataPlatformInstance + - subTypes + - ownership + - container + - globalTags + - glossaryTerms + - institutionalMemory + - browsePaths # unclear if this will be used + - status + - domains +- name: tag + keyAspect: tagKey + aspects: + - tagProperties + - ownership + - deprecation +- name: glossaryTerm + keyAspect: glossaryTermKey + aspects: + - glossaryTermInfo + - institutionalMemory + - ownership + - deprecation + - domains +- name: glossaryNode + keyAspect: glossaryNodeKey + aspects: + - glossaryNodeInfo + - institutionalMemory + - ownership + - status +- name: dataHubIngestionSource + category: internal + keyAspect: dataHubIngestionSourceKey + aspects: + - dataHubIngestionSourceInfo +- name: dataHubSecret + category: internal + keyAspect: dataHubSecretKey + aspects: + - dataHubSecretValue +- name: dataHubExecutionRequest + category: internal + keyAspect: dataHubExecutionRequestKey + aspects: + - dataHubExecutionRequestInput + - dataHubExecutionRequestSignal + - dataHubExecutionRequestResult +- name: assertion + doc: Assertion represents a data quality rule applied on one or more dataset. + category: core + keyAspect: assertionKey + aspects: + - assertionInfo + - dataPlatformInstance + - assertionRunEvent + - status +- name: dataHubRetention + category: internal + keyAspect: dataHubRetentionKey + aspects: + - dataHubRetentionConfig +- name: dataPlatformInstance + category: internal + keyAspect: dataPlatformInstanceKey + aspects: + - dataPlatformInstanceProperties + - ownership + - globalTags + - institutionalMemory + - deprecation + - status +- name: mlModel + category: core + keyAspect: mlModelKey + aspects: + - glossaryTerms + - editableMlModelProperties + - domains +- name: mlModelGroup + category: core + keyAspect: mlModelGroupKey + aspects: + - glossaryTerms + - editableMlModelGroupProperties + - domains +- name: mlFeatureTable + category: core + keyAspect: mlFeatureTableKey + aspects: + - glossaryTerms + - editableMlFeatureTableProperties + - domains +- name: mlFeature + category: core + keyAspect: mlFeatureKey + aspects: + - glossaryTerms + - editableMlFeatureProperties + - domains +- name: mlPrimaryKey + category: core + keyAspect: mlPrimaryKeyKey + aspects: + - glossaryTerms + - editableMlPrimaryKeyProperties + - domains +- name: telemetry + category: internal + keyAspect: telemetryKey + aspects: + - telemetryClientId +- name: dataHubAccessToken + category: internal + keyAspect: dataHubAccessTokenKey + aspects: + - dataHubAccessTokenInfo +- name: test + doc: A DataHub test + category: core + keyAspect: testKey + aspects: + - testInfo +- name: dataHubUpgrade + category: internal + keyAspect: dataHubUpgradeKey + aspects: + - dataHubUpgradeRequest + - dataHubUpgradeResult +- name: inviteToken + category: core + keyAspect: inviteTokenKey + aspects: + - inviteToken +- name: globalSettings + doc: Global settings for an the platform + category: internal + keyAspect: globalSettingsKey + aspects: + - globalSettingsInfo +- name: dataHubRole + category: core + keyAspect: dataHubRoleKey + aspects: + - dataHubRoleInfo +- name: post + category: core + keyAspect: postKey + aspects: + - postInfo +- name: dataHubStepState + category: core + keyAspect: dataHubStepStateKey + aspects: + - dataHubStepStateProperties +- name: dataHubView + category: core + keyAspect: dataHubViewKey + aspects: + - dataHubViewInfo +- name: ownershipType + doc: Ownership Type represents a user-created ownership category for a person or group who is responsible for an asset. + category: core + keyAspect: ownershipTypeKey + aspects: + - ownershipTypeInfo + - status +events: diff --git a/datahub-upgrade/build.gradle b/datahub-upgrade/build.gradle index 625cab6b354c5..5d0edf3ee8427 100644 --- a/datahub-upgrade/build.gradle +++ b/datahub-upgrade/build.gradle @@ -100,7 +100,7 @@ docker { load(true) push(false) } -tasks.getByName("docker").dependsOn([bootJar]) +tasks.getByPath(":datahub-upgrade:docker").dependsOn([bootJar]) task cleanLocalDockerImages { doLast { diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeCleanupConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeCleanupConfig.java index 23ce409c746d0..0fb8b0eb6e20f 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeCleanupConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeCleanupConfig.java @@ -3,7 +3,7 @@ import com.linkedin.datahub.upgrade.nocodecleanup.NoCodeCleanupUpgrade; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; -import io.ebean.EbeanServer; +import io.ebean.Database; import javax.annotation.Nonnull; import org.elasticsearch.client.RestHighLevelClient; import org.springframework.beans.factory.annotation.Autowired; @@ -25,7 +25,7 @@ public class NoCodeCleanupConfig { @DependsOn({"ebeanServer", "graphService", "elasticSearchRestHighLevelClient", INDEX_CONVENTION_BEAN}) @Nonnull public NoCodeCleanupUpgrade createInstance() { - final EbeanServer ebeanServer = applicationContext.getBean(EbeanServer.class); + final Database ebeanServer = applicationContext.getBean(Database.class); final GraphService graphClient = applicationContext.getBean(GraphService.class); final RestHighLevelClient searchClient = applicationContext.getBean(RestHighLevelClient.class); final IndexConvention indexConvention = applicationContext.getBean(IndexConvention.class); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java index 39b3daa73b78f..30175c6fa78c8 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java @@ -5,7 +5,7 @@ import com.linkedin.entity.client.RestliEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.models.registry.EntityRegistry; -import io.ebean.EbeanServer; +import io.ebean.Database; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationContext; @@ -24,7 +24,7 @@ public class NoCodeUpgradeConfig { @DependsOn({"ebeanServer", "entityService", "systemAuthentication", "restliEntityClient", "entityRegistry"}) @Nonnull public NoCodeUpgrade createInstance() { - final EbeanServer ebeanServer = applicationContext.getBean(EbeanServer.class); + final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); final Authentication systemAuthentication = applicationContext.getBean(Authentication.class); final RestliEntityClient entityClient = applicationContext.getBean(RestliEntityClient.class); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java index ebff7f4b899ad..9b0fcf279abf5 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java @@ -7,7 +7,7 @@ import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; -import io.ebean.EbeanServer; +import io.ebean.Database; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationContext; @@ -26,7 +26,7 @@ public class RestoreBackupConfig { "searchService", "entityRegistry"}) @Nonnull public RestoreBackup createInstance() { - final EbeanServer ebeanServer = applicationContext.getBean(EbeanServer.class); + final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); final Authentication systemAuthentication = applicationContext.getBean(Authentication.class); final RestliEntityClient entityClient = applicationContext.getBean(RestliEntityClient.class); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java index ee907005168b8..663cad4a4bff6 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreIndicesConfig.java @@ -5,7 +5,7 @@ import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; -import io.ebean.EbeanServer; +import io.ebean.Database; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationContext; @@ -23,7 +23,7 @@ public class RestoreIndicesConfig { @DependsOn({"ebeanServer", "entityService", "searchService", "graphService", "entityRegistry"}) @Nonnull public RestoreIndices createInstance() { - final EbeanServer ebeanServer = applicationContext.getBean(EbeanServer.class); + final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); final EntitySearchService entitySearchService = applicationContext.getBean(EntitySearchService.class); final GraphService graphService = applicationContext.getBean(GraphService.class); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/CreateAspectTableStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/CreateAspectTableStep.java index 3b78e95a7b751..7ed7169bf20bc 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/CreateAspectTableStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/CreateAspectTableStep.java @@ -4,7 +4,7 @@ import com.linkedin.datahub.upgrade.UpgradeContext; import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.function.Function; public class CreateAspectTableStep implements UpgradeStep { @@ -17,9 +17,9 @@ enum DbType { MARIA } - private final EbeanServer _server; + private final Database _server; - public CreateAspectTableStep(final EbeanServer server) { + public CreateAspectTableStep(final Database server) { _server = server; } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java index 6553bb80bb1fa..1b5770a11ff62 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java @@ -17,10 +17,12 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectV1; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.models.EntitySpec; -import io.ebean.EbeanServer; +import com.linkedin.util.Pair; +import io.ebean.Database; import io.ebean.PagedList; import java.net.URISyntaxException; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -35,13 +37,13 @@ public class DataMigrationStep implements UpgradeStep { private static final String BROWSE_PATHS_ASPECT_NAME = PegasusUtils.getAspectNameFromSchema(new BrowsePaths().schema()); - private final EbeanServer _server; + private final Database _server; private final EntityService _entityService; private final EntityRegistry _entityRegistry; private final Set urnsWithBrowsePath = new HashSet<>(); public DataMigrationStep( - final EbeanServer server, + final Database server, final EntityService entityService, final EntityRegistry entityRegistry) { _server = server; @@ -132,15 +134,11 @@ public Function executable() { // 6. Write the row back using the EntityService boolean emitMae = oldAspect.getKey().getVersion() == 0L; - _entityService.updateAspect( + _entityService.ingestAspects( urn, - entityName, - newAspectName, - aspectSpec, - aspectRecord, + List.of(Pair.of(newAspectName, aspectRecord)), toAuditStamp(oldAspect), - oldAspect.getKey().getVersion(), - emitMae + null ); // 7. If necessary, emit a browse path aspect. @@ -154,7 +152,7 @@ public Function executable() { browsePathsStamp.setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)); browsePathsStamp.setTime(System.currentTimeMillis()); - _entityService.ingestAspect(urn, BROWSE_PATHS_ASPECT_NAME, browsePaths, browsePathsStamp, null); + _entityService.ingestAspects(urn, List.of(Pair.of(BROWSE_PATHS_ASPECT_NAME, browsePaths)), browsePathsStamp, null); urnsWithBrowsePath.add(urn); } catch (URISyntaxException e) { diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java index c12ff201faf22..ee4a3bc504e77 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java @@ -10,7 +10,7 @@ import com.linkedin.entity.client.RestliEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.models.registry.EntityRegistry; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -25,9 +25,9 @@ public class NoCodeUpgrade implements Upgrade { private final List _steps; private final List _cleanupSteps; - // Upgrade requires the EbeanServer. + // Upgrade requires the Database. public NoCodeUpgrade( - final EbeanServer server, + final Database server, final EntityService entityService, final EntityRegistry entityRegistry, final Authentication systemAuthentication, @@ -60,7 +60,7 @@ private List buildCleanupSteps() { } private List buildUpgradeSteps( - final EbeanServer server, + final Database server, final EntityService entityService, final EntityRegistry entityRegistry, final Authentication systemAuthentication, diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/RemoveAspectV2TableStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/RemoveAspectV2TableStep.java index 440884470463d..cf8e848762f14 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/RemoveAspectV2TableStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/RemoveAspectV2TableStep.java @@ -4,7 +4,7 @@ import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.function.Function; @@ -13,9 +13,9 @@ */ public class RemoveAspectV2TableStep implements UpgradeStep { - private final EbeanServer _server; + private final Database _server; - public RemoveAspectV2TableStep(final EbeanServer server) { + public RemoveAspectV2TableStep(final Database server) { _server = server; } @@ -28,7 +28,7 @@ public String id() { public Function executable() { return (context) -> { context.report().addLine("Cleanup requested. Dropping metadata_aspect_v2"); - _server.execute(_server.createSqlUpdate("DROP TABLE IF EXISTS metadata_aspect_v2")); + _server.execute(_server.sqlUpdate("DROP TABLE IF EXISTS metadata_aspect_v2")); return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED); }; } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java index ec05f210f0132..0fe9afa8cc6f8 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java @@ -5,14 +5,14 @@ import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.metadata.entity.ebean.AspectStorageValidationUtil; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.function.Function; public class UpgradeQualificationStep implements UpgradeStep { - private final EbeanServer _server; + private final Database _server; - UpgradeQualificationStep(EbeanServer server) { + UpgradeQualificationStep(Database server) { _server = server; } @@ -52,7 +52,7 @@ public Function executable() { } // Check whether the upgrade is needed - private boolean isQualified(EbeanServer server, UpgradeContext context) { + private boolean isQualified(Database server, UpgradeContext context) { boolean v1TableExists = AspectStorageValidationUtil.checkV1TableExists(server); if (v1TableExists) { context.report().addLine("-- V1 table exists"); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/DeleteAspectTableStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/DeleteAspectTableStep.java index 2d435cdc28a6b..8005e31e01c67 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/DeleteAspectTableStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/DeleteAspectTableStep.java @@ -4,16 +4,16 @@ import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.function.Function; // Do we need SQL-tech specific migration paths? public class DeleteAspectTableStep implements UpgradeStep { - private final EbeanServer _server; + private final Database _server; - public DeleteAspectTableStep(final EbeanServer server) { + public DeleteAspectTableStep(final Database server) { _server = server; } @@ -31,7 +31,7 @@ public int retryCount() { public Function executable() { return (context) -> { try { - _server.execute(_server.createSqlUpdate("DROP TABLE IF EXISTS metadata_aspect;")); + _server.execute(_server.sqlUpdate("DROP TABLE IF EXISTS metadata_aspect;")); } catch (Exception e) { context.report().addLine("Failed to delete data from legacy table metadata_aspect", e); return new DefaultUpgradeStepResult( diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeCleanupUpgrade.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeCleanupUpgrade.java index c9a13c2208a56..2b5e23c5f8269 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeCleanupUpgrade.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeCleanupUpgrade.java @@ -5,7 +5,7 @@ import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -17,8 +17,8 @@ public class NoCodeCleanupUpgrade implements Upgrade { private final List _steps; private final List _cleanupSteps; - // Upgrade requires the EbeanServer. - public NoCodeCleanupUpgrade(final EbeanServer server, final GraphService graphClient, + // Upgrade requires the Database. + public NoCodeCleanupUpgrade(final Database server, final GraphService graphClient, final RestHighLevelClient searchClient, final IndexConvention indexConvention) { _steps = buildUpgradeSteps(server, graphClient, searchClient, indexConvention); _cleanupSteps = buildCleanupSteps(); @@ -43,7 +43,7 @@ private List buildCleanupSteps() { return Collections.emptyList(); } - private List buildUpgradeSteps(final EbeanServer server, final GraphService graphClient, + private List buildUpgradeSteps(final Database server, final GraphService graphClient, final RestHighLevelClient searchClient, final IndexConvention indexConvention) { final List steps = new ArrayList<>(); steps.add(new NoCodeUpgradeQualificationStep(server)); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeUpgradeQualificationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeUpgradeQualificationStep.java index 52e299d68b45a..67a226f8f0676 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeUpgradeQualificationStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocodecleanup/NoCodeUpgradeQualificationStep.java @@ -5,15 +5,15 @@ import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; import com.linkedin.metadata.entity.ebean.AspectStorageValidationUtil; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.function.Function; public class NoCodeUpgradeQualificationStep implements UpgradeStep { - private final EbeanServer _server; + private final Database _server; - NoCodeUpgradeQualificationStep(EbeanServer server) { + NoCodeUpgradeQualificationStep(Database server) { _server = server; } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/ClearAspectV2TableStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/ClearAspectV2TableStep.java index 711cccf742254..0303739e62afe 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/ClearAspectV2TableStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/ClearAspectV2TableStep.java @@ -5,7 +5,7 @@ import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.function.Function; @@ -14,9 +14,9 @@ */ public class ClearAspectV2TableStep implements UpgradeStep { - private final EbeanServer _server; + private final Database _server; - public ClearAspectV2TableStep(final EbeanServer server) { + public ClearAspectV2TableStep(final Database server) { _server = server; } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java index a9dfa948c7873..67718a6739beb 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java @@ -14,7 +14,7 @@ import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.ArrayList; import java.util.List; @@ -24,7 +24,7 @@ public class RestoreBackup implements Upgrade { private final List _steps; public RestoreBackup( - final EbeanServer server, + final Database server, final EntityService entityService, final EntityRegistry entityRegistry, final Authentication systemAuthentication, @@ -45,7 +45,7 @@ public List steps() { } private List buildSteps( - final EbeanServer server, + final Database server, final EntityService entityService, final EntityRegistry entityRegistry, final Authentication systemAuthentication, 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 f22a52c4877f6..42f7f0073e59b 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 @@ -20,6 +20,8 @@ import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.util.Pair; + import java.lang.reflect.InvocationTargetException; import java.net.URISyntaxException; import java.util.ArrayList; @@ -181,8 +183,7 @@ private void readerExecutable(ReaderWrapper reader, UpgradeContext context) { final long version = aspect.getKey().getVersion(); final AuditStamp auditStamp = toAuditStamp(aspect); futureList.add(_gmsThreadPool.submit(() -> - _entityService.updateAspect(urn, entityName, aspectName, aspectSpec, aspectRecord, auditStamp, - version, version == 0L))); + _entityService.ingestAspects(urn, List.of(Pair.of(aspectName, aspectRecord)), auditStamp, null).get(0).getNewValue())); if (numRows % REPORT_BATCH_SIZE == 0) { for (Future future : futureList) { try { diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java index 9e11a953079a5..ee6a5ed6f1536 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/RestoreIndices.java @@ -10,7 +10,7 @@ import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.ArrayList; import java.util.List; @@ -27,7 +27,7 @@ public class RestoreIndices implements Upgrade { private final List _steps; - public RestoreIndices(final EbeanServer server, final EntityService entityService, + public RestoreIndices(final Database server, final EntityService entityService, final EntityRegistry entityRegistry, final EntitySearchService entitySearchService, final GraphService graphService) { _steps = buildSteps(server, entityService, entityRegistry, entitySearchService, graphService); @@ -43,7 +43,7 @@ public List steps() { return _steps; } - private List buildSteps(final EbeanServer server, final EntityService entityService, + private List buildSteps(final Database server, final EntityService entityService, final EntityRegistry entityRegistry, final EntitySearchService entitySearchService, final GraphService graphService) { final List steps = new ArrayList<>(); 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 ac2457732771d..ce39b3fb562af 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 @@ -9,7 +9,7 @@ import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult; import com.linkedin.metadata.models.registry.EntityRegistry; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.ExpressionList; import java.util.ArrayList; @@ -32,7 +32,7 @@ public class SendMAEStep implements UpgradeStep { private static final long DEFAULT_BATCH_DELAY_MS = 250; private static final int DEFAULT_THREADS = 1; - private final EbeanServer _server; + private final Database _server; private final EntityService _entityService; public class KafkaJob implements Callable { @@ -48,7 +48,7 @@ public RestoreIndicesResult call() { } } - public SendMAEStep(final EbeanServer server, final EntityService entityService, final EntityRegistry entityRegistry) { + public SendMAEStep(final Database server, final EntityService entityService, final EntityRegistry entityRegistry) { _server = server; _entityService = entityService; } diff --git a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTestConfiguration.java b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTestConfiguration.java index fefc853be8c0b..b1bdead58a72b 100644 --- a/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTestConfiguration.java +++ b/datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/UpgradeCliApplicationTestConfiguration.java @@ -6,7 +6,7 @@ import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.SearchService; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.springframework.boot.test.context.TestConfiguration; import org.springframework.boot.test.mock.mockito.MockBean; import org.springframework.context.annotation.Import; @@ -19,7 +19,7 @@ public class UpgradeCliApplicationTestConfiguration { private UpgradeCli upgradeCli; @MockBean - private EbeanServer ebeanServer; + private Database ebeanServer; @MockBean private EntityService _entityService; diff --git a/docs/deploy/environment-vars.md b/docs/deploy/environment-vars.md index af4ae09c009fd..a09db7bc6b82e 100644 --- a/docs/deploy/environment-vars.md +++ b/docs/deploy/environment-vars.md @@ -19,7 +19,7 @@ DataHub works. | Variable | Default | Unit/Type | Components | Description | |------------------------------------|---------|-----------|-------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `ASYNC_INGESTION_DEFAULT` | `false` | boolean | [`GMS`] | Asynchronously process ingestProposals by writing the ingestion MCP to Kafka. Typically enabled with standalone consumers. | +| `ASYNC_INGEST_DEFAULT` | `false` | boolean | [`GMS`] | Asynchronously process ingestProposals by writing the ingestion MCP to Kafka. Typically enabled with standalone consumers. | | `MCP_CONSUMER_ENABLED` | `true` | boolean | [`GMS`, `MCE Consumer`] | When running in standalone mode, disabled on `GMS` and enabled on separate `MCE Consumer`. | | `MCL_CONSUMER_ENABLED` | `true` | boolean | [`GMS`, `MAE Consumer`] | When running in standalone mode, disabled on `GMS` and enabled on separate `MAE Consumer`. | | `PE_CONSUMER_ENABLED` | `true` | boolean | [`GMS`, `MAE Consumer`] | When running in standalone mode, disabled on `GMS` and enabled on separate `MAE Consumer`. | diff --git a/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java b/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java index 65bf250200d13..00b5bb75d901b 100644 --- a/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java +++ b/metadata-dao-impl/kafka-producer/src/main/java/com/linkedin/metadata/dao/producer/KafkaEventProducer.java @@ -1,25 +1,18 @@ package com.linkedin.metadata.dao.producer; import com.datahub.util.exception.ModelConversionException; -import com.google.common.annotations.VisibleForTesting; import com.linkedin.common.urn.Urn; import com.linkedin.metadata.EventUtils; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.snapshot.Snapshot; import com.linkedin.mxe.DataHubUpgradeHistoryEvent; -import com.linkedin.mxe.MetadataAuditEvent; -import com.linkedin.mxe.MetadataAuditOperation; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.PlatformEvent; -import com.linkedin.mxe.SystemMetadata; import com.linkedin.mxe.TopicConvention; import com.linkedin.mxe.TopicConventionImpl; -import com.linkedin.mxe.Topics; import io.opentelemetry.extension.annotations.WithSpan; import java.io.IOException; -import java.util.Arrays; import java.util.concurrent.Future; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -55,45 +48,6 @@ public KafkaEventProducer(@Nonnull final Producer produceMetadataChangeLog(@Nonnull final Urn urn, @Nonnull AspectSpec aspectSpec, @@ -120,7 +74,7 @@ record = EventUtils.pegasusToAvroMCL(metadataChangeLog); @Override @WithSpan public Future produceMetadataChangeProposal(@Nonnull final Urn urn, - @Nonnull final MetadataChangeProposal metadataChangeProposal) { + @Nonnull final MetadataChangeProposal metadataChangeProposal) { GenericRecord record; try { @@ -171,9 +125,4 @@ record = EventUtils.pegasusToAvroDUHE(event); _producer.send(new ProducerRecord(topic, event.getVersion(), record), _kafkaHealthChecker .getKafkaCallBack("History Event", "Event Version: " + event.getVersion())); } - - @VisibleForTesting - static boolean isValidAspectSpecificTopic(@Nonnull String topic) { - return Arrays.stream(Topics.class.getFields()).anyMatch(field -> field.getName().equals(topic)); - } } diff --git a/metadata-io/build.gradle b/metadata-io/build.gradle index e8ef0b3d6819d..d2b584ceb6745 100644 --- a/metadata-io/build.gradle +++ b/metadata-io/build.gradle @@ -38,6 +38,7 @@ dependencies { api externalDependency.kafkaClients api externalDependency.ebean enhance externalDependency.ebeanAgent + implementation externalDependency.ebeanDdl implementation externalDependency.opentelemetryAnnotations implementation externalDependency.resilience4j api externalDependency.springContext @@ -71,6 +72,7 @@ dependencies { // logback >=1.3 required due to `testcontainers` only testImplementation 'ch.qos.logback:logback-classic:1.4.7' + testImplementation 'net.datafaker:datafaker:1.9.0' testAnnotationProcessor externalDependency.lombok 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 35ed8c7905169..a69c6008fea47 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 @@ -25,6 +25,9 @@ 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.event.EventProducer; import com.linkedin.metadata.graph.LineageDirection; import com.linkedin.metadata.query.AutoCompleteResult; @@ -60,6 +63,7 @@ import java.util.Set; import java.util.function.Supplier; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.RequiredArgsConstructor; @@ -535,8 +539,16 @@ public String ingestProposal(@Nonnull final MetadataChangeProposal metadataChang final List additionalChanges = AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService); - Urn urn = _entityService.ingestProposal(metadataChangeProposal, auditStamp, async).getUrn(); - additionalChanges.forEach(proposal -> _entityService.ingestProposal(proposal, auditStamp, async)); + Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), + additionalChanges.stream()); + AspectsBatch batch = AspectsBatchImpl.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(); + + IngestResult one = _entityService.ingestProposal(batch, auditStamp, async).stream() + .findFirst().get(); + + Urn urn = one.getUrn(); tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata()); return urn.toString(); } 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 9967df9207ec7..2d5c5e23ae528 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 @@ -3,8 +3,10 @@ import com.linkedin.common.urn.Urn; 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; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -12,6 +14,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Function; import java.util.function.Supplier; /** @@ -45,37 +48,45 @@ public interface AspectDao { List getAspectsInRange(@Nonnull Urn urn, Set aspectNames, long startTimeMillis, long endTimeMillis); @Nullable - EntityAspect getLatestAspect(@Nonnull final String urn, @Nonnull final String aspectName); + default EntityAspect getLatestAspect(@Nonnull final String urn, @Nonnull final String aspectName) { + return getLatestAspects(Map.of(urn, Set.of(aspectName))).getOrDefault(urn, Map.of()) + .getOrDefault(aspectName, null); + } + + @Nonnull + Map> getLatestAspects(Map> urnAspects); void saveAspect( - @Nonnull final String urn, - @Nonnull final String aspectName, - @Nonnull final String aspectMetadata, - @Nonnull final String actor, - @Nullable final String impersonator, - @Nonnull final Timestamp timestamp, - @Nonnull final String systemMetadata, - final long version, - final boolean insert); + @Nullable Transaction tx, + @Nonnull final String urn, + @Nonnull final String aspectName, + @Nonnull final String aspectMetadata, + @Nonnull final String actor, + @Nullable final String impersonator, + @Nonnull final Timestamp timestamp, + @Nonnull final String systemMetadata, + final long version, + final boolean insert); - void saveAspect(@Nonnull final EntityAspect aspect, final boolean insert); + void saveAspect(@Nullable Transaction tx, @Nonnull final EntityAspect aspect, final boolean insert); long saveLatestAspect( - @Nonnull final String urn, - @Nonnull final String aspectName, - @Nullable final String oldAspectMetadata, - @Nullable final String oldActor, - @Nullable final String oldImpersonator, - @Nullable final Timestamp oldTime, - @Nullable final String oldSystemMetadata, - @Nonnull final String newAspectMetadata, - @Nonnull final String newActor, - @Nullable final String newImpersonator, - @Nonnull final Timestamp newTime, - @Nullable final String newSystemMetadata, - final Long nextVersion); - - void deleteAspect(@Nonnull final EntityAspect aspect); + @Nullable Transaction tx, + @Nonnull final String urn, + @Nonnull final String aspectName, + @Nullable final String oldAspectMetadata, + @Nullable final String oldActor, + @Nullable final String oldImpersonator, + @Nullable final Timestamp oldTime, + @Nullable final String oldSystemMetadata, + @Nonnull final String newAspectMetadata, + @Nonnull final String newActor, + @Nullable final String newImpersonator, + @Nonnull final Timestamp newTime, + @Nullable final String newSystemMetadata, + final Long nextVersion); + + void deleteAspect(@Nullable Transaction tx, @Nonnull final EntityAspect aspect); @Nonnull ListResult listUrns( @@ -92,7 +103,7 @@ Integer countAspect( @Nonnull PagedList getPagedAspects(final RestoreIndicesArgs args); - int deleteUrn(@Nonnull final String urn); + int deleteUrn(@Nullable Transaction tx, @Nonnull final String urn); @Nonnull ListResult listLatestAspectMetadata( @@ -109,16 +120,28 @@ ListResult listAspectMetadata( final int start, final int pageSize); - long getNextVersion(@Nonnull final String urn, @Nonnull final String aspectName); + Map> getNextVersions(@Nonnull Map> urnAspectMap); + + default long getNextVersion(@Nonnull final String urn, @Nonnull final String aspectName) { + return getNextVersions(urn, Set.of(aspectName)).get(aspectName); + } - Map getNextVersions(@Nonnull final String urn, @Nonnull final Set aspectNames); + default Map getNextVersions(@Nonnull final String urn, @Nonnull final Set aspectNames) { + return getNextVersions(Map.of(urn, aspectNames)).get(urn); + } long getMaxVersion(@Nonnull final String urn, @Nonnull final String aspectName); void setWritable(boolean canWrite); @Nonnull - T runInTransactionWithRetry(@Nonnull final Supplier block, final int maxTransactionRetry); + T runInTransactionWithRetry(@Nonnull final Function block, final int maxTransactionRetry); + + @Nonnull + default T runInTransactionWithRetry(@Nonnull final Function block, AspectsBatch batch, + final int maxTransactionRetry) { + return runInTransactionWithRetry(block, maxTransactionRetry); + } default void incrementWriteMetrics(String aspectName, long count, long bytes) { MetricUtils.counter(this.getClass(), 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 e070944b49a05..03b81cfc703c4 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 @@ -4,14 +4,6 @@ import com.linkedin.metadata.config.PreProcessHooks; import com.datahub.util.RecordUtils; import com.datahub.util.exception.ModelConversionException; -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.core.StreamReadConstraints; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.github.fge.jsonpatch.JsonPatch; -import com.github.fge.jsonpatch.JsonPatchException; -import com.github.fge.jsonpatch.Patch; -import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterators; @@ -25,10 +17,7 @@ import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.common.urn.VersionedUrnUtils; -import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.schema.TyperefDataSchema; -import com.linkedin.data.schema.validation.ValidationResult; -import com.linkedin.data.schema.validator.Validator; import com.linkedin.data.template.DataTemplateUtil; import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.template.StringArray; @@ -45,19 +34,20 @@ import com.linkedin.metadata.aspect.Aspect; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.transactions.AbstractBatchItem; +import com.linkedin.metadata.entity.ebean.transactions.PatchBatchItem; +import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; -import com.linkedin.metadata.entity.validation.EntityRegistryUrnValidator; -import com.linkedin.metadata.entity.validation.RecordTemplateValidator; -import com.linkedin.metadata.entity.validation.ValidationUtils; +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; import com.linkedin.metadata.models.RelationshipFieldSpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.models.registry.template.AspectTemplateEngine; import com.linkedin.metadata.query.ListUrnsResult; import com.linkedin.metadata.run.AspectRowSummary; import com.linkedin.metadata.search.utils.BrowsePathV2Utils; @@ -74,9 +64,8 @@ import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; import io.ebean.PagedList; -import java.io.IOException; + import java.net.URISyntaxException; -import java.net.URLEncoder; import java.nio.charset.StandardCharsets; import java.sql.Timestamp; import java.util.ArrayList; @@ -84,19 +73,24 @@ import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.concurrent.TimeUnit; -import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import javax.persistence.EntityNotFoundException; + +import io.ebean.Transaction; import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.*; @@ -126,9 +120,9 @@ * will have version 4. The "true" latest version of an aspect is always equal to the highest stored version * of a given aspect + 1. * - * Note that currently, implementations of this interface are responsible for producing Metadata Audit Events on - * ingestion using {@link #produceMetadataChangeLog(Urn, String, String, AspectSpec, RecordTemplate, RecordTemplate, - * SystemMetadata, SystemMetadata, AuditStamp, ChangeType)}. + * Note that currently, implementations of this interface are responsible for producing Metadata Change Log on + * ingestion using {@link #conditionallyProduceMCLAsync(RecordTemplate, SystemMetadata, RecordTemplate, SystemMetadata, + * MetadataChangeProposal, Urn, AuditStamp, AspectSpec)}. * * TODO: Consider whether we can abstract away virtual versioning semantics to subclasses of this class. */ @@ -140,13 +134,7 @@ public class EntityServiceImpl implements EntityService { * monotonically increasing version incrementing as usual once the latest version is replaced. */ - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - static { - int maxSize = Integer.parseInt(System.getenv().getOrDefault(INGESTION_MAX_SERIALIZED_STRING_LENGTH, MAX_JACKSON_STRING_SIZE)); - OBJECT_MAPPER.getFactory().setStreamReadConstraints(StreamReadConstraints.builder().maxStringLength(maxSize).build()); - } - - private static final int DEFAULT_MAX_TRANSACTION_RETRY = 4; + private static final int DEFAULT_MAX_TRANSACTION_RETRY = 3; protected final AspectDao _aspectDao; private final EventProducer _producer; @@ -158,10 +146,6 @@ public class EntityServiceImpl implements EntityService { private final PreProcessHooks _preProcessHooks; protected static final int MAX_KEYS_PER_QUERY = 500; - private static final int URN_NUM_BYTES_LIMIT = 512; - - // TODO(iprentic): Move this to a common utils location once used in other places - private static final String DELIMITER_SEPARATOR = "␟"; private final Integer ebeanMaxTransactionRetry; @@ -194,8 +178,6 @@ public EntityServiceImpl( ebeanMaxTransactionRetry = retry != null ? retry : DEFAULT_MAX_TRANSACTION_RETRY; } - - /** * Retrieves the latest aspects corresponding to a batch of {@link Urn}s based on a provided * set of aspect names. @@ -204,6 +186,7 @@ public EntityServiceImpl( * @param aspectNames aspects to fetch for each urn in urns set * @return a map of provided {@link Urn} to a List containing the requested aspects. */ + @Override public Map> getLatestAspects( @Nonnull final Set urns, @Nonnull final Set aspectNames) { @@ -220,7 +203,7 @@ public Map> getLatestAspects( // Add "key" aspects for each urn. TODO: Replace this with a materialized key aspect. urnToAspects.keySet().forEach(key -> { - final RecordTemplate keyAspect = buildKeyAspect(key); + final RecordTemplate keyAspect = EntityUtils.buildKeyAspect(_entityRegistry, key); urnToAspects.get(key).add(keyAspect); }); @@ -401,8 +384,7 @@ public Map> getVersionedEnvelopedAspects( .map(UrnUtils::getUrn).collect(Collectors.toSet())); } - private Map> getCorrespondingAspects(Set dbKeys, Set urns) - throws URISyntaxException { + private Map> getCorrespondingAspects(Set dbKeys, Set urns) { final Map envelopedAspectMap = getEnvelopedAspects(dbKeys); @@ -446,34 +428,6 @@ public EnvelopedAspect getLatestEnvelopedAspect( .orElse(null); } - /** - * Retrieves the specific version of the aspect for the given urn - * - * @param entityName name of the entity to fetch - * @param urn urn to fetch - * @param aspectName name of the aspect to fetch - * @param version version to fetch - * @return {@link EnvelopedAspect} object, or null if one cannot be found - */ - @Override - public EnvelopedAspect getEnvelopedAspect( - // TODO: entityName is only used for a debug statement, can we remove this as a param? - String entityName, - @Nonnull Urn urn, - @Nonnull String aspectName, - long version) throws Exception { - log.debug(String.format("Invoked getEnvelopedAspect with entityName: %s, urn: %s, aspectName: %s, version: %s", - urn.getEntityType(), - urn, - aspectName, - version)); - - version = calculateVersionNumber(urn, aspectName, version); - - final EntityAspectIdentifier primaryKey = new EntityAspectIdentifier(urn.toString(), aspectName, version); - return getEnvelopedAspects(ImmutableSet.of(primaryKey)).get(primaryKey); - } - /** * Retrieves an {@link VersionedAspect}, or null if one cannot be found. */ @@ -543,227 +497,192 @@ public ListResult listLatestAspects( aspectMetadataList.getPageSize()); } - - @Nonnull - protected UpdateAspectResult wrappedIngestAspectToLocalDB(@Nonnull final Urn urn, @Nonnull final String aspectName, - @Nonnull final Function, RecordTemplate> updateLambda, - @Nonnull final AuditStamp auditStamp, @Nonnull final SystemMetadata systemMetadata) { - validateUrn(urn); - validateAspect(urn, updateLambda.apply(null)); - return ingestAspectToLocalDB(urn, aspectName, updateLambda, auditStamp, systemMetadata); - } - - @Nonnull - private List> wrappedIngestAspectsToLocalDB(@Nonnull final Urn urn, - @Nonnull List> aspectRecordsToIngest, - @Nonnull final AuditStamp auditStamp, @Nonnull final SystemMetadata providedSystemMetadata) { - validateUrn(urn); - aspectRecordsToIngest.forEach(pair -> validateAspect(urn, pair.getSecond())); - return ingestAspectsToLocalDB(urn, aspectRecordsToIngest, auditStamp, providedSystemMetadata); - } - - // Validates urn subfields using EntityRegistryUrnValidator and does basic field validation for type alignment - // due to validator logic which inherently does coercion - private void validateAspect(Urn urn, RecordTemplate aspect) { - EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(_entityRegistry); - validator.setCurrentEntitySpec(_entityRegistry.getEntitySpec(urn.getEntityType())); - validateAspect(urn, aspect, validator); + /** + * Common batch-like pattern used primarily in tests. + * @param entityUrn the entity urn + * @param pairList list of aspects in pairs of aspect name and record template + * @param auditStamp audit stamp + * @param systemMetadata system metadata + * @return update result + */ + @Override + public List ingestAspects(@Nonnull Urn entityUrn, + List> pairList, + @Nonnull final AuditStamp auditStamp, + SystemMetadata systemMetadata) { + List items = pairList.stream() + .map(pair -> UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(pair.getKey()) + .aspect(pair.getValue()) + .systemMetadata(systemMetadata) + .build(_entityRegistry)) + .collect(Collectors.toList()); + return ingestAspects(AspectsBatchImpl.builder().items(items).build(), auditStamp, true, true); } - private void validateAspect(Urn urn, RecordTemplate aspect, Validator validator) { - Consumer resultFunction = validationResult -> { - throw new IllegalArgumentException("Invalid format for aspect: " + aspect + " for entity: " + urn + "\n Cause: " - + validationResult.getMessages()); }; - RecordTemplateValidator.validate(buildKeyAspect(urn), resultFunction, validator); - RecordTemplateValidator.validate(aspect, resultFunction, validator); - } /** - * Checks whether there is an actual update to the aspect by applying the updateLambda - * If there is an update, push the new version into the local DB. - * Otherwise, do not push the new version, but just update the system metadata. - * DO NOT CALL DIRECTLY, USE WRAPPED METHODS TO VALIDATE URN + * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataChangeLog}. * - * @param urn an urn associated with the new aspect - * @param aspectName name of the aspect being inserted - * @param updateLambda 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 * @param providedSystemMetadata - * @return Details about the new and old version of the aspect + * @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 */ - @Nonnull - @Deprecated - protected UpdateAspectResult ingestAspectToLocalDB( - @Nonnull final Urn urn, - @Nonnull final String aspectName, - @Nonnull final Function, RecordTemplate> updateLambda, - @Nonnull final AuditStamp auditStamp, - @Nonnull final SystemMetadata providedSystemMetadata) { + @Override + public List ingestAspects(@Nonnull final AspectsBatch aspectsBatch, + @Nonnull final AuditStamp auditStamp, + boolean emitMCL, + boolean overwrite) { - return _aspectDao.runInTransactionWithRetry(() -> { - final String urnStr = urn.toString(); - final EntityAspect latest = _aspectDao.getLatestAspect(urnStr, aspectName); - long nextVersion = _aspectDao.getNextVersion(urnStr, aspectName); + Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectsToLocalDB").time(); + List ingestResults = ingestAspectsToLocalDB(aspectsBatch, auditStamp, overwrite); + List mclResults = emitMCL(ingestResults, emitMCL); + ingestToLocalDBTimer.stop(); - return ingestAspectToLocalDBNoTransaction(urn, aspectName, updateLambda, auditStamp, providedSystemMetadata, latest, nextVersion); - }, DEFAULT_MAX_TRANSACTION_RETRY); + return mclResults; } /** - * Apply patch update to aspect within a single transaction + * Checks whether there is an actual update to the aspect by applying the updateLambda + * If there is an update, push the new version into the local DB. + * Otherwise, do not push the new version, but just update the system metadata. * - * @param urn an urn associated with the new aspect - * @param aspectSpec AspectSpec of the aspect to update - * @param jsonPatch JsonPatch to apply to the aspect - * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time * @param providedSystemMetadata + * @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 - @Deprecated - protected UpdateAspectResult patchAspectToLocalDB( - @Nonnull final Urn urn, - @Nonnull final AspectSpec aspectSpec, - @Nonnull final Patch jsonPatch, - @Nonnull final AuditStamp auditStamp, - @Nonnull final SystemMetadata providedSystemMetadata) { - - return _aspectDao.runInTransactionWithRetry(() -> { - final String urnStr = urn.toString(); - final String aspectName = aspectSpec.getName(); - final EntityAspect latest = _aspectDao.getLatestAspect(urnStr, aspectName); - final long nextVersion = _aspectDao.getNextVersion(urnStr, aspectName); - try { + private List ingestAspectsToLocalDB(@Nonnull final AspectsBatch aspectsBatch, + @Nonnull final AuditStamp auditStamp, + boolean overwrite) { - final RecordTemplate currentValue = latest != null - ? EntityUtils.toAspectRecord(urn, aspectName, latest.getMetadata(), _entityRegistry) - : _entityRegistry.getAspectTemplateEngine().getDefaultTemplate(aspectSpec.getName()); - - if (latest == null && currentValue == null) { - // Attempting to patch a value to an aspect which has no default value and no existing value. - throw new UnsupportedOperationException(String.format("Patch not supported for aspect with name %s. " - + "Default aspect is required because no aspect currently exists for urn %s.", aspectName, urn)); - } - - final RecordTemplate updatedValue = _entityRegistry.getAspectTemplateEngine().applyPatch(currentValue, jsonPatch, aspectSpec); + if (aspectsBatch.containsDuplicateAspects()) { + log.warn(String.format("Batch contains duplicates: %s", aspectsBatch)); + } - validateAspect(urn, updatedValue); - return ingestAspectToLocalDBNoTransaction(urn, aspectName, ignored -> updatedValue, auditStamp, providedSystemMetadata, - latest, nextVersion); - } catch (JsonProcessingException | JsonPatchException e) { - throw new IllegalStateException(e); + return _aspectDao.runInTransactionWithRetry((tx) -> { + // Read before write is unfortunate, however batch it + Map> urnAspects = aspectsBatch.getUrnAspectsMap(); + // read #1 + Map> latestAspects = _aspectDao.getLatestAspects(urnAspects); + // read #2 + Map> nextVersions = _aspectDao.getNextVersions(urnAspects); + + 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()); + + // Database Upsert results + List upsertResults = items.stream() + .map(item -> { + final String urnStr = item.getUrn().toString(); + final EntityAspect latest = latestAspects.getOrDefault(urnStr, Map.of()).get(item.getAspectName()); + final long nextVersion = nextVersions.getOrDefault(urnStr, Map.of()).getOrDefault(item.getAspectName(), 0L); + + final UpdateAspectResult result; + if (overwrite || latest == null) { + result = ingestAspectToLocalDB(tx, item.getUrn(), item.getAspectName(), item.getAspect(), + auditStamp, item.getSystemMetadata(), latest, nextVersion).toBuilder().request(item).build(); + + // support inner-batch upserts + latestAspects.computeIfAbsent(urnStr, key -> new HashMap<>()).put(item.getAspectName(), item.toLatestEntityAspect(auditStamp)); + 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()); + result = UpdateAspectResult.builder() + .urn(item.getUrn()) + .request(item) + .oldValue(oldValue) + .newValue(oldValue) + .oldSystemMetadata(oldMetadata) + .newSystemMetadata(oldMetadata) + .operation(MetadataAuditOperation.UPDATE) + .auditStamp(auditStamp) + .maxVersion(latest.getVersion()) + .build(); + } + + return result; + }).collect(Collectors.toList()); + + // commit upserts prior to retention or kafka send, if supported by impl + if (tx != null) { + tx.commitAndContinue(); } - }, DEFAULT_MAX_TRANSACTION_RETRY); - } - /** - * Same as ingestAspectToLocalDB but for multiple aspects - * DO NOT CALL DIRECTLY, USE WRAPPED METHODS TO VALIDATE URN - */ - @Nonnull - @Deprecated - protected List> ingestAspectsToLocalDB( - @Nonnull final Urn urn, - @Nonnull List> aspectRecordsToIngest, - @Nonnull final AuditStamp auditStamp, - @Nonnull final SystemMetadata systemMetadata) { - - return _aspectDao.runInTransactionWithRetry(() -> { - - final Set aspectNames = aspectRecordsToIngest - .stream() - .map(Pair::getFirst) - .collect(Collectors.toSet()); - - Map latestAspects = getLatestAspectForUrn(urn, aspectNames); - Map nextVersions = _aspectDao.getNextVersions(urn.toString(), aspectNames); - - List> result = new ArrayList<>(); - for (Pair aspectRecord: aspectRecordsToIngest) { - String aspectName = aspectRecord.getFirst(); - RecordTemplate newValue = aspectRecord.getSecond(); - EntityAspect latest = latestAspects.get(aspectName); - long nextVersion = nextVersions.get(aspectName); - UpdateAspectResult updateResult = ingestAspectToLocalDBNoTransaction(urn, aspectName, ignored -> newValue, auditStamp, systemMetadata, - latest, nextVersion); - result.add(new Pair<>(aspectName, updateResult)); + // Retention optimization and tx + if (_retentionService != null) { + List retentionBatch = upsertResults.stream() + // Only consider retention when there was a previous version + .filter(result -> latestAspects.containsKey(result.getUrn().toString()) + && latestAspects.get(result.getUrn().toString()).containsKey(result.getRequest().getAspectName())) + .filter(result -> { + RecordTemplate oldAspect = result.getOldValue(); + RecordTemplate newAspect = result.getNewValue(); + // Apply retention policies if there was an update to existing aspect value + return oldAspect != newAspect && oldAspect != null && _retentionService != null; + }) + .map(result -> RetentionService.RetentionContext.builder() + .urn(result.getUrn()) + .aspectName(result.getRequest().getAspectName()) + .maxVersion(Optional.of(result.getMaxVersion())) + .build()) + .collect(Collectors.toList()); + _retentionService.applyRetentionWithPolicyDefaults(retentionBatch); + } else { + log.warn("Retention service is missing!"); } - return result; - }, DEFAULT_MAX_TRANSACTION_RETRY); - } - - @Nonnull - protected SystemMetadata generateSystemMetadataIfEmpty(@Nullable SystemMetadata systemMetadata) { - if (systemMetadata == null) { - systemMetadata = new SystemMetadata(); - systemMetadata.setRunId(DEFAULT_RUN_ID); - systemMetadata.setLastObserved(System.currentTimeMillis()); - } - return systemMetadata; - } - - @VisibleForTesting - void validateUrn(@Nonnull final Urn urn) { - EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(_entityRegistry); - validator.setCurrentEntitySpec(_entityRegistry.getEntitySpec(urn.getEntityType())); - RecordTemplateValidator.validate(buildKeyAspect(urn), validationResult -> { - throw new IllegalArgumentException("Invalid urn: " + urn + "\n Cause: " - + validationResult.getMessages()); }, validator); - if (urn.toString().trim().length() != urn.toString().length()) { - throw new IllegalArgumentException("Error: cannot provide an URN with leading or trailing whitespace"); - } - if (URLEncoder.encode(urn.toString()).length() > URN_NUM_BYTES_LIMIT) { - throw new IllegalArgumentException("Error: cannot provide an URN longer than " + Integer.toString(URN_NUM_BYTES_LIMIT) + " bytes (when URL encoded)"); - } - if (urn.toString().contains(DELIMITER_SEPARATOR)) { - throw new IllegalArgumentException("Error: URN cannot contain " + DELIMITER_SEPARATOR + " character"); - } - try { - Urn.createFromString(urn.toString()); - } catch (URISyntaxException e) { - throw new IllegalArgumentException(e); - } + return upsertResults; + }, aspectsBatch, DEFAULT_MAX_TRANSACTION_RETRY); } - @Override - public void ingestAspects(@Nonnull final Urn urn, @Nonnull List> aspectRecordsToIngest, - @Nonnull final AuditStamp auditStamp, @Nullable SystemMetadata systemMetadata) { - - systemMetadata = generateSystemMetadataIfEmpty(systemMetadata); - - Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectsToLocalDB").time(); - List> ingestResults = wrappedIngestAspectsToLocalDB(urn, aspectRecordsToIngest, auditStamp, systemMetadata); - ingestToLocalDBTimer.stop(); - - for (Pair result: ingestResults) { - sendEventForUpdateAspectResult(urn, result.getFirst(), result.getSecond()); + @Nonnull + private List emitMCL(List sqlResults, boolean emitMCL) { + List withEmitMCL = sqlResults.stream() + .map(result -> emitMCL ? conditionallyProduceMCLAsync(result) : result) + .collect(Collectors.toList()); + + // join futures messages, capture error state + List> statusPairs = withEmitMCL.stream() + .filter(result -> result.getMclFuture() != null) + .map(result -> { + try { + result.getMclFuture().get(); + return Pair.of(true, result); + } catch (InterruptedException | ExecutionException e) { + return Pair.of(false, result); + } + }).collect(Collectors.toList()); + + if (statusPairs.stream().anyMatch(p -> !p.getFirst())) { + log.error("Failed to produce MCLs: {}", statusPairs.stream() + .filter(p -> !p.getFirst()) + .map(Pair::getValue) + .map(v -> v.getRequest().toString()) + .collect(Collectors.toList())); + // TODO restoreIndices? + throw new RuntimeException("Failed to produce MCLs"); } - } - /** - * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. - * - * Note that in general, this should not be used externally. It is currently serving upgrade scripts and - * is as such public. - * - * @param urn an urn associated with the new aspect - * @param aspectName name of the aspect being inserted - * @param newValue value of the aspect being inserted - * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time - * @param systemMetadata - * @return the {@link RecordTemplate} representation of the written aspect object - */ - @Override - public RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName, - @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp, @Nullable SystemMetadata systemMetadata) { - - log.debug("Invoked ingestAspect with urn: {}, aspectName: {}, newValue: {}", urn, aspectName, newValue); - - systemMetadata = generateSystemMetadataIfEmpty(systemMetadata); - - Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectToLocalDB").time(); - UpdateAspectResult result = wrappedIngestAspectToLocalDB(urn, aspectName, ignored -> newValue, auditStamp, systemMetadata); - ingestToLocalDBTimer.stop(); - - return sendEventForUpdateAspectResult(urn, aspectName, result); + return withEmitMCL; } /** @@ -783,94 +702,39 @@ public RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String */ @Nullable @Override - public RecordTemplate ingestAspectIfNotPresent(@Nonnull Urn urn, @Nonnull String aspectName, - @Nonnull RecordTemplate newValue, @Nonnull AuditStamp auditStamp, @Nullable SystemMetadata systemMetadata) { + public RecordTemplate ingestAspectIfNotPresent(@Nonnull Urn urn, + @Nonnull String aspectName, + @Nonnull RecordTemplate newValue, + @Nonnull AuditStamp auditStamp, + @Nonnull SystemMetadata systemMetadata) { log.debug("Invoked ingestAspectIfNotPresent with urn: {}, aspectName: {}, newValue: {}", urn, aspectName, newValue); - final SystemMetadata internalSystemMetadata = generateSystemMetadataIfEmpty(systemMetadata); - - Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectToLocalDB").time(); - UpdateAspectResult result = _aspectDao.runInTransactionWithRetry(() -> { - final String urnStr = urn.toString(); - final EntityAspect latest = _aspectDao.getLatestAspect(urnStr, aspectName); - if (latest == null) { - long nextVersion = _aspectDao.getNextVersion(urnStr, aspectName); - - return ingestAspectToLocalDBNoTransaction(urn, aspectName, ignored -> newValue, auditStamp, - internalSystemMetadata, latest, nextVersion); - } - RecordTemplate oldValue = EntityUtils.toAspectRecord(urn, aspectName, latest.getMetadata(), getEntityRegistry()); - SystemMetadata oldMetadata = EntityUtils.parseSystemMetadata(latest.getSystemMetadata()); - return new UpdateAspectResult(urn, oldValue, oldValue, oldMetadata, oldMetadata, MetadataAuditOperation.UPDATE, auditStamp, - latest.getVersion()); - }, DEFAULT_MAX_TRANSACTION_RETRY); - ingestToLocalDBTimer.stop(); + AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() + .one(UpsertBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .aspect(newValue) + .systemMetadata(systemMetadata) + .build(_entityRegistry)) + .build(); + List ingested = ingestAspects(aspectsBatch, auditStamp, true, false); - return sendEventForUpdateAspectResult(urn, aspectName, result); - } - - protected RecordTemplate sendEventForUpdateAspectResult(@Nonnull final Urn urn, @Nonnull final String aspectName, - @Nonnull UpdateAspectResult result) { - - final RecordTemplate oldValue = result.getOldValue(); - final RecordTemplate updatedValue = result.getNewValue(); - final SystemMetadata oldSystemMetadata = result.getOldSystemMetadata(); - final SystemMetadata updatedSystemMetadata = result.getNewSystemMetadata(); - - // Apply retention policies asynchronously if there was an update to existing aspect value - if (oldValue != updatedValue && oldValue != null && _retentionService != null) { - _retentionService.applyRetention(urn, aspectName, - Optional.of(new RetentionService.RetentionContext(Optional.of(result.getMaxVersion())))); - } - - // Produce MCL after a successful update - boolean isNoOp = oldValue == updatedValue; - if (!isNoOp || _alwaysEmitChangeLog || shouldAspectEmitChangeLog(urn, aspectName)) { - log.debug(String.format("Producing MetadataChangeLog for ingested aspect %s, urn %s", aspectName, urn)); - String entityName = urnToEntityName(urn); - EntitySpec entitySpec = getEntityRegistry().getEntitySpec(entityName); - AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); - if (aspectSpec == null) { - throw new RuntimeException(String.format("Unknown aspect %s for entity %s", aspectName, entityName)); - } - - Timer.Context produceMCLTimer = MetricUtils.timer(this.getClass(), "produceMCL").time(); - produceMetadataChangeLog(urn, entityName, aspectName, aspectSpec, oldValue, updatedValue, oldSystemMetadata, - updatedSystemMetadata, result.getAuditStamp(), isNoOp ? ChangeType.RESTATE : ChangeType.UPSERT); - produceMCLTimer.stop(); - - // For legacy reasons, keep producing to the MAE event stream without blocking ingest - try { - Timer.Context produceMAETimer = MetricUtils.timer(this.getClass(), "produceMAE").time(); - produceMetadataAuditEvent(urn, aspectName, oldValue, updatedValue, result.getOldSystemMetadata(), - result.getNewSystemMetadata(), MetadataAuditOperation.UPDATE); - produceMAETimer.stop(); - } catch (Exception e) { - log.warn("Unable to produce legacy MAE, entity may not have legacy Snapshot schema.", e); - } - } else { - log.debug("Skipped producing MetadataAuditEvent for ingested aspect {}, urn {}. Aspect has not changed.", - aspectName, urn); - } - return updatedValue; + return ingested.stream().findFirst().get().getNewValue(); } /** - * Validates that a change type is valid for the given aspect - * @param changeType - * @param aspectSpec - * @return + * Wrapper around batch method for single item + * @param proposal the proposal + * @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 */ - private boolean isValidChangeType(ChangeType changeType, AspectSpec aspectSpec) { - if (aspectSpec.isTimeseries()) { - // Timeseries aspects only support UPSERT - return ChangeType.UPSERT.equals(changeType); - } else { - return (ChangeType.UPSERT.equals(changeType) || ChangeType.PATCH.equals(changeType)); - } + @Override + public IngestResult ingestProposal(MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async) { + return ingestProposal(AspectsBatchImpl.builder().mcps(List.of(proposal), getEntityRegistry()).build(), auditStamp, + async).stream().findFirst().get(); } - /** * Ingest a new {@link MetadataChangeProposal}. Note that this method does NOT include any additional aspects or do any * enrichment, instead it changes only those which are provided inside the metadata change proposal. @@ -878,174 +742,127 @@ private boolean isValidChangeType(ChangeType changeType, AspectSpec aspectSpec) * Do not use this method directly for creating new entities, as it DOES NOT create an Entity Key aspect in the DB. Instead, * use an Entity Client. * - * @param mcp the proposal to ingest + * @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 IngestProposalResult} containing the results + * @return an {@link IngestResult} containing the results */ @Override - public IngestProposalResult ingestProposal(@Nonnull MetadataChangeProposal mcp, - AuditStamp auditStamp, final boolean async) { - - log.debug("entity type = {}", mcp.getEntityType()); - EntitySpec entitySpec = getEntityRegistry().getEntitySpec(mcp.getEntityType()); - log.debug("entity spec = {}", entitySpec); - - Urn entityUrn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); + public Set ingestProposal(AspectsBatch aspectsBatch, AuditStamp auditStamp, final boolean async) { - AspectSpec aspectSpec = validateAspect(mcp, entitySpec); + Stream timeseriesIngestResults = ingestTimeseriesProposal(aspectsBatch, auditStamp); + Stream nonTimeseriesIngestResults = async ? ingestProposalAsync(aspectsBatch) + : ingestProposalSync(aspectsBatch, auditStamp); - log.debug("aspect spec = {}", aspectSpec); - - if (!isValidChangeType(mcp.getChangeType(), aspectSpec)) { - throw new UnsupportedOperationException( - "ChangeType not supported: " + mcp.getChangeType() + " for aspect " + mcp.getAspectName()); - } - - SystemMetadata systemMetadata = generateSystemMetadataIfEmpty(mcp.getSystemMetadata()); - systemMetadata.setRegistryName(aspectSpec.getRegistryName()); - systemMetadata.setRegistryVersion(aspectSpec.getRegistryVersion().toString()); - - RecordTemplate oldAspect = null; - SystemMetadata oldSystemMetadata = null; - RecordTemplate newAspect; - SystemMetadata newSystemMetadata; - - if (!aspectSpec.isTimeseries()) { - if (!async) { - // When async mode is turned off, we write to primary store for non timeseries aspects - UpdateAspectResult result; - switch (mcp.getChangeType()) { - case UPSERT: - result = performUpsert(mcp, aspectSpec, systemMetadata, entityUrn, auditStamp); - break; - case PATCH: - result = performPatch(mcp, aspectSpec, systemMetadata, entityUrn, auditStamp); - break; - default: - // Should never reach since we throw error above - throw new UnsupportedOperationException("ChangeType not supported: " + mcp.getChangeType()); - } - oldAspect = result != null ? result.getOldValue() : null; - oldSystemMetadata = result != null ? result.getOldSystemMetadata() : null; - newAspect = result != null ? result.getNewValue() : null; - newSystemMetadata = result != null ? result.getNewSystemMetadata() : null; - } else { - // When async is turned on, we write to proposal log and return without waiting - _producer.produceMetadataChangeProposal(entityUrn, mcp); - return new IngestProposalResult(entityUrn, false, true); - } - } else { - // For timeseries aspects - newAspect = convertToRecordTemplate(mcp, aspectSpec); - newSystemMetadata = mcp.getSystemMetadata(); - } - - boolean didUpdate = - emitChangeLog(oldAspect, oldSystemMetadata, newAspect, newSystemMetadata, mcp, entityUrn, auditStamp, - aspectSpec); - - return new IngestProposalResult(entityUrn, didUpdate, false); + return Stream.concat(timeseriesIngestResults, nonTimeseriesIngestResults).collect(Collectors.toSet()); } - private 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())); + /** + * 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 = aspectsBatch.getItems().stream() + .filter(item -> item.getAspectSpec().isTimeseries() && item.getChangeType() != ChangeType.UPSERT) + .collect(Collectors.toList()); + if (!unsupported.isEmpty()) { + throw new UnsupportedOperationException("ChangeType not supported: " + unsupported.stream() + .map(AbstractBatchItem::getChangeType).collect(Collectors.toSet())); } - return aspectSpec; - } - - private UpdateAspectResult performUpsert(MetadataChangeProposal mcp, AspectSpec aspectSpec, SystemMetadata - systemMetadata, Urn entityUrn, AuditStamp auditStamp) { - RecordTemplate aspect = convertToRecordTemplate(mcp, aspectSpec); - log.debug("aspect = {}", aspect); - - return upsertAspect(aspect, systemMetadata, mcp, entityUrn, auditStamp, aspectSpec); + List, Boolean>>>> timeseriesResults = aspectsBatch.getItems().stream() + .filter(item -> item.getAspectSpec().isTimeseries()) + .map(item -> (UpsertBatchItem) item) + .map(item -> Pair.of(item, conditionallyProduceMCLAsync(null, null, item.getAspect(), item.getSystemMetadata(), + item.getMetadataChangeProposal(), item.getUrn(), auditStamp, item.getAspectSpec()))) + .collect(Collectors.toList()); + + return timeseriesResults.stream().map(result -> { + Optional, Boolean>> emissionStatus = result.getSecond(); + + emissionStatus.ifPresent(status -> { + try { + status.getFirst().get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + + UpsertBatchItem request = result.getFirst(); + return IngestResult.builder() + .urn(request.getUrn()) + .request(request) + .publishedMCL(emissionStatus.map(status -> status.getFirst() != null).orElse(false)) + .processedMCL(emissionStatus.map(Pair::getSecond).orElse(false)) + .build(); + }); } - private UpdateAspectResult performPatch(MetadataChangeProposal mcp, AspectSpec aspectSpec, SystemMetadata - systemMetadata, Urn entityUrn, AuditStamp auditStamp) { - if (!supportsPatch(aspectSpec)) { - // 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."); - } - Patch jsonPatch = convertToJsonPatch(mcp); - log.debug("patch = {}", jsonPatch); - - return patchAspect(jsonPatch, systemMetadata, mcp, entityUrn, auditStamp, aspectSpec); - } + /** + * For async ingestion of non-timeseries, any change type + * @param aspectsBatch non-timeseries ingest aspects + * @return produced items to the MCP topic + */ + private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { + List nonTimeseries = aspectsBatch.getItems().stream() + .filter(item -> !item.getAspectSpec().isTimeseries()) + .collect(Collectors.toList()); - private boolean supportsPatch(AspectSpec aspectSpec) { - // Limit initial support to defined templates - return AspectTemplateEngine.SUPPORTED_TEMPLATES.contains(aspectSpec.getName()); - } + List> futures = nonTimeseries.stream().map(item -> + // When async is turned on, we write to proposal log and return without waiting + _producer.produceMetadataChangeProposal(item.getUrn(), item.getMetadataChangeProposal())) + .filter(Objects::nonNull) + .collect(Collectors.toList()); - private RecordTemplate convertToRecordTemplate(MetadataChangeProposal mcp, AspectSpec aspectSpec) { - RecordTemplate aspect; try { - aspect = GenericRecordUtils.deserializeAspect(mcp.getAspect().getValue(), - mcp.getAspect().getContentType(), aspectSpec); - ValidationUtils.validateOrThrow(aspect); - } catch (ModelConversionException e) { - throw new RuntimeException( - String.format("Could not deserialize %s for aspect %s", mcp.getAspect().getValue(), - mcp.getAspectName())); + return nonTimeseries.stream().map(item -> + IngestResult.builder() + .urn(item.getUrn()) + .request(item) + .publishedMCP(true) + .build()); + } finally { + futures.forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); } - log.debug("aspect = {}", aspect); - return aspect; } - private Patch convertToJsonPatch(MetadataChangeProposal mcp) { - JsonNode json; - try { - json = OBJECT_MAPPER.readTree(mcp.getAspect().getValue().asString(StandardCharsets.UTF_8)); - return JsonPatch.fromJson(json); - } catch (IOException e) { - throw new IllegalArgumentException("Invalid JSON Patch: " + mcp.getAspect().getValue(), e); + private Stream ingestProposalSync(AspectsBatch aspectsBatch, AuditStamp auditStamp) { + AspectsBatchImpl nonTimeseries = AspectsBatchImpl.builder() + .items(aspectsBatch.getItems().stream() + .filter(item -> !item.getAspectSpec().isTimeseries()) + .collect(Collectors.toList())) + .build(); + + List unsupported = nonTimeseries.getItems().stream() + .filter(item -> item.getMetadataChangeProposal().getChangeType() != ChangeType.PATCH + && item.getMetadataChangeProposal().getChangeType() != ChangeType.UPSERT) + .collect(Collectors.toList()); + if (!unsupported.isEmpty()) { + throw new UnsupportedOperationException("ChangeType not supported: " + unsupported.stream() + .map(item -> item.getMetadataChangeProposal().getChangeType()).collect(Collectors.toSet())); } - } - private UpdateAspectResult upsertAspect(final RecordTemplate aspect, final SystemMetadata systemMetadata, - MetadataChangeProposal mcp, Urn entityUrn, AuditStamp auditStamp, AspectSpec aspectSpec) { - Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestProposalToLocalDB").time(); - UpdateAspectResult result = - wrappedIngestAspectToLocalDB(entityUrn, mcp.getAspectName(), ignored -> aspect, auditStamp, - systemMetadata); - ingestToLocalDBTimer.stop(); - RecordTemplate oldAspect = result.getOldValue(); - RecordTemplate newAspect = result.getNewValue(); - // Apply retention policies asynchronously if there was an update to existing aspect value - if (oldAspect != newAspect && oldAspect != null && _retentionService != null) { - _retentionService.applyRetention(entityUrn, aspectSpec.getName(), - Optional.of(new RetentionService.RetentionContext(Optional.of(result.getMaxVersion())))); - } - return result; - } + List upsertResults = ingestAspects(nonTimeseries, auditStamp, true, true); - private UpdateAspectResult patchAspect(final Patch patch, final SystemMetadata systemMetadata, - MetadataChangeProposal mcp, Urn entityUrn, AuditStamp auditStamp, AspectSpec aspectSpec) { - Timer.Context patchAspectToLocalDBTimer = MetricUtils.timer(this.getClass(), "patchAspect").time(); - UpdateAspectResult result = patchAspectToLocalDB(entityUrn, aspectSpec, patch, auditStamp, systemMetadata); - patchAspectToLocalDBTimer.stop(); - RecordTemplate oldAspect = result.getOldValue(); - RecordTemplate newAspect = result.getNewValue(); - // Apply retention policies asynchronously if there was an update to existing aspect value - if (oldAspect != newAspect && oldAspect != null && _retentionService != null) { - _retentionService.applyRetention(entityUrn, aspectSpec.getName(), - Optional.of(new RetentionService.RetentionContext(Optional.of(result.getMaxVersion())))); - } - return result; + return upsertResults.stream().map(result -> { + AbstractBatchItem item = result.getRequest(); + + return IngestResult.builder() + .urn(item.getUrn()) + .request(item) + .publishedMCL(result.getMclFuture() != null) + .sqlCommitted(true) + .isUpdate(result.getOldValue() != null) + .build(); + }); } @Override @@ -1070,43 +887,19 @@ public String batchApplyRetention(Integer start, Integer count, Integer attemptW return result.toString(); } - private boolean emitChangeLog(@Nullable RecordTemplate oldAspect, @Nullable SystemMetadata oldSystemMetadata, - RecordTemplate newAspect, SystemMetadata newSystemMetadata, - MetadataChangeProposal mcp, Urn entityUrn, - AuditStamp auditStamp, AspectSpec aspectSpec) { - boolean isNoOp = oldAspect == newAspect; - if (!isNoOp || _alwaysEmitChangeLog || shouldAspectEmitChangeLog(aspectSpec)) { - log.debug("Producing MetadataChangeLog for ingested aspect {}, urn {}", mcp.getAspectName(), entityUrn); - - final MetadataChangeLog metadataChangeLog = constructMCL(mcp, urnToEntityName(entityUrn), entityUrn, - isNoOp ? ChangeType.RESTATE : ChangeType.UPSERT, aspectSpec.getName(), auditStamp, newAspect, newSystemMetadata, - oldAspect, oldSystemMetadata); - - log.debug("Serialized MCL event: {}", metadataChangeLog); - - produceMetadataChangeLog(entityUrn, aspectSpec, metadataChangeLog); - preprocessEvent(metadataChangeLog); - - return true; - } else { - log.debug( - "Skipped producing MetadataChangeLog for ingested aspect {}, urn {}. Aspect has not changed.", - mcp.getAspectName(), entityUrn); - return false; - } - } - - private void preprocessEvent(MetadataChangeLog metadataChangeLog) { + private boolean preprocessEvent(MetadataChangeLog metadataChangeLog) { if (_preProcessHooks.isUiEnabled()) { if (metadataChangeLog.getSystemMetadata() != null) { if (metadataChangeLog.getSystemMetadata().getProperties() != null) { if (UI_SOURCE.equals(metadataChangeLog.getSystemMetadata().getProperties().get(APP_SOURCE))) { // Pre-process the update indices hook for UI updates to avoid perceived lag from Kafka _updateIndicesService.handleChangeEvent(metadataChangeLog); + return true; } } } } + return false; } @Override @@ -1130,6 +923,8 @@ public RestoreIndicesResult restoreIndices(@Nonnull RestoreIndicesArgs args, @No logger.accept(String.format( "Reading rows %s through %s from the aspects table completed.", args.start, args.start + args.batchSize)); + LinkedList> futures = new LinkedList<>(); + for (EbeanAspectV2 aspect : rows != null ? rows.getList() : List.of()) { // 1. Extract an Entity type from the entity Urn result.timeGetRowMs = System.currentTimeMillis() - startTime; @@ -1193,14 +988,21 @@ public RestoreIndicesResult restoreIndices(@Nonnull RestoreIndicesArgs args, @No latestSystemMetadata.setProperties(properties); // 5. Produce MAE events for the aspect record - produceMetadataChangeLog(urn, entityName, aspectName, aspectSpec, null, aspectRecord, null, + futures.add(alwaysProduceMCLAsync(urn, entityName, aspectName, aspectSpec, null, aspectRecord, null, latestSystemMetadata, new AuditStamp().setActor(UrnUtils.getUrn(SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); result.sendMessageMs += System.currentTimeMillis() - startTime; rowsMigrated++; } + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); try { TimeUnit.MILLISECONDS.sleep(args.batchDelayMs); } catch (InterruptedException e) { @@ -1211,40 +1013,6 @@ public RestoreIndicesResult restoreIndices(@Nonnull RestoreIndicesArgs args, @No return result; } - /** - * Updates a particular version of an aspect & optionally emits a {@link com.linkedin.mxe.MetadataAuditEvent}. - * - * Note that in general, this should not be used externally. It is currently serving upgrade scripts and - * is as such public. - * - * @param urn an urn associated with the aspect to update - * @param entityName name of the entity being updated - * @param aspectName name of the aspect being updated - * @param aspectSpec spec of the aspect being updated - * @param newValue new value of the aspect being updated - * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time - * @param version specific version of the aspect being requests - * @param emitMae whether a {@link com.linkedin.mxe.MetadataAuditEvent} should be emitted in correspondence upon - * successful update - * @return the {@link RecordTemplate} representation of the requested aspect object - */ - @Override - public RecordTemplate updateAspect( - @Nonnull final Urn urn, - @Nonnull final String entityName, - @Nonnull final String aspectName, - @Nonnull final AspectSpec aspectSpec, - @Nonnull final RecordTemplate newValue, - @Nonnull final AuditStamp auditStamp, - @Nonnull final long version, - @Nonnull final boolean emitMae) { - log.debug( - "Invoked updateAspect with urn: {}, aspectName: {}, newValue: {}, version: {}, emitMae: {}", urn, - aspectName, newValue, version, emitMae); - return updateAspect(urn, entityName, aspectName, aspectSpec, newValue, auditStamp, version, emitMae, - DEFAULT_MAX_TRANSACTION_RETRY); - } - /** * Lists the entity URNs found in storage. * @@ -1309,60 +1077,59 @@ public Map getEntities(@Nonnull final Set urns, @Nonnull Set, Boolean> alwaysProduceMCLAsync(@Nonnull final Urn urn, @Nonnull final AspectSpec aspectSpec, + @Nonnull final MetadataChangeLog metadataChangeLog) { + Future future = _producer.produceMetadataChangeLog(urn, aspectSpec, metadataChangeLog); + return Pair.of(future, preprocessEvent(metadataChangeLog)); } - protected Snapshot buildKeySnapshot(@Nonnull final Urn urn) { - final RecordTemplate keyAspectValue = buildKeyAspect(urn); - return toSnapshotUnion(toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, keyAspectValue)))); + @Override + public Pair, Boolean> alwaysProduceMCLAsync(@Nonnull final Urn urn, @Nonnull String entityName, @Nonnull String aspectName, + @Nonnull final AspectSpec aspectSpec, @Nullable final RecordTemplate oldAspectValue, + @Nullable final RecordTemplate newAspectValue, @Nullable final SystemMetadata oldSystemMetadata, + @Nullable final SystemMetadata newSystemMetadata, @Nonnull AuditStamp auditStamp, + @Nonnull final ChangeType changeType) { + final MetadataChangeLog metadataChangeLog = constructMCL(null, entityName, urn, changeType, aspectName, auditStamp, + newAspectValue, newSystemMetadata, oldAspectValue, oldSystemMetadata); + return alwaysProduceMCLAsync(urn, aspectSpec, metadataChangeLog); } - @Override - public void produceMetadataAuditEventForKey(@Nonnull final Urn urn, - @Nullable final SystemMetadata newSystemMetadata) { + public Optional, Boolean>> conditionallyProduceMCLAsync(@Nullable RecordTemplate oldAspect, + @Nullable SystemMetadata oldSystemMetadata, + RecordTemplate newAspect, SystemMetadata newSystemMetadata, + @Nullable MetadataChangeProposal mcp, Urn entityUrn, + AuditStamp auditStamp, AspectSpec aspectSpec) { + boolean isNoOp = oldAspect == newAspect; + if (!isNoOp || _alwaysEmitChangeLog || shouldAspectEmitChangeLog(aspectSpec)) { + log.debug("Producing MetadataChangeLog for ingested aspect {}, urn {}", aspectSpec.getName(), entityUrn); - final Snapshot newSnapshot = buildKeySnapshot(urn); + final MetadataChangeLog metadataChangeLog = constructMCL(mcp, urnToEntityName(entityUrn), entityUrn, + isNoOp ? ChangeType.RESTATE : ChangeType.UPSERT, aspectSpec.getName(), auditStamp, newAspect, newSystemMetadata, + oldAspect, oldSystemMetadata); - _producer.produceMetadataAuditEvent(urn, null, newSnapshot, null, newSystemMetadata, MetadataAuditOperation.UPDATE); + log.debug("Serialized MCL event: {}", metadataChangeLog); + Pair, Boolean> emissionStatus = alwaysProduceMCLAsync(entityUrn, aspectSpec, metadataChangeLog); + return emissionStatus.getFirst() != null ? Optional.of(emissionStatus) : Optional.empty(); + } else { + log.debug( + "Skipped producing MetadataChangeLog for ingested aspect {}, urn {}. Aspect has not changed.", + aspectSpec.getName(), entityUrn); + return Optional.empty(); + } } - /** - * Produces a {@link com.linkedin.mxe.MetadataChangeLog} from a - * new & previous aspect. - * - * @param urn the urn associated with the entity changed - * @param aspectSpec AspectSpec of the aspect being updated - * @param metadataChangeLog metadata change log to push into MCL kafka topic - */ - @Override - public void produceMetadataChangeLog(@Nonnull final Urn urn, AspectSpec aspectSpec, - @Nonnull final MetadataChangeLog metadataChangeLog) { - _producer.produceMetadataChangeLog(urn, aspectSpec, metadataChangeLog); - } + private UpdateAspectResult conditionallyProduceMCLAsync(UpdateAspectResult result) { + AbstractBatchItem request = result.getRequest(); + Optional, Boolean>> emissionStatus = conditionallyProduceMCLAsync(result.getOldValue(), result.getOldSystemMetadata(), + result.getNewValue(), result.getNewSystemMetadata(), + request.getMetadataChangeProposal(), result.getUrn(), result.getAuditStamp(), request.getAspectSpec()); - @Override - public void produceMetadataChangeLog(@Nonnull final Urn urn, @Nonnull String entityName, @Nonnull String aspectName, - @Nonnull final AspectSpec aspectSpec, @Nullable final RecordTemplate oldAspectValue, - @Nullable final RecordTemplate newAspectValue, @Nullable final SystemMetadata oldSystemMetadata, - @Nullable final SystemMetadata newSystemMetadata, @Nonnull AuditStamp auditStamp, @Nonnull final ChangeType changeType) { - final MetadataChangeLog metadataChangeLog = constructMCL(null, entityName, urn, changeType, aspectName, auditStamp, - newAspectValue, newSystemMetadata, oldAspectValue, oldSystemMetadata); - produceMetadataChangeLog(urn, aspectSpec, metadataChangeLog); + return emissionStatus.map(status -> + result.toBuilder() + .mclFuture(status.getFirst()) + .processedMCL(status.getSecond()) + .build() + ).orElse(result); } @Override @@ -1390,7 +1157,7 @@ public void ingestEntity(Entity entity, AuditStamp auditStamp) { @Override public void ingestEntity(@Nonnull Entity entity, @Nonnull AuditStamp auditStamp, - @Nonnull SystemMetadata systemMetadata) { + @Nonnull SystemMetadata systemMetadata) { log.debug("Invoked ingestEntity with entity {}, audit stamp {} systemMetadata {}", entity, auditStamp, systemMetadata.toString()); ingestSnapshotUnion(entity.getValue(), auditStamp, systemMetadata); } @@ -1432,63 +1199,93 @@ private boolean isAspectMissing(String entityType, String aspectName, Set> generateDefaultAspectsIfMissing(@Nonnull final Urn urn, - Set includedAspects) { + public Pair>> generateDefaultAspectsOnFirstWrite(@Nonnull final Urn urn, + Map includedAspects) { + List> returnAspects = new ArrayList<>(); - Set aspectsToGet = new HashSet<>(); - String entityType = urnToEntityName(urn); + final String keyAspectName = getKeyAspectName(urn); + final Map latestAspects = new HashMap<>(getLatestAspectsForUrn(urn, Set.of(keyAspectName))); - boolean shouldCheckBrowsePath = isAspectMissing(entityType, BROWSE_PATHS_ASPECT_NAME, includedAspects); - if (shouldCheckBrowsePath) { - aspectsToGet.add(BROWSE_PATHS_ASPECT_NAME); - } + // key aspect: does not exist in database && is being written + boolean generateDefaults = !latestAspects.containsKey(keyAspectName) && includedAspects.containsKey(keyAspectName); - boolean shouldCheckBrowsePathV2 = isAspectMissing(entityType, BROWSE_PATHS_V2_ASPECT_NAME, includedAspects); - if (shouldCheckBrowsePathV2) { - aspectsToGet.add(BROWSE_PATHS_V2_ASPECT_NAME); - } + // conditionally generate defaults + if (generateDefaults) { + String entityType = urnToEntityName(urn); + Set aspectsToGet = new HashSet<>(); - boolean shouldCheckDataPlatform = isAspectMissing(entityType, DATA_PLATFORM_INSTANCE_ASPECT_NAME, includedAspects); - if (shouldCheckDataPlatform) { - aspectsToGet.add(DATA_PLATFORM_INSTANCE_ASPECT_NAME); - } + boolean shouldCheckBrowsePath = isAspectMissing(entityType, BROWSE_PATHS_ASPECT_NAME, includedAspects.keySet()); + if (shouldCheckBrowsePath) { + aspectsToGet.add(BROWSE_PATHS_ASPECT_NAME); + } - List> aspects = new ArrayList<>(); - final String keyAspectName = getKeyAspectName(urn); - aspectsToGet.add(keyAspectName); + boolean shouldCheckBrowsePathV2 = isAspectMissing(entityType, BROWSE_PATHS_V2_ASPECT_NAME, includedAspects.keySet()); + if (shouldCheckBrowsePathV2) { + aspectsToGet.add(BROWSE_PATHS_V2_ASPECT_NAME); + } - Map latestAspects = getLatestAspectsForUrn(urn, aspectsToGet); + boolean shouldCheckDataPlatform = isAspectMissing(entityType, DATA_PLATFORM_INSTANCE_ASPECT_NAME, includedAspects.keySet()); + if (shouldCheckDataPlatform) { + aspectsToGet.add(DATA_PLATFORM_INSTANCE_ASPECT_NAME); + } - RecordTemplate keyAspect = latestAspects.get(keyAspectName); - if (keyAspect == null) { - keyAspect = buildKeyAspect(urn); - aspects.add(Pair.of(keyAspectName, keyAspect)); - } + // fetch additional aspects + latestAspects.putAll(getLatestAspectsForUrn(urn, aspectsToGet)); - if (shouldCheckBrowsePath && latestAspects.get(BROWSE_PATHS_ASPECT_NAME) == null) { - try { - BrowsePaths generatedBrowsePath = buildDefaultBrowsePath(urn); - aspects.add(Pair.of(BROWSE_PATHS_ASPECT_NAME, generatedBrowsePath)); - } catch (URISyntaxException e) { - log.error("Failed to parse urn: {}", urn); + if (shouldCheckBrowsePath && latestAspects.get(BROWSE_PATHS_ASPECT_NAME) == null + && !includedAspects.containsKey(BROWSE_PATHS_ASPECT_NAME)) { + try { + BrowsePaths generatedBrowsePath = buildDefaultBrowsePath(urn); + returnAspects.add(Pair.of(BROWSE_PATHS_ASPECT_NAME, generatedBrowsePath)); + } catch (URISyntaxException e) { + log.error("Failed to parse urn: {}", urn); + } } - } - if (shouldCheckBrowsePathV2 && latestAspects.get(BROWSE_PATHS_V2_ASPECT_NAME) == null) { - try { - BrowsePathsV2 generatedBrowsePathV2 = buildDefaultBrowsePathV2(urn, false); - aspects.add(Pair.of(BROWSE_PATHS_V2_ASPECT_NAME, generatedBrowsePathV2)); - } catch (URISyntaxException e) { - log.error("Failed to parse urn: {}", urn); + if (shouldCheckBrowsePathV2 && latestAspects.get(BROWSE_PATHS_V2_ASPECT_NAME) == null + && !includedAspects.containsKey(BROWSE_PATHS_V2_ASPECT_NAME)) { + try { + BrowsePathsV2 generatedBrowsePathV2 = buildDefaultBrowsePathV2(urn, false); + returnAspects.add(Pair.of(BROWSE_PATHS_V2_ASPECT_NAME, generatedBrowsePathV2)); + } catch (URISyntaxException e) { + log.error("Failed to parse urn: {}", urn); + } } - } - if (shouldCheckDataPlatform && latestAspects.get(DATA_PLATFORM_INSTANCE_ASPECT_NAME) == null) { - DataPlatformInstanceUtils.buildDataPlatformInstance(entityType, keyAspect) - .ifPresent(aspect -> aspects.add(Pair.of(DATA_PLATFORM_INSTANCE_ASPECT_NAME, aspect))); + if (shouldCheckDataPlatform && latestAspects.get(DATA_PLATFORM_INSTANCE_ASPECT_NAME) == null + && !includedAspects.containsKey(DATA_PLATFORM_INSTANCE_ASPECT_NAME)) { + RecordTemplate keyAspect = includedAspects.get(keyAspectName); + DataPlatformInstanceUtils.buildDataPlatformInstance(entityType, keyAspect) + .ifPresent(aspect -> returnAspects.add(Pair.of(DATA_PLATFORM_INSTANCE_ASPECT_NAME, aspect))); + } } - return aspects; + return Pair.of(latestAspects.containsKey(keyAspectName), returnAspects); + } + + @Override + public List> generateDefaultAspectsIfMissing(@Nonnull final Urn urn, + Map includedAspects) { + + final String keyAspectName = getKeyAspectName(urn); + + if (includedAspects.containsKey(keyAspectName)) { + return generateDefaultAspectsOnFirstWrite(urn, includedAspects).getValue(); + } else { + // No key aspect being written, generate it and potentially suggest writing it later + HashMap includedWithKeyAspect = new HashMap<>(includedAspects); + Pair keyAspect = Pair.of(keyAspectName, EntityUtils.buildKeyAspect(_entityRegistry, urn)); + includedWithKeyAspect.put(keyAspect.getKey(), keyAspect.getValue()); + + Pair>> returnAspects = generateDefaultAspectsOnFirstWrite(urn, includedWithKeyAspect); + + // missing key aspect in database, add it + if (!returnAspects.getFirst()) { + returnAspects.getValue().add(keyAspect); + } + + return returnAspects.getValue(); + } } private void ingestSnapshotUnion(@Nonnull final Snapshot snapshotUnion, @Nonnull final AuditStamp auditStamp, @@ -1500,28 +1297,18 @@ private void ingestSnapshotUnion(@Nonnull final Snapshot snapshotUnion, @Nonnull log.info("INGEST urn {} with system metadata {}", urn.toString(), systemMetadata.toString()); aspectRecordsToIngest.addAll(generateDefaultAspectsIfMissing(urn, - aspectRecordsToIngest.stream().map(pair -> pair.getFirst()).collect(Collectors.toSet()))); + aspectRecordsToIngest.stream().collect(Collectors.toMap(Pair::getKey, Pair::getValue)))); - ingestAspects(urn, aspectRecordsToIngest, auditStamp, systemMetadata); - } + AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() + .items(aspectRecordsToIngest.stream().map(pair -> UpsertBatchItem.builder() + .urn(urn) + .aspectName(pair.getKey()) + .aspect(pair.getValue()) + .systemMetadata(systemMetadata) + .build(_entityRegistry)).collect(Collectors.toList())) + .build(); - @Override - public Snapshot buildSnapshot(@Nonnull final Urn urn, @Nonnull final RecordTemplate aspectValue) { - // if the aspect value is the key, we do not need to include the key a second time - if (PegasusUtils.getAspectNameFromSchema(aspectValue.schema()).equals(getKeyAspectName(urn))) { - return toSnapshotUnion(toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, aspectValue)))); - } - - final RecordTemplate keyAspectValue = buildKeyAspect(urn); - return toSnapshotUnion( - toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, keyAspectValue), toAspectUnion(urn, aspectValue)))); - } - - protected RecordTemplate buildKeyAspect(@Nonnull final Urn urn) { - final EntitySpec spec = _entityRegistry.getEntitySpec(urnToEntityName(urn)); - final AspectSpec keySpec = spec.getKeyAspectSpec(); - final RecordDataSchema keySchema = keySpec.getPegasusSchema(); - return EntityKeyUtils.convertUrnToEntityKey(urn, keySpec); + ingestAspects(aspectsBatch, auditStamp, true, true); } @Override @@ -1641,24 +1428,33 @@ public RollbackRunResult rollbackWithConditions(List aspectRow List removedAspects = new ArrayList<>(); AtomicInteger rowsDeletedFromEntityDeletion = new AtomicInteger(0); - aspectRows.forEach(aspectToRemove -> { - + List> futures = aspectRows.stream().map(aspectToRemove -> { RollbackResult result = deleteAspect(aspectToRemove.getUrn(), aspectToRemove.getAspectName(), conditions, hardDelete); if (result != null) { Optional aspectSpec = getAspectSpec(result.entityName, result.aspectName); if (!aspectSpec.isPresent()) { log.error("Issue while rolling back: unknown aspect {} for entity {}", result.entityName, result.aspectName); - return; + return null; } rowsDeletedFromEntityDeletion.addAndGet(result.additionalRowsAffected); removedAspects.add(aspectToRemove); - produceMetadataChangeLog(result.getUrn(), result.getEntityName(), result.getAspectName(), aspectSpec.get(), + return alwaysProduceMCLAsync(result.getUrn(), result.getEntityName(), result.getAspectName(), aspectSpec.get(), result.getOldValue(), result.getNewValue(), result.getOldSystemMetadata(), result.getNewSystemMetadata(), // TODO: use properly attributed audit stamp. createSystemAuditStamp(), - result.getChangeType()); + result.getChangeType()).getFirst(); + } + + return null; + }).filter(Objects::nonNull).collect(Collectors.toList()); + + futures.forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); } }); @@ -1697,11 +1493,19 @@ public RollbackRunResult deleteUrn(Urn urn) { rowsDeletedFromEntityDeletion = result.additionalRowsAffected; removedAspects.add(summary); - produceMetadataChangeLog(result.getUrn(), result.getEntityName(), result.getAspectName(), keySpec, + Future future = alwaysProduceMCLAsync(result.getUrn(), result.getEntityName(), result.getAspectName(), keySpec, result.getOldValue(), result.getNewValue(), result.getOldSystemMetadata(), result.getNewSystemMetadata(), // TODO: Use a proper inferred audit stamp createSystemAuditStamp(), - result.getChangeType()); + result.getChangeType()).getFirst(); + + if (future != null) { + try { + future.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + } } return new RollbackRunResult(removedAspects, rowsDeletedFromEntityDeletion); @@ -1752,7 +1556,7 @@ public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map { + final RollbackResult result = _aspectDao.runInTransactionWithRetry((tx) -> { Integer additionalRowsDeleted = 0; // 1. Fetch the latest existing version of the aspect. @@ -1799,7 +1603,7 @@ public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map _aspectDao.deleteAspect(aspect)); + aspectsToDelete.forEach(aspect -> _aspectDao.deleteAspect(tx, aspect)); if (survivingAspect != null) { // if there was a surviving aspect, copy its information into the latest row @@ -1810,15 +1614,15 @@ public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map getEnvelopedAspects(final S private EnvelopedAspect getKeyEnvelopedAspect(final Urn urn) { final EntitySpec spec = getEntityRegistry().getEntitySpec(PegasusUtils.urnToEntityName(urn)); final AspectSpec keySpec = spec.getKeyAspectSpec(); - final RecordDataSchema keySchema = keySpec.getPegasusSchema(); final com.linkedin.entity.Aspect aspect = new com.linkedin.entity.Aspect(EntityKeyUtils.convertUrnToEntityKey(urn, keySpec).data()); @@ -1992,19 +1795,19 @@ private EnvelopedAspect getKeyEnvelopedAspect(final Urn urn) { } @Nonnull - private UpdateAspectResult ingestAspectToLocalDBNoTransaction( - @Nonnull final Urn urn, - @Nonnull final String aspectName, - @Nonnull final Function, RecordTemplate> updateLambda, - @Nonnull final AuditStamp auditStamp, - @Nonnull final SystemMetadata providedSystemMetadata, - @Nullable final EntityAspect latest, - @Nonnull final Long nextVersion) { + private UpdateAspectResult ingestAspectToLocalDB( + @Nullable Transaction tx, + @Nonnull final Urn urn, + @Nonnull final String aspectName, + @Nonnull final RecordTemplate newValue, + @Nonnull final AuditStamp auditStamp, + @Nonnull final SystemMetadata providedSystemMetadata, + @Nullable final EntityAspect latest, + @Nonnull final Long nextVersion) { // 2. Compare the latest existing and new. final RecordTemplate oldValue = latest == null ? null : EntityUtils.toAspectRecord(urn, aspectName, latest.getMetadata(), getEntityRegistry()); - final RecordTemplate newValue = updateLambda.apply(Optional.ofNullable(oldValue)); // 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 @@ -2014,20 +1817,28 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction( latest.setSystemMetadata(RecordUtils.toJsonString(latestSystemMetadata)); - _aspectDao.saveAspect(latest, false); + log.info("Ingesting aspect with name {}, urn {}", aspectName, urn); + _aspectDao.saveAspect(tx, latest, false); // metrics _aspectDao.incrementWriteMetrics(aspectName, 1, latest.getAspect().getBytes(StandardCharsets.UTF_8).length); - return new UpdateAspectResult(urn, oldValue, oldValue, - EntityUtils.parseSystemMetadata(latest.getSystemMetadata()), latestSystemMetadata, - MetadataAuditOperation.UPDATE, auditStamp, 0); + return UpdateAspectResult.builder() + .urn(urn) + .oldValue(oldValue) + .newValue(oldValue) + .oldSystemMetadata(EntityUtils.parseSystemMetadata(latest.getSystemMetadata())) + .newSystemMetadata(latestSystemMetadata) + .operation(MetadataAuditOperation.UPDATE) + .auditStamp(auditStamp) + .maxVersion(0) + .build(); } // 4. Save the newValue as the latest version log.debug("Ingesting aspect with name {}, urn {}", aspectName, urn); String newValueStr = EntityUtils.toJsonAspect(newValue); - long versionOfOld = _aspectDao.saveLatestAspect(urn.toString(), aspectName, latest == null ? null : EntityUtils.toJsonAspect(oldValue), + long versionOfOld = _aspectDao.saveLatestAspect(tx, urn.toString(), aspectName, latest == null ? null : EntityUtils.toJsonAspect(oldValue), latest == null ? null : latest.getCreatedBy(), latest == null ? null : latest.getCreatedFor(), latest == null ? null : latest.getCreatedOn(), latest == null ? null : latest.getSystemMetadata(), newValueStr, auditStamp.getActor().toString(), @@ -2037,75 +1848,16 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction( // metrics _aspectDao.incrementWriteMetrics(aspectName, 1, newValueStr.getBytes(StandardCharsets.UTF_8).length); - return new UpdateAspectResult(urn, oldValue, newValue, - latest == null ? null : EntityUtils.parseSystemMetadata(latest.getSystemMetadata()), providedSystemMetadata, - MetadataAuditOperation.UPDATE, auditStamp, versionOfOld); - } - - @Nonnull - private Map getLatestAspectForUrn(@Nonnull final Urn urn, @Nonnull final Set aspectNames) { - Set urns = new HashSet<>(); - urns.add(urn); - - Map result = new HashMap<>(); - getLatestAspect(urns, aspectNames).forEach((key, aspectEntry) -> { - final String aspectName = key.getAspect(); - result.put(aspectName, aspectEntry); - }); - return result; - } - - @Nonnull - private RecordTemplate updateAspect( - @Nonnull final Urn urn, - @Nonnull final String entityName, - @Nonnull final String aspectName, - @Nonnull final AspectSpec aspectSpec, - @Nonnull final RecordTemplate value, - @Nonnull final AuditStamp auditStamp, - @Nonnull final long version, - @Nonnull final boolean emitMae, - final int maxTransactionRetry) { - - final UpdateAspectResult result = _aspectDao.runInTransactionWithRetry(() -> { - - final EntityAspect oldAspect = _aspectDao.getAspect(urn.toString(), aspectName, version); - final RecordTemplate oldValue = - oldAspect == null ? null : EntityUtils.toAspectRecord(urn, aspectName, oldAspect.getMetadata(), getEntityRegistry()); - - SystemMetadata oldSystemMetadata = - oldAspect == null ? new SystemMetadata() : EntityUtils.parseSystemMetadata(oldAspect.getSystemMetadata()); - // create a duplicate of the old system metadata to update and write back - SystemMetadata newSystemMetadata = - oldAspect == null ? new SystemMetadata() : EntityUtils.parseSystemMetadata(oldAspect.getSystemMetadata()); - newSystemMetadata.setLastObserved(System.currentTimeMillis()); - - log.debug("Updating aspect with name {}, urn {}", aspectName, urn); - String aspectStr = EntityUtils.toJsonAspect(value); - _aspectDao.saveAspect(urn.toString(), aspectName, aspectStr, auditStamp.getActor().toString(), - auditStamp.hasImpersonator() ? auditStamp.getImpersonator().toString() : null, - new Timestamp(auditStamp.getTime()), EntityUtils.toJsonAspect(newSystemMetadata), version, oldAspect == null); - - // metrics - _aspectDao.incrementWriteMetrics(aspectName, 1, aspectStr.getBytes(StandardCharsets.UTF_8).length); - - return new UpdateAspectResult(urn, oldValue, value, oldSystemMetadata, newSystemMetadata, - MetadataAuditOperation.UPDATE, auditStamp, version); - }, maxTransactionRetry); - - final RecordTemplate oldValue = result.getOldValue(); - final RecordTemplate newValue = result.getNewValue(); - - if (emitMae) { - log.debug("Producing MetadataAuditEvent for updated aspect {}, urn {}", aspectName, urn); - produceMetadataChangeLog(urn, entityName, aspectName, aspectSpec, oldValue, newValue, - result.getOldSystemMetadata(), result.getNewSystemMetadata(), auditStamp, ChangeType.UPSERT); - } else { - log.debug("Skipped producing MetadataAuditEvent for updated aspect {}, urn {}. emitMAE is false.", - aspectName, urn); - } - - return newValue; + return UpdateAspectResult.builder() + .urn(urn) + .oldValue(oldValue) + .newValue(newValue) + .oldSystemMetadata(latest == null ? null : EntityUtils.parseSystemMetadata(latest.getSystemMetadata())) + .newSystemMetadata(providedSystemMetadata) + .operation(MetadataAuditOperation.UPDATE) + .auditStamp(auditStamp) + .maxVersion(versionOfOld) + .build(); } /** @@ -2174,14 +1926,7 @@ private DataPlatformInfo getDataPlatformInfo(Urn urn) { return null; } - private boolean shouldAspectEmitChangeLog(@Nonnull final Urn urn, @Nonnull final String aspectName) { - final String entityName = urnToEntityName(urn); - final EntitySpec entitySpec = getEntityRegistry().getEntitySpec(entityName); - final AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); - return shouldAspectEmitChangeLog(aspectSpec); - } - - private boolean shouldAspectEmitChangeLog(@Nonnull final AspectSpec aspectSpec) { + private static boolean shouldAspectEmitChangeLog(@Nonnull final AspectSpec aspectSpec) { final List relationshipFieldSpecs = aspectSpec.getRelationshipFieldSpecs(); return relationshipFieldSpecs.stream().anyMatch(RelationshipFieldSpec::isLineageRelationship); } 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 a989f519910d6..ffd63479589bc 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 @@ -8,10 +8,13 @@ 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.validation.EntityRegistryUrnValidator; import com.linkedin.metadata.entity.validation.RecordTemplateValidator; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.EntityRegistry; +import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.PegasusUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; @@ -21,9 +24,11 @@ import lombok.extern.slf4j.Slf4j; import java.net.URISyntaxException; +import java.net.URLEncoder; import java.util.List; import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.utils.PegasusUtils.urnToEntityName; @Slf4j @@ -32,6 +37,9 @@ public class EntityUtils { private EntityUtils() { } + public static final int URN_NUM_BYTES_LIMIT = 512; + public static final String URN_DELIMITER_SEPARATOR = "␟"; + @Nonnull public static String toJsonAspect(@Nonnull final RecordTemplate aspectRecord) { return RecordUtils.toJsonString(aspectRecord); @@ -60,10 +68,8 @@ public static void ingestChangeProposals( @Nonnull Urn actor, @Nonnull Boolean async ) { - // TODO: Replace this with a batch ingest proposals endpoint. - for (MetadataChangeProposal change : changes) { - entityService.ingestProposal(change, EntityUtils.getAuditStamp(actor), async); - } + entityService.ingestProposal(AspectsBatchImpl.builder() + .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), async); } /** @@ -167,4 +173,34 @@ public static boolean checkIfRemoved(EntityService entityService, Urn entityUrn) return false; } } + + public static RecordTemplate buildKeyAspect(@Nonnull EntityRegistry entityRegistry, @Nonnull final Urn urn) { + final EntitySpec spec = entityRegistry.getEntitySpec(urnToEntityName(urn)); + final AspectSpec keySpec = spec.getKeyAspectSpec(); + return EntityKeyUtils.convertUrnToEntityKey(urn, keySpec); + } + + public static void validateUrn(@Nonnull EntityRegistry entityRegistry, @Nonnull final Urn urn) { + EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(entityRegistry); + validator.setCurrentEntitySpec(entityRegistry.getEntitySpec(urn.getEntityType())); + RecordTemplateValidator.validate(EntityUtils.buildKeyAspect(entityRegistry, urn), validationResult -> { + throw new IllegalArgumentException("Invalid urn: " + urn + "\n Cause: " + + validationResult.getMessages()); }, validator); + + if (urn.toString().trim().length() != urn.toString().length()) { + throw new IllegalArgumentException("Error: cannot provide an URN with leading or trailing whitespace"); + } + if (URLEncoder.encode(urn.toString()).length() > URN_NUM_BYTES_LIMIT) { + throw new IllegalArgumentException("Error: cannot provide an URN longer than " + Integer.toString(URN_NUM_BYTES_LIMIT) + " bytes (when URL encoded)"); + } + if (urn.toString().contains(URN_DELIMITER_SEPARATOR)) { + throw new IllegalArgumentException("Error: URN cannot contain " + URN_DELIMITER_SEPARATOR + " character"); + } + try { + Urn.createFromString(urn.toString()); + } catch (URISyntaxException e) { + throw new IllegalArgumentException(e); + } + } + } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java index dcb58ee324a23..b215dd4a5d1ed 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 @@ -39,12 +39,13 @@ import java.util.Map; import java.util.Objects; import java.util.Set; -import java.util.function.Supplier; +import java.util.function.Function; import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; import io.ebean.PagedList; +import io.ebean.Transaction; import lombok.extern.slf4j.Slf4j; import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.*; @@ -85,6 +86,19 @@ public EntityAspect getLatestAspect(@Nonnull String urn, @Nonnull String aspectN return getAspect(urn, aspectName, ASPECT_LATEST_VERSION); } + @Override + public Map> getLatestAspects(Map> urnAspects) { + return urnAspects.entrySet().stream() + .map(entry -> Map.entry(entry.getKey(), entry.getValue().stream() + .map(aspectName -> { + EntityAspect aspect = getLatestAspect(entry.getKey(), aspectName); + return aspect != null ? Map.entry(aspectName, aspect) : null; + }) + .filter(Objects::nonNull) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + @Override public long getMaxVersion(@Nonnull final String urn, @Nonnull final String aspectName) { validateConnection(); @@ -149,7 +163,7 @@ private Map getMaxVersions(@Nonnull final String urn, @Nonnull fin } @Override - public void saveAspect(@Nonnull EntityAspect aspect, final boolean insert) { + public void saveAspect(@Nullable Transaction tx, @Nonnull EntityAspect aspect, final boolean insert) { validateConnection(); SimpleStatement statement = generateSaveStatement(aspect, insert); _cqlSession.execute(statement); @@ -241,7 +255,7 @@ public ListResult listAspectMetadata( @Override @Nonnull - public T runInTransactionWithRetry(@Nonnull final Supplier block, final int maxTransactionRetry) { + public T runInTransactionWithRetry(@Nonnull final Function block, final int maxTransactionRetry) { validateConnection(); int retryCount = 0; Exception lastException; @@ -249,7 +263,7 @@ public T runInTransactionWithRetry(@Nonnull final Supplier block, final i do { try { // TODO: Try to bend this code to make use of Cassandra batches. This method is called from single-urn operations, so perf should not suffer much - return block.get(); + return block.apply(null); } catch (DriverException exception) { lastException = exception; } @@ -319,7 +333,7 @@ private static AuditStamp toAuditStamp(@Nonnull final EntityAspect aspect) { } @Override - public void deleteAspect(@Nonnull final EntityAspect aspect) { + public void deleteAspect(@Nullable Transaction tx, @Nonnull final EntityAspect aspect) { validateConnection(); SimpleStatement ss = deleteFrom(CassandraAspect.TABLE_NAME) .whereColumn(CassandraAspect.URN_COLUMN).isEqualTo(literal(aspect.getUrn())) @@ -332,7 +346,7 @@ public void deleteAspect(@Nonnull final EntityAspect aspect) { } @Override - public int deleteUrn(@Nonnull final String urn) { + public int deleteUrn(@Nullable Transaction tx, @Nonnull final String urn) { validateConnection(); SimpleStatement ss = deleteFrom(CassandraAspect.TABLE_NAME) .whereColumn(CassandraAspect.URN_COLUMN).isEqualTo(literal(urn)) @@ -453,42 +467,42 @@ public Iterable listAllUrns(int start, int pageSize) { } @Override - public long getNextVersion(@Nonnull final String urn, @Nonnull final String aspectName) { + public Map> getNextVersions(Map> urnAspectMap) { validateConnection(); - Map versions = getNextVersions(urn, ImmutableSet.of(aspectName)); - return versions.get(aspectName); - } + Map> result = new HashMap<>(); - @Override - public Map getNextVersions(@Nonnull final String urn, @Nonnull final Set aspectNames) { - validateConnection(); - Map maxVersions = getMaxVersions(urn, aspectNames); - Map nextVersions = new HashMap<>(); + for (Map.Entry> aspectNames : urnAspectMap.entrySet()) { + Map maxVersions = getMaxVersions(aspectNames.getKey(), aspectNames.getValue()); + Map nextVersions = new HashMap<>(); - for (String aspectName: aspectNames) { - long latestVersion = maxVersions.get(aspectName); - long nextVal = latestVersion < 0 ? ASPECT_LATEST_VERSION : latestVersion + 1L; - nextVersions.put(aspectName, nextVal); + for (String aspectName : aspectNames.getValue()) { + long latestVersion = maxVersions.get(aspectName); + long nextVal = latestVersion < 0 ? ASPECT_LATEST_VERSION : latestVersion + 1L; + nextVersions.put(aspectName, nextVal); + } + + result.put(aspectNames.getKey(), nextVersions); } - return nextVersions; + return result; } @Override public long saveLatestAspect( - @Nonnull final String urn, - @Nonnull final String aspectName, - @Nullable final String oldAspectMetadata, - @Nullable final String oldActor, - @Nullable final String oldImpersonator, - @Nullable final Timestamp oldTime, - @Nullable final String oldSystemMetadata, - @Nonnull final String newAspectMetadata, - @Nonnull final String newActor, - @Nullable final String newImpersonator, - @Nonnull final Timestamp newTime, - @Nullable final String newSystemMetadata, - final Long nextVersion + @Nullable Transaction tx, + @Nonnull final String urn, + @Nonnull final String aspectName, + @Nullable final String oldAspectMetadata, + @Nullable final String oldActor, + @Nullable final String oldImpersonator, + @Nullable final Timestamp oldTime, + @Nullable final String oldSystemMetadata, + @Nonnull final String newAspectMetadata, + @Nonnull final String newActor, + @Nullable final String newImpersonator, + @Nonnull final Timestamp newTime, + @Nullable final String newSystemMetadata, + final Long nextVersion ) { validateConnection(); @@ -574,15 +588,16 @@ public void setWritable(boolean canWrite) { @Override public void saveAspect( - @Nonnull final String urn, - @Nonnull final String aspectName, - @Nonnull final String aspectMetadata, - @Nonnull final String actor, - @Nullable final String impersonator, - @Nonnull final Timestamp timestamp, - @Nonnull final String systemMetadata, - final long version, - final boolean insert) { + @Nullable Transaction tx, + @Nonnull final String urn, + @Nonnull final String aspectName, + @Nonnull final String aspectMetadata, + @Nonnull final String actor, + @Nullable final String impersonator, + @Nonnull final Timestamp timestamp, + @Nonnull final String systemMetadata, + final long version, + final boolean insert) { validateConnection(); final EntityAspect aspect = new EntityAspect( @@ -596,7 +611,7 @@ public void saveAspect( impersonator ); - saveAspect(aspect, insert); + saveAspect(tx, aspect, insert); // metrics incrementWriteMetrics(aspectName, 1, aspectMetadata.getBytes(StandardCharsets.UTF_8).length); 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 3d8245b324ce5..9ebb6b26fc43d 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 @@ -14,8 +14,11 @@ import com.linkedin.metadata.entity.RetentionService; import com.linkedin.metadata.entity.EntityAspectIdentifier; import com.linkedin.metadata.entity.EntityAspect; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; 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; import com.linkedin.retention.TimeBasedRetention; @@ -53,22 +56,33 @@ public EntityService getEntityService() { return _entityService; } + @Override + protected AspectsBatch buildAspectsBatch(List mcps) { + return AspectsBatchImpl.builder() + .mcps(mcps, _entityService.getEntityRegistry()) + .build(); + } + @Override @WithSpan - public void applyRetention(@Nonnull Urn urn, @Nonnull String aspectName, Retention retentionPolicy, Optional retentionContext) { - log.debug("Applying retention to urn {}, aspectName {}", urn, aspectName); - // If no policies are set or has indefinite policy set, do not apply any retention - if (retentionPolicy.data().isEmpty()) { - return; - } + protected void applyRetention(List retentionContexts) { - if (retentionPolicy.hasVersion()) { - applyVersionBasedRetention(urn, aspectName, retentionPolicy.getVersion(), retentionContext.flatMap(RetentionService.RetentionContext::getMaxVersion)); - } + List nonEmptyContexts = retentionContexts.stream() + .filter(context -> context.getRetentionPolicy().isPresent() + && !context.getRetentionPolicy().get().data().isEmpty()) + .collect(Collectors.toList()); - if (retentionPolicy.hasTime()) { - applyTimeBasedRetention(urn, aspectName, retentionPolicy.getTime()); - } + nonEmptyContexts.forEach(context -> { + if (context.getRetentionPolicy().map(Retention::hasVersion).orElse(false)) { + Retention retentionPolicy = context.getRetentionPolicy().get(); + applyVersionBasedRetention(context.getUrn(), context.getAspectName(), retentionPolicy.getVersion(), context.getMaxVersion()); + } + + if (context.getRetentionPolicy().map(Retention::hasTime).orElse(false)) { + Retention retentionPolicy = context.getRetentionPolicy().get(); + applyTimeBasedRetention(context.getUrn(), context.getAspectName(), retentionPolicy.getTime()); + } + }); } @Override @@ -103,7 +117,12 @@ public void batchApplyRetention(@Nullable String entityName, @Nullable String as .findFirst() .map(DataHubRetentionConfig::getRetention); retentionPolicy.ifPresent(retention -> - applyRetention(urn, aspectNameFromRecord, retention, Optional.of(new RetentionContext(Optional.of(id.getVersion()))))); + applyRetention(List.of(RetentionContext.builder() + .urn(urn) + .aspectName(aspectNameFromRecord) + .retentionPolicy(retentionPolicy) + .maxVersion(Optional.of(id.getVersion())) + .build()))); i += 1; if (i % _batchSize == 0) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/AspectStorageValidationUtil.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/AspectStorageValidationUtil.java index f12e2ba521b15..c0aef268e14c9 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/AspectStorageValidationUtil.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/AspectStorageValidationUtil.java @@ -1,7 +1,7 @@ package com.linkedin.metadata.entity.ebean; import com.linkedin.metadata.Constants; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.SqlQuery; import io.ebean.SqlRow; @@ -16,33 +16,33 @@ private AspectStorageValidationUtil() { } - public static long getV1RowCount(EbeanServer server) { + public static long getV1RowCount(Database server) { return server.find(EbeanAspectV1.class).findCount(); } /** * Get the number of rows created not by the DataHub system actor (urn:li:corpuser:__datahub_system) */ - public static long getV2NonSystemRowCount(EbeanServer server) { + public static long getV2NonSystemRowCount(Database server) { return server.find(EbeanAspectV2.class).where(ne("createdby", Constants.SYSTEM_ACTOR)).findCount(); } - public static boolean checkV2TableExists(EbeanServer server) { + public static boolean checkV2TableExists(Database server) { final String queryStr = "SELECT * FROM INFORMATION_SCHEMA.TABLES \n" - + "WHERE TABLE_NAME = 'metadata_aspect_v2'"; + + "WHERE lower(TABLE_NAME) = 'metadata_aspect_v2'"; - final SqlQuery query = server.createSqlQuery(queryStr); + final SqlQuery query = server.sqlQuery(queryStr); final List rows = query.findList(); return rows.size() > 0; } - public static boolean checkV1TableExists(EbeanServer server) { + public static boolean checkV1TableExists(Database server) { final String queryStr = "SELECT * FROM INFORMATION_SCHEMA.TABLES \n" - + "WHERE TABLE_NAME = 'metadata_aspect'"; + + "WHERE lower(TABLE_NAME) = 'metadata_aspect'"; - final SqlQuery query = server.createSqlQuery(queryStr); + final SqlQuery query = server.sqlQuery(queryStr); final List rows = query.findList(); return rows.size() > 0; } 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 9e65c752b531b..30886db264994 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 @@ -11,13 +11,16 @@ 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; import com.linkedin.metadata.query.ExtraInfoArray; import com.linkedin.metadata.query.ListResultMetadata; import com.linkedin.metadata.search.utils.QueryUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; import io.ebean.DuplicateKeyException; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.ExpressionList; import io.ebean.Junction; import io.ebean.PagedList; @@ -27,26 +30,24 @@ import io.ebean.Transaction; import io.ebean.TxScope; import io.ebean.annotation.TxIsolation; -import io.ebean.annotation.Platform; -import io.ebean.config.dbplatform.DatabasePlatform; -import io.ebean.plugin.SpiServer; import java.net.URISyntaxException; -import java.sql.SQLException; import java.sql.Timestamp; import java.time.Clock; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.function.Supplier; +import java.util.function.Function; import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import javax.persistence.RollbackException; + import javax.persistence.PersistenceException; import javax.persistence.Table; + import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; @@ -54,7 +55,7 @@ @Slf4j public class EbeanAspectDao implements AspectDao, AspectMigrationsDao { - private final EbeanServer _server; + private final Database _server; private boolean _connectionValidated = false; private final Clock _clock = Clock.systemUTC(); @@ -67,7 +68,7 @@ public class EbeanAspectDao implements AspectDao, AspectMigrationsDao { // more testing. private int _queryKeysCount = 375; // 0 means no pagination on keys - public EbeanAspectDao(@Nonnull final EbeanServer server) { + public EbeanAspectDao(@Nonnull final Database server) { _server = server; } @@ -77,10 +78,10 @@ public void setWritable(boolean canWrite) { } /** - * Return the {@link EbeanServer} server instance used for customized queries. + * Return the {@link Database} server instance used for customized queries. * Only used in tests. */ - public EbeanServer getServer() { + public Database getServer() { return _server; } @@ -107,19 +108,20 @@ private boolean validateConnection() { @Override public long saveLatestAspect( - @Nonnull final String urn, - @Nonnull final String aspectName, - @Nullable final String oldAspectMetadata, - @Nullable final String oldActor, - @Nullable final String oldImpersonator, - @Nullable final Timestamp oldTime, - @Nullable final String oldSystemMetadata, - @Nonnull final String newAspectMetadata, - @Nonnull final String newActor, - @Nullable final String newImpersonator, - @Nonnull final Timestamp newTime, - @Nullable final String newSystemMetadata, - final Long nextVersion + @Nullable Transaction tx, + @Nonnull final String urn, + @Nonnull final String aspectName, + @Nullable final String oldAspectMetadata, + @Nullable final String oldActor, + @Nullable final String oldImpersonator, + @Nullable final Timestamp oldTime, + @Nullable final String oldSystemMetadata, + @Nonnull final String newAspectMetadata, + @Nonnull final String newActor, + @Nullable final String newImpersonator, + @Nonnull final Timestamp newTime, + @Nullable final String newSystemMetadata, + final Long nextVersion ) { validateConnection(); @@ -130,26 +132,27 @@ public long saveLatestAspect( long largestVersion = ASPECT_LATEST_VERSION; if (oldAspectMetadata != null && oldTime != null) { largestVersion = nextVersion; - saveAspect(urn, aspectName, oldAspectMetadata, oldActor, oldImpersonator, oldTime, oldSystemMetadata, largestVersion, true); + saveAspect(tx, urn, aspectName, oldAspectMetadata, oldActor, oldImpersonator, oldTime, oldSystemMetadata, largestVersion, true); } // Save newValue as the latest version (v0) - saveAspect(urn, aspectName, newAspectMetadata, newActor, newImpersonator, newTime, newSystemMetadata, ASPECT_LATEST_VERSION, oldAspectMetadata == null); + saveAspect(tx, urn, aspectName, newAspectMetadata, newActor, newImpersonator, newTime, newSystemMetadata, ASPECT_LATEST_VERSION, oldAspectMetadata == null); return largestVersion; } @Override public void saveAspect( - @Nonnull final String urn, - @Nonnull final String aspectName, - @Nonnull final String aspectMetadata, - @Nonnull final String actor, - @Nullable final String impersonator, - @Nonnull final Timestamp timestamp, - @Nonnull final String systemMetadata, - final long version, - final boolean insert) { + @Nullable Transaction tx, + @Nonnull final String urn, + @Nonnull final String aspectName, + @Nonnull final String aspectMetadata, + @Nonnull final String actor, + @Nullable final String impersonator, + @Nonnull final Timestamp timestamp, + @Nonnull final String systemMetadata, + final long version, + final boolean insert) { validateConnection(); @@ -163,47 +166,38 @@ public void saveAspect( aspect.setCreatedFor(impersonator); } - saveEbeanAspect(aspect, insert); + saveEbeanAspect(tx, aspect, insert); } @Override - public void saveAspect(@Nonnull final EntityAspect aspect, final boolean insert) { + public void saveAspect(@Nullable Transaction tx, @Nonnull final EntityAspect aspect, final boolean insert) { EbeanAspectV2 ebeanAspect = EbeanAspectV2.fromEntityAspect(aspect); - saveEbeanAspect(ebeanAspect, insert); + saveEbeanAspect(tx, ebeanAspect, insert); } - private void saveEbeanAspect(@Nonnull final EbeanAspectV2 ebeanAspect, final boolean insert) { + private void saveEbeanAspect(@Nullable Transaction tx, @Nonnull final EbeanAspectV2 ebeanAspect, final boolean insert) { validateConnection(); if (insert) { - _server.insert(ebeanAspect); + _server.insert(ebeanAspect, tx); } else { - _server.update(ebeanAspect); + _server.update(ebeanAspect, tx); } } @Override - @Nullable - public EntityAspect getLatestAspect(@Nonnull final String urn, @Nonnull final String aspectName) { + public Map> getLatestAspects(@Nonnull Map> urnAspects) { validateConnection(); - final EbeanAspectV2.PrimaryKey key = new EbeanAspectV2.PrimaryKey(urn, aspectName, ASPECT_LATEST_VERSION); - EbeanAspectV2 ebeanAspect = _server.find(EbeanAspectV2.class, key); - return ebeanAspect == null ? null : ebeanAspect.toEntityAspect(); - } - @Override - public long getMaxVersion(@Nonnull final String urn, @Nonnull final String aspectName) { - validateConnection(); - List result = _server.find(EbeanAspectV2.class) - .where() - .eq("urn", urn) - .eq("aspect", aspectName) - .orderBy() - .desc("version") - .findList(); - if (result.size() == 0) { - return -1; - } - return result.get(0).getKey().getVersion(); + List keys = urnAspects.entrySet().stream() + .flatMap(entry -> entry.getValue().stream() + .map(aspect -> new EbeanAspectV2.PrimaryKey(entry.getKey(), aspect, ASPECT_LATEST_VERSION)) + ).collect(Collectors.toList()); + + List results = _server.find(EbeanAspectV2.class) + .where().idIn(keys) + .findList(); + + return toUrnAspectMap(results); } @Override @@ -240,16 +234,16 @@ public EntityAspect getAspect(@Nonnull final EntityAspectIdentifier key) { } @Override - public void deleteAspect(@Nonnull final EntityAspect aspect) { + public void deleteAspect(@Nullable Transaction tx, @Nonnull final EntityAspect aspect) { validateConnection(); EbeanAspectV2 ebeanAspect = EbeanAspectV2.fromEntityAspect(aspect); - _server.delete(ebeanAspect); + _server.delete(ebeanAspect, tx); } @Override - public int deleteUrn(@Nonnull final String urn) { + public int deleteUrn(@Nullable Transaction tx, @Nonnull final String urn) { validateConnection(); - return _server.createQuery(EbeanAspectV2.class).where().eq(EbeanAspectV2.URN_COLUMN, urn).delete(); + return _server.createQuery(EbeanAspectV2.class).where().eq(EbeanAspectV2.URN_COLUMN, urn).delete(tx); } @Override @@ -497,49 +491,37 @@ public ListResult listLatestAspectMetadata( @Override @Nonnull - public T runInTransactionWithRetry(@Nonnull final Supplier block, final int maxTransactionRetry) { + public T runInTransactionWithRetry(@Nonnull final Function block, final int maxTransactionRetry) { + return runInTransactionWithRetry(block, null, maxTransactionRetry); + } + + @Override + @Nonnull + public T runInTransactionWithRetry(@Nonnull final Function block, @Nullable AspectsBatch batch, + final int maxTransactionRetry) { validateConnection(); int retryCount = 0; - Exception lastException; + Exception lastException = null; T result = null; do { try (Transaction transaction = _server.beginTransaction(TxScope.requiresNew().setIsolation(TxIsolation.REPEATABLE_READ))) { transaction.setBatchMode(true); - result = block.get(); + result = block.apply(transaction); transaction.commit(); lastException = null; break; - } catch (RollbackException | DuplicateKeyException exception) { - MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); - lastException = exception; } catch (PersistenceException exception) { - MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); - // TODO: replace this logic by catching SerializableConflictException above once the exception is available - SpiServer pluginApi = _server.getPluginApi(); - DatabasePlatform databasePlatform = pluginApi.getDatabasePlatform(); - - if (databasePlatform.isPlatform(Platform.POSTGRES)) { - Throwable cause = exception.getCause(); - if (cause instanceof SQLException) { - SQLException sqlException = (SQLException) cause; - String sqlState = sqlException.getSQLState(); - while (sqlState == null && sqlException.getCause() instanceof SQLException) { - sqlException = (SQLException) sqlException.getCause(); - sqlState = sqlException.getSQLState(); - } - - // version 11.33.3 of io.ebean does not have a SerializableConflictException (will be available with version 11.44.1), - // therefore when using a PostgreSQL database we have to check the SQL state 40001 here to retry the transactions - // also in case of serialization errors ("could not serialize access due to concurrent update") - if (sqlState.equals("40001")) { - lastException = exception; - continue; - } + if (exception instanceof DuplicateKeyException) { + if (batch != null && batch.getItems().stream().allMatch(a -> a.getAspectName().equals(a.getEntitySpec().getKeyAspectSpec().getName()))) { + log.warn("Skipping DuplicateKeyException retry since aspect is the key aspect. {}", batch.getUrnAspectsMap().keySet()); + continue; } } - throw exception; + MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); + log.warn("Retryable PersistenceException: {}", exception.getMessage()); + lastException = exception; } } while (++retryCount <= maxTransactionRetry); @@ -552,57 +534,64 @@ public T runInTransactionWithRetry(@Nonnull final Supplier block, final i } @Override - public long getNextVersion(@Nonnull final String urn, @Nonnull final String aspectName) { + public long getMaxVersion(@Nonnull final String urn, @Nonnull final String aspectName) { validateConnection(); final List result = _server.find(EbeanAspectV2.class) - .where() - .eq(EbeanAspectV2.URN_COLUMN, urn.toString()) - .eq(EbeanAspectV2.ASPECT_COLUMN, aspectName) - .orderBy() - .desc(EbeanAspectV2.VERSION_COLUMN) - .setMaxRows(1) - .findIds(); + .where() + .eq(EbeanAspectV2.URN_COLUMN, urn.toString()) + .eq(EbeanAspectV2.ASPECT_COLUMN, aspectName) + .orderBy() + .desc(EbeanAspectV2.VERSION_COLUMN) + .setMaxRows(1) + .findIds(); - return result.isEmpty() ? 0 : result.get(0).getVersion() + 1L; + return result.isEmpty() ? -1 : result.get(0).getVersion(); } - @Override - public Map getNextVersions(@Nonnull final String urn, @Nonnull final Set aspectNames) { + public Map> getNextVersions(@Nonnull Map> urnAspects) { validateConnection(); - Map result = new HashMap<>(); + Junction queryJunction = _server.find(EbeanAspectV2.class) - .select("aspect, max(version)") - .where() - .eq("urn", urn) - .or(); + .select("urn, aspect, max(version)") + .where() + .in("urn", urnAspects.keySet()) + .or(); ExpressionList exp = null; - for (String aspectName: aspectNames) { + for (Map.Entry> entry: urnAspects.entrySet()) { if (exp == null) { - exp = queryJunction.eq("aspect", aspectName); + exp = queryJunction.and() + .eq("urn", entry.getKey()) + .in("aspect", entry.getValue()) + .endAnd(); } else { - exp = exp.eq("aspect", aspectName); + exp = exp.and() + .eq("urn", entry.getKey()) + .in("aspect", entry.getValue()) + .endAnd(); } } + + Map> result = new HashMap<>(); + // Default next version 0 + urnAspects.forEach((key, value) -> { + Map defaultNextVersion = new HashMap<>(); + value.forEach(aspectName -> defaultNextVersion.put(aspectName, 0L)); + result.put(key, defaultNextVersion); + }); + if (exp == null) { return result; } - // Order by ascending version so that the results are correctly populated. - // TODO: Improve the below logic to be more explicit. - exp.orderBy().asc(EbeanAspectV2.VERSION_COLUMN); + List dbResults = exp.endOr().findIds(); for (EbeanAspectV2.PrimaryKey key: dbResults) { - result.put(key.getAspect(), key.getVersion()); - } - - for (String aspectName: aspectNames) { - long nextVal = ASPECT_LATEST_VERSION; - if (result.containsKey(aspectName)) { - nextVal = result.get(aspectName) + 1L; + if (result.get(key.getUrn()).get(key.getAspect()) <= key.getVersion()) { + result.get(key.getUrn()).put(key.getAspect(), key.getVersion() + 1L); } - result.put(aspectName, nextVal); } + return result; } @@ -676,4 +665,21 @@ public List getAspectsInRange(@Nonnull Urn urn, Set aspect .findList(); return ebeanAspects.stream().map(EbeanAspectV2::toEntityAspect).collect(Collectors.toList()); } + + private static Map toAspectMap(Set beans) { + return beans.stream().map(bean -> Map.entry(bean.getAspect(), bean)) + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().toEntityAspect())); + } + + private static Map> toUrnAspectMap(Collection beans) { + return beans.stream() + .collect(Collectors.groupingBy(EbeanAspectV2::getUrn, Collectors.toSet())) + .entrySet().stream() + .map(e -> Map.entry(e.getKey(), toAspectMap(e.getValue()))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + private static String buildMetricName(EntitySpec entitySpec, AspectSpec aspectSpec, String status) { + return String.join(MetricUtils.DELIMITER, List.of(entitySpec.getName(), aspectSpec.getName(), status.toLowerCase())); + } } 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 1c0729b7c27e4..d94ec1fa7ae2b 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 @@ -4,25 +4,28 @@ import com.datahub.util.RecordUtils; 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.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; import com.linkedin.retention.TimeBasedRetention; import com.linkedin.retention.VersionBasedRetention; import com.linkedin.metadata.Constants; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.Expression; import io.ebean.ExpressionList; import io.ebean.PagedList; import io.ebean.Query; import io.ebean.Transaction; +import io.ebean.TxScope; import io.ebeaninternal.server.expression.Op; import io.ebeaninternal.server.expression.SimpleExpression; import io.opentelemetry.extension.annotations.WithSpan; import java.sql.Timestamp; import java.time.Clock; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Objects; @@ -38,7 +41,7 @@ @RequiredArgsConstructor public class EbeanRetentionService extends RetentionService { private final EntityService _entityService; - private final EbeanServer _server; + private final Database _server; private final int _batchSize; private final Clock _clock = Clock.systemUTC(); @@ -48,35 +51,59 @@ public EntityService getEntityService() { return _entityService; } + @Override + protected AspectsBatch buildAspectsBatch(List mcps) { + return AspectsBatchImpl.builder() + .mcps(mcps, _entityService.getEntityRegistry()) + .build(); + } + @Override @WithSpan - public void applyRetention(@Nonnull Urn urn, @Nonnull String aspectName, Retention retentionPolicy, - Optional retentionContext) { - log.debug("Applying retention to urn {}, aspectName {}", urn, aspectName); - // If no policies are set or has indefinite policy set, do not apply any retention - if (retentionPolicy.data().isEmpty()) { - return; - } - ExpressionList deleteQuery = _server.find(EbeanAspectV2.class) - .where() - .eq(EbeanAspectV2.URN_COLUMN, urn.toString()) - .eq(EbeanAspectV2.ASPECT_COLUMN, aspectName) - .ne(EbeanAspectV2.VERSION_COLUMN, Constants.ASPECT_LATEST_VERSION) - .or(); - - List filterList = new ArrayList<>(); - if (retentionPolicy.hasVersion()) { - getVersionBasedRetentionQuery(urn, aspectName, retentionPolicy.getVersion(), - retentionContext.flatMap(RetentionService.RetentionContext::getMaxVersion)).ifPresent(filterList::add); - } - if (retentionPolicy.hasTime()) { - filterList.add(getTimeBasedRetentionQuery(retentionPolicy.getTime())); - } + protected void applyRetention(List retentionContexts) { + + List nonEmptyContexts = retentionContexts.stream() + .filter(context -> context.getRetentionPolicy().isPresent() + && !context.getRetentionPolicy().get().data().isEmpty()).collect(Collectors.toList()); // Only run delete if at least one of the retention policies are applicable - if (!filterList.isEmpty()) { - filterList.forEach(deleteQuery::add); - deleteQuery.endOr().delete(); + if (!nonEmptyContexts.isEmpty()) { + ExpressionList deleteQuery = _server.find(EbeanAspectV2.class) + .where() + .ne(EbeanAspectV2.VERSION_COLUMN, Constants.ASPECT_LATEST_VERSION) + .or(); + + boolean applied = false; + for (RetentionContext context : nonEmptyContexts) { + Retention retentionPolicy = context.getRetentionPolicy().get(); + + if (retentionPolicy.hasVersion()) { + boolean appliedVersion = getVersionBasedRetentionQuery(context.getUrn(), context.getAspectName(), + retentionPolicy.getVersion(), context.getMaxVersion()) + .map(expr -> + deleteQuery.and() + .eq(EbeanAspectV2.URN_COLUMN, context.getUrn().toString()) + .eq(EbeanAspectV2.ASPECT_COLUMN, context.getAspectName()) + .add(expr) + .endAnd() + ).isPresent(); + + applied = appliedVersion || applied; + } + + if (retentionPolicy.hasTime()) { + deleteQuery.and() + .eq(EbeanAspectV2.URN_COLUMN, context.getUrn().toString()) + .eq(EbeanAspectV2.ASPECT_COLUMN, context.getAspectName()) + .add(getTimeBasedRetentionQuery(retentionPolicy.getTime())) + .endAnd(); + applied = true; + } + } + + if (applied) { + deleteQuery.endOr().delete(); + } } } @@ -118,36 +145,46 @@ private void applyRetention( Map retentionPolicyMap, BulkApplyRetentionResult applyRetentionResult ) { - try (Transaction transaction = _server.beginTransaction()) { + try (Transaction transaction = _server.beginTransaction(TxScope.required())) { transaction.setBatchMode(true); transaction.setBatchSize(_batchSize); - for (EbeanAspectV2 row : rows.getList()) { - // Only run for cases where there's multiple versions of the aspect - if (row.getVersion() == 0) { - continue; - } - // 1. Extract an Entity type from the entity Urn - Urn urn; - try { - urn = Urn.createFromString(row.getUrn()); - } catch (Exception e) { - log.error("Failed to serialize urn {}", row.getUrn(), e); - continue; - } - final String aspectNameFromRecord = row.getAspect(); - log.debug("Handling urn {} aspect {}", row.getUrn(), row.getAspect()); - // Get the retention policies to apply from the local retention policy map - Optional retentionPolicy = getRetentionKeys(urn.getEntityType(), aspectNameFromRecord).stream() - .map(key -> retentionPolicyMap.get(key.toString())) - .filter(Objects::nonNull) - .findFirst() - .map(DataHubRetentionConfig::getRetention); - retentionPolicy.ifPresent(retention -> applyRetention(urn, aspectNameFromRecord, retention, - Optional.of(new RetentionContext(Optional.of(row.getVersion()))))); - if (applyRetentionResult != null) { - applyRetentionResult.rowsHandled += 1; - } + + List retentionContexts = rows.getList().stream() + .filter(row -> row.getVersion() != 0) + .map(row -> { + // 1. Extract an Entity type from the entity Urn + Urn urn; + try { + urn = Urn.createFromString(row.getUrn()); + } catch (Exception e) { + log.error("Failed to serialize urn {}", row.getUrn(), e); + return null; + } + + final String aspectNameFromRecord = row.getAspect(); + log.debug("Handling urn {} aspect {}", row.getUrn(), row.getAspect()); + // Get the retention policies to apply from the local retention policy map + Optional retentionPolicy = getRetentionKeys(urn.getEntityType(), aspectNameFromRecord).stream() + .map(key -> retentionPolicyMap.get(key.toString())) + .filter(Objects::nonNull) + .findFirst() + .map(DataHubRetentionConfig::getRetention); + + return RetentionService.RetentionContext.builder() + .urn(urn) + .aspectName(aspectNameFromRecord) + .retentionPolicy(retentionPolicy) + .maxVersion(Optional.of(row.getVersion())) + .build(); + }) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + + applyRetention(retentionContexts); + if (applyRetentionResult != null) { + applyRetentionResult.rowsHandled += retentionContexts.size(); } + transaction.commit(); } } 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 new file mode 100644 index 0000000000000..ca5e070bc5ca7 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java @@ -0,0 +1,67 @@ +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 lombok.Builder; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + + +@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/ebean/transactions/PatchBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java new file mode 100644 index 0000000000000..cc0b3d915b407 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java @@ -0,0 +1,188 @@ +package com.linkedin.metadata.entity.ebean.transactions; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.StreamReadConstraints; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.github.fge.jsonpatch.JsonPatch; +import com.github.fge.jsonpatch.JsonPatchException; +import com.github.fge.jsonpatch.Patch; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.events.metadata.ChangeType; +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.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; +import lombok.Builder; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Objects; + +import static com.linkedin.metadata.Constants.*; + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class PatchBatchItem extends AbstractBatchItem { + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + static { + int maxSize = Integer.parseInt(System.getenv().getOrDefault(INGESTION_MAX_SERIALIZED_STRING_LENGTH, MAX_JACKSON_STRING_SIZE)); + OBJECT_MAPPER.getFactory().setStreamReadConstraints(StreamReadConstraints.builder().maxStringLength(maxSize).build()); + } + + // urn an urn associated with the new aspect + private final Urn urn; + // aspectName name of the aspect being inserted + private final String aspectName; + private final SystemMetadata systemMetadata; + + private final Patch patch; + + private final MetadataChangeProposal metadataChangeProposal; + + // derived + private final EntitySpec entitySpec; + private final AspectSpec aspectSpec; + + @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() + .urn(getUrn()) + .aspectName(getAspectName()) + .metadataChangeProposal(getMetadataChangeProposal()) + .systemMetadata(getSystemMetadata()); + + AspectTemplateEngine aspectTemplateEngine = entityRegistry.getAspectTemplateEngine(); + + RecordTemplate currentValue = recordTemplate != null ? recordTemplate + : aspectTemplateEngine.getDefaultTemplate(getAspectName()); + + if (currentValue == null) { + // Attempting to patch a value to an aspect which has no default value and no existing value. + throw new UnsupportedOperationException(String.format("Patch not supported for aspect with name %s. " + + "Default aspect is required because no aspect currently exists for urn %s.", getAspectName(), getUrn())); + } + + try { + builder.aspect(aspectTemplateEngine.applyPatch(currentValue, getPatch(), getAspectSpec())); + } catch (JsonProcessingException | JsonPatchException e) { + throw new RuntimeException(e); + } + + return builder.build(entityRegistry); + } + + public static class PatchBatchItemBuilder { + + public PatchBatchItem build(EntityRegistry entityRegistry) { + EntityUtils.validateUrn(entityRegistry, this.urn); + log.debug("entity type = {}", this.urn.getEntityType()); + + entitySpec(entityRegistry.getEntitySpec(this.urn.getEntityType())); + log.debug("entity spec = {}", this.entitySpec); + + aspectSpec(ValidationUtils.validate(this.entitySpec, this.aspectName)); + log.debug("aspect spec = {}", this.aspectSpec); + + if (this.patch == null) { + throw new IllegalArgumentException(String.format("Missing patch to apply. Aspect: %s", + this.aspectSpec.getName())); + } + + return new PatchBatchItem(this.urn, this.aspectName, generateSystemMetadataIfEmpty(this.systemMetadata), + this.patch, this.metadataChangeProposal, this.entitySpec, this.aspectSpec); + } + + public static PatchBatchItem build(MetadataChangeProposal mcp, EntityRegistry entityRegistry) { + log.debug("entity type = {}", mcp.getEntityType()); + EntitySpec entitySpec = entityRegistry.getEntitySpec(mcp.getEntityType()); + AspectSpec aspectSpec = validateAspect(mcp, entitySpec); + + if (!isValidChangeType(ChangeType.PATCH, aspectSpec)) { + throw new UnsupportedOperationException("ChangeType not supported: " + mcp.getChangeType() + + " for aspect " + mcp.getAspectName()); + } + + Urn urn = mcp.getEntityUrn(); + if (urn == null) { + 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); + } + + private PatchBatchItemBuilder entitySpec(EntitySpec entitySpec) { + this.entitySpec = entitySpec; + return this; + } + + private PatchBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { + this.aspectSpec = aspectSpec; + return this; + } + + private static Patch convertToJsonPatch(MetadataChangeProposal mcp) { + JsonNode json; + try { + json = OBJECT_MAPPER.readTree(mcp.getAspect().getValue().asString(StandardCharsets.UTF_8)); + return JsonPatch.fromJson(json); + } catch (IOException e) { + throw new IllegalArgumentException("Invalid JSON Patch: " + mcp.getAspect().getValue(), e); + } + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + PatchBatchItem that = (PatchBatchItem) o; + return urn.equals(that.urn) && aspectName.equals(that.aspectName) && Objects.equals(systemMetadata, that.systemMetadata) && patch.equals(that.patch); + } + + @Override + public int hashCode() { + return Objects.hash(urn, aspectName, systemMetadata, patch); + } + + @Override + public String toString() { + return "PatchBatchItem{" + + "urn=" + urn + + ", aspectName='" + aspectName + + '\'' + + ", systemMetadata=" + systemMetadata + + ", patch=" + patch + + '}'; + } +} 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/transactions/UpsertBatchItem.java new file mode 100644 index 0000000000000..bd58d267a8308 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java @@ -0,0 +1,173 @@ +package com.linkedin.metadata.entity.ebean.transactions; + +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.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.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; +import lombok.Builder; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +import java.sql.Timestamp; +import java.util.Objects; + +import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; + + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class UpsertBatchItem extends AbstractBatchItem { + + // urn an urn associated with the new aspect + private final Urn urn; + // aspectName name of the aspect being inserted + private final String aspectName; + private final SystemMetadata systemMetadata; + + private final RecordTemplate aspect; + + private final MetadataChangeProposal metadataChangeProposal; + + // derived + private final EntitySpec entitySpec; + private final AspectSpec aspectSpec; + + @Override + public ChangeType getChangeType() { + return ChangeType.UPSERT; + } + + @Override + public void validateUrn(EntityRegistry entityRegistry, Urn urn) { + EntityUtils.validateUrn(entityRegistry, urn); + } + + public EntityAspect toLatestEntityAspect(AuditStamp auditStamp) { + EntityAspect latest = new EntityAspect(); + latest.setAspect(getAspectName()); + latest.setMetadata(EntityUtils.toJsonAspect(getAspect())); + latest.setUrn(getUrn().toString()); + latest.setVersion(ASPECT_LATEST_VERSION); + latest.setCreatedOn(new Timestamp(auditStamp.getTime())); + latest.setCreatedBy(auditStamp.getActor().toString()); + return latest; + } + + public static class UpsertBatchItemBuilder { + + public UpsertBatchItem build(EntityRegistry entityRegistry) { + EntityUtils.validateUrn(entityRegistry, this.urn); + log.debug("entity type = {}", this.urn.getEntityType()); + + entitySpec(entityRegistry.getEntitySpec(this.urn.getEntityType())); + log.debug("entity spec = {}", this.entitySpec); + + aspectSpec(ValidationUtils.validate(this.entitySpec, this.aspectName)); + log.debug("aspect spec = {}", this.aspectSpec); + + ValidationUtils.validateRecordTemplate(entityRegistry, this.entitySpec, this.urn, this.aspect); + + return new UpsertBatchItem(this.urn, this.aspectName, AbstractBatchItem.generateSystemMetadataIfEmpty(this.systemMetadata), + this.aspect, this.metadataChangeProposal, this.entitySpec, this.aspectSpec); + } + + public static UpsertBatchItem build(MetadataChangeProposal mcp, EntityRegistry entityRegistry) { + if (!mcp.getChangeType().equals(ChangeType.UPSERT)) { + throw new IllegalArgumentException("Invalid MCP, this class only supports change type of UPSERT."); + } + + log.debug("entity type = {}", mcp.getEntityType()); + EntitySpec entitySpec = entityRegistry.getEntitySpec(mcp.getEntityType()); + AspectSpec aspectSpec = validateAspect(mcp, entitySpec); + + if (!isValidChangeType(ChangeType.UPSERT, aspectSpec)) { + throw new UnsupportedOperationException("ChangeType not supported: " + mcp.getChangeType() + + " for aspect " + mcp.getAspectName()); + } + + Urn urn = mcp.getEntityUrn(); + if (urn == null) { + 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); + } + + private UpsertBatchItemBuilder entitySpec(EntitySpec entitySpec) { + this.entitySpec = entitySpec; + return this; + } + + private UpsertBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { + this.aspectSpec = aspectSpec; + return this; + } + + private static RecordTemplate convertToRecordTemplate(MetadataChangeProposal mcp, AspectSpec aspectSpec) { + RecordTemplate aspect; + try { + aspect = GenericRecordUtils.deserializeAspect(mcp.getAspect().getValue(), + mcp.getAspect().getContentType(), aspectSpec); + ValidationUtils.validateOrThrow(aspect); + } catch (ModelConversionException e) { + throw new RuntimeException( + String.format("Could not deserialize %s for aspect %s", mcp.getAspect().getValue(), + mcp.getAspectName())); + } + log.debug("aspect = {}", aspect); + return aspect; + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + UpsertBatchItem that = (UpsertBatchItem) o; + return urn.equals(that.urn) && aspectName.equals(that.aspectName) && Objects.equals(systemMetadata, that.systemMetadata) && aspect.equals(that.aspect); + } + + @Override + public int hashCode() { + return Objects.hash(urn, aspectName, systemMetadata, aspect); + } + + @Override + public String toString() { + return "UpsertBatchItem{" + + "urn=" + + urn + + ", aspectName='" + + aspectName + + '\'' + + ", systemMetadata=" + + systemMetadata + + ", aspect=" + + aspect + + '}'; + } +} 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 99bb323e51ecb..6182b27333cbb 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 @@ -1,8 +1,16 @@ package com.linkedin.metadata.entity.validation; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.schema.validation.ValidationResult; import com.linkedin.data.template.RecordTemplate; +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 lombok.extern.slf4j.Slf4j; +import java.util.function.Consumer; + @Slf4j public class ValidationUtils { @@ -33,6 +41,36 @@ public static void validateOrWarn(RecordTemplate record) { }); } + public static AspectSpec validate(EntitySpec entitySpec, String aspectName) { + if (aspectName == null || aspectName.isEmpty()) { + throw new UnsupportedOperationException("Aspect name is required for create and update operations"); + } + + AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); + + if (aspectSpec == null) { + throw new RuntimeException( + String.format("Unknown aspect %s for entity %s", aspectName, entitySpec.getName())); + } + + return aspectSpec; + } + + public static void validateRecordTemplate(EntityRegistry entityRegistry, EntitySpec entitySpec, Urn urn, RecordTemplate aspect) { + EntityRegistryUrnValidator validator = new EntityRegistryUrnValidator(entityRegistry); + validator.setCurrentEntitySpec(entitySpec); + Consumer resultFunction = validationResult -> { + throw new IllegalArgumentException("Invalid format for aspect: " + entitySpec.getName() + "\n Cause: " + + validationResult.getMessages()); }; + 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); + } + private ValidationUtils() { } } \ No newline at end of file diff --git a/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java b/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java index c83764284c0c4..ffadc07124727 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/event/EventProducer.java @@ -1,16 +1,12 @@ package com.linkedin.metadata.event; import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; -import com.linkedin.metadata.snapshot.Snapshot; import com.linkedin.mxe.DataHubUpgradeHistoryEvent; import com.linkedin.mxe.MetadataChangeLog; -import com.linkedin.mxe.MetadataAuditOperation; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.PlatformEvent; -import com.linkedin.mxe.SystemMetadata; import io.opentelemetry.extension.annotations.WithSpan; import java.util.concurrent.Future; import javax.annotation.Nonnull; @@ -22,27 +18,6 @@ */ public interface EventProducer { - /** - * Deprecated! Replaced by {@link #produceMetadataChangeLog(Urn, AspectSpec, MetadataChangeLog)} - * - * Produces a {@link com.linkedin.mxe.MetadataAuditEvent} from a - * new & previous Entity {@link Snapshot}. - * @param urn the urn associated with the entity changed - * @param oldSnapshot a {@link RecordTemplate} corresponding to the old snapshot. - * @param newSnapshot a {@link RecordTemplate} corresponding to the new snapshot. - * @param oldSystemMetadata - * @param newSystemMetadata - */ - @Deprecated - void produceMetadataAuditEvent( - @Nonnull final Urn urn, - @Nullable final Snapshot oldSnapshot, - @Nonnull final Snapshot newSnapshot, - @Nullable SystemMetadata oldSystemMetadata, - @Nullable SystemMetadata newSystemMetadata, - MetadataAuditOperation operation - ); - /** * Produces a {@link com.linkedin.mxe.MetadataChangeLog} from a * new & previous aspect. 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 2361bcc22780a..e95378a616d97 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,12 @@ 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.key.CorpUserKey; import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; import java.util.Map; import javax.annotation.Nonnull; @@ -25,12 +29,19 @@ public static Map ingestCorpUserKeyAspects(EntityService entit public static Map ingestCorpUserKeyAspects(EntityService entityService, int aspectCount, int startIndex) { String aspectName = AspectGenerationUtils.getAspectName(new CorpUserKey()); Map aspects = new HashMap<>(); + 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); - entityService.ingestAspect(urn, aspectName, aspect, AspectGenerationUtils.createAuditStamp(), AspectGenerationUtils.createSystemMetadata()); + items.add(UpsertBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .aspect(aspect) + .systemMetadata(AspectGenerationUtils.createSystemMetadata()) + .build(entityService.getEntityRegistry())); } + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } @@ -43,13 +54,20 @@ public static Map ingestCorpUserInfoAspects(@Nonnull final En 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<>(); 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); - entityService.ingestAspect(urn, aspectName, aspect, AspectGenerationUtils.createAuditStamp(), AspectGenerationUtils.createSystemMetadata()); + items.add(UpsertBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .aspect(aspect) + .systemMetadata(AspectGenerationUtils.createSystemMetadata()) + .build(entityService.getEntityRegistry())); } + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } @@ -62,14 +80,21 @@ public static Map ingestChartInfoAspects(@Nonnull final EntitySe 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<>(); 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); String description = String.format("Test description %d", i); ChartInfo aspect = AspectGenerationUtils.createChartInfo(title, description); aspects.put(urn, aspect); - entityService.ingestAspect(urn, aspectName, aspect, AspectGenerationUtils.createAuditStamp(), AspectGenerationUtils.createSystemMetadata()); + items.add(UpsertBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .aspect(aspect) + .systemMetadata(AspectGenerationUtils.createSystemMetadata()) + .build(entityService.getEntityRegistry())); } + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java b/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java index 36ebec5a42849..46d08bc8887b9 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java @@ -18,7 +18,7 @@ import com.linkedin.metadata.snapshot.Snapshot; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; -import io.ebean.EbeanServer; +import io.ebean.Database; import java.util.List; import org.testng.Assert; import org.testng.annotations.Test; @@ -39,7 +39,7 @@ public AspectUtilsTest() throws EntityRegistryException { @Test public void testAdditionalChanges() { - EbeanServer server = EbeanTestUtils.createTestServer(); + Database server = EbeanTestUtils.createTestServer(); EbeanAspectDao aspectDao = new EbeanAspectDao(server); aspectDao.setConnectionValidated(true); EventProducer mockProducer = mock(EventProducer.class); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java b/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java index 45c4c16864b07..7e9605cbe3db0 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/ESTestUtils.java @@ -34,8 +34,9 @@ import org.testcontainers.elasticsearch.ElasticsearchContainer; import org.testcontainers.utility.DockerImageName; -import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.*; -import static com.linkedin.metadata.DockerTestUtils.*; +import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.AUTO_COMPLETE_ENTITY_TYPES; +import static com.linkedin.datahub.graphql.resolvers.search.SearchUtils.SEARCHABLE_ENTITY_TYPES; +import static com.linkedin.metadata.DockerTestUtils.checkContainerEngine; public class ESTestUtils { private ESTestUtils() { diff --git a/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java b/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java index d8d7efeff87d4..180166e963fca 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java @@ -1,7 +1,7 @@ package com.linkedin.metadata; -import io.ebean.EbeanServer; -import io.ebean.EbeanServerFactory; +import io.ebean.Database; +import io.ebean.DatabaseFactory; import io.ebean.config.ServerConfig; import io.ebean.datasource.DataSourceConfig; @@ -13,8 +13,8 @@ private EbeanTestUtils() { } @Nonnull - public static EbeanServer createTestServer() { - return EbeanServerFactory.create(createTestingH2ServerConfig()); + public static Database createTestServer() { + return DatabaseFactory.create(createTestingH2ServerConfig()); } @Nonnull @@ -22,7 +22,7 @@ private static ServerConfig createTestingH2ServerConfig() { DataSourceConfig dataSourceConfig = new DataSourceConfig(); dataSourceConfig.setUsername("tester"); dataSourceConfig.setPassword(""); - dataSourceConfig.setUrl("jdbc:h2:mem:;IGNORECASE=TRUE;"); + dataSourceConfig.setUrl("jdbc:h2:mem:test;IGNORECASE=TRUE;mode=mysql;"); dataSourceConfig.setDriver("org.h2.Driver"); ServerConfig serverConfig = new ServerConfig(); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java index 62f8827b574b8..9e453e6e75677 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java @@ -7,7 +7,7 @@ import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.registry.EntityRegistryException; import com.linkedin.metadata.service.UpdateIndicesService; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -22,7 +22,7 @@ public EbeanAspectMigrationsDaoTest() throws EntityRegistryException { @BeforeMethod public void setupTest() { - EbeanServer server = EbeanTestUtils.createTestServer(); + Database server = EbeanTestUtils.createTestServer(); _mockProducer = mock(EventProducer.class); EbeanAspectDao dao = new EbeanAspectDao(server); dao.setConnectionValidated(true); 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 9126aad62895d..90f9baa4ca4c2 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 @@ -1,5 +1,7 @@ package com.linkedin.metadata.entity; +import com.linkedin.common.AuditStamp; +import com.linkedin.metadata.Constants; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; @@ -10,21 +12,36 @@ import com.linkedin.metadata.EbeanTestUtils; 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.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; import com.linkedin.metadata.models.registry.EntityRegistryException; import com.linkedin.metadata.query.ListUrnsResult; import com.linkedin.metadata.service.UpdateIndicesService; import com.linkedin.metadata.utils.PegasusUtils; +import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; -import io.ebean.EbeanServer; +import io.datahub.test.DataGenerator; +import io.ebean.Database; import io.ebean.Transaction; import io.ebean.TxScope; import io.ebean.annotation.TxIsolation; +import org.apache.commons.lang3.tuple.Triple; import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import java.net.URISyntaxException; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; @@ -44,10 +61,10 @@ public EbeanEntityServiceTest() throws EntityRegistryException { @BeforeMethod public void setupTest() { - EbeanServer server = EbeanTestUtils.createTestServer(); + Database server = EbeanTestUtils.createTestServer(); _mockProducer = mock(EventProducer.class); _aspectDao = new EbeanAspectDao(server); - _aspectDao.setConnectionValidated(true); + _mockUpdateIndicesService = mock(UpdateIndicesService.class); PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(true); @@ -86,15 +103,34 @@ public void testIngestListLatestAspects() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); - _entityServiceImpl.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - _entityServiceImpl.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); + + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .aspect(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // List aspects ListResult batch1 = _entityServiceImpl.listLatestAspects(entityUrn1.getEntityType(), aspectName, 0, 2); @@ -131,15 +167,34 @@ public void testIngestListUrns() throws AssertionError { // Ingest CorpUserInfo Aspect #1 RecordTemplate writeAspect1 = AspectGenerationUtils.createCorpUserKey(entityUrn1); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #2 RecordTemplate writeAspect2 = AspectGenerationUtils.createCorpUserKey(entityUrn2); - _entityServiceImpl.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 RecordTemplate writeAspect3 = AspectGenerationUtils.createCorpUserKey(entityUrn3); - _entityServiceImpl.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); + + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .aspect(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // List aspects urns ListUrnsResult batch1 = _entityServiceImpl.listUrns(entityUrn1.getEntityType(), 0, 2); @@ -163,7 +218,7 @@ public void testIngestListUrns() throws AssertionError { @Override @Test public void testNestedTransactions() throws AssertionError { - EbeanServer server = _aspectDao.getServer(); + Database server = _aspectDao.getServer(); try (Transaction transaction = server.beginTransaction(TxScope.requiresNew() .setIsolation(TxIsolation.REPEATABLE_READ))) { @@ -183,4 +238,178 @@ public void testNestedTransactions() throws AssertionError { } System.out.println("done"); } + + @Test + public void dataGeneratorThreadingTest() { + DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl); + List aspects = List.of("status", "globalTags", "glossaryTerms"); + List> testData = dataGenerator.generateMCPs("dataset", 25, aspects) + .collect(Collectors.toList()); + + // Expected no duplicates aspects + List duplicates = testData.stream() + .flatMap(Collection::stream) + .map(mcp -> Triple.of(mcp.getEntityUrn().toString(), mcp.getAspectName(), 0L)) + .collect(Collectors.groupingBy(Triple::toString)) + .entrySet().stream() + .filter(e -> e.getValue().size() > 1) + .map(Map.Entry::getKey) + .collect(Collectors.toList()); + assertEquals(duplicates.size(), 0, duplicates.toString()); + } + + /** + * This test is designed to detect multi-threading persistence exceptions like duplicate key, + * exceptions that exceed retry limits or unnecessary versions. + */ + @Test + public void multiThreadingTest() { + DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl); + Database server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); + + // Add data + List aspects = List.of("status", "globalTags", "glossaryTerms"); + List> testData = dataGenerator.generateMCPs("dataset", 25, aspects) + .collect(Collectors.toList()); + + executeThreadingTest(_entityServiceImpl, testData, 15); + + // Expected aspects + Set> generatedAspectIds = testData.stream() + .flatMap(Collection::stream) + .map(mcp -> Triple.of(mcp.getEntityUrn().toString(), mcp.getAspectName(), 0L)) + .collect(Collectors.toSet()); + + // Actual inserts + Set> actualAspectIds = server.sqlQuery( + "select urn, aspect, version from metadata_aspect_v2").findList().stream() + .map(row -> Triple.of(row.getString("urn"), row.getString("aspect"), row.getLong("version"))) + .collect(Collectors.toSet()); + + // Assert State + Set> additions = actualAspectIds.stream() + .filter(id -> !generatedAspectIds.contains(id)) + .collect(Collectors.toSet()); + assertEquals(additions.size(), 0, String.format("Expected no additional aspects. Found: %s", additions)); + + Set> missing = generatedAspectIds.stream() + .filter(id -> !actualAspectIds.contains(id)) + .collect(Collectors.toSet()); + assertEquals(missing.size(), 0, String.format("Expected all generated aspects to be inserted. Missing: %s", missing)); + } + + /** + * Don't blame multi-threading for what might not be a threading issue. + * Perform the multi-threading test with 1 thread. + */ + @Test + public void singleThreadingTest() { + DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl); + Database server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); + + // Add data + List aspects = List.of("status", "globalTags", "glossaryTerms"); + List> testData = dataGenerator.generateMCPs("dataset", 25, aspects) + .collect(Collectors.toList()); + + executeThreadingTest(_entityServiceImpl, testData, 1); + + // Expected aspects + Set> generatedAspectIds = testData.stream() + .flatMap(Collection::stream) + .map(mcp -> Triple.of(mcp.getEntityUrn().toString(), mcp.getAspectName(), 0L)) + .collect(Collectors.toSet()); + + // Actual inserts + Set> actualAspectIds = server.sqlQuery( + "select urn, aspect, version from metadata_aspect_v2").findList().stream() + .map(row -> Triple.of(row.getString("urn"), row.getString("aspect"), row.getLong("version"))) + .collect(Collectors.toSet()); + + // Assert State + Set> additions = actualAspectIds.stream() + .filter(id -> !generatedAspectIds.contains(id)) + .collect(Collectors.toSet()); + assertEquals(additions.size(), 0, String.format("Expected no additional aspects. Found: %s", additions)); + + Set> missing = generatedAspectIds.stream() + .filter(id -> !actualAspectIds.contains(id)) + .collect(Collectors.toSet()); + assertEquals(missing.size(), 0, String.format("Expected all generated aspects to be inserted. Missing: %s", missing)); + } + + private static void executeThreadingTest(EntityServiceImpl entityService, List> testData, + int threadCount) { + Database server = ((EbeanAspectDao) entityService._aspectDao).getServer(); + server.sqlUpdate("truncate metadata_aspect_v2"); + + int count = Objects.requireNonNull(server.sqlQuery( + "select count(*) as cnt from metadata_aspect_v2").findOne()).getInteger("cnt"); + assertEquals(count, 0, "Expected exactly 0 rows at the start."); + + // Create ingest proposals in parallel, mimic the smoke-test ingestion + final LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(threadCount * 2); + + // Spin up workers + List writeThreads = IntStream.range(0, threadCount) + .mapToObj(threadId -> new Thread(new MultiThreadTestWorker(queue, entityService))) + .collect(Collectors.toList()); + writeThreads.forEach(Thread::start); + + testData.forEach(mcps -> { + try { + queue.put(mcps); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + + // Terminate workers with empty mcp + IntStream.range(0, threadCount).forEach(threadId -> { + try { + queue.put(List.of()); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + + // Wait for threads to finish + writeThreads.forEach(thread -> { + try { + thread.join(10000); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + } + + private static class MultiThreadTestWorker implements Runnable { + private final EntityServiceImpl entityService; + private final LinkedBlockingQueue> queue; + + public MultiThreadTestWorker(LinkedBlockingQueue> queue, EntityServiceImpl entityService) { + this.queue = queue; + this.entityService = entityService; + } + + public void run() { + try { + while (true) { + List mcps = queue.take(); + if (mcps.isEmpty()) { + break; + } + final AuditStamp auditStamp = new AuditStamp(); + 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); + } + } 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 d485981f32a07..c0d2a3783c0a7 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 @@ -32,6 +32,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.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; @@ -46,7 +48,6 @@ import com.linkedin.metadata.snapshot.Snapshot; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.GenericAspect; -import com.linkedin.mxe.MetadataAuditOperation; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; @@ -146,9 +147,6 @@ public void testIngestGetEntity() throws Exception { assertNull(mcl.getPreviousSystemMetadata()); assertEquals(mcl.getChangeType(), ChangeType.UPSERT); - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); } @@ -183,9 +181,6 @@ public void testAddKey() throws Exception { assertNull(mcl.getPreviousSystemMetadata()); assertEquals(mcl.getChangeType(), ChangeType.UPSERT); - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); } @@ -254,12 +249,6 @@ public void testIngestGetEntities() throws Exception { assertNull(mcl.getPreviousSystemMetadata()); assertEquals(mcl.getChangeType(), ChangeType.UPSERT); - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn1), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn2), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); } @@ -314,12 +303,6 @@ public void testIngestGetEntitiesV2() throws Exception { EnvelopedAspect envelopedKey2 = readEntityResponse2.getAspects().get(keyName); assertTrue(DataTemplateUtil.areEqual(expectedKey2, new CorpUserKey(envelopedKey2.getValue().data()))); - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn1), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn2), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn1), Mockito.any(), Mockito.any()); @@ -382,12 +365,6 @@ public void testIngestGetEntitiesVersionedV2() throws Exception { EnvelopedAspect envelopedKey2 = readEntityResponse2.getAspects().get(keyName); assertTrue(DataTemplateUtil.areEqual(expectedKey2, new CorpUserKey(envelopedKey2.getValue().data()))); - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn1), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn2), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn1), Mockito.any(), Mockito.any()); @@ -413,7 +390,6 @@ public void testIngestAspectsGetLatestAspects() throws Exception { pairToIngest.add(getAspectRecordPair(writeAspect2, CorpUserInfo.class)); SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(); - _entityServiceImpl.ingestAspects(entityUrn, pairToIngest, TEST_AUDIT_STAMP, metadata1); Map latestAspects = _entityServiceImpl.getLatestAspectsForUrn( @@ -425,8 +401,6 @@ public void testIngestAspectsGetLatestAspects() throws Exception { verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.any()); - verify(_mockProducer, times(2)).produceMetadataAuditEvent(Mockito.eq(entityUrn), - Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); verifyNoMoreInteractions(_mockProducer); } @@ -444,7 +418,6 @@ public void testReingestAspectsGetLatestAspects() throws Exception { pairToIngest.add(getAspectRecordPair(writeAspect1, CorpUserInfo.class)); SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(); - _entityServiceImpl.ingestAspects(entityUrn, pairToIngest, TEST_AUDIT_STAMP, metadata1); final MetadataChangeLog initialChangeLog = new MetadataChangeLog(); @@ -478,8 +451,6 @@ public void testReingestAspectsGetLatestAspects() throws Exception { verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(initialChangeLog)); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), - Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); // Mockito detects the previous invocation and throws an error in verifying the second call unless invocations are cleared clearInvocations(_mockProducer); @@ -488,8 +459,6 @@ public void testReingestAspectsGetLatestAspects() throws Exception { verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(restateChangeLog)); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), - Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); verifyNoMoreInteractions(_mockProducer); @@ -541,8 +510,6 @@ public void testReingestLineageAspect() throws Exception { verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(initialChangeLog)); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), - Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); // Mockito detects the previous invocation and throws an error in verifying the second call unless invocations are cleared clearInvocations(_mockProducer); @@ -551,8 +518,6 @@ public void testReingestLineageAspect() throws Exception { verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(restateChangeLog)); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), - Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); verifyNoMoreInteractions(_mockProducer); @@ -704,9 +669,9 @@ public void testUpdateGetAspect() throws AssertionError { CorpUserInfo writeAspect = AspectGenerationUtils.createCorpUserInfo("email@test.com"); // Validate retrieval of CorpUserInfo Aspect #1 - _entityServiceImpl.updateAspect(entityUrn, "corpuser", aspectName, corpUserInfoSpec, writeAspect, TEST_AUDIT_STAMP, 1, - true); - RecordTemplate readAspect1 = _entityServiceImpl.getAspect(entityUrn, aspectName, 1); + _entityServiceImpl.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect)), TEST_AUDIT_STAMP, null); + + RecordTemplate readAspect1 = _entityServiceImpl.getAspect(entityUrn, aspectName, 0); assertTrue(DataTemplateUtil.areEqual(writeAspect, readAspect1)); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), Mockito.any()); @@ -715,10 +680,13 @@ public void testUpdateGetAspect() throws AssertionError { writeAspect.setEmail("newemail@test.com"); // Validate retrieval of CorpUserInfo Aspect #2 - _entityServiceImpl.updateAspect(entityUrn, "corpuser", aspectName, corpUserInfoSpec, writeAspect, TEST_AUDIT_STAMP, 1, - false); - RecordTemplate readAspect2 = _entityServiceImpl.getAspect(entityUrn, aspectName, 1); + _entityServiceImpl.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect)), TEST_AUDIT_STAMP, null); + + RecordTemplate readAspect2 = _entityServiceImpl.getAspect(entityUrn, aspectName, 0); assertTrue(DataTemplateUtil.areEqual(writeAspect, readAspect2)); + verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), + Mockito.any()); + verifyNoMoreInteractions(_mockProducer); } @@ -731,26 +699,39 @@ public void testGetAspectAtVersion() throws AssertionError { AspectSpec corpUserInfoSpec = _testEntityRegistry.getEntitySpec("corpuser").getAspectSpec("corpUserInfo"); // Ingest CorpUserInfo Aspect #1 - CorpUserInfo writeAspect = AspectGenerationUtils.createCorpUserInfo("email@test.com"); + CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); + CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); // Validate retrieval of CorpUserInfo Aspect #1 - _entityServiceImpl.updateAspect(entityUrn, "corpuser", aspectName, corpUserInfoSpec, writeAspect, TEST_AUDIT_STAMP, 1, - true); + _entityServiceImpl.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect1)), TEST_AUDIT_STAMP, null); - VersionedAspect writtenVersionedAspect = new VersionedAspect(); - writtenVersionedAspect.setAspect(Aspect.create(writeAspect)); - writtenVersionedAspect.setVersion(1); + VersionedAspect writtenVersionedAspect1 = new VersionedAspect(); + writtenVersionedAspect1.setAspect(Aspect.create(writeAspect1)); + writtenVersionedAspect1.setVersion(0); - VersionedAspect readAspect1 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, 1); - assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect, readAspect1)); + VersionedAspect readAspect1 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, 0); + assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), Mockito.any()); - VersionedAspect readAspect2 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, -1); - assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect, readAspect2)); + readAspect1 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, -1); + assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); + + // Validate retrieval of CorpUserInfo Aspect #2 + _entityServiceImpl.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect2)), TEST_AUDIT_STAMP, null); + + VersionedAspect writtenVersionedAspect2 = new VersionedAspect(); + writtenVersionedAspect2.setAspect(Aspect.create(writeAspect2)); + writtenVersionedAspect2.setVersion(0); + + VersionedAspect readAspectVersion2 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, 0); + assertFalse(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspectVersion2)); + assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect2, readAspectVersion2)); + verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), + Mockito.any()); - VersionedAspect readAspectVersion0 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, 0); - assertFalse(DataTemplateUtil.areEqual(writtenVersionedAspect, readAspectVersion0)); + readAspect1 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, -1); + assertFalse(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); verifyNoMoreInteractions(_mockProducer); } @@ -768,19 +749,43 @@ public void testRollbackAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); - _entityServiceImpl.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - _entityServiceImpl.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1Overwrite, TEST_AUDIT_STAMP, metadata2); + + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .aspect(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1Overwrite) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // this should no-op since this run has been overwritten AspectRowSummary rollbackOverwrittenAspect = new AspectRowSummary(); @@ -822,14 +827,33 @@ public void testRollbackKey() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); - RecordTemplate writeKey1 = _entityServiceImpl.buildKeyAspect(entityUrn1); - _entityServiceImpl.ingestAspect(entityUrn1, keyAspectName, writeKey1, TEST_AUDIT_STAMP, metadata1); + RecordTemplate writeKey1 = EntityUtils.buildKeyAspect(_testEntityRegistry, entityUrn1); // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1Overwrite, TEST_AUDIT_STAMP, metadata2); + + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(keyAspectName) + .aspect(writeKey1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1Overwrite) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // this should no-op since the key should have been written in the furst run AspectRowSummary rollbackKeyWithWrongRunId = new AspectRowSummary(); @@ -873,22 +897,51 @@ public void testRollbackUrn() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); - RecordTemplate writeKey1 = _entityServiceImpl.buildKeyAspect(entityUrn1); - _entityServiceImpl.ingestAspect(entityUrn1, keyAspectName, writeKey1, TEST_AUDIT_STAMP, metadata1); + RecordTemplate writeKey1 = EntityUtils.buildKeyAspect(_testEntityRegistry, entityUrn1); // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); - _entityServiceImpl.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - _entityServiceImpl.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - _entityServiceImpl.ingestAspect(entityUrn1, aspectName, writeAspect1Overwrite, TEST_AUDIT_STAMP, metadata2); + + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(keyAspectName) + .aspect(writeKey1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .aspect(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .aspect(writeAspect1Overwrite) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // this should no-op since the key should have been written in the furst run AspectRowSummary rollbackKeyWithWrongRunId = new AspectRowSummary(); @@ -918,8 +971,17 @@ public void testIngestGetLatestAspect() throws AssertionError { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #1 - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, readAspect1)); @@ -931,9 +993,6 @@ public void testIngestGetLatestAspect() throws AssertionError { assertNull(mcl.getPreviousSystemMetadata()); assertEquals(mcl.getChangeType(), ChangeType.UPSERT); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.any(), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); reset(_mockProducer); @@ -941,8 +1000,17 @@ public void testIngestGetLatestAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); + items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #2 - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata2); RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); EntityAspect readAspectDao1 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 1); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 0); @@ -958,9 +1026,6 @@ public void testIngestGetLatestAspect() throws AssertionError { assertNotNull(mcl.getPreviousSystemMetadata()); assertEquals(mcl.getChangeType(), ChangeType.UPSERT); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.notNull(), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); } @@ -975,16 +1040,34 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #1 - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); EnvelopedAspect readAspect1 = _entityServiceImpl.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, new CorpUserInfo(readAspect1.getValue().data()))); // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); + items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #2 - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata2); EnvelopedAspect readAspect2 = _entityServiceImpl.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); EntityAspect readAspectDao1 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 1); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 0); @@ -993,12 +1076,6 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { assertTrue(DataTemplateUtil.areEqual(EntityUtils.parseSystemMetadata(readAspectDao2.getSystemMetadata()), metadata2)); assertTrue(DataTemplateUtil.areEqual(EntityUtils.parseSystemMetadata(readAspectDao1.getSystemMetadata()), metadata1)); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.notNull(), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.any()); @@ -1017,8 +1094,17 @@ public void testIngestSameAspect() throws AssertionError { SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); SystemMetadata metadata3 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-123"); + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #1 - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, readAspect1)); @@ -1030,9 +1116,6 @@ public void testIngestSameAspect() throws AssertionError { assertNull(mcl.getPreviousSystemMetadata()); assertEquals(mcl.getChangeType(), ChangeType.UPSERT); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.eq(null), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); reset(_mockProducer); @@ -1040,8 +1123,17 @@ public void testIngestSameAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); + items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect2) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #2 - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata2); RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, ASPECT_LATEST_VERSION); @@ -1053,9 +1145,6 @@ public void testIngestSameAspect() throws AssertionError { verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), mclCaptor.capture()); - verify(_mockProducer, times(1)).produceMetadataAuditEvent(Mockito.eq(entityUrn), Mockito.notNull(), Mockito.any(), - Mockito.any(), Mockito.any(), Mockito.eq(MetadataAuditOperation.UPDATE)); - verifyNoMoreInteractions(_mockProducer); } @@ -1069,20 +1158,54 @@ public void testRetention() throws AssertionError { // Ingest CorpUserInfo Aspect CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); CorpUserInfo writeAspect1a = AspectGenerationUtils.createCorpUserInfo("email_a@test.com"); - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1a, TEST_AUDIT_STAMP, metadata1); CorpUserInfo writeAspect1b = AspectGenerationUtils.createCorpUserInfo("email_b@test.com"); - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1b, TEST_AUDIT_STAMP, metadata1); String aspectName2 = AspectGenerationUtils.getAspectName(new Status()); // Ingest Status Aspect Status writeAspect2 = new Status().setRemoved(true); - _entityServiceImpl.ingestAspect(entityUrn, aspectName2, writeAspect2, TEST_AUDIT_STAMP, metadata1); Status writeAspect2a = new Status().setRemoved(false); - _entityServiceImpl.ingestAspect(entityUrn, aspectName2, writeAspect2a, TEST_AUDIT_STAMP, metadata1); Status writeAspect2b = new Status().setRemoved(true); - _entityServiceImpl.ingestAspect(entityUrn, aspectName2, writeAspect2b, TEST_AUDIT_STAMP, metadata1); + + List items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1a) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1b) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .aspect(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .aspect(writeAspect2a) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .aspect(writeAspect2b) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName, 1), writeAspect1); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1094,10 +1217,24 @@ public void testRetention() throws AssertionError { // Ingest CorpUserInfo Aspect again CorpUserInfo writeAspect1c = AspectGenerationUtils.createCorpUserInfo("email_c@test.com"); - _entityServiceImpl.ingestAspect(entityUrn, aspectName, writeAspect1c, TEST_AUDIT_STAMP, metadata1); // Ingest Status Aspect again Status writeAspect2c = new Status().setRemoved(false); - _entityServiceImpl.ingestAspect(entityUrn, aspectName2, writeAspect2c, TEST_AUDIT_STAMP, metadata1); + + items = List.of( + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .aspect(writeAspect1c) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + UpsertBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .aspect(writeAspect2c) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); assertNull(_entityServiceImpl.getAspect(entityUrn, aspectName, 1)); assertEquals(_entityServiceImpl.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1203,12 +1340,12 @@ public void testRestoreIndices() throws Exception { public void testValidateUrn() throws Exception { // Valid URN Urn validTestUrn = new Urn("li", "corpuser", new TupleKey("testKey")); - _entityServiceImpl.validateUrn(validTestUrn); + EntityUtils.validateUrn(_testEntityRegistry, validTestUrn); // URN with trailing whitespace Urn testUrnWithTrailingWhitespace = new Urn("li", "corpuser", new TupleKey("testKey ")); try { - _entityServiceImpl.validateUrn(testUrnWithTrailingWhitespace); + EntityUtils.validateUrn(_testEntityRegistry, testUrnWithTrailingWhitespace); Assert.fail("Should have raised IllegalArgumentException for URN with trailing whitespace"); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "Error: cannot provide an URN with leading or trailing whitespace"); @@ -1219,7 +1356,7 @@ public void testValidateUrn() throws Exception { Urn testUrnTooLong = new Urn("li", "corpuser", new TupleKey(stringTooLong)); try { - _entityServiceImpl.validateUrn(testUrnTooLong); + EntityUtils.validateUrn(_testEntityRegistry, testUrnTooLong); Assert.fail("Should have raised IllegalArgumentException for URN too long"); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "Error: cannot provide an URN longer than 512 bytes (when URL encoded)"); @@ -1235,9 +1372,9 @@ public void testValidateUrn() throws Exception { Urn testUrnTooLongWhenEncoded = new Urn("li", "corpUser", new TupleKey(buildStringTooLongWhenEncoded.toString())); Urn testUrnSameLengthWhenEncoded = new Urn("li", "corpUser", new TupleKey(buildStringSameLengthWhenEncoded.toString())); // Same length when encoded should be allowed, the encoded one should not be - _entityServiceImpl.validateUrn(testUrnSameLengthWhenEncoded); + EntityUtils.validateUrn(_testEntityRegistry, testUrnSameLengthWhenEncoded); try { - _entityServiceImpl.validateUrn(testUrnTooLongWhenEncoded); + EntityUtils.validateUrn(_testEntityRegistry, testUrnTooLongWhenEncoded); Assert.fail("Should have raised IllegalArgumentException for URN too long"); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "Error: cannot provide an URN longer than 512 bytes (when URL encoded)"); @@ -1246,9 +1383,9 @@ public void testValidateUrn() throws Exception { // Urn containing disallowed character Urn testUrnSpecialCharValid = new Urn("li", "corpUser", new TupleKey("bob␇")); Urn testUrnSpecialCharInvalid = new Urn("li", "corpUser", new TupleKey("bob␟")); - _entityServiceImpl.validateUrn(testUrnSpecialCharValid); + EntityUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharValid); try { - _entityServiceImpl.validateUrn(testUrnSpecialCharInvalid); + EntityUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharInvalid); Assert.fail("Should have raised IllegalArgumentException for URN containing the illegal char"); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "Error: URN cannot contain ␟ character"); @@ -1256,7 +1393,7 @@ public void testValidateUrn() throws Exception { Urn urnWithMismatchedParens = new Urn("li", "corpuser", new TupleKey("test(Key")); try { - _entityServiceImpl.validateUrn(urnWithMismatchedParens); + EntityUtils.validateUrn(_testEntityRegistry, urnWithMismatchedParens); Assert.fail("Should have raised IllegalArgumentException for URN with mismatched parens"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("mismatched paren nesting")); @@ -1264,18 +1401,18 @@ public void testValidateUrn() throws Exception { Urn invalidType = new Urn("li", "fakeMadeUpType", new TupleKey("testKey")); try { - _entityServiceImpl.validateUrn(invalidType); + EntityUtils.validateUrn(_testEntityRegistry, invalidType); Assert.fail("Should have raised IllegalArgumentException for URN with non-existent entity type"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("Failed to find entity with name fakeMadeUpType")); } Urn validFabricType = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "PROD")); - _entityServiceImpl.validateUrn(validFabricType); + EntityUtils.validateUrn(_testEntityRegistry, validFabricType); Urn invalidFabricType = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "prod")); try { - _entityServiceImpl.validateUrn(invalidFabricType); + EntityUtils.validateUrn(_testEntityRegistry, invalidFabricType); Assert.fail("Should have raised IllegalArgumentException for URN with invalid fabric type"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains(invalidFabricType.toString())); @@ -1283,7 +1420,7 @@ public void testValidateUrn() throws Exception { Urn urnEndingInComma = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "PROD", "")); try { - _entityServiceImpl.validateUrn(urnEndingInComma); + EntityUtils.validateUrn(_testEntityRegistry, urnEndingInComma); Assert.fail("Should have raised IllegalArgumentException for URN ending in comma"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains(urnEndingInComma.toString())); diff --git a/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java index b431f786cd50a..2703dd7fe6cbe 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/timeline/EbeanTimelineServiceTest.java @@ -6,7 +6,7 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectDao; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.registry.EntityRegistryException; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -27,7 +27,7 @@ public EbeanTimelineServiceTest() throws EntityRegistryException { @BeforeMethod public void setupTest() { - EbeanServer server = EbeanTestUtils.createTestServer(); + Database server = EbeanTestUtils.createTestServer(); _aspectDao = new EbeanAspectDao(server); _aspectDao.setConnectionValidated(true); _entityTimelineService = new TimelineServiceImpl(_aspectDao, _testEntityRegistry); diff --git a/metadata-io/src/test/java/io/datahub/test/DataGenerator.java b/metadata-io/src/test/java/io/datahub/test/DataGenerator.java new file mode 100644 index 0000000000000..3b374993cde16 --- /dev/null +++ b/metadata-io/src/test/java/io/datahub/test/DataGenerator.java @@ -0,0 +1,359 @@ +package io.datahub.test; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.GlossaryTermAssociation; +import com.linkedin.common.GlossaryTermAssociationArray; +import com.linkedin.common.TagAssociation; +import com.linkedin.common.TagAssociationArray; +import com.linkedin.common.urn.GlossaryTermUrn; +import com.linkedin.common.urn.TagUrn; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.glossary.GlossaryTermInfo; +import com.linkedin.metadata.Constants; +import com.linkedin.metadata.entity.AspectUtils; +import com.linkedin.metadata.entity.EntityService; +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 net.datafaker.Faker; +import com.linkedin.mxe.MetadataChangeProposal; +import net.datafaker.providers.base.Animal; +import net.datafaker.providers.base.Cat; +import org.apache.commons.lang3.NotImplementedException; + +import javax.annotation.Nonnull; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.net.URISyntaxException; +import java.util.Arrays; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.function.BiFunction; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.LongStream; +import java.util.stream.Stream; + +public class DataGenerator { + private final static Faker FAKER = new Faker(); + private final EntityRegistry entityRegistry; + private final EntityService entityService; + + public DataGenerator(EntityService entityService) { + this.entityService = entityService; + this.entityRegistry = entityService.getEntityRegistry(); + } + + public Stream> generateDatasets() { + return generateMCPs("dataset", 10, List.of()); + } + + public Stream> generateMCPs(String entityName, long count, List aspects) { + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + + // Prevent duplicate tags and terms generated as secondary entities + Set secondaryUrns = new HashSet<>(); + + return LongStream.range(0, count).mapToObj(idx -> { + RecordTemplate key = randomKeyAspect(entitySpec); + MetadataChangeProposal mcp = new MetadataChangeProposal(); + mcp.setEntityType(entitySpec.getName()); + mcp.setAspectName(entitySpec.getKeyAspectName()); + mcp.setAspect(GenericRecordUtils.serializeAspect(key)); + mcp.setEntityUrn(EntityKeyUtils.convertEntityKeyToUrn(key, entityName)); + mcp.setChangeType(ChangeType.UPSERT); + return mcp; + }).flatMap(mcp -> { + // Expand with additional random aspects + List additionalMCPs = new LinkedList<>(); + + for (String aspectName : aspects) { + AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); + if (aspectSpec == null) { + throw new IllegalStateException("Aspect " + aspectName + " not found for entity " + entityName); + } + + RecordTemplate aspect = randomAspectGenerators.getOrDefault(aspectName, + DataGenerator::defaultRandomAspect).apply(entitySpec, aspectSpec); + + // Maybe generate nested entities at the same time, like globalTags/glossaryTerms + List secondaryEntities = nestedRandomAspectGenerators.getOrDefault(aspectSpec.getName(), + (a, c) -> List.of()).apply(aspect, 5).stream() + .filter(secondaryMCP -> { + if (!secondaryUrns.contains(secondaryMCP.getEntityUrn())) { + secondaryUrns.add(secondaryMCP.getEntityUrn()); + return true; + } + return false; + }) + .collect(Collectors.toList()); + additionalMCPs.addAll(secondaryEntities); + + MetadataChangeProposal additionalMCP = new MetadataChangeProposal(); + additionalMCP.setEntityType(entitySpec.getName()); + additionalMCP.setAspectName(aspectName); + additionalMCP.setAspect(GenericRecordUtils.serializeAspect(aspect)); + additionalMCP.setEntityUrn(mcp.getEntityUrn()); + additionalMCP.setChangeType(ChangeType.UPSERT); + + additionalMCPs.add(additionalMCP); + } + + return Stream.concat(Stream.of(mcp), additionalMCPs.stream()); + }).map(mcp -> { + // Expand with default aspects per normal + return Stream.concat(Stream.of(mcp), + AspectUtils.getAdditionalChanges(mcp, entityService, true).stream()).collect(Collectors.toList()); + }); + } + + public static Map> randomAspectGenerators = Map.of( + "glossaryTermInfo", (e, a) -> { + GlossaryTermInfo glossaryTermInfo = (GlossaryTermInfo) defaultRandomAspect(e, a); + glossaryTermInfo.setName(normalize(FAKER.company().buzzword())); + return glossaryTermInfo; + } + ); + + public Map>> nestedRandomAspectGenerators = Map.of( + "globalTags", (aspect, count) -> { + try { + List tags = generateMCPs("tag", count, List.of()) + .map(mcps -> mcps.get(0)) + .collect(Collectors.toList()); + Method setTagsMethod = aspect.getClass().getMethod("setTags", TagAssociationArray.class); + TagAssociationArray tagAssociations = new TagAssociationArray(); + tagAssociations.addAll(tags.stream().map( + tagMCP -> { + try { + return new TagAssociation().setTag(TagUrn.createFromUrn(tagMCP.getEntityUrn())); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + ).collect(Collectors.toList())); + setTagsMethod.invoke(aspect, tagAssociations); + return tags; + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + "glossaryTerms", (aspect, count) -> { + try { + List terms = generateMCPs("glossaryTerm", count, + List.of("glossaryTermInfo")) + .map(mcps -> mcps.get(0)) + .collect(Collectors.toList()); + Method setTermsMethod = aspect.getClass().getMethod("setTerms", GlossaryTermAssociationArray.class); + GlossaryTermAssociationArray termAssociations = new GlossaryTermAssociationArray(); + termAssociations.addAll(terms.stream().map( + termMCP -> { + try { + return new GlossaryTermAssociation() + .setUrn(GlossaryTermUrn.createFromUrn(termMCP.getEntityUrn())); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + ).collect(Collectors.toList())); + setTermsMethod.invoke(aspect, termAssociations); + return terms; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + ); + + private static RecordTemplate defaultRandomAspect(@Nonnull EntitySpec entitySpec, @Nonnull AspectSpec aspectSpec) { + Class aspectClass = aspectSpec.getDataTemplateClass(); + try { + Object aspect = aspectClass.getDeclaredConstructor().newInstance(); + + List booleanMethods = Arrays.stream(aspectClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0] == Boolean.class) + .collect(Collectors.toList()); + + for (Method boolMethod : booleanMethods) { + boolMethod.invoke(aspect, FAKER.random().nextBoolean()); + } + + List stringMethods = Arrays.stream(aspectClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0] == String.class) + .collect(Collectors.toList()); + + for (Method stringMethod : stringMethods) { + String value; + switch (aspectSpec.getName() + "_" + stringMethod.getName()) { + default: + value = FAKER.lorem().characters(8, 16, false); + break; + } + + // global + if (stringMethod.getName().toLowerCase().contains("description") + || stringMethod.getName().toLowerCase().contains("definition")) { + value = FAKER.lorem().paragraph(); + } + + stringMethod.invoke(aspect, value); + } + + List enumMethods = Arrays.stream(aspectClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0].isEnum()) + .collect(Collectors.toList()); + + for (Method enumMethod : enumMethods) { + Object[] enumClass = enumMethod.getParameterTypes()[0].getEnumConstants(); + // Excluding $UNKNOWNs + enumMethod.invoke(aspect, enumClass[FAKER.random().nextInt(0, enumClass.length - 2)]); + } + + // auditStamp + Arrays.stream(aspectClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0] == AuditStamp.class) + .findFirst().ifPresent(auditStampMethod -> { + try { + AuditStamp auditStamp = new AuditStamp() + .setActor(Urn.createFromString(Constants.DATAHUB_ACTOR)) + .setTime(System.currentTimeMillis()); + auditStampMethod.invoke(aspect, auditStamp); + } catch (URISyntaxException | IllegalAccessException | InvocationTargetException e) { + throw new RuntimeException(e); + } + }); + + return aspectClass.cast(aspect); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static RecordTemplate randomKeyAspect(EntitySpec entitySpec) { + Class keyClass = entitySpec.getKeyAspectSpec().getDataTemplateClass(); + try { + Object key = keyClass.getDeclaredConstructor().newInstance(); + + List stringMethods = Arrays.stream(keyClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0] == String.class) + .collect(Collectors.toList()); + + switch (entitySpec.getName()) { + case "tag": + stringMethods.get(0).invoke(key, normalize(FAKER.marketing().buzzwords())); + break; + case "glossaryTerm": + stringMethods.get(0).invoke(key, normalize(UUID.randomUUID().toString())); + break; + case "container": + stringMethods.get(0).invoke(key, FAKER.examplify("b5e95fce839e7d78151ed7e0a7420d84")); + break; + default: + switch (stringMethods.size()) { + case 1: + stringMethods.get(0).invoke(key, String.join(".", multiName(3))); + break; + case 2: + Cat cat = FAKER.cat(); + stringMethods.get(0).invoke(key, cat.breed().toLowerCase()); + stringMethods.get(1).invoke(key, cat.name().toLowerCase()); + break; + default: + Animal animal = FAKER.animal(); + stringMethods.get(0).invoke(key, animal.genus().toLowerCase()); + stringMethods.get(1).invoke(key, animal.species().toLowerCase()); + stringMethods.get(2).invoke(key, animal.name().toLowerCase()); + break; + } + break; + } + + List urnMethods = Arrays.stream(keyClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0] == Urn.class) + .collect(Collectors.toList()); + + for (Method urnMethod : urnMethods) { + switch (entitySpec.getName()) { + case "dataset": + urnMethod.invoke(key, randomUrnLowerCase("dataPlatform", + List.of(randomDataPlatform()))); + break; + default: + throw new NotImplementedException(entitySpec.getName()); + } + } + + List enumMethods = Arrays.stream(keyClass.getMethods()) + .filter(m -> m.getName().startsWith("set") + && m.getParameterCount() == 1 + && m.getParameterTypes()[0].isEnum()) + .collect(Collectors.toList()); + + for (Method enumMethod : enumMethods) { + Object[] enumClass = enumMethod.getParameterTypes()[0].getEnumConstants(); + // Excluding $UNKNOWNs + enumMethod.invoke(key, enumClass[FAKER.random().nextInt(0, enumClass.length - 2)]); + } + + return keyClass.cast(key); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static List multiName(int size) { + switch (size) { + case 1: + return Stream.of(FAKER.marketing().buzzwords()) + .map(String::toLowerCase).collect(Collectors.toList()); + case 2: + Cat cat = FAKER.cat(); + return Stream.of(cat.breed(), cat.name()) + .map(String::toLowerCase).collect(Collectors.toList()); + case 3: + Animal animal = FAKER.animal(); + return Stream.of(animal.genus(), animal.species(), animal.name()) + .map(String::toLowerCase).collect(Collectors.toList()); + default: + return IntStream.range(0, size).mapToObj(i -> FAKER.expression("#{numerify 'test####'}")).collect(Collectors.toList()); + } + } + + private static Urn randomUrnLowerCase(String entityType, List tuple) { + return Urn.createFromTuple(entityType, + tuple.stream().map(DataGenerator::normalize).collect(Collectors.toList())); + } + + private static String normalize(String input) { + return input.toLowerCase().replaceAll("\\W+", "_"); + } + + private static String randomDataPlatform() { + String[] platforms = { + "ambry", "bigquery", "couchbase", "druid", "external", "feast", "glue", "hdfs", "hive", "kafka", "kusto", + "looker", "mongodb", "mssql", "mysql", "oracle", "pinot", "postgres", "presto", "redshift", "s3", + "sagemaker", "snowflake", "teradata", "voldemort" + }; + + return platforms[FAKER.random().nextInt(0, platforms.length - 1)]; + } +} diff --git a/metadata-jobs/mae-consumer-job/build.gradle b/metadata-jobs/mae-consumer-job/build.gradle index ca099eea5a8a3..51c758f434328 100644 --- a/metadata-jobs/mae-consumer-job/build.gradle +++ b/metadata-jobs/mae-consumer-job/build.gradle @@ -58,7 +58,7 @@ docker { load(true) push(false) } -tasks.getByName("docker").dependsOn([bootJar]) +tasks.getByPath(":metadata-jobs:mae-consumer-job:docker").dependsOn([bootJar]) task cleanLocalDockerImages { doLast { diff --git a/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java b/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java index 72665ffa0b76e..3b44ede0f1d43 100644 --- a/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java +++ b/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java @@ -8,7 +8,7 @@ import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.systemmetadata.ElasticSearchSystemMetadataService; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.springframework.boot.test.context.TestConfiguration; import org.springframework.boot.test.mock.mockito.MockBean; import org.springframework.context.annotation.Import; @@ -27,7 +27,7 @@ public class MaeConsumerApplicationTestConfiguration { private RestliEntityClient restliEntityClient; @MockBean - private EbeanServer ebeanServer; + private Database ebeanServer; @MockBean private EntityRegistry entityRegistry; diff --git a/metadata-jobs/mce-consumer-job/build.gradle b/metadata-jobs/mce-consumer-job/build.gradle index b72d4baff23d6..daf41a1e0303e 100644 --- a/metadata-jobs/mce-consumer-job/build.gradle +++ b/metadata-jobs/mce-consumer-job/build.gradle @@ -69,7 +69,7 @@ docker { load(true) push(false) } -tasks.getByName("docker").dependsOn([bootJar]) +tasks.getByPath(":metadata-jobs:mce-consumer-job:docker").dependsOn([bootJar]) task cleanLocalDockerImages { doLast { diff --git a/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java b/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java index 2d09cf2043575..558a7b9d90ccb 100644 --- a/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java +++ b/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java @@ -11,7 +11,7 @@ import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.restli.client.Client; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.TestConfiguration; import org.springframework.boot.test.mock.mockito.MockBean; @@ -44,7 +44,7 @@ public RestliEntityClient restliEntityClient() { } @MockBean - public EbeanServer ebeanServer; + public Database ebeanServer; @MockBean protected TimeseriesAspectService timeseriesAspectService; 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 e37a351e0365a..125bba7ec3280 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,6 +12,7 @@ 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.key.DataHubAccessTokenKey; import com.linkedin.metadata.utils.AuditStampUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -19,11 +20,12 @@ import java.util.Base64; import java.util.Date; import java.util.HashMap; -import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; @@ -124,11 +126,12 @@ public String generateAccessToken(@Nonnull final TokenType type, @Nonnull final log.info("About to ingest access token metadata {}", proposal); final AuditStamp auditStamp = AuditStampUtils.createDefaultAuditStamp().setActor(UrnUtils.getUrn(actorUrn)); - // Need this to write key aspect - final List additionalChanges = AspectUtils.getAdditionalChanges(proposal, _entityService); + Stream proposalStream = Stream.concat(Stream.of(proposal), + AspectUtils.getAdditionalChanges(proposal, _entityService).stream()); - _entityService.ingestProposal(proposal, auditStamp, false); - additionalChanges.forEach(mcp -> _entityService.ingestProposal(mcp, auditStamp, false)); + _entityService.ingestProposal(AspectsBatchImpl.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(), auditStamp, false); return accessToken; } diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java index 44673f693b555..f5ce938c411c6 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java @@ -106,6 +106,7 @@ public void testAuthenticateSuccess() throws Exception { final AspectSpec keyAspectSpec = configEntityRegistry.getEntitySpec(Constants.ACCESS_TOKEN_ENTITY_NAME).getKeyAspectSpec(); Mockito.when(mockService.getKeyAspectSpec(Mockito.eq(Constants.ACCESS_TOKEN_ENTITY_NAME))).thenReturn(keyAspectSpec); Mockito.when(mockService.exists(Mockito.any(Urn.class))).thenReturn(true); + Mockito.when(mockService.getEntityRegistry()).thenReturn(configEntityRegistry); final DataHubTokenAuthenticator authenticator = new DataHubTokenAuthenticator(); authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY, SALT_CONFIG_NAME, diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java index 75a9114529d7a..1c46e864a559e 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java @@ -13,6 +13,7 @@ import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import java.util.Date; import java.util.Map; + import org.mockito.Mockito; import org.testng.annotations.Test; @@ -157,6 +158,7 @@ public void generateRevokeToken() throws TokenException { DataHubTokenAuthenticatorTest.class.getClassLoader().getResourceAsStream("test-entity-registry.yaml")); final AspectSpec keyAspectSpec = configEntityRegistry.getEntitySpec(Constants.ACCESS_TOKEN_ENTITY_NAME).getKeyAspectSpec(); + Mockito.when(mockService.getEntityRegistry()).thenReturn(configEntityRegistry); Mockito.when(mockService.getKeyAspectSpec(Mockito.eq(Constants.ACCESS_TOKEN_ENTITY_NAME))).thenReturn(keyAspectSpec); Mockito.when(mockService.exists(Mockito.any(Urn.class))).thenReturn(true); final RollbackRunResult result = new RollbackRunResult(ImmutableList.of(), 0); @@ -174,4 +176,8 @@ public void generateRevokeToken() throws TokenException { // Validation should fail. assertThrows(TokenException.class, () -> tokenService.validateAccessToken(token)); } + + private void mockStateful() { + + } } diff --git a/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml b/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml index acdc5ead92fb1..48c1b9450bd7d 100644 --- a/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml +++ b/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml @@ -4,4 +4,14 @@ entities: keyAspect: dataHubAccessTokenKey aspects: - dataHubAccessTokenInfo + - name: corpuser + keyAspect: corpUserKey + aspects: + - corpUserInfo + - corpUserEditableInfo + - corpUserStatus + - groupMembership + - status + - corpUserCredentials + - corpUserSettings events: diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EbeanServerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EbeanServerFactory.java index b7759d906f5b4..9feb7e469d018 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EbeanServerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EbeanServerFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.entity; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.config.ServerConfig; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -25,7 +25,7 @@ public class EbeanServerFactory { @DependsOn({"gmsEbeanServiceConfig"}) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @Nonnull - protected EbeanServer createServer() { + protected Database createServer() { ServerConfig serverConfig = applicationContext.getBean(ServerConfig.class); // Make sure that the serverConfig includes the package that contains DAO's Ebean model. if (!serverConfig.getPackages().contains(EBEAN_MODEL_PACKAGE)) { @@ -33,7 +33,7 @@ protected EbeanServer createServer() { } // TODO: Consider supporting SCSI try { - return io.ebean.EbeanServerFactory.create(serverConfig); + return io.ebean.DatabaseFactory.create(serverConfig); } catch (NullPointerException ne) { log.error("Failed to connect to the server. Is it up?"); throw ne; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectDaoFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectDaoFactory.java index c4f0dae4c5fd5..925689c8609db 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectDaoFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectDaoFactory.java @@ -4,7 +4,7 @@ import com.linkedin.metadata.entity.AspectDao; import com.linkedin.metadata.entity.cassandra.CassandraAspectDao; import com.linkedin.metadata.entity.ebean.EbeanAspectDao; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -19,7 +19,7 @@ public class EntityAspectDaoFactory { @DependsOn({"gmsEbeanServiceConfig"}) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @Nonnull - protected AspectDao createEbeanInstance(EbeanServer server) { + protected AspectDao createEbeanInstance(Database server) { return new EbeanAspectDao(server); } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectMigrationsDaoFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectMigrationsDaoFactory.java index 0e83c1af9c66a..4000f7d6ed058 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectMigrationsDaoFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/EntityAspectMigrationsDaoFactory.java @@ -4,7 +4,7 @@ import com.linkedin.metadata.entity.AspectMigrationsDao; import com.linkedin.metadata.entity.cassandra.CassandraAspectDao; import com.linkedin.metadata.entity.ebean.EbeanAspectDao; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -19,7 +19,7 @@ public class EntityAspectMigrationsDaoFactory { @DependsOn({"gmsEbeanServiceConfig"}) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @Nonnull - protected AspectMigrationsDao createEbeanInstance(EbeanServer server) { + protected AspectMigrationsDao createEbeanInstance(Database server) { return new EbeanAspectDao(server); } 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 737773d0972e2..b13bf5813d47e 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 @@ -6,7 +6,7 @@ import com.linkedin.metadata.entity.RetentionService; import com.linkedin.metadata.entity.cassandra.CassandraRetentionService; import com.linkedin.metadata.entity.ebean.EbeanRetentionService; -import io.ebean.EbeanServer; +import io.ebean.Database; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; @@ -46,7 +46,7 @@ protected RetentionService createCassandraInstance(CqlSession session) { @DependsOn({"ebeanServer", "entityService"}) @ConditionalOnProperty(name = "entityService.impl", havingValue = "ebean", matchIfMissing = true) @Nonnull - protected RetentionService createEbeanInstance(EbeanServer server) { + protected RetentionService createEbeanInstance(Database server) { RetentionService retentionService = new EbeanRetentionService(_entityService, server, _batchSize); _entityService.setRetentionService(retentionService); return retentionService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IndexDataPlatformsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IndexDataPlatformsStep.java index 43b71d36e0e38..b26eb67465c0d 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IndexDataPlatformsStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IndexDataPlatformsStep.java @@ -13,10 +13,15 @@ import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; import com.linkedin.metadata.search.EntitySearchService; + import java.util.Collections; import java.util.HashSet; +import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -73,6 +78,7 @@ private int getAndReIndexDataPlatforms(AuditStamp auditStamp, AspectSpec dataPla ); // Loop over Data platforms and produce changelog + List> futures = new LinkedList<>(); for (Urn dpUrn : dataPlatformUrns) { EntityResponse dataPlatformEntityResponse = dataPlatformInfoResponses.get(dpUrn); if (dataPlatformEntityResponse == null) { @@ -86,7 +92,7 @@ private int getAndReIndexDataPlatforms(AuditStamp auditStamp, AspectSpec dataPla continue; } - _entityService.produceMetadataChangeLog( + futures.add(_entityService.alwaysProduceMCLAsync( dpUrn, Constants.DATA_PLATFORM_ENTITY_NAME, Constants.DATA_PLATFORM_INFO_ASPECT_NAME, @@ -96,9 +102,17 @@ private int getAndReIndexDataPlatforms(AuditStamp auditStamp, AspectSpec dataPla null, null, auditStamp, - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); } + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + return listResult.getTotal(); } 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 d5165713ab193..30608e984a0f2 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 @@ -8,6 +8,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.models.AspectSpec; import com.linkedin.metadata.utils.DataPlatformInstanceUtils; import com.linkedin.metadata.utils.EntityKeyUtils; @@ -15,6 +17,8 @@ import lombok.extern.slf4j.Slf4j; import javax.annotation.Nonnull; +import java.util.LinkedList; +import java.util.List; import java.util.Optional; import static com.linkedin.metadata.Constants.*; @@ -60,19 +64,25 @@ public void execute() throws Exception { while (start < numEntities) { log.info("Reading urns {} to {} from the aspects table to generate dataplatform instance aspects", start, start + BATCH_SIZE); - Iterable urns = _migrationsDao.listAllUrns(start, start + BATCH_SIZE); - for (String urnStr : urns) { + + List items = new LinkedList<>(); + + for (String urnStr : _migrationsDao.listAllUrns(start, start + BATCH_SIZE)) { Urn urn = Urn.createFromString(urnStr); Optional dataPlatformInstance = getDataPlatformInstance(urn); - if (!dataPlatformInstance.isPresent()) { - continue; + if (dataPlatformInstance.isPresent()) { + items.add(UpsertBatchItem.builder() + .urn(urn) + .aspectName(DATA_PLATFORM_INSTANCE_ASPECT_NAME) + .aspect(dataPlatformInstance.get()) + .build(_entityService.getEntityRegistry())); } + } - final AuditStamp aspectAuditStamp = - new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); + 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.ingestAspect(urn, DATA_PLATFORM_INSTANCE_ASPECT_NAME, dataPlatformInstance.get(), aspectAuditStamp, null); - } log.info("Finished ingesting DataPlatformInstance for urn {} to {}", start, start + BATCH_SIZE); 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 11a45c4960b05..e4ad215eec864 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,6 +12,14 @@ import com.linkedin.metadata.entity.EntityService; import java.io.IOException; import java.net.URISyntaxException; +import java.util.List; +import java.util.Spliterator; +import java.util.Spliterators; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; + +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; +import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.core.io.ClassPathResource; @@ -49,32 +57,32 @@ public void execute() throws IOException, URISyntaxException { } // 2. For each JSON object, cast into a DataPlatformSnapshot object. - for (final JsonNode dataPlatform : dataPlatforms) { - final String urnString; - final Urn urn; - try { - urnString = dataPlatform.get("urn").asText(); - urn = Urn.createFromString(urnString); - } catch (URISyntaxException e) { - log.error("Malformed urn: {}", dataPlatform.get("urn").asText()); - throw new RuntimeException("Malformed urn", e); - } - - final DataPlatformInfo existingInfo = - (DataPlatformInfo) _entityService.getLatestAspect(urn, PLATFORM_ASPECT_NAME); - // Skip ingesting for this JSON object if info already exists. - if (existingInfo != null) { - log.debug(String.format("%s already exists for %s. Skipping...", PLATFORM_ASPECT_NAME, urnString)); - continue; - } - - final DataPlatformInfo info = - RecordUtils.toRecordTemplate(DataPlatformInfo.class, dataPlatform.get("aspect").toString()); - - final AuditStamp aspectAuditStamp = - new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); - - _entityService.ingestAspect(urn, PLATFORM_ASPECT_NAME, info, aspectAuditStamp, null); - } + List dataPlatformAspects = StreamSupport.stream( + Spliterators.spliteratorUnknownSize(dataPlatforms.iterator(), Spliterator.ORDERED), false) + .map(dataPlatform -> { + final String urnString; + final Urn urn; + try { + urnString = dataPlatform.get("urn").asText(); + urn = Urn.createFromString(urnString); + } catch (URISyntaxException e) { + log.error("Malformed urn: {}", dataPlatform.get("urn").asText()); + throw new RuntimeException("Malformed urn", e); + } + + final DataPlatformInfo info = + RecordUtils.toRecordTemplate(DataPlatformInfo.class, dataPlatform.get("aspect").toString()); + + return UpsertBatchItem.builder() + .urn(urn) + .aspectName(PLATFORM_ASPECT_NAME) + .aspect(info) + .build(_entityService.getEntityRegistry()); + }).collect(Collectors.toList()); + + _entityService.ingestAspects(AspectsBatchImpl.builder().items(dataPlatformAspects).build(), + new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), + 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 08a867d710419..55d612618ff9f 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 @@ -9,6 +9,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.boot.UpgradeStep; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.utils.EntityKeyUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -19,6 +20,8 @@ import lombok.extern.slf4j.Slf4j; import org.springframework.core.io.ClassPathResource; +import java.util.List; + import static com.linkedin.metadata.Constants.*; @@ -96,8 +99,6 @@ private void ingestOwnershipType(final Urn ownershipTypeUrn, final OwnershipType keyAspectProposal.setChangeType(ChangeType.UPSERT); keyAspectProposal.setEntityUrn(ownershipTypeUrn); - _entityService.ingestProposal(keyAspectProposal, auditStamp, false); - final MetadataChangeProposal proposal = new MetadataChangeProposal(); proposal.setEntityUrn(ownershipTypeUrn); proposal.setEntityType(OWNERSHIP_TYPE_ENTITY_NAME); @@ -107,7 +108,9 @@ private void ingestOwnershipType(final Urn ownershipTypeUrn, final OwnershipType proposal.setAspect(GenericRecordUtils.serializeAspect(info)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(proposal, auditStamp, false); + _entityService.ingestProposal(AspectsBatchImpl.builder() + .mcps(List.of(keyAspectProposal, proposal), _entityService.getEntityRegistry()).build(), auditStamp, + false); } @Nonnull 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 1025cacb3685c..87dcfd736da40 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 @@ -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.transactions.AspectsBatchImpl; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; @@ -28,8 +29,10 @@ import java.net.URISyntaxException; import java.util.Collections; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Optional; + import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.core.io.ClassPathResource; @@ -172,9 +175,6 @@ private void ingestPolicy(final Urn urn, final DataHubPolicyInfo info) throws UR keyAspectProposal.setChangeType(ChangeType.UPSERT); keyAspectProposal.setEntityUrn(urn); - _entityService.ingestProposal(keyAspectProposal, - new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); - final MetadataChangeProposal proposal = new MetadataChangeProposal(); proposal.setEntityUrn(urn); proposal.setEntityType(POLICY_ENTITY_NAME); @@ -182,8 +182,11 @@ private void ingestPolicy(final Urn urn, final DataHubPolicyInfo info) throws UR proposal.setAspect(GenericRecordUtils.serializeAspect(info)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(proposal, - new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); + _entityService.ingestProposal(AspectsBatchImpl.builder() + .mcps(List.of(keyAspectProposal, proposal), _entityRegistry) + .build(), + new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), + false); } private boolean hasPolicy(Urn policyUrn) { 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 b9f43cbf898a7..99be185113968 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 @@ -10,6 +10,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.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.utils.EntityKeyUtils; @@ -18,6 +19,7 @@ import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.policy.DataHubRoleInfo; import java.net.URISyntaxException; +import java.util.List; import javax.annotation.Nonnull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -99,9 +101,6 @@ private void ingestRole(final Urn roleUrn, final DataHubRoleInfo dataHubRoleInfo keyAspectProposal.setChangeType(ChangeType.UPSERT); keyAspectProposal.setEntityUrn(roleUrn); - _entityService.ingestProposal(keyAspectProposal, - new AuditStamp().setActor(Urn.createFromString(SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); - final MetadataChangeProposal proposal = new MetadataChangeProposal(); proposal.setEntityUrn(roleUrn); proposal.setEntityType(DATAHUB_ROLE_ENTITY_NAME); @@ -109,10 +108,12 @@ private void ingestRole(final Urn roleUrn, final DataHubRoleInfo dataHubRoleInfo proposal.setAspect(GenericRecordUtils.serializeAspect(dataHubRoleInfo)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(proposal, - new AuditStamp().setActor(Urn.createFromString(SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); + _entityService.ingestProposal(AspectsBatchImpl.builder() + .mcps(List.of(keyAspectProposal, proposal), _entityRegistry).build(), + new AuditStamp().setActor(Urn.createFromString(SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), + false); - _entityService.produceMetadataChangeLog(roleUrn, DATAHUB_ROLE_ENTITY_NAME, DATAHUB_ROLE_INFO_ASPECT_NAME, + _entityService.alwaysProduceMCLAsync(roleUrn, DATAHUB_ROLE_ENTITY_NAME, DATAHUB_ROLE_INFO_ASPECT_NAME, roleInfoAspectSpec, null, dataHubRoleInfo, null, null, auditStamp, ChangeType.RESTATE); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRootUserStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRootUserStep.java index b322afb809d2b..febcb9d4ec8a4 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRootUserStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestRootUserStep.java @@ -16,7 +16,9 @@ import com.linkedin.metadata.utils.EntityKeyUtils; import java.io.IOException; import java.net.URISyntaxException; +import java.util.List; +import com.linkedin.util.Pair; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.core.io.ClassPathResource; @@ -67,8 +69,11 @@ public void execute() throws IOException, URISyntaxException { final CorpUserKey key = (CorpUserKey) EntityKeyUtils.convertUrnToEntityKey(urn, getUserKeyAspectSpec()); final AuditStamp aspectAuditStamp = new AuditStamp().setActor(Urn.createFromString(SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); - _entityService.ingestAspect(urn, CORP_USER_KEY_ASPECT_NAME, key, aspectAuditStamp, null); - _entityService.ingestAspect(urn, USER_INFO_ASPECT_NAME, info, aspectAuditStamp, null); + + _entityService.ingestAspects(urn, List.of( + Pair.of(CORP_USER_KEY_ASPECT_NAME, key), + Pair.of(USER_INFO_ASPECT_NAME, info) + ), aspectAuditStamp, null); } private AspectSpec getUserKeyAspectSpec() { 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 6e1522051bfab..1f5f7f26ed89b 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 @@ -16,7 +16,11 @@ import lombok.extern.slf4j.Slf4j; import javax.annotation.Nonnull; +import java.util.LinkedList; +import java.util.List; import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; @Slf4j public class RestoreColumnLineageIndices extends UpgradeStep { @@ -89,6 +93,7 @@ private int getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStamp return latestAspects.getTotalCount(); } + List> futures = new LinkedList<>(); for (int i = 0; i < latestAspects.getValues().size(); i++) { ExtraInfo info = latestAspects.getMetadata().getExtraInfos().get(i); RecordTemplate upstreamLineageRecord = latestAspects.getValues().get(i); @@ -99,7 +104,7 @@ private int getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStamp continue; } - _entityService.produceMetadataChangeLog( + futures.add(_entityService.alwaysProduceMCLAsync( urn, Constants.DATASET_ENTITY_NAME, Constants.UPSTREAM_LINEAGE_ASPECT_NAME, @@ -109,9 +114,17 @@ private int getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStamp null, null, auditStamp, - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); } + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + return latestAspects.getTotalCount(); } @@ -140,6 +153,7 @@ private int getAndRestoreInputFieldsIndices(String entityName, int start, AuditS return latestAspects.getTotalCount(); } + List> futures = new LinkedList<>(); for (int i = 0; i < latestAspects.getValues().size(); i++) { ExtraInfo info = latestAspects.getMetadata().getExtraInfos().get(i); RecordTemplate inputFieldsRecord = latestAspects.getValues().get(i); @@ -150,7 +164,7 @@ private int getAndRestoreInputFieldsIndices(String entityName, int start, AuditS continue; } - _entityService.produceMetadataChangeLog( + futures.add(_entityService.alwaysProduceMCLAsync( urn, entityName, Constants.INPUT_FIELDS_ASPECT_NAME, @@ -160,9 +174,17 @@ private int getAndRestoreInputFieldsIndices(String entityName, int start, AuditS null, null, auditStamp, - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); } + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + return latestAspects.getTotalCount(); } } 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 989ee1a39b169..355936fe1994c 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 @@ -23,8 +23,12 @@ import java.net.URISyntaxException; import java.util.Collections; import java.util.HashSet; +import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; import javax.annotation.Nonnull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -126,6 +130,7 @@ private void getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStam } // Loop over datasets and produce changelog + List> futures = new LinkedList<>(); for (Urn datasetUrn : datasetUrns) { EntityResponse response = upstreamLineageResponse.get(datasetUrn); if (response == null) { @@ -137,7 +142,7 @@ private void getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStam continue; } - _entityService.produceMetadataChangeLog( + futures.add(_entityService.alwaysProduceMCLAsync( datasetUrn, DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, @@ -147,8 +152,16 @@ private void getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStam null, null, auditStamp, - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); } + + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); } private UpstreamLineage getUpstreamLineage(EntityResponse entityResponse) { diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java index 097dcfdfdf52e..4de2bea9a76a9 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndices.java @@ -16,10 +16,15 @@ import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchEntity; import com.linkedin.metadata.search.SearchResult; + import java.util.Collections; import java.util.HashSet; +import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; import java.util.stream.Collectors; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -87,6 +92,7 @@ null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) ); // Loop over Terms and produce changelog + List> futures = new LinkedList<>(); for (Urn termUrn : termUrns) { EntityResponse termEntityResponse = termInfoResponses.get(termUrn); if (termEntityResponse == null) { @@ -99,7 +105,7 @@ null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) continue; } - _entityService.produceMetadataChangeLog( + futures.add(_entityService.alwaysProduceMCLAsync( termUrn, Constants.GLOSSARY_TERM_ENTITY_NAME, Constants.GLOSSARY_TERM_INFO_ASPECT_NAME, @@ -109,9 +115,17 @@ null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) null, null, auditStamp, - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); } + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + return termsResult.getNumEntities(); } @@ -130,6 +144,7 @@ null, null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) ); // Loop over Nodes and produce changelog + List> futures = new LinkedList<>(); for (Urn nodeUrn : nodeUrns) { EntityResponse nodeEntityResponse = nodeInfoResponses.get(nodeUrn); if (nodeEntityResponse == null) { @@ -142,7 +157,7 @@ null, null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) continue; } - _entityService.produceMetadataChangeLog( + futures.add(_entityService.alwaysProduceMCLAsync( nodeUrn, Constants.GLOSSARY_NODE_ENTITY_NAME, Constants.GLOSSARY_NODE_INFO_ASPECT_NAME, @@ -152,9 +167,17 @@ null, null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) null, null, auditStamp, - ChangeType.RESTATE); + ChangeType.RESTATE).getFirst()); } + futures.stream().filter(Objects::nonNull).forEach(f -> { + try { + f.get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + }); + return nodesResult.getNumEntities(); } 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 82f3a82c135ce..0ae8eb2cba808 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 @@ -5,6 +5,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.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; @@ -95,15 +96,24 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw final IngestDataPlatformInstancesStep step = new IngestDataPlatformInstancesStep(entityService, migrationsDao); step.execute(); - verify(entityService, times(countOfChartEntities)) - .ingestAspect( - argThat(arg -> arg.getEntityType().equals("chart")), - eq(DATA_PLATFORM_INSTANCE_ASPECT_NAME), - any(DataPlatformInstance.class), + verify(entityService, times(1)) + .ingestAspects( + argThat(arg -> + arg.getItems().stream() + .allMatch(item -> item.getUrn().getEntityType().equals("chart") + && item.getAspectName().equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) + && ((UpsertBatchItem) item).getAspect() instanceof DataPlatformInstance) + ), any(), - any()); + anyBoolean(), + anyBoolean()); verify(entityService, times(0)) - .ingestAspect(argThat(arg -> !arg.getEntityType().equals("chart")), anyString(), any(), any(), any()); + .ingestAspects(argThat(arg -> + !arg.getItems().stream() + .allMatch(item -> item.getUrn().getEntityType().equals("chart") + && item.getAspectName().equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) + && ((UpsertBatchItem) item).getAspect() instanceof DataPlatformInstance) + ), any(), anyBoolean(), anyBoolean()); } @NotNull @@ -135,6 +145,7 @@ private void mockDBWithWorkToDo( when(migrationsDao.checkIfAspectExists(DATA_PLATFORM_INSTANCE_ASPECT_NAME)).thenReturn(false); when(migrationsDao.countEntities()).thenReturn((long) allUrnsInDB.size()); when(migrationsDao.listAllUrns(anyInt(), anyInt())).thenReturn(allUrnsInDB); + when(entityService.getEntityRegistry()).thenReturn(entityRegistry); } private List insertMockEntities(int count, String entity, String urnTemplate, EntityRegistry entityRegistry, EntityService entityService) { diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndicesTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndicesTest.java index b73e749142863..aca5e322567d8 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndicesTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndicesTest.java @@ -21,6 +21,7 @@ import com.linkedin.metadata.query.ExtraInfoArray; import com.linkedin.metadata.query.ListResultMetadata; import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.util.Pair; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -29,6 +30,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.Future; public class RestoreColumnLineageIndicesTest { @@ -64,7 +66,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.any(AuditStamp.class), Mockito.eq(false) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(datasetUrn), Mockito.eq(Constants.DATASET_ENTITY_NAME), Mockito.eq(Constants.UPSTREAM_LINEAGE_ASPECT_NAME), @@ -76,7 +78,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(chartUrn), Mockito.eq(Constants.CHART_ENTITY_NAME), Mockito.eq(Constants.INPUT_FIELDS_ASPECT_NAME), @@ -88,7 +90,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(dashboardUrn), Mockito.eq(Constants.DASHBOARD_ENTITY_NAME), Mockito.eq(Constants.INPUT_FIELDS_ASPECT_NAME), @@ -126,7 +128,7 @@ public void testExecuteWithNewVersion() throws Exception { Mockito.any(AuditStamp.class), Mockito.eq(false) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(datasetUrn), Mockito.eq(Constants.DATASET_ENTITY_NAME), Mockito.eq(Constants.UPSTREAM_LINEAGE_ASPECT_NAME), @@ -138,7 +140,7 @@ public void testExecuteWithNewVersion() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(chartUrn), Mockito.eq(Constants.CHART_ENTITY_NAME), Mockito.eq(Constants.INPUT_FIELDS_ASPECT_NAME), @@ -150,7 +152,7 @@ public void testExecuteWithNewVersion() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(dashboardUrn), Mockito.eq(Constants.DASHBOARD_ENTITY_NAME), Mockito.eq(Constants.INPUT_FIELDS_ASPECT_NAME), @@ -188,7 +190,7 @@ public void testDoesNotExecuteWithSameVersion() throws Exception { Mockito.any(AuditStamp.class), Mockito.eq(false) ); - Mockito.verify(mockService, Mockito.times(0)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(0)).alwaysProduceMCLAsync( Mockito.eq(datasetUrn), Mockito.eq(Constants.DATASET_ENTITY_NAME), Mockito.eq(Constants.UPSTREAM_LINEAGE_ASPECT_NAME), @@ -200,7 +202,7 @@ public void testDoesNotExecuteWithSameVersion() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(0)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(0)).alwaysProduceMCLAsync( Mockito.eq(chartUrn), Mockito.eq(Constants.CHART_ENTITY_NAME), Mockito.eq(Constants.INPUT_FIELDS_ASPECT_NAME), @@ -212,7 +214,7 @@ public void testDoesNotExecuteWithSameVersion() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(0)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(0)).alwaysProduceMCLAsync( Mockito.eq(dashboardUrn), Mockito.eq(Constants.DASHBOARD_ENTITY_NAME), Mockito.eq(Constants.INPUT_FIELDS_ASPECT_NAME), @@ -234,6 +236,12 @@ private void mockGetUpstreamLineage(@Nonnull Urn datasetUrn, @Nonnull EntityServ .setAudit(new AuditStamp().setActor(UrnUtils.getUrn("urn:li:corpuser:test")).setTime(0L)) ); + Mockito.when(mockService.alwaysProduceMCLAsync( + Mockito.any(Urn.class), Mockito.anyString(), Mockito.anyString(), Mockito.any(AspectSpec.class), + Mockito.eq(null), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.any(ChangeType.class) + )).thenReturn(Pair.of(Mockito.mock(Future.class), false)); + Mockito.when(mockService.listLatestAspects( Mockito.eq(Constants.DATASET_ENTITY_NAME), Mockito.eq(Constants.UPSTREAM_LINEAGE_ASPECT_NAME), diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java index d56fbed07f890..3753904053256 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreGlossaryIndicesTest.java @@ -21,6 +21,7 @@ import com.linkedin.metadata.search.SearchResult; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.util.Pair; import java.util.List; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -29,6 +30,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.Map; +import java.util.concurrent.Future; public class RestoreGlossaryIndicesTest { @@ -94,6 +96,11 @@ public void testExecuteFirstTime() throws Exception { upgradeEntityUrn, Collections.singleton(Constants.DATA_HUB_UPGRADE_REQUEST_ASPECT_NAME) )).thenReturn(null); + Mockito.when(mockService.alwaysProduceMCLAsync( + Mockito.any(Urn.class), Mockito.anyString(), Mockito.anyString(), Mockito.any(AspectSpec.class), + Mockito.eq(null), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.any(ChangeType.class) + )).thenReturn(Pair.of(Mockito.mock(Future.class), false)); mockGetTermInfo(glossaryTermUrn, mockSearchService, mockService); mockGetNodeInfo(glossaryNodeUrn, mockSearchService, mockService); @@ -111,7 +118,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.any(AuditStamp.class), Mockito.eq(false) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(glossaryTermUrn), Mockito.eq(Constants.GLOSSARY_TERM_ENTITY_NAME), Mockito.eq(Constants.GLOSSARY_TERM_INFO_ASPECT_NAME), @@ -123,7 +130,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(glossaryNodeUrn), Mockito.eq(Constants.GLOSSARY_NODE_ENTITY_NAME), Mockito.eq(Constants.GLOSSARY_NODE_INFO_ASPECT_NAME), @@ -155,6 +162,11 @@ public void testExecutesWithNewVersion() throws Exception { upgradeEntityUrn, Collections.singleton(Constants.DATA_HUB_UPGRADE_REQUEST_ASPECT_NAME) )).thenReturn(response); + Mockito.when(mockService.alwaysProduceMCLAsync( + Mockito.any(Urn.class), Mockito.anyString(), Mockito.anyString(), Mockito.any(AspectSpec.class), + Mockito.eq(null), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.any(ChangeType.class) + )).thenReturn(Pair.of(Mockito.mock(Future.class), false)); mockGetTermInfo(glossaryTermUrn, mockSearchService, mockService); mockGetNodeInfo(glossaryNodeUrn, mockSearchService, mockService); @@ -164,7 +176,6 @@ public void testExecutesWithNewVersion() throws Exception { RestoreGlossaryIndices restoreIndicesStep = new RestoreGlossaryIndices(mockService, mockSearchService, mockRegistry); restoreIndicesStep.execute(); - Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.GLOSSARY_TERM_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.GLOSSARY_NODE_ENTITY_NAME); Mockito.verify(mockService, Mockito.times(2)).ingestProposal( @@ -172,7 +183,7 @@ public void testExecutesWithNewVersion() throws Exception { Mockito.any(AuditStamp.class), Mockito.eq(false) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(glossaryTermUrn), Mockito.eq(Constants.GLOSSARY_TERM_ENTITY_NAME), Mockito.eq(Constants.GLOSSARY_TERM_INFO_ASPECT_NAME), @@ -184,7 +195,7 @@ public void testExecutesWithNewVersion() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(1)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(1)).alwaysProduceMCLAsync( Mockito.eq(glossaryNodeUrn), Mockito.eq(Constants.GLOSSARY_NODE_ENTITY_NAME), Mockito.eq(Constants.GLOSSARY_NODE_INFO_ASPECT_NAME), @@ -233,7 +244,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { Mockito.any(AuditStamp.class), Mockito.anyBoolean() ); - Mockito.verify(mockService, Mockito.times(0)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(0)).alwaysProduceMCLAsync( Mockito.eq(glossaryTermUrn), Mockito.eq(Constants.GLOSSARY_TERM_ENTITY_NAME), Mockito.eq(Constants.GLOSSARY_TERM_INFO_ASPECT_NAME), @@ -245,7 +256,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { Mockito.any(), Mockito.eq(ChangeType.RESTATE) ); - Mockito.verify(mockService, Mockito.times(0)).produceMetadataChangeLog( + Mockito.verify(mockService, Mockito.times(0)).alwaysProduceMCLAsync( Mockito.eq(glossaryNodeUrn), Mockito.eq(Constants.GLOSSARY_NODE_ENTITY_NAME), Mockito.eq(Constants.GLOSSARY_NODE_INFO_ASPECT_NAME), diff --git a/metadata-service/factories/src/test/resources/test-entity-registry.yaml b/metadata-service/factories/src/test/resources/test-entity-registry.yaml index 45aa9b9554fb4..fe32b413751e6 100644 --- a/metadata-service/factories/src/test/resources/test-entity-registry.yaml +++ b/metadata-service/factories/src/test/resources/test-entity-registry.yaml @@ -8,3 +8,9 @@ entities: keyAspect: chartKey aspects: - domains + - dataPlatformInstance + - name: dataPlatform + category: core + keyAspect: dataPlatformKey + aspects: + - dataPlatformInfo \ No newline at end of file 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 52842d2f32b2a..4d0e5e7df29d5 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 @@ -19,8 +19,10 @@ import com.linkedin.entity.Aspect; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.IngestProposalResult; +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.validation.ValidationException; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.entity.AspectUtils; @@ -49,6 +51,7 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; import org.reflections.Reflections; @@ -266,10 +269,17 @@ public static Pair ingestProposal(com.linkedin.mxe.MetadataChan log.info("Proposal: {}", serviceProposal); Throwable exceptionally = null; try { - IngestProposalResult proposalResult = entityService.ingestProposal(serviceProposal, auditStamp, false); - Urn urn = proposalResult.getUrn(); - additionalChanges.forEach(proposal -> entityService.ingestProposal(proposal, auditStamp, false)); - return new Pair<>(urn.toString(), proposalResult.isDidUpdate()); + Stream proposalStream = Stream.concat(Stream.of(serviceProposal), + AspectUtils.getAdditionalChanges(serviceProposal, entityService).stream()); + + AspectsBatch batch = AspectsBatchImpl.builder().mcps(proposalStream.collect(Collectors.toList()), + entityService.getEntityRegistry()).build(); + + Set proposalResult = + entityService.ingestProposal(batch, auditStamp, false); + + Urn urn = proposalResult.stream().findFirst().get().getUrn(); + return new Pair<>(urn.toString(), proposalResult.stream().anyMatch(IngestResult::isSqlCommitted)); } catch (ValidationException ve) { exceptionally = ve; throw HttpClientErrorException.create(HttpStatus.UNPROCESSABLE_ENTITY, ve.getMessage(), null, null, null); 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 80cc80067fba2..229e71168557d 100644 --- a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java +++ b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java @@ -9,6 +9,7 @@ import com.linkedin.metadata.config.PreProcessHooks; import com.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.metadata.entity.AspectDao; +import com.linkedin.metadata.entity.UpdateAspectResult; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.service.UpdateIndicesService; @@ -35,8 +36,12 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.function.Function; + +import io.ebean.Transaction; import mock.MockEntityRegistry; import mock.MockEntityService; +import org.mockito.ArgumentMatchers; import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -60,6 +65,11 @@ public void setup() throws NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException { EntityRegistry mockEntityRegistry = new MockEntityRegistry(); AspectDao aspectDao = Mockito.mock(AspectDao.class); + Mockito.when(aspectDao.runInTransactionWithRetry( + ArgumentMatchers.>any(), any(), anyInt())).thenAnswer(i -> + ((Function) i.getArgument(0)).apply(Mockito.mock(Transaction.class)) + ); + EventProducer mockEntityEventProducer = Mockito.mock(EventProducer.class); UpdateIndicesService mockUpdateIndicesService = mock(UpdateIndicesService.class); PreProcessHooks preProcessHooks = new PreProcessHooks(); 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 cf7985a7c26a9..852b6cfcb4b22 100644 --- a/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java +++ b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java @@ -24,9 +24,7 @@ import com.linkedin.metadata.entity.EntityServiceImpl; import com.linkedin.metadata.entity.ListResult; import com.linkedin.metadata.entity.RollbackRunResult; -import com.linkedin.metadata.entity.UpdateAspectResult; import com.linkedin.metadata.event.EventProducer; -import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; import com.linkedin.metadata.run.AspectRowSummary; @@ -40,16 +38,13 @@ import com.linkedin.schema.SchemaFieldDataType; import com.linkedin.schema.SchemaMetadata; import com.linkedin.schema.StringType; -import com.linkedin.util.Pair; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; -import java.util.function.Function; import javax.annotation.Nonnull; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -136,12 +131,6 @@ public EnvelopedAspect getLatestEnvelopedAspect(@Nonnull String entityName, @Non return null; } - @Override - public EnvelopedAspect getEnvelopedAspect(@Nonnull String entityName, @Nonnull Urn urn, @Nonnull String aspectName, - long version) throws Exception { - return null; - } - @Override public VersionedAspect getVersionedAspect(@Nonnull Urn urn, @Nonnull String aspectName, long version) { return null; @@ -153,7 +142,7 @@ public ListResult listLatestAspects(@Nonnull String entityName, return null; } - @Nonnull +/* @Nonnull @Override protected UpdateAspectResult ingestAspectToLocalDB(@Nonnull Urn urn, @Nonnull String aspectName, @Nonnull Function, RecordTemplate> updateLambda, @Nonnull AuditStamp auditStamp, @@ -168,7 +157,7 @@ protected List> ingestAspectsToLocalDB(@Nonnull @Nonnull List> aspectRecordsToIngest, @Nonnull AuditStamp auditStamp, @Nonnull SystemMetadata providedSystemMetadata) { return Collections.emptyList(); - } + }*/ @Nullable @Override @@ -177,13 +166,6 @@ public RecordTemplate ingestAspectIfNotPresent(@NotNull Urn urn, @NotNull String return null; } - @Override - public RecordTemplate updateAspect(@Nonnull Urn urn, @Nonnull String entityName, @Nonnull String aspectName, - @Nonnull AspectSpec aspectSpec, @Nonnull RecordTemplate newValue, @Nonnull AuditStamp auditStamp, - @Nonnull long version, @Nonnull boolean emitMae) { - return null; - } - @Override public ListUrnsResult listUrns(@Nonnull String entityName, int start, int count) { return null; 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 5c06dbfc8b25a..936c8bb67e645 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 @@ -8,6 +8,9 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.linkedin.aspect.GetTimeseriesAspectValuesResponse; +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.resources.operations.Utils; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; @@ -16,7 +19,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.IngestProposalResult; import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.query.filter.Filter; @@ -41,6 +43,10 @@ import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import javax.inject.Inject; @@ -179,7 +185,7 @@ public Task ingestProposal( @ActionParam(PARAM_ASYNC) @Optional(UNSET) String async) throws URISyntaxException { log.info("INGEST PROPOSAL proposal: {}", metadataChangeProposal); - boolean asyncBool; + final boolean asyncBool; if (UNSET.equals(async)) { asyncBool = Boolean.parseBoolean(System.getenv(ASYNC_INGEST_DEFAULT_NAME)); } else { @@ -200,18 +206,34 @@ public Task ingestProposal( return RestliUtil.toTask(() -> { log.debug("Proposal: {}", metadataChangeProposal); try { - IngestProposalResult result = _entityService.ingestProposal(metadataChangeProposal, auditStamp, asyncBool); - Urn responseUrn = result.getUrn(); + 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()); - if (!asyncBool) { - AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService) - .forEach(proposal -> _entityService.ingestProposal(proposal, auditStamp, asyncBool)); + batch = AspectsBatchImpl.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(); } - if (!result.isQueued()) { - tryIndexRunId(responseUrn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); + Set results = + _entityService.ingestProposal(batch, auditStamp, asyncBool); + + IngestResult one = results.stream() + .findFirst() + .get(); + + // Update runIds, only works for existing documents, so ES document must exist + Urn resultUrn = one.getUrn(); + if (one.isProcessedMCL() || one.isUpdate()) { + tryIndexRunId(resultUrn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); } - return responseUrn.toString(); + return resultUrn.toString(); } catch (ValidationException e) { throw new RestLiServiceException(HttpStatus.S_422_UNPROCESSABLE_ENTITY, e.getMessage()); } 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 c6d36a6e29f10..351a3d8f24e36 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 @@ -5,6 +5,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authentication.AuthenticationContext; import com.datahub.plugins.auth.authorization.Authorizer; +import com.linkedin.common.AuditStamp; import com.linkedin.common.FabricType; import com.linkedin.common.urn.DataPlatformUrn; import com.linkedin.common.urn.DatasetUrn; @@ -14,8 +15,9 @@ import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.AspectDao; 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.EntityServiceImpl; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -24,6 +26,8 @@ import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import java.net.URISyntaxException; +import java.util.List; + import mock.MockEntityRegistry; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -78,8 +82,34 @@ public void testAsyncDefaultAspects() throws URISyntaxException { reset(_producer, _aspectDao); - when(_aspectDao.runInTransactionWithRetry(any(), anyInt())) - .thenReturn(new UpdateAspectResult(urn, null, properties, null, null, null, null, 0)); + UpsertBatchItem req = UpsertBatchItem.builder() + .urn(urn) + .aspectName(mcp.getAspectName()) + .aspect(mcp.getAspect()) + .metadataChangeProposal(mcp) + .build(_entityRegistry); + when(_aspectDao.runInTransactionWithRetry(any(), any(), anyInt())) + .thenReturn(List.of( + UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name1")) + .auditStamp(new AuditStamp()) + .request(req).build(), + UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name2")) + .auditStamp(new AuditStamp()) + .request(req).build(), + UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name3")) + .auditStamp(new AuditStamp()) + .request(req).build(), + UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name4")) + .auditStamp(new AuditStamp()) + .request(req).build(), + UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name5")) + .auditStamp(new AuditStamp()) + .request(req).build())); _aspectResource.ingestProposal(mcp, "false"); verify(_producer, times(5)).produceMetadataChangeLog(eq(urn), any(AspectSpec.class), any(MetadataChangeLog.class)); verifyNoMoreInteractions(_producer); 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 e062d55254f90..40a5e3a07ae6d 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 @@ -1,6 +1,7 @@ package com.linkedin.metadata.entity; 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.urn.Urn; @@ -31,22 +32,52 @@ public class AspectUtils { private AspectUtils() { } + public static final Set SUPPORTED_TYPES = Set.of(ChangeType.UPSERT, ChangeType.CREATE, ChangeType.PATCH); + public static List getAdditionalChanges( - @Nonnull MetadataChangeProposal metadataChangeProposal, - @Nonnull EntityService entityService) { - // No additional changes for delete operation - if (metadataChangeProposal.getChangeType() == ChangeType.DELETE) { + @Nonnull MetadataChangeProposal metadataChangeProposal, + @Nonnull EntityService entityService, + boolean onPrimaryKeyInsertOnly) { + + // No additional changes for unsupported operations + if (!SUPPORTED_TYPES.contains(metadataChangeProposal.getChangeType())) { return Collections.emptyList(); } final Urn urn = EntityKeyUtils.getUrnFromProposal(metadataChangeProposal, - entityService.getKeyAspectSpec(metadataChangeProposal.getEntityType())); + entityService.getKeyAspectSpec(metadataChangeProposal.getEntityType())); + + final Map includedAspects; + if (metadataChangeProposal.getChangeType() != ChangeType.PATCH) { + RecordTemplate aspectRecord = GenericRecordUtils.deserializeAspect(metadataChangeProposal.getAspect().getValue(), + metadataChangeProposal.getAspect().getContentType(), entityService.getEntityRegistry() + .getEntitySpec(urn.getEntityType()).getAspectSpec(metadataChangeProposal.getAspectName())); + includedAspects = ImmutableMap.of(metadataChangeProposal.getAspectName(), aspectRecord); + } else { + includedAspects = ImmutableMap.of(); + } + + if (onPrimaryKeyInsertOnly) { + return entityService.generateDefaultAspectsOnFirstWrite(urn, includedAspects) + .getValue() + .stream() + .map(entry -> getProposalFromAspect(entry.getKey(), entry.getValue(), metadataChangeProposal)) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } else { + return entityService.generateDefaultAspectsIfMissing(urn, includedAspects) + .stream() + .map(entry -> getProposalFromAspect(entry.getKey(), entry.getValue(), metadataChangeProposal)) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } + } + + public static List getAdditionalChanges( + @Nonnull MetadataChangeProposal metadataChangeProposal, + @Nonnull EntityService entityService) { - return entityService.generateDefaultAspectsIfMissing(urn, ImmutableSet.of(metadataChangeProposal.getAspectName())) - .stream() - .map(entry -> getProposalFromAspect(entry.getKey(), entry.getValue(), metadataChangeProposal)) - .filter(Objects::nonNull) - .collect(Collectors.toList()); + return getAdditionalChanges(metadataChangeProposal, entityService, false); } public static Map batchGetLatestAspect( 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 35b7bc4589b32..40284efe7ac82 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 @@ -270,9 +270,9 @@ private void updateAspect(Urn urn, String aspectName, RecordTemplate prevAspect, proposal.setAspect(GenericRecordUtils.serializeAspect(newAspect)); final AuditStamp auditStamp = new AuditStamp().setActor(UrnUtils.getUrn(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); - final IngestProposalResult ingestProposalResult = _entityService.ingestProposal(proposal, auditStamp, false); + final IngestResult ingestProposalResult = _entityService.ingestProposal(proposal, auditStamp, false); - if (!ingestProposalResult.isDidUpdate()) { + if (!ingestProposalResult.isSqlCommitted()) { log.error("Failed to ingest aspect with references removed. Before {}, after: {}, please check MCP processor" + " logs for more information", prevAspect, newAspect); handleError(new DeleteEntityServiceError("Failed to ingest new aspect", 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 25edff740037e..86043f4b7cd27 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 @@ -13,12 +13,11 @@ import com.linkedin.metadata.aspect.VersionedAspect; 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; import com.linkedin.metadata.run.AspectRowSummary; -import com.linkedin.metadata.snapshot.Snapshot; -import com.linkedin.mxe.MetadataAuditOperation; import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; @@ -28,6 +27,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.Future; import java.util.function.Consumer; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -140,22 +140,6 @@ EnvelopedAspect getLatestEnvelopedAspect( @Nonnull final Urn urn, @Nonnull final String aspectName) throws Exception; - /** - * Retrieves the specific version of the aspect for the given urn - * - * @param entityName name of the entity to fetch - * @param urn urn to fetch - * @param aspectName name of the aspect to fetch - * @param version version to fetch - * @return {@link EnvelopedAspect} object, or null if one cannot be found - */ - EnvelopedAspect getEnvelopedAspect( - // TODO: entityName is only used for a debug statement, can we remove this as a param? - String entityName, - @Nonnull Urn urn, - @Nonnull String aspectName, - long version) throws Exception; - @Deprecated VersionedAspect getVersionedAspect(@Nonnull Urn urn, @Nonnull String aspectName, long version); @@ -165,24 +149,11 @@ ListResult listLatestAspects( final int start, final int count); - void ingestAspects(@Nonnull final Urn urn, @Nonnull List> aspectRecordsToIngest, + List ingestAspects(@Nonnull final Urn urn, @Nonnull List> aspectRecordsToIngest, @Nonnull final AuditStamp auditStamp, @Nullable SystemMetadata systemMetadata); - /** - * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. - * - * Note that in general, this should not be used externally. It is currently serving upgrade scripts and - * is as such public. - * - * @param urn an urn associated with the new aspect - * @param aspectName name of the aspect being inserted - * @param newValue value of the aspect being inserted - * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time - * @param systemMetadata - * @return the {@link RecordTemplate} representation of the written aspect object - */ - RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName, - @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp, @Nullable SystemMetadata systemMetadata); + List ingestAspects(@Nonnull final AspectsBatch aspectsBatch, @Nonnull final AuditStamp auditStamp, + boolean emitMCL, boolean overwrite); /** * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. @@ -211,17 +182,6 @@ String batchApplyRetention(Integer start, Integer count, Integer attemptWithVers // TODO: Extract this to a different service, doesn't need to be here RestoreIndicesResult restoreIndices(@Nonnull RestoreIndicesArgs args, @Nonnull Consumer logger); - @Deprecated - RecordTemplate updateAspect( - @Nonnull final Urn urn, - @Nonnull final String entityName, - @Nonnull final String aspectName, - @Nonnull final AspectSpec aspectSpec, - @Nonnull final RecordTemplate newValue, - @Nonnull final AuditStamp auditStamp, - @Nonnull final long version, - @Nonnull final boolean emitMae); - ListUrnsResult listUrns(@Nonnull final String entityName, final int start, final int count); @Deprecated @@ -230,23 +190,14 @@ RecordTemplate updateAspect( @Deprecated Map getEntities(@Nonnull final Set urns, @Nonnull Set aspectNames); - @Deprecated - void produceMetadataAuditEvent(@Nonnull final Urn urn, @Nonnull final String aspectName, - @Nullable final RecordTemplate oldAspectValue, @Nullable final RecordTemplate newAspectValue, - @Nullable final SystemMetadata oldSystemMetadata, @Nullable final SystemMetadata newSystemMetadata, - @Nullable final MetadataAuditOperation operation); - - @Deprecated - void produceMetadataAuditEventForKey(@Nonnull final Urn urn, - @Nullable final SystemMetadata newSystemMetadata); - - void produceMetadataChangeLog(@Nonnull final Urn urn, AspectSpec aspectSpec, + Pair, Boolean> alwaysProduceMCLAsync(@Nonnull final Urn urn, AspectSpec aspectSpec, @Nonnull final MetadataChangeLog metadataChangeLog); - void produceMetadataChangeLog(@Nonnull final Urn urn, @Nonnull String entityName, @Nonnull String aspectName, - @Nonnull final AspectSpec aspectSpec, @Nullable final RecordTemplate oldAspectValue, - @Nullable final RecordTemplate newAspectValue, @Nullable final SystemMetadata oldSystemMetadata, - @Nullable final SystemMetadata newSystemMetadata, @Nonnull AuditStamp auditStamp, @Nonnull final ChangeType changeType); + Pair, Boolean> alwaysProduceMCLAsync(@Nonnull final Urn urn, @Nonnull String entityName, @Nonnull String aspectName, + @Nonnull final AspectSpec aspectSpec, @Nullable final RecordTemplate oldAspectValue, + @Nullable final RecordTemplate newAspectValue, @Nullable final SystemMetadata oldSystemMetadata, + @Nullable final SystemMetadata newSystemMetadata, @Nonnull AuditStamp auditStamp, + @Nonnull final ChangeType changeType); RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName); @@ -261,9 +212,6 @@ void ingestEntities(@Nonnull final List entities, @Nonnull final AuditSt void ingestEntity(@Nonnull Entity entity, @Nonnull AuditStamp auditStamp, @Nonnull SystemMetadata systemMetadata); - @Deprecated - Snapshot buildSnapshot(@Nonnull final Urn urn, @Nonnull final RecordTemplate aspectValue); - void setRetentionService(RetentionService retentionService); AspectSpec getKeyAspectSpec(@Nonnull final Urn urn); @@ -272,8 +220,33 @@ void ingestEntity(@Nonnull Entity entity, @Nonnull AuditStamp auditStamp, String getKeyAspectName(@Nonnull final Urn urn); + /** + * Generate default aspects if not present in the database. + * @param urn entity urn + * @param includedAspects aspects being written + * @return additional aspects to be written + */ List> generateDefaultAspectsIfMissing(@Nonnull final Urn urn, - Set includedAspects); + Map includedAspects); + + /** + * Generate default aspects if the entity key aspect is NOT in the database **AND** + * the key aspect is being written, present in `includedAspects`. + * + * Does not automatically create key aspects. + * @see EntityService#generateDefaultAspectsIfMissing if key aspects need autogeneration + * + * This version is more efficient in that it only generates additional writes + * when a new entity is being minted for the first time. The drawback is that it will not automatically + * add key aspects, in case the producer is not bothering to ensure that the entity exists + * before writing non-key aspects. + * + * @param urn entity urn + * @param includedAspects aspects being written + * @return whether key aspect exists in database and the additional aspects to be written + */ + Pair>> generateDefaultAspectsOnFirstWrite(@Nonnull final Urn urn, + Map includedAspects); AspectSpec getKeyAspectSpec(@Nonnull final String entityName); @@ -289,8 +262,16 @@ List> generateDefaultAspectsIfMissing(@Nonnull fina RollbackRunResult rollbackWithConditions(List aspectRows, Map conditions, boolean hardDelete); - IngestProposalResult ingestProposal(@Nonnull MetadataChangeProposal mcp, - AuditStamp auditStamp, final boolean async); + Set ingestProposal(AspectsBatch aspectsBatch, AuditStamp auditStamp, final boolean async); + + /** + * If you have more than 1 proposal use the {AspectsBatch} method + * @param proposal the metadata proposal to ingest + * @param auditStamp audit information + * @param async async ingestion or sync ingestion + * @return ingestion result + */ + IngestResult ingestProposal(MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async); Boolean exists(Urn urn); 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 new file mode 100644 index 0000000000000..5e4ed6259a7f7 --- /dev/null +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java @@ -0,0 +1,18 @@ +package com.linkedin.metadata.entity; + +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.entity.transactions.AbstractBatchItem; +import lombok.Builder; +import lombok.Value; + +@Builder(toBuilder = true) +@Value +public class IngestResult { + Urn urn; + AbstractBatchItem request; + boolean publishedMCL; + boolean processedMCL; + boolean publishedMCP; + boolean sqlCommitted; + boolean isUpdate; // update else insert +} 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 a27cb8076721a..1cdd9965c4bfc 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 @@ -9,6 +9,7 @@ import com.linkedin.metadata.Constants; 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; @@ -21,10 +22,11 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.CompletableFuture; import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; + +import lombok.Builder; import lombok.SneakyThrows; import lombok.Value; @@ -65,7 +67,7 @@ public Retention getRetention(@Nonnull String entityName, @Nonnull String aspect } // Get list of datahub retention keys that match the input entity name and aspect name - protected List getRetentionKeys(@Nonnull String entityName, @Nonnull String aspectName) { + protected static List getRetentionKeys(@Nonnull String entityName, @Nonnull String aspectName) { return ImmutableList.of( new DataHubRetentionKey().setEntityName(entityName).setAspectName(aspectName), new DataHubRetentionKey().setEntityName(entityName).setAspectName(ALL), @@ -87,12 +89,13 @@ protected List getRetentionKeys(@Nonnull String entityName, @Nonnull String */ @SneakyThrows public boolean setRetention(@Nullable String entityName, @Nullable String aspectName, - @Nonnull DataHubRetentionConfig retentionConfig) { + @Nonnull DataHubRetentionConfig retentionConfig) { validateRetention(retentionConfig.getRetention()); DataHubRetentionKey retentionKey = new DataHubRetentionKey(); retentionKey.setEntityName(entityName != null ? entityName : ALL); retentionKey.setAspectName(aspectName != null ? aspectName : ALL); Urn retentionUrn = EntityKeyUtils.convertEntityKeyToUrn(retentionKey, Constants.DATAHUB_RETENTION_ENTITY); + MetadataChangeProposal keyProposal = new MetadataChangeProposal(); GenericAspect keyAspect = GenericRecordUtils.serializeAspect(retentionKey); keyProposal.setAspect(keyAspect); @@ -100,17 +103,22 @@ public boolean setRetention(@Nullable String entityName, @Nullable String aspect keyProposal.setEntityType(Constants.DATAHUB_RETENTION_ENTITY); keyProposal.setChangeType(ChangeType.UPSERT); keyProposal.setEntityUrn(retentionUrn); - AuditStamp auditStamp = - new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); - getEntityService().ingestProposal(keyProposal, auditStamp, false); + MetadataChangeProposal aspectProposal = keyProposal.clone(); GenericAspect retentionAspect = GenericRecordUtils.serializeAspect(retentionConfig); aspectProposal.setAspect(retentionAspect); aspectProposal.setAspectName(Constants.DATAHUB_RETENTION_ASPECT); - aspectProposal.setChangeType(ChangeType.UPSERT); - return getEntityService().ingestProposal(aspectProposal, auditStamp, false).isDidUpdate(); + + AuditStamp auditStamp = + new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); + AspectsBatch batch = buildAspectsBatch(List.of(keyProposal, aspectProposal)); + + return getEntityService().ingestProposal(batch, auditStamp, false).stream() + .anyMatch(IngestResult::isSqlCommitted); } + protected abstract AspectsBatch buildAspectsBatch(List mcps); + /** * Delete the retention policy set for given entity and aspect. * @@ -140,41 +148,40 @@ private void validateRetention(Retention retention) { } } - /** - * Apply retention policies given the urn and aspect name asynchronously - * - * @param urn Urn of the entity - * @param aspectName Name of the aspect - * @param context Additional context that could be used to apply retention - */ - public void applyRetentionAsync(@Nonnull Urn urn, @Nonnull String aspectName, Optional context) { - CompletableFuture.runAsync(() -> applyRetention(urn, aspectName, context)); - } - /** * Apply retention policies given the urn and aspect name * - * @param urn Urn of the entity - * @param aspectName Name of the aspect - * @param context Additional context that could be used to apply retention + * @param retentionContexts urn, aspect name, and additional context that could be used to apply retention */ - public void applyRetention(@Nonnull Urn urn, @Nonnull String aspectName, Optional context) { - Retention retentionPolicy = getRetention(urn.getEntityType(), aspectName); - if (retentionPolicy.data().isEmpty()) { - return; - } - applyRetention(urn, aspectName, retentionPolicy, context); + public void applyRetentionWithPolicyDefaults(@Nonnull List retentionContexts) { + List withDefaults = retentionContexts.stream() + .map(context -> { + if (context.getRetentionPolicy().isEmpty()) { + Retention retentionPolicy = getRetention(context.getUrn().getEntityType(), context.getAspectName()); + return context.toBuilder() + .retentionPolicy(Optional.of(retentionPolicy)) + .build(); + } else { + return context; + } + }) + .filter(context -> context.getRetentionPolicy().isPresent() + && !context.getRetentionPolicy().get().data().isEmpty()) + .collect(Collectors.toList()); + + applyRetention(withDefaults); } /** - * Apply retention policies given the urn and aspect name and policies - * @param urn Urn of the entity - * @param aspectName Name of the aspect - * @param retentionPolicy Retention policies to apply - * @param retentionContext Additional context that could be used to apply retention + * Apply retention policies given the urn and aspect name and policies. This protected + * method assumes that the policy is provided, however we likely need to fetch these + * from system configuration. + * + * Users of this should use {@link #applyRetentionWithPolicyDefaults(List)}) + * + * @param retentionContexts Additional context that could be used to apply retention */ - public abstract void applyRetention(@Nonnull Urn urn, @Nonnull String aspectName, Retention retentionPolicy, - Optional retentionContext); + protected abstract void applyRetention(List retentionContexts); /** * Batch apply retention to all records that match the input entityName and aspectName @@ -189,9 +196,16 @@ public abstract void applyRetention(@Nonnull Urn urn, @Nonnull String aspectName */ public abstract BulkApplyRetentionResult batchApplyRetentionEntities(@Nonnull BulkApplyRetentionArgs args); - @Value + @Builder(toBuilder = true) public static class RetentionContext { - Optional maxVersion; + @Nonnull + Urn urn; + @Nonnull + String aspectName; + @Builder.Default + Optional retentionPolicy = Optional.empty(); + @Builder.Default + Optional maxVersion = Optional.empty(); } } 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 68ecdbd87dd16..06199814d30dd 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,14 +3,20 @@ 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.mxe.MetadataAuditOperation; import com.linkedin.mxe.SystemMetadata; +import lombok.Builder; import lombok.Value; +import java.util.concurrent.Future; + +@Builder(toBuilder = true) @Value public class UpdateAspectResult { Urn urn; + AbstractBatchItem request; RecordTemplate oldValue; RecordTemplate newValue; SystemMetadata oldSystemMetadata; @@ -18,4 +24,6 @@ public class UpdateAspectResult { MetadataAuditOperation operation; AuditStamp auditStamp; long maxVersion; + boolean processedMCL; + Future mclFuture; } 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 new file mode 100644 index 0000000000000..03a2b4e2a7f73 --- /dev/null +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java @@ -0,0 +1,92 @@ +package com.linkedin.metadata.entity.transactions; + +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; + +import static com.linkedin.metadata.Constants.*; + + +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 new file mode 100644 index 0000000000000..1d3da08130071 --- /dev/null +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java @@ -0,0 +1,22 @@ +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-service/war/build.gradle b/metadata-service/war/build.gradle index ae207e0260e60..3bd2695c927a7 100644 --- a/metadata-service/war/build.gradle +++ b/metadata-service/war/build.gradle @@ -82,7 +82,7 @@ docker { load(true) push(false) } -tasks.getByName("docker").dependsOn([build, war]) +tasks.getByPath(":metadata-service:war:docker").dependsOn([build, war]) task cleanLocalDockerImages { doLast { diff --git a/mock-entity-registry/src/main/java/mock/MockEntitySpec.java b/mock-entity-registry/src/main/java/mock/MockEntitySpec.java index f43c1f7fd6613..d740fff29e258 100644 --- a/mock-entity-registry/src/main/java/mock/MockEntitySpec.java +++ b/mock-entity-registry/src/main/java/mock/MockEntitySpec.java @@ -57,7 +57,7 @@ public EntityAnnotation getEntityAnnotation() { @Override public String getKeyAspectName() { - return null; + return _name + "Key"; } @Override diff --git a/settings.gradle b/settings.gradle index f326cbf10b96b..270672e929e88 100644 --- a/settings.gradle +++ b/settings.gradle @@ -34,7 +34,6 @@ include ':docker:postgres-setup' include ':docker:kafka-setup' include ':docker:datahub-ingestion' include ':docker:datahub-ingestion-base' -include ':docker:datahub-ingestion-slim' include 'metadata-models' include 'metadata-models-validator' include 'metadata-testing:metadata-models-test-utils' diff --git a/smoke-test/run-quickstart.sh b/smoke-test/run-quickstart.sh index d40e4a5e7a4aa..050b5d2db95c9 100755 --- a/smoke-test/run-quickstart.sh +++ b/smoke-test/run-quickstart.sh @@ -15,4 +15,4 @@ echo "test_user:test_pass" >> ~/.datahub/plugins/frontend/auth/user.props echo "DATAHUB_VERSION = $DATAHUB_VERSION" DATAHUB_TELEMETRY_ENABLED=false \ DOCKER_COMPOSE_BASE="file://$( dirname "$DIR" )" \ -datahub docker quickstart --version ${DATAHUB_VERSION} --standalone_consumers --dump-logs-on-failure --kafka-setup \ No newline at end of file +datahub docker quickstart --version ${DATAHUB_VERSION} --standalone_consumers --dump-logs-on-failure --kafka-setup diff --git a/smoke-test/tests/tags-and-terms/data.json b/smoke-test/tests/tags-and-terms/data.json index d018061796296..349400f099339 100644 --- a/smoke-test/tests/tags-and-terms/data.json +++ b/smoke-test/tests/tags-and-terms/data.json @@ -180,5 +180,44 @@ } }, "proposedDelta": null + }, + { + "auditHeader": null, + "proposedSnapshot": { + "com.linkedin.pegasus2avro.metadata.snapshot.TagSnapshot": { + "urn": "urn:li:tag:Legacy", + "aspects": [ + { + "com.linkedin.pegasus2avro.tag.TagProperties": { + "name": "Legacy", + "description": "Indicates the dataset is no longer supported" + } + } + ] + } + }, + "proposedDelta": null + }, + { + "auditHeader": null, + "proposedSnapshot": { + "com.linkedin.pegasus2avro.metadata.snapshot.GlossaryTermSnapshot": { + "urn": "urn:li:glossaryTerm:SavingAccount", + "aspects": [ + { + "com.linkedin.pegasus2avro.glossary.GlossaryTermInfo": { + "definition": "a product provided to consumers and businesses by a bank or similar depository institution such as a checking account, savings account, certificate of deposit, debit or pre-paid card, or credit card", + "sourceRef": "FIBO", + "termSource": "EXTERNAL", + "sourceUrl": "https://spec.edmcouncil.org/fibo/ontology/FBC/FunctionalEntities/FinancialServicesEntities/BankingProduct", + "customProperties": { + "FQDN": "SavingAccount" + } + } + } + ] + } + }, + "proposedDelta": null } ] \ No newline at end of file From d30898637d8f803ca5165459c67f56d4b5070278 Mon Sep 17 00:00:00 2001 From: Ne Enot Date: Sun, 3 Sep 2023 04:37:26 +0300 Subject: [PATCH 2/2] fix(frontend): Fix"Logout with OIDC not working" (#8773) --- .../app/controllers/CentralLogoutController.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datahub-frontend/app/controllers/CentralLogoutController.java b/datahub-frontend/app/controllers/CentralLogoutController.java index b1220c7f3f111..5e24fe9f8220c 100644 --- a/datahub-frontend/app/controllers/CentralLogoutController.java +++ b/datahub-frontend/app/controllers/CentralLogoutController.java @@ -16,7 +16,8 @@ */ @Slf4j public class CentralLogoutController extends LogoutController { - private static final String DEFAULT_BASE_URL_PATH = "/login"; + private static final String AUTH_URL_CONFIG_PATH = "/login"; + private static final String DEFAULT_BASE_URL_PATH = "/"; private static Boolean _isOidcEnabled = false; @Inject @@ -36,8 +37,7 @@ public CentralLogoutController(Config config) { public Result executeLogout(Http.Request request) { if (_isOidcEnabled) { try { - return Results.redirect(DEFAULT_BASE_URL_PATH) - .removingFromSession(request); + return logout(request).toCompletableFuture().get().withNewSession(); } catch (Exception e) { log.error("Caught exception while attempting to perform SSO logout! It's likely that SSO integration is mis-configured.", e); return redirect( @@ -47,7 +47,7 @@ public Result executeLogout(Http.Request request) { .withNewSession(); } } - return Results.redirect(DEFAULT_BASE_URL_PATH) + return Results.redirect(AUTH_URL_CONFIG_PATH) .withNewSession(); } }