From 0e37907b3e3912321317a03221b771ba6db3c4a9 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Mon, 17 Jul 2023 18:33:14 -0500 Subject: [PATCH 01/27] refactor(rentention): refactor retention --- .../metadata/entity/EntityService.java | 23 +++- .../metadata/entity/RetentionService.java | 71 +++++----- .../cassandra/CassandraRetentionService.java | 35 +++-- .../entity/ebean/EbeanRetentionService.java | 130 +++++++++++------- 4 files changed, 156 insertions(+), 103 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java index e5d549b95754d..c1bda73dd7f75 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -812,8 +812,12 @@ protected RecordTemplate sendEventForUpdateAspectResult(@Nonnull final Urn urn, // 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.maxVersion)))); + _retentionService.applyRetentionWithPolicyDefaults(List.of( + RetentionService.RetentionContext.builder() + .urn(urn) + .aspectName(aspectName) + .maxVersion(Optional.of(result.maxVersion)) + .build())); } // Produce MCL after a successful update @@ -1019,8 +1023,12 @@ private UpdateAspectResult upsertAspect(final RecordTemplate aspect, final Syste 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.maxVersion)))); + _retentionService.applyRetentionWithPolicyDefaults(List.of( + RetentionService.RetentionContext.builder() + .urn(entityUrn) + .aspectName(aspectSpec.getName()) + .maxVersion(Optional.of(result.maxVersion)) + .build())); } return result; } @@ -1034,8 +1042,11 @@ private UpdateAspectResult patchAspect(final Patch patch, final SystemMetadata s 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.maxVersion)))); + _retentionService.applyRetentionWithPolicyDefaults(List.of(RetentionService.RetentionContext.builder() + .urn(entityUrn) + .aspectName(aspectSpec.getName()) + .maxVersion(Optional.of(result.maxVersion)) + .build())); } return result; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java index e4cdb5f531b93..d02512a9a338e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java @@ -21,10 +21,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 +66,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,7 +88,7 @@ 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); @@ -140,41 +141,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(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 +189,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-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..1da8e3375035f 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 @@ -55,20 +55,24 @@ public EntityService getEntityService() { @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 +107,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/EbeanRetentionService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanRetentionService.java index 1c0729b7c27e4..3c62ded6ce5bc 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 @@ -22,7 +22,6 @@ 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; @@ -50,33 +49,50 @@ public EntityService getEntityService() { @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(); + } } } @@ -121,33 +137,43 @@ private void applyRetention( try (Transaction transaction = _server.beginTransaction()) { 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(); } } From 70c5b5e5241e4250012dba4f6f5cbd0f31831384 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 18 Jul 2023 19:19:16 -0500 Subject: [PATCH 02/27] feat(sql): ebean transaction batches --- .../resolvers/embed/UpdateEmbedResolver.java | 2 +- .../resolvers/mutate/MutationUtils.java | 2 +- .../mutate/UpdateUserSettingResolver.java | 2 +- .../resolvers/mutate/util/DeleteUtils.java | 7 +- .../mutate/util/DeprecationUtils.java | 7 +- .../resolvers/mutate/util/DomainUtils.java | 7 +- .../resolvers/mutate/util/LabelUtils.java | 7 +- .../resolvers/mutate/util/OwnerUtils.java | 7 +- .../linkedin/datahub/graphql/TestUtils.java | 45 +- .../DeleteAssertionResolverTest.java | 10 +- .../BatchUpdateSoftDeletedResolverTest.java | 22 +- .../BatchUpdateDeprecationResolverTest.java | 23 +- .../UpdateDeprecationResolverTest.java | 2 +- .../domain/BatchSetDomainResolverTest.java | 33 +- .../domain/CreateDomainResolverTest.java | 6 +- .../domain/SetDomainResolverTest.java | 10 +- .../domain/UnsetDomainResolverTest.java | 8 +- .../embed/UpdateEmbedResolverTest.java | 22 +- .../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 | 12 +- .../owner/BatchAddOwnersResolverTest.java | 16 +- .../owner/BatchRemoveOwnersResolverTest.java | 14 +- .../resolvers/tag/AddTagsResolverTest.java | 12 +- .../tag/BatchAddTagsResolverTest.java | 24 +- .../tag/BatchRemoveTagsResolverTest.java | 27 +- .../resolvers/tag/CreateTagResolverTest.java | 6 +- .../tag/SetTagColorResolverTest.java | 8 +- .../resolvers/term/AddTermsResolverTest.java | 17 +- .../term/BatchAddTermsResolverTest.java | 16 +- .../term/BatchRemoveTermsResolverTest.java | 14 +- .../test/resources/test-entity-registry.yaml | 295 +++++++++ .../upgrade/nocode/DataMigrationStep.java | 8 +- .../metadata/client/JavaEntityClient.java | 15 +- .../linkedin/metadata/entity/AspectDao.java | 22 +- .../linkedin/metadata/entity/AspectUtils.java | 37 ++ .../metadata/entity/DeleteEntityService.java | 2 +- .../metadata/entity/EntityService.java | 621 +++++++----------- .../linkedin/metadata/entity/EntityUtils.java | 40 ++ .../metadata/entity/RetentionService.java | 19 +- .../entity/cassandra/CassandraAspectDao.java | 47 +- .../cassandra/CassandraRetentionService.java | 1 + .../metadata/entity/ebean/EbeanAspectDao.java | 118 ++-- .../ebean/transactions/AspectsBatch.java | 150 +++++ .../ebean/transactions/AspectsBatchItem.java | 123 ++++ .../metadata/AspectIngestionUtils.java | 31 +- .../entity/EbeanEntityServiceTest.java | 54 +- .../metadata/entity/EntityServiceTest.java | 271 ++++++-- .../token/StatefulTokenService.java | 12 +- .../token/StatefulTokenServiceTest.java | 6 + .../linkedin/metadata/boot/BootstrapStep.java | 2 +- .../linkedin/metadata/boot/UpgradeStep.java | 4 +- .../boot/steps/BackfillBrowsePathsV2Step.java | 2 +- .../IngestDataPlatformInstancesStep.java | 26 +- .../boot/steps/IngestDataPlatformsStep.java | 62 +- .../IngestDefaultGlobalSettingsStep.java | 2 +- .../boot/steps/IngestOwnershipTypesStep.java | 9 +- .../boot/steps/IngestPoliciesStep.java | 17 +- .../metadata/boot/steps/IngestRolesStep.java | 11 +- .../boot/steps/IngestRootUserStep.java | 9 +- .../boot/steps/RestoreDbtSiblingsIndices.java | 2 +- .../steps/UpgradeDefaultBrowsePathsStep.java | 2 +- .../steps/BackfillBrowsePathsV2StepTest.java | 4 +- .../IngestDataPlatformInstancesStepTest.java | 23 +- .../IngestDefaultGlobalSettingsStepTest.java | 4 +- .../RestoreColumnLineageIndicesTest.java | 6 +- .../steps/RestoreGlossaryIndicesTest.java | 6 +- .../UpgradeDefaultBrowsePathsStepTest.java | 8 +- .../openapi/util/MappingUtil.java | 18 +- .../src/test/java/mock/MockEntityService.java | 10 +- .../resources/entity/AspectResource.java | 32 +- .../entity/BatchIngestionRunResource.java | 2 +- 78 files changed, 1728 insertions(+), 864 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/AspectsBatch.java create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java index 86b8eb5564152..8fc931310570a 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java @@ -58,7 +58,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw updateEmbed(embed, input); final MetadataChangeProposal proposal = buildMetadataChangeProposalWithUrn(entityUrn, EMBED_ASPECT_NAME, embed); - _entityService.ingestProposal( + _entityService.ingestSingleProposal( proposal, new AuditStamp().setActor(UrnUtils.getUrn(context.getActorUrn())).setTime(System.currentTimeMillis()), false diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java index 0cf9acd62f736..524080bc94edb 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java @@ -29,7 +29,7 @@ private MutationUtils() { } public static void persistAspect(Urn urn, String aspectName, RecordTemplate aspect, Urn actor, EntityService entityService) { final MetadataChangeProposal proposal = buildMetadataChangeProposalWithUrn(urn, aspectName, aspect); - entityService.ingestProposal(proposal, getAuditStamp(actor), false); + entityService.ingestSingleProposal(proposal, getAuditStamp(actor), false); } /** diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java index 8c1d32c470f44..e97c20b77434e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java @@ -56,7 +56,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw MetadataChangeProposal proposal = buildMetadataChangeProposalWithUrn(actor, CORP_USER_SETTINGS_ASPECT_NAME, newSettings); - _entityService.ingestProposal(proposal, getAuditStamp(actor), false); + _entityService.ingestSingleProposal(proposal, getAuditStamp(actor), false); return true; } catch (Exception e) { diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java index 480357b89bb14..9ec91b4784f08 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java @@ -12,6 +12,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -72,9 +73,7 @@ private static MetadataChangeProposal buildSoftDeleteProposal( } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - // TODO: Replace this with a batch ingest proposals endpoint. - for (MetadataChangeProposal change : changes) { - entityService.ingestProposal(change, getAuditStamp(actor), false); - } + entityService.ingestProposal(AspectsBatch.builder() + .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } \ No newline at end of file diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java index 615d76c650286..4d3b0770bd363 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java @@ -14,6 +14,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -87,9 +88,7 @@ private static MetadataChangeProposal buildUpdateDeprecationProposal( } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - // TODO: Replace this with a batch ingest proposals endpoint. - for (MetadataChangeProposal change : changes) { - entityService.ingestProposal(change, getAuditStamp(actor), false); - } + entityService.ingestProposal(AspectsBatch.builder() + .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } \ No newline at end of file diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java index 213c08454d278..4da6ee5c2db07 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java @@ -14,6 +14,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -86,9 +87,7 @@ public static void validateDomain(Urn domainUrn, EntityService entityService) { } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - // TODO: Replace this with a batch ingest proposals endpoint. - for (MetadataChangeProposal change : changes) { - entityService.ingestProposal(change, getAuditStamp(actor), false); - } + entityService.ingestProposal(AspectsBatch.builder() + .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } \ No newline at end of file diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java index 1922a02fc1ca0..ae673f15a6cf0 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java @@ -20,6 +20,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.schema.EditableSchemaFieldInfo; import com.linkedin.schema.EditableSchemaMetadata; @@ -554,9 +555,7 @@ private static GlossaryTermAssociationArray removeTermsIfExists(GlossaryTerms te } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - // TODO: Replace this with a batch ingest proposals endpoint. - for (MetadataChangeProposal change : changes) { - entityService.ingestProposal(change, getAuditStamp(actor), false); - } + entityService.ingestProposal(AspectsBatch.builder() + .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java index a08419b5226b4..fd86fa377b723 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java @@ -21,6 +21,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -270,10 +271,8 @@ public static Boolean validateRemoveInput( } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - // TODO: Replace this with a batch ingest proposals endpoint. - for (MetadataChangeProposal change : changes) { - entityService.ingestProposal(change, getAuditStamp(actor), false); - } + entityService.ingestProposal(AspectsBatch.builder() + .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } public static void addCreatorAsOwner( 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..67afd4d3a0cc8 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.AspectsBatch; +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,21 +100,44 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ } public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { - Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( - Mockito.eq(proposal), - Mockito.any(AuditStamp.class), - Mockito.eq(false) + verifyIngestProposal(mockService, numberOfInvocations, List.of(proposal)); + } + + public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, List proposals) { + AspectsBatch batch = AspectsBatch.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)).ingestSingleProposal( + 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(AspectsBatch.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); } + public static void verifySingleIngestProposal(EntityService mockService, int numberOfInvocations) { + Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestSingleProposal( + 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(), 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..6703d8726fa5d 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 @@ -13,6 +13,8 @@ import com.linkedin.metadata.entity.EntityService; 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 +31,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 +63,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 +104,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 +144,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,7 +162,7 @@ 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(), 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..444d7481a23b2 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 @@ -14,6 +14,8 @@ import com.linkedin.metadata.entity.EntityService; 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 +32,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 +70,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 +83,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 +120,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 +161,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,7 +180,7 @@ 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(), 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..138fd1d9bf024 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 @@ -17,6 +17,8 @@ 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 +37,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 +76,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 +91,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 +134,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 +149,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 +185,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 +221,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 +255,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,7 +274,7 @@ 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(), 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..3cfb1c55b3275 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 @@ -41,7 +41,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 +64,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 +76,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 +99,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 +122,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); @@ -153,7 +145,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 @@ -173,7 +165,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..e1324edaf8a7f 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 @@ -30,7 +30,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 +75,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 +120,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 +148,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 +176,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,7 +195,7 @@ 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(), 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..ff119451dd859 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 @@ -35,7 +35,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 +82,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 +98,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 +150,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 +163,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 +202,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 +247,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,7 +277,7 @@ 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(), 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..c76f2ec4acb12 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 @@ -32,7 +32,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 +66,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 +112,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 +152,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,7 +172,7 @@ 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(), 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..0596924b54a25 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 @@ -32,7 +32,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 +82,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 +128,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 +156,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 +184,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,7 +203,7 @@ 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(), 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..6009912f498ec 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 @@ -17,6 +17,8 @@ import com.linkedin.metadata.entity.EntityService; 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 +37,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 +80,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 +100,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 +142,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 +158,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)), @@ -194,7 +190,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)), @@ -233,7 +229,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchAddTagsResolver resolver = new BatchAddTagsResolver(mockService); @@ -257,7 +253,7 @@ 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(), 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..2e4314ef27e54 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 @@ -20,6 +20,7 @@ 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 +39,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 +77,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 +85,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 +138,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)), @@ -196,7 +185,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); BatchRemoveTagsResolver resolver = new BatchRemoveTagsResolver(mockService); @@ -220,7 +209,7 @@ 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(), 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..a2d2693f716dd 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,6 +13,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.AddTermsResolver; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; @@ -31,7 +32,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 +59,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(AspectsBatch.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -77,7 +78,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 +105,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(AspectsBatch.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -119,7 +120,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)), @@ -149,7 +150,7 @@ public void testGetFailureTermDoesNotExist() 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)), @@ -179,7 +180,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception { @Test public void testGetUnauthorized() throws Exception { - EntityService mockService = Mockito.mock(EntityService.class); + EntityService mockService = getMockEntityService(); AddTermsResolver resolver = new AddTermsResolver(mockService); @@ -200,7 +201,7 @@ 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(), 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..77200c8b72a22 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 @@ -32,7 +32,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 +66,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 +83,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 +118,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 +131,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 +160,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 +196,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,7 +217,7 @@ 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(), 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..c0f1c06ba0556 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 @@ -32,7 +32,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 +66,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 +115,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 +157,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,7 +179,7 @@ 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(), 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/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java index 689b1fb997f38..fb9272f68c37a 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,13 +17,11 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectV1; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.util.Pair; import io.ebean.EbeanServer; import io.ebean.PagedList; import java.net.URISyntaxException; -import java.util.HashSet; -import java.util.Map; -import java.util.Optional; -import java.util.Set; +import java.util.*; import java.util.concurrent.TimeUnit; import java.util.function.Function; @@ -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/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java index d541127d74093..d861fdbe46bc5 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,7 @@ import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.DeleteEntityService; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.graph.LineageDirection; import com.linkedin.metadata.query.AutoCompleteResult; @@ -50,6 +51,7 @@ import com.linkedin.parseq.retry.backoff.BackoffPolicy; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.r2.RemoteInvocationException; +import com.linkedin.util.Pair; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; @@ -60,6 +62,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; @@ -534,8 +537,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 = AspectsBatch.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(); + + EntityService.IngestProposalResult one = _entityService.ingestProposal(batch, auditStamp, async).stream() + .findFirst().map(Pair::getSecond).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 ebb5139d35cae..c82f06fae0c5e 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 @@ -4,6 +4,7 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import io.ebean.PagedList; +import io.ebean.Transaction; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -11,6 +12,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Function; import java.util.function.Supplier; /** @@ -42,7 +44,13 @@ 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, @@ -106,14 +114,20 @@ 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); } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectUtils.java index e062d55254f90..58e24825d1f2b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectUtils.java @@ -4,11 +4,17 @@ import com.google.common.collect.ImmutableSet; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; +import com.linkedin.data.schema.validation.ValidationResult; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.events.metadata.ChangeType; +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.GenericRecordUtils; import com.linkedin.mxe.GenericAspect; @@ -19,6 +25,7 @@ import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.function.Consumer; import java.util.stream.Collectors; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -131,4 +138,34 @@ public static AuditStamp getAuditStamp(Urn actor) { auditStamp.setActor(actor); return auditStamp; } + + 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); + } } \ No newline at end of file diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java index bf077dfe0eb21..a411dbb19524f 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java @@ -270,7 +270,7 @@ 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 EntityService.IngestProposalResult ingestProposalResult = _entityService.ingestProposal(proposal, auditStamp, false); + final EntityService.IngestProposalResult ingestProposalResult = _entityService.ingestSingleProposal(proposal, auditStamp, false); if (!ingestProposalResult.isDidUpdate()) { log.error("Failed to ingest aspect with references removed. Before {}, after: {}, please check MCP processor" diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java index c1bda73dd7f75..7245cc59c3b37 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -10,7 +10,6 @@ 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; @@ -24,10 +23,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,12 +41,12 @@ import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; 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.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; @@ -76,7 +72,6 @@ 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; @@ -94,6 +89,7 @@ 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; @@ -181,10 +177,6 @@ public static class IngestProposalResult { public static final String DATA_PLATFORM_INSTANCE = "dataPlatformInstance"; 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 = "␟"; public EntityService( @Nonnull final AspectDao aspectDao, @@ -227,7 +219,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 +393,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); @@ -445,33 +436,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 - */ - 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. */ @@ -539,67 +503,135 @@ 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); + /** + * 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 + */ + public List ingestAspects(Urn entityUrn, + List> pairList, + @Nonnull final AuditStamp auditStamp, + SystemMetadata systemMetadata) { + List items = pairList.stream() + .map(pair -> AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(pair.getKey()) + .value(pair.getValue()) + .systemMetadata(systemMetadata) + .build(_entityRegistry)) + .collect(Collectors.toList()); + return ingestAspects(AspectsBatch.builder().items(items).build(), auditStamp, true, true); } - @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); - } + /** + * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. + * + * @param aspectsBatch aspects to write + * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time + * @param emitMae whether a {@link com.linkedin.mxe.MetadataAuditEvent} should be emitted in correspondence upon + * successful update + * @return the {@link RecordTemplate} representation of the written aspect object + */ + public List ingestAspects(@Nonnull final AspectsBatch aspectsBatch, + @Nonnull final AuditStamp auditStamp, + boolean emitMae, + boolean overwrite) { - // 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); - } + Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectsToLocalDB").time(); + List ingestResults = ingestAspectsToLocalDB(aspectsBatch, auditStamp, overwrite); + ingestToLocalDBTimer.stop(); + + if (emitMae) { + Streams.zip(aspectsBatch.getItems().stream(), ingestResults.stream(), (aspect, result) -> + sendEventForUpdateAspectResult(aspect.getUrn(), aspect.getAspectName(), result) + ).collect(Collectors.toList()); + } - 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); + return ingestResults; } + /** * 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 * - * @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 + * @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 ingestAspectToLocalDB( - @Nonnull final Urn urn, - @Nonnull final String aspectName, - @Nonnull final Function, RecordTemplate> updateLambda, - @Nonnull final AuditStamp auditStamp, - @Nonnull final SystemMetadata providedSystemMetadata) { + private List ingestAspectsToLocalDB(@Nonnull final AspectsBatch aspectsBatch, + @Nonnull final AuditStamp auditStamp, + boolean overwrite) { + + 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); + + // Upsert results + List> upsertResults = aspectsBatch.getItems().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) { + SystemMetadata systemMetadata = generateSystemMetadataIfEmpty(item.getSystemMetadata()); + result = ingestAspectToLocalDBNoTransaction(item.getUrn(), item.getAspectName(), item.getLambda(), + auditStamp, systemMetadata, latest, nextVersion); + + // 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 = new UpdateAspectResult(item.getUrn(), oldValue, oldValue, oldMetadata, oldMetadata, MetadataAuditOperation.UPDATE, auditStamp, + latest.getVersion()); + } + + return Pair.of(item, result); + }).collect(Collectors.toList()); + + // commit upserts prior to retention, if supported by impl + if (tx != null) { + tx.commitAndContinue(); + } - return _aspectDao.runInTransactionWithRetry(() -> { - final String urnStr = urn.toString(); - final EntityAspect latest = _aspectDao.getLatestAspect(urnStr, aspectName); - long nextVersion = _aspectDao.getNextVersion(urnStr, aspectName); + // Retention optimization and tx + if (_retentionService != null) { + List retentionBatch = upsertResults.stream() + // Only consider retention when there was a previous version + .filter(resultPair -> latestAspects.containsKey(resultPair.getKey().getUrn().toString()) + && latestAspects.get(resultPair.getKey().getUrn().toString()).containsKey(resultPair.getKey().getAspectName())) + .filter(resultPair -> { + RecordTemplate oldAspect = resultPair.getSecond().getOldValue(); + RecordTemplate newAspect = resultPair.getSecond().getNewValue(); + // Apply retention policies if there was an update to existing aspect value + return oldAspect != newAspect && oldAspect != null && _retentionService != null; + }) + .map(resultPair -> RetentionService.RetentionContext.builder() + .urn(resultPair.getKey().getUrn()) + .aspectName(resultPair.getKey().getAspectName()) + .maxVersion(Optional.of(resultPair.getValue().getMaxVersion())) + .build()) + .collect(Collectors.toList()); + _retentionService.applyRetentionWithPolicyDefaults(retentionBatch); + } else { + log.warn("Retention service is missing!"); + } - return ingestAspectToLocalDBNoTransaction(urn, aspectName, updateLambda, auditStamp, providedSystemMetadata, latest, nextVersion); + return upsertResults.stream().map(Pair::getValue).collect(Collectors.toList()); }, DEFAULT_MAX_TRANSACTION_RETRY); } @@ -621,7 +653,7 @@ protected UpdateAspectResult patchAspectToLocalDB( @Nonnull final AuditStamp auditStamp, @Nonnull final SystemMetadata providedSystemMetadata) { - return _aspectDao.runInTransactionWithRetry(() -> { + return _aspectDao.runInTransactionWithRetry((tx) -> { final String urnStr = urn.toString(); final String aspectName = aspectSpec.getName(); final EntityAspect latest = _aspectDao.getLatestAspect(urnStr, aspectName); @@ -640,7 +672,7 @@ protected UpdateAspectResult patchAspectToLocalDB( final RecordTemplate updatedValue = _entityRegistry.getAspectTemplateEngine().applyPatch(currentValue, jsonPatch, aspectSpec); - validateAspect(urn, updatedValue); + AspectUtils.validateRecordTemplate(_entityRegistry, urn, updatedValue); return ingestAspectToLocalDBNoTransaction(urn, aspectName, ignored -> updatedValue, auditStamp, providedSystemMetadata, latest, nextVersion); } catch (JsonProcessingException | JsonPatchException e) { @@ -649,42 +681,6 @@ protected UpdateAspectResult patchAspectToLocalDB( }, 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)); - } - return result; - }, DEFAULT_MAX_TRANSACTION_RETRY); - } - @Nonnull protected SystemMetadata generateSystemMetadataIfEmpty(@Nullable SystemMetadata systemMetadata) { if (systemMetadata == null) { @@ -695,71 +691,6 @@ protected SystemMetadata generateSystemMetadataIfEmpty(@Nullable SystemMetadata 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); - } - } - - 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()); - } - } - - /** - * 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 - */ - 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); - } - /** * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. * @@ -776,30 +707,24 @@ public RecordTemplate ingestAspect(@Nonnull final Urn urn, @Nonnull final String * @return the {@link RecordTemplate} representation of the written aspect object */ @Nullable - 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(); + AspectsBatch aspectsBatch = AspectsBatch.builder() + .aspect(AspectsBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .value(newValue) + .systemMetadata(systemMetadata) + .build(_entityRegistry)) + .build(); + List ingested = ingestAspects(aspectsBatch, auditStamp, true, false); - return sendEventForUpdateAspectResult(urn, aspectName, result); + return ingested.stream().findFirst().get().getNewValue(); } protected RecordTemplate sendEventForUpdateAspectResult(@Nonnull final Urn urn, @Nonnull final String aspectName, @@ -853,119 +778,103 @@ protected RecordTemplate sendEventForUpdateAspectResult(@Nonnull final Urn urn, } /** - * Validates that a change type is valid for the given aspect - * @param changeType - * @param aspectSpec - * @return - */ - 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)); - } - } - - - /** - * 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. - * - * 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 + * 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 IngestProposalResult} containing the results */ - 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()); - - AspectSpec aspectSpec = validateAspect(mcp, entitySpec); - - log.debug("aspect spec = {}", aspectSpec); - - if (!isValidChangeType(mcp.getChangeType(), aspectSpec)) { - throw new UnsupportedOperationException( - "ChangeType not supported: " + mcp.getChangeType() + " for aspect " + mcp.getAspectName()); + public IngestProposalResult ingestSingleProposal(MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async) { + return ingestProposal(AspectsBatch.builder() + .mcps(List.of(proposal), getEntityRegistry()) + .build(), auditStamp, async).stream().findFirst().get().getValue(); + } + + /** + * 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. + * + * 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 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 + */ + public Set> ingestProposal(AspectsBatch aspectsBatch, + AuditStamp auditStamp, + final boolean async) { + + Stream> timeseriesIngestResults = ingestTimeseriesProposal(aspectsBatch, auditStamp); + Stream> nonTimeseriesIngestResults = async ? ingestProposalAsync(aspectsBatch) + : ingestProposalSync(aspectsBatch, auditStamp); + + return Stream.concat(timeseriesIngestResults, nonTimeseriesIngestResults).collect(Collectors.toSet()); + } + + private Stream> ingestTimeseriesProposal(AspectsBatch aspectsBatch, AuditStamp auditStamp) { + List unsupported = aspectsBatch.getItems().stream() + .filter(item -> item.getAspectSpec().isTimeseries() && item.getMcp().getChangeType() != ChangeType.UPSERT) + .collect(Collectors.toList()); + if (!unsupported.isEmpty()) { + throw new UnsupportedOperationException("ChangeType not supported: " + unsupported.stream() + .map(item -> item.getMcp().getChangeType()).collect(Collectors.toSet())); } - SystemMetadata systemMetadata = generateSystemMetadataIfEmpty(mcp.getSystemMetadata()); - systemMetadata.setRegistryName(aspectSpec.getRegistryName()); - systemMetadata.setRegistryVersion(aspectSpec.getRegistryVersion().toString()); + List timeseries = aspectsBatch.getItems().stream() + .filter(item -> item.getAspectSpec().isTimeseries()) + .collect(Collectors.toList()); - 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 timeseries.stream().map(item -> { + boolean didUpdate = emitChangeLog(null, null, item.getAspect(), item.getSystemMetadata(), + item.getMcp(), item.getUrn(), auditStamp, item.getAspectSpec()); + return Pair.of(item, new IngestProposalResult(item.getUrn(), didUpdate, false)); + }); } - 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()); + private Stream> ingestProposalAsync(AspectsBatch aspectsBatch) { + List nonTimeseries = aspectsBatch.getItems().stream() + .filter(item -> !item.getAspectSpec().isTimeseries()) + .collect(Collectors.toList()); - if (aspectSpec == null) { - throw new RuntimeException( - String.format("Unknown aspect %s for entity %s", mcp.getAspectName(), - mcp.getEntityType())); - } - - return aspectSpec; + return nonTimeseries.stream().map(item -> { + // When async is turned on, we write to proposal log and return without waiting + _producer.produceMetadataChangeProposal(item.getUrn(), item.getMcp()); + return Pair.of(item, new IngestProposalResult(item.getUrn(), false, true)); + }); } - private UpdateAspectResult performUpsert(MetadataChangeProposal mcp, AspectSpec aspectSpec, SystemMetadata - systemMetadata, Urn entityUrn, AuditStamp auditStamp) { - RecordTemplate aspect = convertToRecordTemplate(mcp, aspectSpec); - log.debug("aspect = {}", aspect); + private Stream> ingestProposalSync(AspectsBatch aspectsBatch, AuditStamp auditStamp) { + List unsupported = aspectsBatch.getItems().stream() + .filter(item -> item.getMcp().getChangeType() != ChangeType.PATCH + && item.getMcp().getChangeType() != ChangeType.UPSERT) + .collect(Collectors.toList()); + if (!unsupported.isEmpty()) { + throw new UnsupportedOperationException("ChangeType not supported: " + unsupported.stream() + .map(item -> item.getMcp().getChangeType()).collect(Collectors.toSet())); + } - return upsertAspect(aspect, systemMetadata, mcp, entityUrn, auditStamp, aspectSpec); + Stream> patches = aspectsBatch.getItems().stream() + .filter(item -> item.getMcp().getChangeType() == ChangeType.PATCH) + .map(item -> Pair.of(item, performPatch(item.getMcp(), item.getAspectSpec(), item.getSystemMetadata(), + item.getUrn(), auditStamp))); + + List upsertItems = aspectsBatch.getItems() + .stream().filter(item -> item.getMcp().getChangeType() == ChangeType.UPSERT) + .collect(Collectors.toList()); + List upsertResults = ingestAspects(AspectsBatch.builder().items(upsertItems).build(), + auditStamp, false, true); + Stream> upserts = Streams.zip(upsertItems.stream(), upsertResults.stream(), Pair::of); + + return Stream.concat(patches, upserts).map(resultPair -> { + AspectsBatchItem item = resultPair.getFirst(); + UpdateAspectResult result = resultPair.getSecond(); + boolean didUpdate = emitChangeLog(result.getOldValue(), result.getOldSystemMetadata(), + result.getNewValue(), result.getNewSystemMetadata(), item.getMcp(), item.getUrn(), + auditStamp, item.getAspectSpec()); + return Pair.of(item, new IngestProposalResult(item.getUrn(), didUpdate, false)); + }); } private UpdateAspectResult performPatch(MetadataChangeProposal mcp, AspectSpec aspectSpec, SystemMetadata @@ -979,7 +888,7 @@ private UpdateAspectResult performPatch(MetadataChangeProposal mcp, AspectSpec a Patch jsonPatch = convertToJsonPatch(mcp); log.debug("patch = {}", jsonPatch); - return patchAspect(jsonPatch, systemMetadata, mcp, entityUrn, auditStamp, aspectSpec); + return patchAspect(jsonPatch, systemMetadata, entityUrn, auditStamp, aspectSpec); } private boolean supportsPatch(AspectSpec aspectSpec) { @@ -987,21 +896,6 @@ private boolean supportsPatch(AspectSpec aspectSpec) { return AspectTemplateEngine.SUPPORTED_TEMPLATES.contains(aspectSpec.getName()); } - 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())); - } - log.debug("aspect = {}", aspect); - return aspect; - } - private Patch convertToJsonPatch(MetadataChangeProposal mcp) { JsonNode json; try { @@ -1012,29 +906,8 @@ private Patch convertToJsonPatch(MetadataChangeProposal mcp) { } } - 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.applyRetentionWithPolicyDefaults(List.of( - RetentionService.RetentionContext.builder() - .urn(entityUrn) - .aspectName(aspectSpec.getName()) - .maxVersion(Optional.of(result.maxVersion)) - .build())); - } - return result; - } - - private UpdateAspectResult patchAspect(final Patch patch, final SystemMetadata systemMetadata, - MetadataChangeProposal mcp, Urn entityUrn, AuditStamp auditStamp, AspectSpec aspectSpec) { + private UpdateAspectResult patchAspect(final Patch patch, final SystemMetadata systemMetadata, 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(); @@ -1323,7 +1196,7 @@ public void produceMetadataAuditEvent(@Nonnull final Urn urn, @Nonnull final Str } protected Snapshot buildKeySnapshot(@Nonnull final Urn urn) { - final RecordTemplate keyAspectValue = buildKeyAspect(urn); + final RecordTemplate keyAspectValue = EntityUtils.buildKeyAspect(_entityRegistry, urn); return toSnapshotUnion(toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, keyAspectValue)))); } @@ -1448,7 +1321,7 @@ public List> generateDefaultAspectsIfMissing(@Nonnu RecordTemplate keyAspect = latestAspects.get(keyAspectName); if (keyAspect == null) { - keyAspect = buildKeyAspect(urn); + keyAspect = EntityUtils.buildKeyAspect(_entityRegistry, urn); aspects.add(Pair.of(keyAspectName, keyAspect)); } @@ -1487,9 +1360,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().map(Pair::getFirst).collect(Collectors.toSet()))); - ingestAspects(urn, aspectRecordsToIngest, auditStamp, systemMetadata); + AspectsBatch aspectsBatch = AspectsBatch.builder() + .items(aspectRecordsToIngest.stream().map(pair -> AspectsBatchItem.builder() + .urn(urn) + .aspectName(pair.getKey()) + .value(pair.getValue()) + .systemMetadata(systemMetadata) + .build(_entityRegistry)).collect(Collectors.toList())) + .build(); + + ingestAspects(aspectsBatch, auditStamp, true, true); } public Snapshot buildSnapshot(@Nonnull final Urn urn, @Nonnull final RecordTemplate aspectValue) { @@ -1498,18 +1380,11 @@ public Snapshot buildSnapshot(@Nonnull final Urn urn, @Nonnull final RecordTempl return toSnapshotUnion(toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, aspectValue)))); } - final RecordTemplate keyAspectValue = buildKeyAspect(urn); + final RecordTemplate keyAspectValue = EntityUtils.buildKeyAspect(_entityRegistry, 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); - } - public AspectSpec getKeyAspectSpec(@Nonnull final Urn urn) { return getKeyAspectSpec(urnToEntityName(urn)); } @@ -1724,7 +1599,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. @@ -1802,7 +1677,7 @@ 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()); @@ -2005,19 +1879,6 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction( 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, @@ -2030,7 +1891,7 @@ private RecordTemplate updateAspect( @Nonnull final boolean emitMae, final int maxTransactionRetry) { - final UpdateAspectResult result = _aspectDao.runInTransactionWithRetry(() -> { + final UpdateAspectResult result = _aspectDao.runInTransactionWithRetry((tx) -> { final EntityAspect oldAspect = _aspectDao.getAspect(urn.toString(), aspectName, version); final RecordTemplate oldValue = @@ -2138,7 +1999,7 @@ private boolean shouldAspectEmitChangeLog(@Nonnull final Urn urn, @Nonnull final 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 53099e900ee4d..2becb4778a5b1 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 @@ -6,17 +6,24 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.template.RecordTemplate; +import com.linkedin.entity.Entity; import com.linkedin.entity.EnvelopedAspect; +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.SystemMetadata; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; +import java.net.URISyntaxException; +import java.net.URLEncoder; + import static com.linkedin.metadata.entity.EntityService.*; +import static com.linkedin.metadata.utils.PegasusUtils.urnToEntityName; @Slf4j @@ -25,6 +32,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); @@ -95,4 +105,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/RetentionService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java index d02512a9a338e..5627b53a49a0a 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java @@ -7,6 +7,7 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; import com.linkedin.metadata.key.DataHubRetentionKey; @@ -94,6 +95,7 @@ public boolean setRetention(@Nullable String entityName, @Nullable String aspect 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); @@ -101,15 +103,20 @@ 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 = AspectsBatch.builder() + .mcps(List.of(keyProposal, aspectProposal), getEntityService().getEntityRegistry()) + .build(); + + return getEntityService().ingestProposal(batch, auditStamp, false).stream() + .anyMatch(resultPair -> resultPair.getSecond().isDidUpdate()); } /** @@ -146,7 +153,7 @@ private void validateRetention(Retention retention) { * * @param retentionContexts urn, aspect name, and additional context that could be used to apply retention */ - public void applyRetentionWithPolicyDefaults(List retentionContexts) { + public void applyRetentionWithPolicyDefaults(@Nonnull List retentionContexts) { List withDefaults = retentionContexts.stream() .map(context -> { if (context.getRetentionPolicy().isEmpty()) { 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 d75a61d6246ea..eac1482e5ee12 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 @@ -38,12 +38,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.*; @@ -84,6 +85,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(); @@ -240,7 +254,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; @@ -248,7 +262,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; } @@ -452,25 +466,24 @@ 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.getValue()) { + long latestVersion = maxVersions.get(aspectName); + long nextVal = latestVersion < 0 ? ASPECT_LATEST_VERSION : latestVersion + 1L; + nextVersions.put(aspectName, nextVal); + } - for (String aspectName: aspectNames) { - 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 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 1da8e3375035f..ad77e6f361a4c 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 @@ -16,6 +16,7 @@ import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.retention.DataHubRetentionConfig; import com.linkedin.retention.Retention; import com.linkedin.retention.TimeBasedRetention; 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 8f32be42e6d95..cca85c1204b97 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 @@ -33,12 +33,13 @@ 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; @@ -180,28 +181,19 @@ private void saveEbeanAspect(@Nonnull final EbeanAspectV2 ebeanAspect, final boo } @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 @@ -495,7 +487,7 @@ 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) { validateConnection(); int retryCount = 0; Exception lastException; @@ -504,7 +496,7 @@ public T runInTransactionWithRetry(@Nonnull final Supplier block, final i 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; @@ -547,57 +539,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; } @@ -671,4 +670,17 @@ 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)); + } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java new file mode 100644 index 0000000000000..7c5416f99563e --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java @@ -0,0 +1,150 @@ +package com.linkedin.metadata.entity.ebean.transactions; + +import com.datahub.util.exception.ModelConversionException; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +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.events.metadata.ChangeType; +import lombok.Builder; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; + +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class AspectsBatch { + private final List items; + + public Map> getUrnAspectsMap() { + return items.stream() + .map(aspect -> Map.entry(aspect.getUrn().toString(), aspect.getAspectName())) + .collect(Collectors.groupingBy(Map.Entry::getKey, Collectors.mapping(Map.Entry::getValue, Collectors.toSet()))); + } + + public static class AspectsBatchBuilder { + + /** + * Just one aspect record template + * @param data aspect data + * @return builder + */ + public AspectsBatchBuilder aspect(AspectsBatchItem data) { + this.items = List.of(data); + return this; + } + + public AspectsBatchBuilder mcps(List mcps, EntityRegistry entityRegistry) { + this.items = mcps.stream().map(mcp -> toAspectBatchItem(mcp, entityRegistry)).collect(Collectors.toList()); + return this; + } + + private static AspectsBatchItem toAspectBatchItem(MetadataChangeProposal mcp, EntityRegistry entityRegistry) { + log.debug("entity type = {}", mcp.getEntityType()); + EntitySpec entitySpec = entityRegistry.getEntitySpec(mcp.getEntityType()); + AspectSpec aspectSpec = validateAspect(mcp, entitySpec); + RecordTemplate aspect = convertToRecordTemplate(mcp, aspectSpec); + log.debug("aspect = {}", aspect); + + if (!isValidChangeType(mcp.getChangeType(), aspectSpec)) { + throw new UnsupportedOperationException("ChangeType not supported: " + mcp.getChangeType() + + " for aspect " + mcp.getAspectName()); + } + + Urn urn = mcp.getEntityUrn(); + if (urn == null) { + urn = EntityKeyUtils.getUrnFromProposal(mcp, aspectSpec); + } + + return AspectsBatchItem.builder() + .urn(urn) + .aspectName(mcp.getAspectName()) + .systemMetadata(mcp.getSystemMetadata()) + .mcp(mcp) + .value(aspect) + .build(entityRegistry); + } + + 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; + } + + private 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 + */ + private static 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 boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AspectsBatch that = (AspectsBatch) o; + return items.equals(that.items); + } + + @Override + public int hashCode() { + return Objects.hash(items); + } + + @Override + public String toString() { + return "AspectsBatch{" + + "items=" + + items + + '}'; + } +} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java new file mode 100644 index 0000000000000..426845642e5ba --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java @@ -0,0 +1,123 @@ +package com.linkedin.metadata.entity.ebean.transactions; + +import com.linkedin.common.AuditStamp; +import com.linkedin.common.urn.Urn; +import com.linkedin.data.template.RecordTemplate; +import com.linkedin.metadata.entity.AspectUtils; +import com.linkedin.metadata.entity.EntityAspect; +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 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 java.util.Optional; +import java.util.function.Function; + +import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; + +@Slf4j +@Getter +@Builder(toBuilder = true) +public class AspectsBatchItem { + // 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; + // updateLambda Function to apply to the latest version of the aspect to get the updated version + private final Function, RecordTemplate> lambda; + + private final MetadataChangeProposal mcp; + + // derived + private final EntitySpec entitySpec; + private final AspectSpec aspectSpec; + + public RecordTemplate getAspect() { + return lambda.apply(null); + } + + public EntityAspect toLatestEntityAspect(AuditStamp auditStamp) { + EntityAspect latest = new EntityAspect(); + latest.setAspect(aspectName); + latest.setMetadata(EntityUtils.toJsonAspect(getAspect())); + latest.setUrn(urn.toString()); + latest.setVersion(ASPECT_LATEST_VERSION); + latest.setCreatedOn(new Timestamp(auditStamp.getTime())); + latest.setCreatedBy(auditStamp.getActor().toString()); + return latest; + } + + public static class AspectsBatchItemBuilder { + public AspectsBatchItemBuilder value(final RecordTemplate recordTemplate) { + this.lambda = ignored -> recordTemplate; + return this; + } + + public AspectsBatchItem 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(AspectUtils.validate(this.entitySpec, this.aspectName)); + log.debug("aspect spec = {}", this.aspectSpec); + + RecordTemplate aspect = this.lambda.apply(null); + AspectUtils.validateRecordTemplate(entityRegistry, this.entitySpec, this.urn, aspect); + + return new AspectsBatchItem(this.urn, this.aspectName, this.systemMetadata, this.lambda, this.mcp, + this.entitySpec, this.aspectSpec); + } + + private AspectsBatchItemBuilder entitySpec(EntitySpec entitySpec) { + this.entitySpec = entitySpec; + return this; + } + + private AspectsBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { + this.aspectSpec = aspectSpec; + return this; + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AspectsBatchItem that = (AspectsBatchItem) o; + return urn.equals(that.urn) && aspectName.equals(that.aspectName) + && Objects.equals(systemMetadata, that.systemMetadata) + && lambda.apply(null).equals(that.lambda.apply(null)) + && Objects.equals(mcp, that.mcp); + } + + @Override + public int hashCode() { + return Objects.hash(urn, aspectName, systemMetadata, lambda.apply(null), mcp); + } + + @Override + public String toString() { + return "AspectsBatchItem{" + + "urn=" + urn + + ", aspectName='" + + aspectName + '\'' + + ", systemMetadata=" + systemMetadata + + ", lambda=" + lambda + + ", mcp=" + mcp + + '}'; + } +} 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..d37eb39a81076 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.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; 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(AspectsBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .value(aspect) + .systemMetadata(AspectGenerationUtils.createSystemMetadata()) + .build(entityService.getEntityRegistry())); } + entityService.ingestAspects(AspectsBatch.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(AspectsBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .value(aspect) + .systemMetadata(AspectGenerationUtils.createSystemMetadata()) + .build(entityService.getEntityRegistry())); } + entityService.ingestAspects(AspectsBatch.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(AspectsBatchItem.builder() + .urn(urn) + .aspectName(aspectName) + .value(aspect) + .systemMetadata(AspectGenerationUtils.createSystemMetadata()) + .build(entityService.getEntityRegistry())); } + entityService.ingestAspects(AspectsBatch.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } } 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 3d5b8c6cfaa2b..83f0fc723bb3e 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 @@ -10,6 +10,8 @@ import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.ebean.EbeanAspectDao; import com.linkedin.metadata.entity.ebean.EbeanRetentionService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; import com.linkedin.metadata.models.registry.EntityRegistryException; @@ -25,6 +27,8 @@ import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import java.util.List; + import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; @@ -86,15 +90,34 @@ public void testIngestListLatestAspects() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); - _entityService.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - _entityService.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); + + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .value(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // List aspects ListResult batch1 = _entityService.listLatestAspects(entityUrn1.getEntityType(), aspectName, 0, 2); @@ -131,15 +154,34 @@ public void testIngestListUrns() throws AssertionError { // Ingest CorpUserInfo Aspect #1 RecordTemplate writeAspect1 = AspectGenerationUtils.createCorpUserKey(entityUrn1); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #2 RecordTemplate writeAspect2 = AspectGenerationUtils.createCorpUserKey(entityUrn2); - _entityService.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 RecordTemplate writeAspect3 = AspectGenerationUtils.createCorpUserKey(entityUrn3); - _entityService.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); + + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .value(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // List aspects urns ListUrnsResult batch1 = _entityService.listUrns(entityUrn1.getEntityType(), 0, 2); 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 1d11bd786c4cf..c1e3c852fab9c 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.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.key.CorpUserKey; @@ -61,6 +63,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import org.junit.Assert; import org.mockito.ArgumentCaptor; @@ -413,7 +416,6 @@ public void testIngestAspectsGetLatestAspects() throws Exception { pairToIngest.add(getAspectRecordPair(writeAspect2, CorpUserInfo.class)); SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(); - _entityService.ingestAspects(entityUrn, pairToIngest, TEST_AUDIT_STAMP, metadata1); Map latestAspects = _entityService.getLatestAspectsForUrn( @@ -444,7 +446,6 @@ public void testReingestAspectsGetLatestAspects() throws Exception { pairToIngest.add(getAspectRecordPair(writeAspect1, CorpUserInfo.class)); SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(); - _entityService.ingestAspects(entityUrn, pairToIngest, TEST_AUDIT_STAMP, metadata1); final MetadataChangeLog initialChangeLog = new MetadataChangeLog(); @@ -578,7 +579,7 @@ public void testReingestLineageProposal() throws Exception { mcp1.setSystemMetadata(metadata1); mcp1.setAspectName(UPSTREAM_LINEAGE_ASPECT_NAME); - _entityService.ingestProposal(mcp1, TEST_AUDIT_STAMP, false); + _entityService.ingestSingleProposal(mcp1, TEST_AUDIT_STAMP, false); final MetadataChangeLog initialChangeLog = new MetadataChangeLog(); initialChangeLog.setEntityType(entityUrn.getEntityType()); @@ -613,7 +614,7 @@ public void testReingestLineageProposal() throws Exception { // Mockito detects the previous invocation and throws an error in verifying the second call unless invocations are cleared clearInvocations(_mockProducer); - _entityService.ingestProposal(mcp1, TEST_AUDIT_STAMP, false); + _entityService.ingestSingleProposal(mcp1, TEST_AUDIT_STAMP, false); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(restateChangeLog)); @@ -640,7 +641,7 @@ public void testIngestTimeseriesAspect() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetProfileSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityService.ingestProposal(gmce, TEST_AUDIT_STAMP, false); + _entityService.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, false); } @Test @@ -659,7 +660,7 @@ public void testAsyncProposalVersioned() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetPropertiesSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityService.ingestProposal(gmce, TEST_AUDIT_STAMP, true); + _entityService.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, true); verify(_mockProducer, times(0)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.any()); verify(_mockProducer, times(1)).produceMetadataChangeProposal(Mockito.eq(entityUrn), @@ -685,7 +686,7 @@ public void testAsyncProposalTimeseries() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetProfileSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityService.ingestProposal(gmce, TEST_AUDIT_STAMP, true); + _entityService.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, true); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.any()); verify(_mockProducer, times(0)).produceMetadataChangeProposal(Mockito.eq(entityUrn), @@ -768,19 +769,43 @@ public void testRollbackAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); - _entityService.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - _entityService.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1Overwrite, TEST_AUDIT_STAMP, metadata2); + + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .value(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1Overwrite) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.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 +847,33 @@ public void testRollbackKey() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); - RecordTemplate writeKey1 = _entityService.buildKeyAspect(entityUrn1); - _entityService.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"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1Overwrite, TEST_AUDIT_STAMP, metadata2); + + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(keyAspectName) + .value(writeKey1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1Overwrite) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.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 +917,51 @@ public void testRollbackUrn() throws AssertionError { // Ingest CorpUserInfo Aspect #1 CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); - RecordTemplate writeKey1 = _entityService.buildKeyAspect(entityUrn1); - _entityService.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"); - _entityService.ingestAspect(entityUrn2, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - _entityService.ingestAspect(entityUrn3, aspectName, writeAspect3, TEST_AUDIT_STAMP, metadata1); // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - _entityService.ingestAspect(entityUrn1, aspectName, writeAspect1Overwrite, TEST_AUDIT_STAMP, metadata2); + + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(keyAspectName) + .value(writeKey1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn2) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn3) + .aspectName(aspectName) + .value(writeAspect3) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn1) + .aspectName(aspectName) + .value(writeAspect1Overwrite) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.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 +991,17 @@ public void testIngestGetLatestAspect() throws AssertionError { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #1 - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); RecordTemplate readAspect1 = _entityService.getLatestAspect(entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, readAspect1)); @@ -941,8 +1023,17 @@ public void testIngestGetLatestAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); + items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #2 - _entityService.ingestAspect(entityUrn, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata2); RecordTemplate readAspect2 = _entityService.getLatestAspect(entityUrn, aspectName); EntityAspect readAspectDao1 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 1); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 0); @@ -975,16 +1066,34 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #1 - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); EnvelopedAspect readAspect1 = _entityService.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( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #2 - _entityService.ingestAspect(entityUrn, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata2); EnvelopedAspect readAspect2 = _entityService.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); EntityAspect readAspectDao1 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 1); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 0); @@ -1017,8 +1126,17 @@ public void testIngestSameAspect() throws AssertionError { SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); SystemMetadata metadata3 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-123"); + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #1 - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); RecordTemplate readAspect1 = _entityService.getLatestAspect(entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, readAspect1)); @@ -1040,8 +1158,17 @@ public void testIngestSameAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #2 CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); + items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect2) + .systemMetadata(metadata2) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + // Validate retrieval of CorpUserInfo Aspect #2 - _entityService.ingestAspect(entityUrn, aspectName, writeAspect2, TEST_AUDIT_STAMP, metadata2); RecordTemplate readAspect2 = _entityService.getLatestAspect(entityUrn, aspectName); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, ASPECT_LATEST_VERSION); @@ -1069,20 +1196,54 @@ public void testRetention() throws AssertionError { // Ingest CorpUserInfo Aspect CorpUserInfo writeAspect1 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1, TEST_AUDIT_STAMP, metadata1); CorpUserInfo writeAspect1a = AspectGenerationUtils.createCorpUserInfo("email_a@test.com"); - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1a, TEST_AUDIT_STAMP, metadata1); CorpUserInfo writeAspect1b = AspectGenerationUtils.createCorpUserInfo("email_b@test.com"); - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1b, TEST_AUDIT_STAMP, metadata1); String aspectName2 = AspectGenerationUtils.getAspectName(new Status()); // Ingest Status Aspect Status writeAspect2 = new Status().setRemoved(true); - _entityService.ingestAspect(entityUrn, aspectName2, writeAspect2, TEST_AUDIT_STAMP, metadata1); Status writeAspect2a = new Status().setRemoved(false); - _entityService.ingestAspect(entityUrn, aspectName2, writeAspect2a, TEST_AUDIT_STAMP, metadata1); Status writeAspect2b = new Status().setRemoved(true); - _entityService.ingestAspect(entityUrn, aspectName2, writeAspect2b, TEST_AUDIT_STAMP, metadata1); + + List items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1a) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1b) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .value(writeAspect2) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .value(writeAspect2a) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .value(writeAspect2b) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); assertEquals(_entityService.getAspect(entityUrn, aspectName, 1), writeAspect1); assertEquals(_entityService.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1094,10 +1255,24 @@ public void testRetention() throws AssertionError { // Ingest CorpUserInfo Aspect again CorpUserInfo writeAspect1c = AspectGenerationUtils.createCorpUserInfo("email_c@test.com"); - _entityService.ingestAspect(entityUrn, aspectName, writeAspect1c, TEST_AUDIT_STAMP, metadata1); // Ingest Status Aspect again Status writeAspect2c = new Status().setRemoved(false); - _entityService.ingestAspect(entityUrn, aspectName2, writeAspect2c, TEST_AUDIT_STAMP, metadata1); + + items = List.of( + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName) + .value(writeAspect1c) + .systemMetadata(metadata1) + .build(_testEntityRegistry), + AspectsBatchItem.builder() + .urn(entityUrn) + .aspectName(aspectName2) + .value(writeAspect2c) + .systemMetadata(metadata1) + .build(_testEntityRegistry) + ); + _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); assertNull(_entityService.getAspect(entityUrn, aspectName, 1)); assertEquals(_entityService.getAspect(entityUrn, aspectName2, 1), writeAspect2); @@ -1203,12 +1378,12 @@ public void testRestoreIndices() throws Exception { public void testValidateUrn() throws Exception { // Valid URN Urn validTestUrn = new Urn("li", "corpuser", new TupleKey("testKey")); - _entityService.validateUrn(validTestUrn); + EntityUtils.validateUrn(_testEntityRegistry, validTestUrn); // URN with trailing whitespace Urn testUrnWithTrailingWhitespace = new Urn("li", "corpuser", new TupleKey("testKey ")); try { - _entityService.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 +1394,7 @@ public void testValidateUrn() throws Exception { Urn testUrnTooLong = new Urn("li", "corpuser", new TupleKey(stringTooLong)); try { - _entityService.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 +1410,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 - _entityService.validateUrn(testUrnSameLengthWhenEncoded); + EntityUtils.validateUrn(_testEntityRegistry, testUrnSameLengthWhenEncoded); try { - _entityService.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 +1421,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␟")); - _entityService.validateUrn(testUrnSpecialCharValid); + EntityUtils.validateUrn(_testEntityRegistry, testUrnSpecialCharValid); try { - _entityService.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 +1431,7 @@ public void testValidateUrn() throws Exception { Urn urnWithMismatchedParens = new Urn("li", "corpuser", new TupleKey("test(Key")); try { - _entityService.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 +1439,18 @@ public void testValidateUrn() throws Exception { Urn invalidType = new Urn("li", "fakeMadeUpType", new TupleKey("testKey")); try { - _entityService.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")); - _entityService.validateUrn(validFabricType); + EntityUtils.validateUrn(_testEntityRegistry, validFabricType); Urn invalidFabricType = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "prod")); try { - _entityService.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 +1458,7 @@ public void testValidateUrn() throws Exception { Urn urnEndingInComma = new Urn("li", "dataset", new TupleKey("urn:li:dataPlatform:foo", "bar", "PROD", "")); try { - _entityService.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())); @@ -1312,7 +1487,7 @@ public void testUIPreProcessedProposal() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetPropertiesSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityService.ingestProposal(gmce, TEST_AUDIT_STAMP, false); + _entityService.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, false); ArgumentCaptor captor = ArgumentCaptor.forClass(MetadataChangeLog.class); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), captor.capture()); 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..6c73c009db9bf 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.AspectsBatch; import com.linkedin.metadata.key.DataHubAccessTokenKey; import com.linkedin.metadata.utils.AuditStampUtils; import com.linkedin.metadata.utils.GenericRecordUtils; @@ -24,6 +25,8 @@ 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 +127,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(AspectsBatch.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/token/StatefulTokenServiceTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java index 75a9114529d7a..99f242bac6b42 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,8 @@ import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import java.util.Date; import java.util.Map; + +import com.linkedin.metadata.models.registry.EntityRegistry; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -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/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java index 876a0871fa4cb..b7477fdcc0af9 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java @@ -63,6 +63,6 @@ static void setUpgradeResult(Urn urn, EntityService entityService) throws URISyn upgradeProposal.setAspectName(Constants.DATA_HUB_UPGRADE_RESULT_ASPECT_NAME); upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(upgradeResult)); upgradeProposal.setChangeType(ChangeType.UPSERT); - entityService.ingestProposal(upgradeProposal, auditStamp, false); + entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java index dbbcf3a139bf1..b408af8b66720 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java @@ -92,7 +92,7 @@ private void ingestUpgradeRequestAspect() throws URISyntaxException { upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(upgradeRequest)); upgradeProposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(upgradeProposal, auditStamp, false); + _entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); } private void ingestUpgradeResultAspect() throws URISyntaxException { @@ -107,7 +107,7 @@ private void ingestUpgradeResultAspect() throws URISyntaxException { upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(upgradeResult)); upgradeProposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(upgradeProposal, auditStamp, false); + _entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); } private void cleanUpgradeAfterError(Exception e, String errorMessage) { diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java index 627c5934140a8..a18cd5362b8e2 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java @@ -135,7 +135,7 @@ private void ingestBrowsePathsV2(Urn urn, AuditStamp auditStamp) throws Exceptio proposal.setChangeType(ChangeType.UPSERT); proposal.setSystemMetadata(new SystemMetadata().setRunId(EntityService.DEFAULT_RUN_ID).setLastObserved(System.currentTimeMillis())); proposal.setAspect(GenericRecordUtils.serializeAspect(browsePathsV2)); - _entityService.ingestProposal( + _entityService.ingestSingleProposal( proposal, auditStamp, false 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..35816e824d31f 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.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; 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(AspectsBatchItem.builder() + .urn(urn) + .aspectName(DATA_PLATFORM_INSTANCE_ASPECT_NAME) + .value(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(AspectsBatch.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..800217c1d420a 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.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; 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 AspectsBatchItem.builder() + .urn(urn) + .aspectName(PLATFORM_ASPECT_NAME) + .value(info) + .build(_entityService.getEntityRegistry()); + }).collect(Collectors.toList()); + + _entityService.ingestAspects(AspectsBatch.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/IngestDefaultGlobalSettingsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java index 5bc80f46e6478..1541ca40d8911 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java @@ -116,7 +116,7 @@ public void execute() throws IOException, URISyntaxException { proposal.setAspect(GenericRecordUtils.serializeAspect(newSettings)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal( + _entityService.ingestSingleProposal( proposal, new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java index 08a867d710419..596af9cb857f6 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.AspectsBatch; 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(AspectsBatch.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..1de96218332ec 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.AspectsBatch; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.query.ListUrnsResult; @@ -26,10 +27,8 @@ import com.linkedin.policy.DataHubPolicyInfo; import java.io.IOException; import java.net.URISyntaxException; -import java.util.Collections; -import java.util.HashSet; -import java.util.Map; -import java.util.Optional; +import java.util.*; + import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.core.io.ClassPathResource; @@ -172,9 +171,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 +178,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(AspectsBatch.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..0475e85157fa1 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.AspectsBatch; 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,8 +108,10 @@ 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(AspectsBatch.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, 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/RestoreDbtSiblingsIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreDbtSiblingsIndices.java index 989ee1a39b169..617d8043c106b 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 @@ -168,6 +168,6 @@ private void ingestUpgradeAspect(String aspectName, RecordTemplate aspect, Audit upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(aspect)); upgradeProposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(upgradeProposal, auditStamp, false); + _entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java index b990400b38491..9def1e544dc85 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java @@ -126,7 +126,7 @@ private void migrateBrowsePath(Urn urn, AuditStamp auditStamp) throws Exception proposal.setChangeType(ChangeType.UPSERT); proposal.setSystemMetadata(new SystemMetadata().setRunId(EntityService.DEFAULT_RUN_ID).setLastObserved(System.currentTimeMillis())); proposal.setAspect(GenericRecordUtils.serializeAspect(newPaths)); - _entityService.ingestProposal( + _entityService.ingestSingleProposal( proposal, auditStamp, false diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java index 49fce75ab7c61..1fb979cba2d3c 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java @@ -97,7 +97,7 @@ public void testExecuteNoExistingBrowsePaths() throws Exception { Mockito.eq(null) ); // Verify that 11 aspects are ingested, 2 for the upgrade request / result, 9 for ingesting 1 of each entity type - Mockito.verify(mockService, Mockito.times(11)).ingestProposal( + Mockito.verify(mockService, Mockito.times(11)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -124,7 +124,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { BackfillBrowsePathsV2Step backfillBrowsePathsV2Step = new BackfillBrowsePathsV2Step(mockService, mockSearchService); backfillBrowsePathsV2Step.execute(); - Mockito.verify(mockService, Mockito.times(0)).ingestProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean() 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 fa351d7539e02..68618a9e87a79 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 @@ -95,15 +95,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) + && 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) + && item.getAspect() instanceof DataPlatformInstance) + ), any(), anyBoolean(), anyBoolean()); } @NotNull diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java index 24bdd193a39c8..50ab1f53455f1 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java @@ -38,7 +38,7 @@ public void testExecuteValidSettingsNoExistingSettings() throws Exception { GlobalSettingsInfo expectedResult = new GlobalSettingsInfo(); expectedResult.setViews(new GlobalViewsSettings().setDefaultView(UrnUtils.getUrn("urn:li:dataHubView:test"))); - Mockito.verify(entityService, times(1)).ingestProposal( + Mockito.verify(entityService, times(1)).ingestSingleProposal( Mockito.eq(buildUpdateSettingsProposal(expectedResult)), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -65,7 +65,7 @@ public void testExecuteValidSettingsExistingSettings() throws Exception { GlobalSettingsInfo expectedResult = new GlobalSettingsInfo(); expectedResult.setViews(new GlobalViewsSettings().setDefaultView(UrnUtils.getUrn("urn:li:dataHubView:custom"))); - Mockito.verify(entityService, times(1)).ingestProposal( + Mockito.verify(entityService, times(1)).ingestSingleProposal( Mockito.eq(buildUpdateSettingsProposal(expectedResult)), Mockito.any(AuditStamp.class), Mockito.eq(false) 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..52be8a5d364ee 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 @@ -59,7 +59,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.CHART_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.DASHBOARD_ENTITY_NAME); // creates upgradeRequest and upgradeResult aspects - Mockito.verify(mockService, Mockito.times(2)).ingestProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -121,7 +121,7 @@ public void testExecuteWithNewVersion() throws Exception { Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.CHART_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.DASHBOARD_ENTITY_NAME); // creates upgradeRequest and upgradeResult aspects - Mockito.verify(mockService, Mockito.times(2)).ingestProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -183,7 +183,7 @@ public void testDoesNotExecuteWithSameVersion() throws Exception { Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.CHART_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.DASHBOARD_ENTITY_NAME); // creates upgradeRequest and upgradeResult aspects - Mockito.verify(mockService, Mockito.times(0)).ingestProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) 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 88c63110ee63e..1429da55940d5 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 @@ -105,7 +105,7 @@ public void testExecuteFirstTime() throws Exception { 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( + Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -166,7 +166,7 @@ public void testExecutesWithNewVersion() throws Exception { 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( + Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -227,7 +227,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { Mockito.verify(mockSearchService, Mockito.times(0)).search(Constants.GLOSSARY_NODE_ENTITY_NAME, "", null, null, 0, 1000, new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)); - Mockito.verify(mockService, Mockito.times(0)).ingestProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean() diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java index 5e4ad6e7fe880..a9410fb2522f2 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java @@ -87,7 +87,7 @@ public void testExecuteNoExistingBrowsePaths() throws Exception { Mockito.eq(5000) ); // Verify that 4 aspects are ingested, 2 for the upgrade request / result, but none for ingesting - Mockito.verify(mockService, Mockito.times(2)).ingestProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -155,7 +155,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.eq(5000) ); // Verify that 4 aspects are ingested, 2 for the upgrade request / result and 2 for the browse pahts - Mockito.verify(mockService, Mockito.times(4)).ingestProposal( + Mockito.verify(mockService, Mockito.times(4)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -223,7 +223,7 @@ public void testDoesNotRunWhenBrowsePathIsNotQualified() throws Exception { Mockito.eq(5000) ); // Verify that 2 aspects are ingested, only those for the upgrade step - Mockito.verify(mockService, Mockito.times(2)).ingestProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -249,7 +249,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { UpgradeDefaultBrowsePathsStep step = new UpgradeDefaultBrowsePathsStep(mockService); step.execute(); - Mockito.verify(mockService, Mockito.times(0)).ingestProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean() 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 7ce41924fe92d..b89f9d75e7ee8 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 @@ -20,6 +20,8 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RollbackRunResult; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.entity.AspectUtils; @@ -48,6 +50,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; @@ -265,10 +268,17 @@ public static Pair ingestProposal(com.linkedin.mxe.MetadataChan log.info("Proposal: {}", serviceProposal); Throwable exceptionally = null; try { - EntityService.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 = AspectsBatch.builder().mcps(proposalStream.collect(Collectors.toList()), + entityService.getEntityRegistry()).build(); + + Set> proposalResult = + entityService.ingestProposal(batch, auditStamp, false); + + Urn urn = proposalResult.stream().findFirst().get().getSecond().getUrn(); + return new Pair<>(urn.toString(), proposalResult.stream().anyMatch(resultPair -> resultPair.getSecond().isDidUpdate())); } 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/mock/MockEntityService.java b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java index 9838b00abd191..d9fd22b753340 100644 --- a/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java +++ b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java @@ -135,12 +135,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; @@ -152,7 +146,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, @@ -167,7 +161,7 @@ protected List> ingestAspectsToLocalDB(@Nonnull @Nonnull List> aspectRecordsToIngest, @Nonnull AuditStamp auditStamp, @Nonnull SystemMetadata providedSystemMetadata) { return Collections.emptyList(); - } + }*/ @Nullable @Override 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 0004c82b43a6f..a90c95767c85f 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,8 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.linkedin.aspect.GetTimeseriesAspectValuesResponse; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; import com.linkedin.metadata.resources.operations.Utils; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; @@ -37,9 +39,13 @@ import com.linkedin.restli.server.annotations.RestLiCollection; import com.linkedin.restli.server.annotations.RestMethod; import com.linkedin.restli.server.resources.CollectionResourceTaskTemplate; +import com.linkedin.util.Pair; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; +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; @@ -199,18 +205,26 @@ public Task ingestProposal( return RestliUtil.toTask(() -> { log.debug("Proposal: {}", metadataChangeProposal); try { - EntityService.IngestProposalResult result = _entityService.ingestProposal(metadataChangeProposal, auditStamp, asyncBool); - Urn responseUrn = result.getUrn(); + 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)); - } + AspectsBatch batch = AspectsBatch.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(); + + Set> results = + _entityService.ingestProposal(batch, auditStamp, asyncBool); + + EntityService.IngestProposalResult one = results.stream() + .map(Pair::getSecond) + .findFirst() + .get(); - if (!result.isQueued()) { - tryIndexRunId(responseUrn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); + // Update runIds + if (!one.isQueued()) { + tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); } - return responseUrn.toString(); + return urn.toString(); } catch (ValidationException e) { throw new RestLiServiceException(HttpStatus.S_422_UNPROCESSABLE_ENTITY, e.getMessage()); } diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java index 674609c7f67df..b9321c460b824 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java @@ -294,7 +294,7 @@ private void updateExecutionRequestStatus(String runId, String status) { proposal.setAspect(GenericRecordUtils.serializeAspect(requestResult)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(proposal, + _entityService.ingestSingleProposal(proposal, new AuditStamp().setActor(UrnUtils.getUrn(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); } } catch (Exception e) { From 3f13b99552503e866ae2b519038bfda6b8c736f2 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Wed, 19 Jul 2023 09:09:08 -0500 Subject: [PATCH 03/27] lint --- .../main/java/com/linkedin/metadata/entity/EntityService.java | 1 - .../src/main/java/com/linkedin/metadata/entity/EntityUtils.java | 1 - .../metadata/entity/cassandra/CassandraRetentionService.java | 1 - 3 files changed, 3 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java index 7245cc59c3b37..ceec91f97a33b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -47,7 +47,6 @@ 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.ValidationUtils; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; 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 2becb4778a5b1..5da71a79a3317 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 @@ -6,7 +6,6 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.entity.Entity; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.metadata.entity.validation.EntityRegistryUrnValidator; import com.linkedin.metadata.entity.validation.RecordTemplateValidator; 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 ad77e6f361a4c..1da8e3375035f 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 @@ -16,7 +16,6 @@ import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.entity.retention.BulkApplyRetentionArgs; import com.linkedin.metadata.entity.retention.BulkApplyRetentionResult; -import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.retention.DataHubRetentionConfig; import com.linkedin.retention.Retention; import com.linkedin.retention.TimeBasedRetention; From 496362db472ff9b75f11173b079262051adc964f Mon Sep 17 00:00:00 2001 From: David Leifker Date: Wed, 19 Jul 2023 11:07:19 -0500 Subject: [PATCH 04/27] fixing tests --- .../resolvers/term/AddTermsResolverTest.java | 1 - .../upgrade/nocode/DataMigrationStep.java | 6 ++++- .../ebean/transactions/AspectsBatch.java | 2 +- .../token/StatefulTokenService.java | 1 - .../DataHubTokenAuthenticatorTest.java | 1 + .../token/StatefulTokenServiceTest.java | 1 + .../test/resources/test-entity-registry.yaml | 10 +++++++ .../boot/steps/IngestPoliciesStep.java | 6 ++++- .../IngestDataPlatformInstancesStepTest.java | 1 + .../test/resources/test-entity-registry.yaml | 6 +++++ .../java/entities/EntitiesControllerTest.java | 10 +++++++ .../resources/entity/AspectResource.java | 26 +++++++++++++------ .../resources/entity/AspectResourceTest.java | 9 ++++++- 13 files changed, 66 insertions(+), 14 deletions(-) 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 a2d2693f716dd..94584c3750d2c 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 @@ -14,7 +14,6 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; 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 fb9272f68c37a..a3c92c0486415 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 @@ -21,7 +21,11 @@ import io.ebean.EbeanServer; import io.ebean.PagedList; import java.net.URISyntaxException; -import java.util.*; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.function.Function; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java index 7c5416f99563e..042564fc7eb87 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java @@ -64,7 +64,7 @@ private static AspectsBatchItem toAspectBatchItem(MetadataChangeProposal mcp, En Urn urn = mcp.getEntityUrn(); if (urn == null) { - urn = EntityKeyUtils.getUrnFromProposal(mcp, aspectSpec); + urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } return AspectsBatchItem.builder() 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 6c73c009db9bf..40cf3ae1d5b9a 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 @@ -20,7 +20,6 @@ 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; 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 99f242bac6b42..12c1307f996d6 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 @@ -159,6 +159,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); 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/metadata/boot/steps/IngestPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java index 1de96218332ec..67291946f83dd 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 @@ -27,7 +27,11 @@ import com.linkedin.policy.DataHubPolicyInfo; import java.io.IOException; import java.net.URISyntaxException; -import java.util.*; +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; 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 68618a9e87a79..f1950d3b798d9 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 @@ -142,6 +142,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/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/test/java/entities/EntitiesControllerTest.java b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java index 2dff636de2f27..4cf2758d86f70 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.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.AspectDao; +import com.linkedin.metadata.entity.EntityService; 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(), 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/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 a90c95767c85f..cf798b1461ade 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 @@ -43,6 +43,7 @@ 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; @@ -184,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 { @@ -205,12 +206,20 @@ public Task ingestProposal( return RestliUtil.toTask(() -> { log.debug("Proposal: {}", metadataChangeProposal); try { - Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), - AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService).stream()); + final AspectsBatch batch; + if (asyncBool) { + // if async we'll expand the additional changes later, no need to do this early + batch = AspectsBatch.builder() + .mcps(List.of(metadataChangeProposal), _entityService.getEntityRegistry()) + .build(); + } else { + Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), + AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService).stream()); - AspectsBatch batch = AspectsBatch.builder() - .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) - .build(); + batch = AspectsBatch.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(); + } Set> results = _entityService.ingestProposal(batch, auditStamp, asyncBool); @@ -221,10 +230,11 @@ public Task ingestProposal( .get(); // Update runIds + Urn resultUrn = one.getUrn(); if (!one.isQueued()) { - tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); + tryIndexRunId(resultUrn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); } - return urn.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 0d3642b91758a..12bd0c1628c13 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 @@ -22,6 +22,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; @@ -77,7 +79,12 @@ public void testAsyncDefaultAspects() throws URISyntaxException { reset(_producer, _aspectDao); when(_aspectDao.runInTransactionWithRetry(any(), anyInt())) - .thenReturn(new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0)); + .thenReturn(List.of( + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0))); _aspectResource.ingestProposal(mcp, "false"); verify(_producer, times(5)).produceMetadataChangeLog(eq(urn), any(AspectSpec.class), any(MetadataChangeLog.class)); verifyNoMoreInteractions(_producer); From dd55590e720d2bbdd07b8b4c104f9af50455861f Mon Sep 17 00:00:00 2001 From: David Leifker Date: Wed, 19 Jul 2023 15:35:15 -0500 Subject: [PATCH 05/27] fix-up patch --- .../metadata/entity/EntityService.java | 13 +---- .../ebean/transactions/AspectsBatch.java | 14 +++--- .../ebean/transactions/AspectsBatchItem.java | 49 +++++++++++++++---- 3 files changed, 49 insertions(+), 27 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java index ceec91f97a33b..7593619972244 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -584,9 +584,8 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa final UpdateAspectResult result; if (overwrite || latest == null) { - SystemMetadata systemMetadata = generateSystemMetadataIfEmpty(item.getSystemMetadata()); result = ingestAspectToLocalDBNoTransaction(item.getUrn(), item.getAspectName(), item.getLambda(), - auditStamp, systemMetadata, latest, nextVersion); + auditStamp, item.getSystemMetadata(), latest, nextVersion); // support inner-batch upserts latestAspects.computeIfAbsent(urnStr, key -> new HashMap<>()).put(item.getAspectName(), item.toLatestEntityAspect(auditStamp)); @@ -680,16 +679,6 @@ protected UpdateAspectResult patchAspectToLocalDB( }, 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; - } - /** * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. * diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java index 042564fc7eb87..a9505d8e0120b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java @@ -54,8 +54,6 @@ private static AspectsBatchItem toAspectBatchItem(MetadataChangeProposal mcp, En log.debug("entity type = {}", mcp.getEntityType()); EntitySpec entitySpec = entityRegistry.getEntitySpec(mcp.getEntityType()); AspectSpec aspectSpec = validateAspect(mcp, entitySpec); - RecordTemplate aspect = convertToRecordTemplate(mcp, aspectSpec); - log.debug("aspect = {}", aspect); if (!isValidChangeType(mcp.getChangeType(), aspectSpec)) { throw new UnsupportedOperationException("ChangeType not supported: " + mcp.getChangeType() @@ -67,13 +65,17 @@ private static AspectsBatchItem toAspectBatchItem(MetadataChangeProposal mcp, En urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } - return AspectsBatchItem.builder() + AspectsBatchItem.AspectsBatchItemBuilder builder = AspectsBatchItem.builder() .urn(urn) .aspectName(mcp.getAspectName()) .systemMetadata(mcp.getSystemMetadata()) - .mcp(mcp) - .value(aspect) - .build(entityRegistry); + .mcp(mcp); + + if (!mcp.getChangeType().equals(ChangeType.PATCH)) { + builder.value(convertToRecordTemplate(mcp, aspectSpec)); + } + + return builder.build(entityRegistry); } private static RecordTemplate convertToRecordTemplate(MetadataChangeProposal mcp, AspectSpec aspectSpec) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java index 426845642e5ba..f7f9d21b7a35b 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java @@ -5,6 +5,7 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.metadata.entity.AspectUtils; import com.linkedin.metadata.entity.EntityAspect; +import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -15,6 +16,8 @@ import lombok.Getter; import lombok.extern.slf4j.Slf4j; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.sql.Timestamp; import java.util.Objects; import java.util.Optional; @@ -71,11 +74,13 @@ public AspectsBatchItem build(EntityRegistry entityRegistry) { aspectSpec(AspectUtils.validate(this.entitySpec, this.aspectName)); log.debug("aspect spec = {}", this.aspectSpec); - RecordTemplate aspect = this.lambda.apply(null); - AspectUtils.validateRecordTemplate(entityRegistry, this.entitySpec, this.urn, aspect); + if (this.lambda != null) { + RecordTemplate aspect = this.lambda.apply(null); + AspectUtils.validateRecordTemplate(entityRegistry, this.entitySpec, this.urn, aspect); + } - return new AspectsBatchItem(this.urn, this.aspectName, this.systemMetadata, this.lambda, this.mcp, - this.entitySpec, this.aspectSpec); + return new AspectsBatchItem(this.urn, this.aspectName, generateSystemMetadataIfEmpty(this.systemMetadata), + this.lambda, this.mcp, this.entitySpec, this.aspectSpec); } private AspectsBatchItemBuilder entitySpec(EntitySpec entitySpec) { @@ -87,6 +92,16 @@ private AspectsBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { this.aspectSpec = aspectSpec; return this; } + + @Nonnull + private SystemMetadata generateSystemMetadataIfEmpty(@Nullable SystemMetadata systemMetadata) { + if (systemMetadata == null) { + systemMetadata = new SystemMetadata(); + systemMetadata.setRunId(EntityService.DEFAULT_RUN_ID); + systemMetadata.setLastObserved(System.currentTimeMillis()); + } + return systemMetadata; + } } @Override @@ -97,16 +112,32 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { return false; } + AspectsBatchItem that = (AspectsBatchItem) o; - return urn.equals(that.urn) && aspectName.equals(that.aspectName) - && Objects.equals(systemMetadata, that.systemMetadata) - && lambda.apply(null).equals(that.lambda.apply(null)) - && Objects.equals(mcp, that.mcp); + + if (!urn.equals(that.urn)) { + return false; + } + if (!aspectName.equals(that.aspectName)) { + return false; + } + if (!systemMetadata.equals(that.systemMetadata)) { + return false; + } + if (!Objects.equals(lambda, that.lambda)) { + return false; + } + return mcp.equals(that.mcp); } @Override public int hashCode() { - return Objects.hash(urn, aspectName, systemMetadata, lambda.apply(null), mcp); + int result = urn.hashCode(); + result = 31 * result + aspectName.hashCode(); + result = 31 * result + systemMetadata.hashCode(); + result = 31 * result + (lambda != null ? lambda.hashCode() : 0); + result = 31 * result + mcp.hashCode(); + return result; } @Override From e65d5403c41e8aa0af9d8a8c13c7929643a06f3f Mon Sep 17 00:00:00 2001 From: David Leifker Date: Sat, 22 Jul 2023 11:06:16 -0500 Subject: [PATCH 06/27] EntityService refactor --- .../upgrade/nocode/DataMigrationStep.java | 10 +- .../restorebackup/RestoreStorageStep.java | 5 +- docs/deploy/environment-vars.md | 2 +- .../dao/producer/KafkaEventProducer.java | 53 +- .../metadata/client/JavaEntityClient.java | 5 +- .../metadata/entity/DeleteEntityService.java | 4 +- .../metadata/entity/EntityService.java | 709 +++++++----------- .../metadata/entity/RetentionService.java | 2 +- .../metadata/entity/ebean/EbeanAspectDao.java | 20 +- .../entity/ebean/EbeanRetentionService.java | 3 +- .../ebean/transactions/AbstractBatchItem.java | 88 +++ .../ebean/transactions/AspectsBatch.java | 102 +-- .../ebean/transactions/AspectsBatchItem.java | 154 ---- .../ebean/transactions/PatchBatchItem.java | 182 +++++ .../ebean/transactions/UpsertBatchItem.java | 168 +++++ .../metadata/event/EventProducer.java | 25 - .../metadata/AspectIngestionUtils.java | 20 +- .../entity/EbeanEntityServiceTest.java | 30 +- .../metadata/entity/EntityServiceTest.java | 225 +++--- .../boot/steps/IndexDataPlatformsStep.java | 2 +- .../IngestDataPlatformInstancesStep.java | 8 +- .../boot/steps/IngestDataPlatformsStep.java | 8 +- .../metadata/boot/steps/IngestRolesStep.java | 2 +- .../steps/RestoreColumnLineageIndices.java | 4 +- .../boot/steps/RestoreDbtSiblingsIndices.java | 2 +- .../boot/steps/RestoreGlossaryIndices.java | 4 +- .../IngestDataPlatformInstancesStepTest.java | 5 +- .../RestoreColumnLineageIndicesTest.java | 18 +- .../steps/RestoreGlossaryIndicesTest.java | 12 +- .../openapi/util/MappingUtil.java | 7 +- .../src/test/java/mock/MockEntityService.java | 7 - .../resources/entity/AspectResource.java | 11 +- .../resources/entity/AspectResourceTest.java | 33 +- .../cypress/cypress/e2e/siblings/siblings.js | 2 +- smoke-test/tests/tags-and-terms/data.json | 39 + 35 files changed, 982 insertions(+), 989 deletions(-) create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java delete mode 100644 metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.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 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 a3c92c0486415..2133a2f36ee29 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 @@ -134,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. 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/docs/deploy/environment-vars.md b/docs/deploy/environment-vars.md index 0721aaa7964b7..addb47908f720 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/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java index d861fdbe46bc5..1b2a13844bf5e 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 @@ -51,7 +51,6 @@ import com.linkedin.parseq.retry.backoff.BackoffPolicy; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.r2.RemoteInvocationException; -import com.linkedin.util.Pair; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; @@ -543,8 +542,8 @@ public String ingestProposal(@Nonnull final MetadataChangeProposal metadataChang .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) .build(); - EntityService.IngestProposalResult one = _entityService.ingestProposal(batch, auditStamp, async).stream() - .findFirst().map(Pair::getSecond).get(); + EntityService.IngestResult one = _entityService.ingestProposal(batch, auditStamp, async).stream() + .findFirst().get(); Urn urn = one.getUrn(); tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata()); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java index a411dbb19524f..f95227373f556 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java +++ b/metadata-io/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 EntityService.IngestProposalResult ingestProposalResult = _entityService.ingestSingleProposal(proposal, auditStamp, false); + final EntityService.IngestResult ingestProposalResult = _entityService.ingestSingleProposal(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-io/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java index 7593619972244..2b943d88a5623 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java @@ -3,13 +3,6 @@ import com.codahale.metrics.Timer; 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.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterators; @@ -41,8 +34,10 @@ import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; +import com.linkedin.metadata.entity.ebean.transactions.AbstractBatchItem; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; +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; @@ -52,7 +47,6 @@ 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; @@ -69,29 +63,32 @@ import com.linkedin.mxe.SystemMetadata; import com.linkedin.util.Pair; import io.ebean.PagedList; -import java.io.IOException; + import java.net.URISyntaxException; -import java.nio.charset.StandardCharsets; import java.sql.Timestamp; import java.util.ArrayList; import java.util.Arrays; 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 lombok.Builder; import lombok.Value; import lombok.extern.slf4j.Slf4j; @@ -122,9 +119,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. */ @@ -136,15 +133,11 @@ public class 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()); - } - + @Builder(toBuilder = true) @Value public static class UpdateAspectResult { Urn urn; + UpsertBatchItem request; RecordTemplate oldValue; RecordTemplate newValue; SystemMetadata oldSystemMetadata; @@ -152,13 +145,20 @@ public static class UpdateAspectResult { MetadataAuditOperation operation; AuditStamp auditStamp; long maxVersion; + boolean processedMCL; + Future mclFuture; } + @Builder(toBuilder = true) @Value - public static class IngestProposalResult { + public static class IngestResult { Urn urn; - boolean didUpdate; - boolean queued; + AbstractBatchItem request; + boolean publishedMCL; + boolean processedMCL; + boolean publishedMCP; + boolean sqlCommitted; + boolean isUpdate; // update else insert } private static final int DEFAULT_MAX_TRANSACTION_RETRY = 3; @@ -514,11 +514,11 @@ public List ingestAspects(Urn entityUrn, List> pairList, @Nonnull final AuditStamp auditStamp, SystemMetadata systemMetadata) { - List items = pairList.stream() - .map(pair -> AspectsBatchItem.builder() + List items = pairList.stream() + .map(pair -> UpsertBatchItem.builder() .urn(entityUrn) .aspectName(pair.getKey()) - .value(pair.getValue()) + .aspect(pair.getValue()) .systemMetadata(systemMetadata) .build(_entityRegistry)) .collect(Collectors.toList()); @@ -526,30 +526,25 @@ public List ingestAspects(Urn entityUrn, } /** - * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataAuditEvent}. + * Ingests (inserts) a new version of an entity aspect & emits a {@link com.linkedin.mxe.MetadataChangeLog}. * * @param aspectsBatch aspects to write * @param auditStamp an {@link AuditStamp} containing metadata about the writer & current time - * @param emitMae whether a {@link com.linkedin.mxe.MetadataAuditEvent} should be emitted in correspondence upon + * @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 */ public List ingestAspects(@Nonnull final AspectsBatch aspectsBatch, @Nonnull final AuditStamp auditStamp, - boolean emitMae, + boolean emitMCL, boolean overwrite) { Timer.Context ingestToLocalDBTimer = MetricUtils.timer(this.getClass(), "ingestAspectsToLocalDB").time(); List ingestResults = ingestAspectsToLocalDB(aspectsBatch, auditStamp, overwrite); + List mclResults = emitMCL(ingestResults, emitMCL); ingestToLocalDBTimer.stop(); - if (emitMae) { - Streams.zip(aspectsBatch.getItems().stream(), ingestResults.stream(), (aspect, result) -> - sendEventForUpdateAspectResult(aspect.getUrn(), aspect.getAspectName(), result) - ).collect(Collectors.toList()); - } - - return ingestResults; + return mclResults; } /** @@ -567,6 +562,10 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa @Nonnull final AuditStamp auditStamp, boolean overwrite) { + if (aspectsBatch.containsDuplicateAspects()) { + log.warn(String.format("Batch contains duplicates: %s", aspectsBatch)); + } + return _aspectDao.runInTransactionWithRetry((tx) -> { // Read before write is unfortunate, however batch it Map> urnAspects = aspectsBatch.getUrnAspectsMap(); @@ -575,8 +574,24 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa // read #2 Map> nextVersions = _aspectDao.getNextVersions(urnAspects); - // Upsert results - List> upsertResults = aspectsBatch.getItems().stream() + 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()); @@ -584,8 +599,8 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa final UpdateAspectResult result; if (overwrite || latest == null) { - result = ingestAspectToLocalDBNoTransaction(item.getUrn(), item.getAspectName(), item.getLambda(), - auditStamp, item.getSystemMetadata(), latest, nextVersion); + result = ingestAspectToLocalDBNoTransaction(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)); @@ -594,14 +609,23 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa RecordTemplate oldValue = EntityUtils.toAspectRecord(item.getUrn().getEntityType(), item.getAspectName(), latest.getMetadata(), getEntityRegistry()); SystemMetadata oldMetadata = EntityUtils.parseSystemMetadata(latest.getSystemMetadata()); - result = new UpdateAspectResult(item.getUrn(), oldValue, oldValue, oldMetadata, oldMetadata, MetadataAuditOperation.UPDATE, auditStamp, - latest.getVersion()); + 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 Pair.of(item, result); + return result; }).collect(Collectors.toList()); - // commit upserts prior to retention, if supported by impl + // commit upserts prior to retention or kafka send, if supported by impl if (tx != null) { tx.commitAndContinue(); } @@ -610,18 +634,18 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa if (_retentionService != null) { List retentionBatch = upsertResults.stream() // Only consider retention when there was a previous version - .filter(resultPair -> latestAspects.containsKey(resultPair.getKey().getUrn().toString()) - && latestAspects.get(resultPair.getKey().getUrn().toString()).containsKey(resultPair.getKey().getAspectName())) - .filter(resultPair -> { - RecordTemplate oldAspect = resultPair.getSecond().getOldValue(); - RecordTemplate newAspect = resultPair.getSecond().getNewValue(); + .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(resultPair -> RetentionService.RetentionContext.builder() - .urn(resultPair.getKey().getUrn()) - .aspectName(resultPair.getKey().getAspectName()) - .maxVersion(Optional.of(resultPair.getValue().getMaxVersion())) + .map(result -> RetentionService.RetentionContext.builder() + .urn(result.getUrn()) + .aspectName(result.getRequest().getAspectName()) + .maxVersion(Optional.of(result.getMaxVersion())) .build()) .collect(Collectors.toList()); _retentionService.applyRetentionWithPolicyDefaults(retentionBatch); @@ -629,54 +653,39 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa log.warn("Retention service is missing!"); } - return upsertResults.stream().map(Pair::getValue).collect(Collectors.toList()); + return upsertResults; }, DEFAULT_MAX_TRANSACTION_RETRY); } - /** - * Apply patch update to aspect within a single transaction - * - * @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 - * @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((tx) -> { - 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 { - - 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)); - } + private List emitMCL(List sqlResults, boolean emitMCL) { + List withEmitMCL = sqlResults.stream() + .map(result -> emitMCL ? conditionallyProduceMCLAsync(result) : result) + .collect(Collectors.toList()); - final RecordTemplate updatedValue = _entityRegistry.getAspectTemplateEngine().applyPatch(currentValue, jsonPatch, aspectSpec); + // 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"); + } - AspectUtils.validateRecordTemplate(_entityRegistry, urn, updatedValue); - return ingestAspectToLocalDBNoTransaction(urn, aspectName, ignored -> updatedValue, auditStamp, providedSystemMetadata, - latest, nextVersion); - } catch (JsonProcessingException | JsonPatchException e) { - throw new IllegalStateException(e); - } - }, DEFAULT_MAX_TRANSACTION_RETRY); + return withEmitMCL; } /** @@ -703,10 +712,10 @@ public RecordTemplate ingestAspectIfNotPresent(@Nonnull Urn urn, log.debug("Invoked ingestAspectIfNotPresent with urn: {}, aspectName: {}, newValue: {}", urn, aspectName, newValue); AspectsBatch aspectsBatch = AspectsBatch.builder() - .aspect(AspectsBatchItem.builder() + .one(UpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) - .value(newValue) + .aspect(newValue) .systemMetadata(systemMetadata) .build(_entityRegistry)) .build(); @@ -715,67 +724,16 @@ public RecordTemplate ingestAspectIfNotPresent(@Nonnull Urn urn, return ingested.stream().findFirst().get().getNewValue(); } - 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.applyRetentionWithPolicyDefaults(List.of( - RetentionService.RetentionContext.builder() - .urn(urn) - .aspectName(aspectName) - .maxVersion(Optional.of(result.maxVersion)) - .build())); - } - - // 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; - } - /** * 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 IngestProposalResult} containing the results + * @return an {@link IngestResult} containing the results */ - public IngestProposalResult ingestSingleProposal(MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async) { - return ingestProposal(AspectsBatch.builder() - .mcps(List.of(proposal), getEntityRegistry()) - .build(), auditStamp, async).stream().findFirst().get().getValue(); + public IngestResult ingestSingleProposal(MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async) { + return ingestProposal(AspectsBatch.builder().mcps(List.of(proposal), getEntityRegistry()).build(), auditStamp, async) + .stream().findFirst().get(); } /** @@ -788,128 +746,123 @@ public IngestProposalResult ingestSingleProposal(MetadataChangeProposal proposal * @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 */ - public Set> ingestProposal(AspectsBatch aspectsBatch, - AuditStamp auditStamp, - final boolean async) { + public Set ingestProposal(AspectsBatch aspectsBatch, AuditStamp auditStamp, final boolean async) { - Stream> timeseriesIngestResults = ingestTimeseriesProposal(aspectsBatch, auditStamp); - Stream> nonTimeseriesIngestResults = async ? ingestProposalAsync(aspectsBatch) + Stream timeseriesIngestResults = ingestTimeseriesProposal(aspectsBatch, auditStamp); + Stream nonTimeseriesIngestResults = async ? ingestProposalAsync(aspectsBatch) : ingestProposalSync(aspectsBatch, auditStamp); return Stream.concat(timeseriesIngestResults, nonTimeseriesIngestResults).collect(Collectors.toSet()); } - private Stream> ingestTimeseriesProposal(AspectsBatch aspectsBatch, AuditStamp auditStamp) { - List unsupported = aspectsBatch.getItems().stream() - .filter(item -> item.getAspectSpec().isTimeseries() && item.getMcp().getChangeType() != ChangeType.UPSERT) + /** + * 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(item -> item.getMcp().getChangeType()).collect(Collectors.toSet())); + .map(AbstractBatchItem::getChangeType).collect(Collectors.toSet())); } - List timeseries = aspectsBatch.getItems().stream() + 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 timeseries.stream().map(item -> { - boolean didUpdate = emitChangeLog(null, null, item.getAspect(), item.getSystemMetadata(), - item.getMcp(), item.getUrn(), auditStamp, item.getAspectSpec()); - return Pair.of(item, new IngestProposalResult(item.getUrn(), didUpdate, false)); + 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 Stream> ingestProposalAsync(AspectsBatch aspectsBatch) { - List nonTimeseries = aspectsBatch.getItems().stream() + /** + * 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()); - return nonTimeseries.stream().map(item -> { + 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.getMcp()); - return Pair.of(item, new IngestProposalResult(item.getUrn(), false, true)); - }); - } - - private Stream> ingestProposalSync(AspectsBatch aspectsBatch, AuditStamp auditStamp) { - List unsupported = aspectsBatch.getItems().stream() - .filter(item -> item.getMcp().getChangeType() != ChangeType.PATCH - && item.getMcp().getChangeType() != ChangeType.UPSERT) + _producer.produceMetadataChangeProposal(item.getUrn(), item.getMetadataChangeProposal())) + .filter(Objects::nonNull) .collect(Collectors.toList()); - if (!unsupported.isEmpty()) { - throw new UnsupportedOperationException("ChangeType not supported: " + unsupported.stream() - .map(item -> item.getMcp().getChangeType()).collect(Collectors.toSet())); + + try { + 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); + } + }); } + } - Stream> patches = aspectsBatch.getItems().stream() - .filter(item -> item.getMcp().getChangeType() == ChangeType.PATCH) - .map(item -> Pair.of(item, performPatch(item.getMcp(), item.getAspectSpec(), item.getSystemMetadata(), - item.getUrn(), auditStamp))); + private Stream ingestProposalSync(AspectsBatch aspectsBatch, AuditStamp auditStamp) { + AspectsBatch nonTimeseries = AspectsBatch.builder() + .items(aspectsBatch.getItems().stream() + .filter(item -> !item.getAspectSpec().isTimeseries()) + .collect(Collectors.toList())) + .build(); - List upsertItems = aspectsBatch.getItems() - .stream().filter(item -> item.getMcp().getChangeType() == ChangeType.UPSERT) + List unsupported = nonTimeseries.getItems().stream() + .filter(item -> item.getMetadataChangeProposal().getChangeType() != ChangeType.PATCH + && item.getMetadataChangeProposal().getChangeType() != ChangeType.UPSERT) .collect(Collectors.toList()); - List upsertResults = ingestAspects(AspectsBatch.builder().items(upsertItems).build(), - auditStamp, false, true); - Stream> upserts = Streams.zip(upsertItems.stream(), upsertResults.stream(), Pair::of); - - return Stream.concat(patches, upserts).map(resultPair -> { - AspectsBatchItem item = resultPair.getFirst(); - UpdateAspectResult result = resultPair.getSecond(); - boolean didUpdate = emitChangeLog(result.getOldValue(), result.getOldSystemMetadata(), - result.getNewValue(), result.getNewSystemMetadata(), item.getMcp(), item.getUrn(), - auditStamp, item.getAspectSpec()); - return Pair.of(item, new IngestProposalResult(item.getUrn(), didUpdate, false)); - }); - } - - 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."); + if (!unsupported.isEmpty()) { + throw new UnsupportedOperationException("ChangeType not supported: " + unsupported.stream() + .map(item -> item.getMetadataChangeProposal().getChangeType()).collect(Collectors.toSet())); } - Patch jsonPatch = convertToJsonPatch(mcp); - log.debug("patch = {}", jsonPatch); - return patchAspect(jsonPatch, systemMetadata, entityUrn, auditStamp, aspectSpec); - } + List upsertResults = ingestAspects(nonTimeseries, auditStamp, true, true); - private boolean supportsPatch(AspectSpec aspectSpec) { - // Limit initial support to defined templates - return AspectTemplateEngine.SUPPORTED_TEMPLATES.contains(aspectSpec.getName()); - } - - 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); - } - } + return upsertResults.stream().map(result -> { + AbstractBatchItem item = result.getRequest(); - private UpdateAspectResult patchAspect(final Patch patch, final SystemMetadata systemMetadata, 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.applyRetentionWithPolicyDefaults(List.of(RetentionService.RetentionContext.builder() - .urn(entityUrn) - .aspectName(aspectSpec.getName()) - .maxVersion(Optional.of(result.maxVersion)) - .build())); - } - return result; + return IngestResult.builder() + .urn(item.getUrn()) + .request(item) + .publishedMCL(result.getMclFuture() != null) + .sqlCommitted(true) + .isUpdate(result.getOldValue() != null) + .build(); + }); } public String batchApplyRetention(Integer start, Integer count, Integer attemptWithVersion, String aspectName, @@ -933,43 +886,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; } public Integer getCountAspect(@Nonnull String aspectName, @Nullable String urnLike) { @@ -991,6 +920,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; @@ -1054,14 +985,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) { @@ -1072,39 +1010,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 - */ - 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. * @@ -1165,57 +1070,58 @@ public Map getEntities(@Nonnull final Set urns, @Nonnull Set toEntity(entry.getValue()))); } - public 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) { - log.debug(String.format("Producing MetadataAuditEvent for ingested aspect %s, urn %s", aspectName, urn)); - if (aspectName.equals(getKeyAspectName(urn))) { - produceMetadataAuditEventForKey(urn, newSystemMetadata); - } else { - final Snapshot newSnapshot = buildSnapshot(urn, newAspectValue); - Snapshot oldSnapshot = null; - if (oldAspectValue != null) { - oldSnapshot = buildSnapshot(urn, oldAspectValue); - } - _producer.produceMetadataAuditEvent(urn, oldSnapshot, newSnapshot, oldSystemMetadata, newSystemMetadata, - operation); - } + public Pair, 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 = EntityUtils.buildKeyAspect(_entityRegistry, urn); - return toSnapshotUnion(toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, keyAspectValue)))); + 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); } - 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 - */ - 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()); - 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); } public RecordTemplate getLatestAspect(@Nonnull final Urn urn, @Nonnull final String aspectName) { @@ -1351,10 +1257,10 @@ private void ingestSnapshotUnion(@Nonnull final Snapshot snapshotUnion, @Nonnull aspectRecordsToIngest.stream().map(Pair::getFirst).collect(Collectors.toSet()))); AspectsBatch aspectsBatch = AspectsBatch.builder() - .items(aspectRecordsToIngest.stream().map(pair -> AspectsBatchItem.builder() + .items(aspectRecordsToIngest.stream().map(pair -> UpsertBatchItem.builder() .urn(urn) .aspectName(pair.getKey()) - .value(pair.getValue()) + .aspect(pair.getValue()) .systemMetadata(systemMetadata) .build(_entityRegistry)).collect(Collectors.toList())) .build(); @@ -1362,17 +1268,6 @@ private void ingestSnapshotUnion(@Nonnull final Snapshot snapshotUnion, @Nonnull ingestAspects(aspectsBatch, auditStamp, true, true); } - 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 = EntityUtils.buildKeyAspect(_entityRegistry, urn); - return toSnapshotUnion( - toSnapshotRecord(urn, ImmutableList.of(toAspectUnion(urn, keyAspectValue), toAspectUnion(urn, aspectValue)))); - } - public AspectSpec getKeyAspectSpec(@Nonnull final Urn urn) { return getKeyAspectSpec(urnToEntityName(urn)); } @@ -1493,7 +1388,7 @@ public RollbackRunResult rollbackWithConditions(List aspectRow rowsDeletedFromEntityDeletion.addAndGet(result.additionalRowsAffected); removedAspects.add(aspectToRemove); - produceMetadataChangeLog(result.getUrn(), result.getEntityName(), result.getAspectName(), aspectSpec.get(), + alwaysProduceMCLAsync(result.getUrn(), result.getEntityName(), result.getAspectName(), aspectSpec.get(), result.getOldValue(), result.getNewValue(), result.getOldSystemMetadata(), result.getNewSystemMetadata(), // TODO: use properly attributed audit stamp. createSystemAuditStamp(), @@ -1535,7 +1430,7 @@ public RollbackRunResult deleteUrn(Urn urn) { rowsDeletedFromEntityDeletion = result.additionalRowsAffected; removedAspects.add(summary); - produceMetadataChangeLog(result.getUrn(), result.getEntityName(), result.getAspectName(), keySpec, + alwaysProduceMCLAsync(result.getUrn(), result.getEntityName(), result.getAspectName(), keySpec, result.getOldValue(), result.getNewValue(), result.getOldSystemMetadata(), result.getNewSystemMetadata(), // TODO: Use a proper inferred audit stamp createSystemAuditStamp(), @@ -1827,7 +1722,7 @@ private EnvelopedAspect getKeyEnvelopedAspect(final Urn urn) { private UpdateAspectResult ingestAspectToLocalDBNoTransaction( @Nonnull final Urn urn, @Nonnull final String aspectName, - @Nonnull final Function, RecordTemplate> updateLambda, + @Nonnull final RecordTemplate newValue, @Nonnull final AuditStamp auditStamp, @Nonnull final SystemMetadata providedSystemMetadata, @Nullable final EntityAspect latest, @@ -1836,7 +1731,6 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction( // 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 @@ -1846,15 +1740,23 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction( latest.setSystemMetadata(RecordUtils.toJsonString(latestSystemMetadata)); + log.info("Ingesting aspect with name {}, urn {}", aspectName, urn); _aspectDao.saveAspect(latest, false); - 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); + log.info("Ingesting aspect with name {}, urn {}", aspectName, urn); long versionOfOld = _aspectDao.saveLatestAspect(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(), @@ -1862,58 +1764,16 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction( auditStamp.hasImpersonator() ? auditStamp.getImpersonator().toString() : null, new Timestamp(auditStamp.getTime()), EntityUtils.toJsonAspect(providedSystemMetadata), nextVersion); - return new UpdateAspectResult(urn, oldValue, newValue, - latest == null ? null : EntityUtils.parseSystemMetadata(latest.getSystemMetadata()), providedSystemMetadata, - MetadataAuditOperation.UPDATE, auditStamp, versionOfOld); - } - - @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((tx) -> { - - 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); - _aspectDao.saveAspect(urn.toString(), aspectName, EntityUtils.toJsonAspect(value), auditStamp.getActor().toString(), - auditStamp.hasImpersonator() ? auditStamp.getImpersonator().toString() : null, - new Timestamp(auditStamp.getTime()), EntityUtils.toJsonAspect(newSystemMetadata), version, oldAspect == null); - - 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(); } /** @@ -1980,13 +1840,6 @@ 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 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/RetentionService.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java index 5627b53a49a0a..5a6ac176bcc77 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/RetentionService.java @@ -116,7 +116,7 @@ public boolean setRetention(@Nullable String entityName, @Nullable String aspect .build(); return getEntityService().ingestProposal(batch, auditStamp, false).stream() - .anyMatch(resultPair -> resultPair.getSecond().isDidUpdate()); + .anyMatch(EntityService.IngestResult::isSqlCommitted); } /** 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 cca85c1204b97..2cfebaeab9839 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 @@ -14,20 +14,9 @@ import com.linkedin.metadata.query.ExtraInfoArray; import com.linkedin.metadata.query.ListResultMetadata; import com.linkedin.metadata.search.utils.QueryUtils; -import io.ebean.DuplicateKeyException; -import io.ebean.EbeanServer; -import io.ebean.ExpressionList; -import io.ebean.Junction; -import io.ebean.PagedList; -import io.ebean.Query; -import io.ebean.RawSql; -import io.ebean.RawSqlBuilder; -import io.ebean.Transaction; -import io.ebean.TxScope; -import io.ebean.annotation.TxIsolation; +import io.ebean.*; import io.ebean.annotation.Platform; -import io.ebean.config.dbplatform.DatabasePlatform; -import io.ebean.plugin.SpiServer; +import io.ebean.annotation.TxIsolation; import java.net.URISyntaxException; import java.sql.SQLException; import java.sql.Timestamp; @@ -43,9 +32,12 @@ import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import javax.persistence.RollbackException; import javax.persistence.PersistenceException; +import javax.persistence.RollbackException; import javax.persistence.Table; + +import io.ebean.config.dbplatform.DatabasePlatform; +import io.ebean.plugin.SpiServer; import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; 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 3c62ded6ce5bc..21bbc38430436 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 @@ -17,6 +17,7 @@ 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; @@ -134,7 +135,7 @@ 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); diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java new file mode 100644 index 0000000000000..a51bf0c943687 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java @@ -0,0 +1,88 @@ +package com.linkedin.metadata.entity.ebean.transactions; + + +import com.linkedin.common.urn.Urn; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.registry.template.AspectTemplateEngine; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.mxe.SystemMetadata; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +public abstract class AbstractBatchItem { + // urn an urn associated with the new aspect + public abstract Urn getUrn(); + + // aspectName name of the aspect being inserted + public abstract String getAspectName(); + + public abstract SystemMetadata getSystemMetadata(); + + public abstract ChangeType getChangeType(); + + public abstract EntitySpec getEntitySpec(); + + public abstract AspectSpec getAspectSpec(); + + public abstract MetadataChangeProposal getMetadataChangeProposal(); + + @Nonnull + protected static SystemMetadata generateSystemMetadataIfEmpty(@Nullable SystemMetadata systemMetadata) { + if (systemMetadata == null) { + systemMetadata = new SystemMetadata(); + systemMetadata.setRunId(EntityService.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-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java index a9505d8e0120b..17037a4bced68 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java @@ -1,16 +1,8 @@ package com.linkedin.metadata.entity.ebean.transactions; -import com.datahub.util.exception.ModelConversionException; -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.entity.validation.ValidationUtils; -import com.linkedin.metadata.models.AspectSpec; -import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.events.metadata.ChangeType; 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.events.metadata.ChangeType; import lombok.Builder; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -25,7 +17,7 @@ @Getter @Builder(toBuilder = true) public class AspectsBatch { - private final List items; + private final List items; public Map> getUrnAspectsMap() { return items.stream() @@ -33,96 +25,32 @@ public Map> getUrnAspectsMap() { .collect(Collectors.groupingBy(Map.Entry::getKey, Collectors.mapping(Map.Entry::getValue, Collectors.toSet()))); } - public static class AspectsBatchBuilder { + public boolean containsDuplicateAspects() { + return items.stream().map(i -> String.format("%s_%s", i.getClass().getName(), i.hashCode())) + .distinct().count() != items.size(); + } + public static class AspectsBatchBuilder { /** * Just one aspect record template * @param data aspect data * @return builder */ - public AspectsBatchBuilder aspect(AspectsBatchItem data) { + public AspectsBatchBuilder one(AbstractBatchItem data) { this.items = List.of(data); return this; } public AspectsBatchBuilder mcps(List mcps, EntityRegistry entityRegistry) { - this.items = mcps.stream().map(mcp -> toAspectBatchItem(mcp, entityRegistry)).collect(Collectors.toList()); + 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; } - - private static AspectsBatchItem toAspectBatchItem(MetadataChangeProposal mcp, EntityRegistry entityRegistry) { - log.debug("entity type = {}", mcp.getEntityType()); - EntitySpec entitySpec = entityRegistry.getEntitySpec(mcp.getEntityType()); - AspectSpec aspectSpec = validateAspect(mcp, entitySpec); - - if (!isValidChangeType(mcp.getChangeType(), 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()); - } - - AspectsBatchItem.AspectsBatchItemBuilder builder = AspectsBatchItem.builder() - .urn(urn) - .aspectName(mcp.getAspectName()) - .systemMetadata(mcp.getSystemMetadata()) - .mcp(mcp); - - if (!mcp.getChangeType().equals(ChangeType.PATCH)) { - builder.value(convertToRecordTemplate(mcp, aspectSpec)); - } - - return builder.build(entityRegistry); - } - - 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; - } - - private 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 - */ - private static 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 diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java deleted file mode 100644 index f7f9d21b7a35b..0000000000000 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchItem.java +++ /dev/null @@ -1,154 +0,0 @@ -package com.linkedin.metadata.entity.ebean.transactions; - -import com.linkedin.common.AuditStamp; -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.entity.AspectUtils; -import com.linkedin.metadata.entity.EntityAspect; -import com.linkedin.metadata.entity.EntityService; -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 com.linkedin.mxe.MetadataChangeProposal; -import com.linkedin.mxe.SystemMetadata; -import lombok.Builder; -import lombok.Getter; -import lombok.extern.slf4j.Slf4j; - -import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import java.sql.Timestamp; -import java.util.Objects; -import java.util.Optional; -import java.util.function.Function; - -import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; - -@Slf4j -@Getter -@Builder(toBuilder = true) -public class AspectsBatchItem { - // 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; - // updateLambda Function to apply to the latest version of the aspect to get the updated version - private final Function, RecordTemplate> lambda; - - private final MetadataChangeProposal mcp; - - // derived - private final EntitySpec entitySpec; - private final AspectSpec aspectSpec; - - public RecordTemplate getAspect() { - return lambda.apply(null); - } - - public EntityAspect toLatestEntityAspect(AuditStamp auditStamp) { - EntityAspect latest = new EntityAspect(); - latest.setAspect(aspectName); - latest.setMetadata(EntityUtils.toJsonAspect(getAspect())); - latest.setUrn(urn.toString()); - latest.setVersion(ASPECT_LATEST_VERSION); - latest.setCreatedOn(new Timestamp(auditStamp.getTime())); - latest.setCreatedBy(auditStamp.getActor().toString()); - return latest; - } - - public static class AspectsBatchItemBuilder { - public AspectsBatchItemBuilder value(final RecordTemplate recordTemplate) { - this.lambda = ignored -> recordTemplate; - return this; - } - - public AspectsBatchItem 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(AspectUtils.validate(this.entitySpec, this.aspectName)); - log.debug("aspect spec = {}", this.aspectSpec); - - if (this.lambda != null) { - RecordTemplate aspect = this.lambda.apply(null); - AspectUtils.validateRecordTemplate(entityRegistry, this.entitySpec, this.urn, aspect); - } - - return new AspectsBatchItem(this.urn, this.aspectName, generateSystemMetadataIfEmpty(this.systemMetadata), - this.lambda, this.mcp, this.entitySpec, this.aspectSpec); - } - - private AspectsBatchItemBuilder entitySpec(EntitySpec entitySpec) { - this.entitySpec = entitySpec; - return this; - } - - private AspectsBatchItemBuilder aspectSpec(AspectSpec aspectSpec) { - this.aspectSpec = aspectSpec; - return this; - } - - @Nonnull - private SystemMetadata generateSystemMetadataIfEmpty(@Nullable SystemMetadata systemMetadata) { - if (systemMetadata == null) { - systemMetadata = new SystemMetadata(); - systemMetadata.setRunId(EntityService.DEFAULT_RUN_ID); - systemMetadata.setLastObserved(System.currentTimeMillis()); - } - return systemMetadata; - } - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - - AspectsBatchItem that = (AspectsBatchItem) o; - - if (!urn.equals(that.urn)) { - return false; - } - if (!aspectName.equals(that.aspectName)) { - return false; - } - if (!systemMetadata.equals(that.systemMetadata)) { - return false; - } - if (!Objects.equals(lambda, that.lambda)) { - return false; - } - return mcp.equals(that.mcp); - } - - @Override - public int hashCode() { - int result = urn.hashCode(); - result = 31 * result + aspectName.hashCode(); - result = 31 * result + systemMetadata.hashCode(); - result = 31 * result + (lambda != null ? lambda.hashCode() : 0); - result = 31 * result + mcp.hashCode(); - return result; - } - - @Override - public String toString() { - return "AspectsBatchItem{" - + "urn=" + urn - + ", aspectName='" - + aspectName + '\'' - + ", systemMetadata=" + systemMetadata - + ", lambda=" + lambda - + ", mcp=" + mcp - + '}'; - } -} 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..b5edccbc57fd8 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java @@ -0,0 +1,182 @@ +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.AspectUtils; +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 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; + } + + 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(AspectUtils.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..9c89858f33c31 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java @@ -0,0 +1,168 @@ +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.AspectUtils; +import com.linkedin.metadata.entity.EntityAspect; +import com.linkedin.metadata.entity.EntityUtils; +import com.linkedin.metadata.entity.validation.ValidationUtils; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.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; + } + + 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(AspectUtils.validate(this.entitySpec, this.aspectName)); + log.debug("aspect spec = {}", this.aspectSpec); + + AspectUtils.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/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 d37eb39a81076..cf1e262a04343 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java @@ -6,7 +6,7 @@ import com.linkedin.identity.CorpUserInfo; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; +import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import com.linkedin.metadata.key.CorpUserKey; import java.util.HashMap; import java.util.LinkedList; @@ -29,15 +29,15 @@ 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<>(); + List items = new LinkedList<>(); for (int i = startIndex; i < startIndex + aspectCount; i++) { Urn urn = UrnUtils.getUrn(String.format("urn:li:corpuser:tester%d", i)); CorpUserKey aspect = AspectGenerationUtils.createCorpUserKey(urn); aspects.put(urn, aspect); - items.add(AspectsBatchItem.builder() + items.add(UpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) - .value(aspect) + .aspect(aspect) .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService.getEntityRegistry())); } @@ -54,16 +54,16 @@ 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<>(); + List items = new LinkedList<>(); for (int i = startIndex; i < startIndex + aspectCount; i++) { Urn urn = UrnUtils.getUrn(String.format("urn:li:corpuser:tester%d", i)); String email = String.format("email%d@test.com", i); CorpUserInfo aspect = AspectGenerationUtils.createCorpUserInfo(email); aspects.put(urn, aspect); - items.add(AspectsBatchItem.builder() + items.add(UpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) - .value(aspect) + .aspect(aspect) .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService.getEntityRegistry())); } @@ -80,17 +80,17 @@ 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<>(); + 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); - items.add(AspectsBatchItem.builder() + items.add(UpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) - .value(aspect) + .aspect(aspect) .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService.getEntityRegistry())); } 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 83f0fc723bb3e..82642fb1ce1fe 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 @@ -11,7 +11,7 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectDao; import com.linkedin.metadata.entity.ebean.EbeanRetentionService; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; +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; @@ -97,23 +97,23 @@ public void testIngestListLatestAspects() throws AssertionError { // Ingest CorpUserInfo Aspect #3 CorpUserInfo writeAspect3 = AspectGenerationUtils.createCorpUserInfo("email3@test.com"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) - .value(writeAspect3) + .aspect(writeAspect3) .systemMetadata(metadata1) .build(_testEntityRegistry) ); @@ -161,23 +161,23 @@ public void testIngestListUrns() throws AssertionError { // Ingest CorpUserInfo Aspect #3 RecordTemplate writeAspect3 = AspectGenerationUtils.createCorpUserKey(entityUrn3); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) - .value(writeAspect3) + .aspect(writeAspect3) .systemMetadata(metadata1) .build(_testEntityRegistry) ); 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 c1e3c852fab9c..e2835a74d6c9a 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 @@ -33,7 +33,7 @@ import com.linkedin.metadata.aspect.CorpUserAspectArray; import com.linkedin.metadata.aspect.VersionedAspect; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; +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; @@ -63,7 +63,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.stream.Collectors; import javax.annotation.Nonnull; import org.junit.Assert; import org.mockito.ArgumentCaptor; @@ -149,9 +148,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); } @@ -186,9 +182,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); } @@ -257,12 +250,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); } @@ -317,12 +304,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()); @@ -385,12 +366,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()); @@ -427,8 +402,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); } @@ -479,8 +452,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); @@ -489,8 +460,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); @@ -542,8 +511,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); @@ -552,8 +519,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); @@ -705,9 +670,9 @@ public void testUpdateGetAspect() throws AssertionError { CorpUserInfo writeAspect = AspectGenerationUtils.createCorpUserInfo("email@test.com"); // Validate retrieval of CorpUserInfo Aspect #1 - _entityService.updateAspect(entityUrn, "corpuser", aspectName, corpUserInfoSpec, writeAspect, TEST_AUDIT_STAMP, 1, - true); - RecordTemplate readAspect1 = _entityService.getAspect(entityUrn, aspectName, 1); + _entityService.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect)), TEST_AUDIT_STAMP, null); + + RecordTemplate readAspect1 = _entityService.getAspect(entityUrn, aspectName, 0); assertTrue(DataTemplateUtil.areEqual(writeAspect, readAspect1)); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), Mockito.any()); @@ -716,10 +681,13 @@ public void testUpdateGetAspect() throws AssertionError { writeAspect.setEmail("newemail@test.com"); // Validate retrieval of CorpUserInfo Aspect #2 - _entityService.updateAspect(entityUrn, "corpuser", aspectName, corpUserInfoSpec, writeAspect, TEST_AUDIT_STAMP, 1, - false); - RecordTemplate readAspect2 = _entityService.getAspect(entityUrn, aspectName, 1); + _entityService.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect)), TEST_AUDIT_STAMP, null); + + RecordTemplate readAspect2 = _entityService.getAspect(entityUrn, aspectName, 0); assertTrue(DataTemplateUtil.areEqual(writeAspect, readAspect2)); + verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), + Mockito.any()); + verifyNoMoreInteractions(_mockProducer); } @@ -732,26 +700,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 - _entityService.updateAspect(entityUrn, "corpuser", aspectName, corpUserInfoSpec, writeAspect, TEST_AUDIT_STAMP, 1, - true); + _entityService.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 = _entityService.getVersionedAspect(entityUrn, aspectName, 1); - assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect, readAspect1)); + VersionedAspect readAspect1 = _entityService.getVersionedAspect(entityUrn, aspectName, 0); + assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), Mockito.any()); - VersionedAspect readAspect2 = _entityService.getVersionedAspect(entityUrn, aspectName, -1); - assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect, readAspect2)); + readAspect1 = _entityService.getVersionedAspect(entityUrn, aspectName, -1); + assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); + + // Validate retrieval of CorpUserInfo Aspect #2 + _entityService.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 = _entityService.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 = _entityService.getVersionedAspect(entityUrn, aspectName, 0); - assertFalse(DataTemplateUtil.areEqual(writtenVersionedAspect, readAspectVersion0)); + readAspect1 = _entityService.getVersionedAspect(entityUrn, aspectName, -1); + assertFalse(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); verifyNoMoreInteractions(_mockProducer); } @@ -779,29 +760,29 @@ public void testRollbackAspect() throws AssertionError { // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) - .value(writeAspect3) + .aspect(writeAspect3) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1Overwrite) + .aspect(writeAspect1Overwrite) .systemMetadata(metadata2) .build(_testEntityRegistry) ); @@ -853,23 +834,23 @@ public void testRollbackKey() throws AssertionError { // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(keyAspectName) - .value(writeKey1) + .aspect(writeKey1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1Overwrite) + .aspect(writeAspect1Overwrite) .systemMetadata(metadata2) .build(_testEntityRegistry) ); @@ -929,35 +910,35 @@ public void testRollbackUrn() throws AssertionError { // Ingest CorpUserInfo Aspect #1 Overwrite CorpUserInfo writeAspect1Overwrite = AspectGenerationUtils.createCorpUserInfo("email1.overwrite@test.com"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(keyAspectName) - .value(writeKey1) + .aspect(writeKey1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn2) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn3) .aspectName(aspectName) - .value(writeAspect3) + .aspect(writeAspect3) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn1) .aspectName(aspectName) - .value(writeAspect1Overwrite) + .aspect(writeAspect1Overwrite) .systemMetadata(metadata2) .build(_testEntityRegistry) ); @@ -991,11 +972,11 @@ public void testIngestGetLatestAspect() throws AssertionError { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry) ); @@ -1013,9 +994,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); @@ -1024,10 +1002,10 @@ public void testIngestGetLatestAspect() throws AssertionError { CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); items = List.of( - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata2) .build(_testEntityRegistry) ); @@ -1049,9 +1027,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); } @@ -1066,11 +1041,11 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { SystemMetadata metadata1 = AspectGenerationUtils.createSystemMetadata(1625792689, "run-123"); SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry) ); @@ -1084,10 +1059,10 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email2@test.com"); items = List.of( - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata2) .build(_testEntityRegistry) ); @@ -1102,12 +1077,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()); @@ -1126,11 +1095,11 @@ public void testIngestSameAspect() throws AssertionError { SystemMetadata metadata2 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-456"); SystemMetadata metadata3 = AspectGenerationUtils.createSystemMetadata(1635792689, "run-123"); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry) ); @@ -1148,9 +1117,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); @@ -1159,10 +1125,10 @@ public void testIngestSameAspect() throws AssertionError { CorpUserInfo writeAspect2 = AspectGenerationUtils.createCorpUserInfo("email@test.com"); items = List.of( - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata2) .build(_testEntityRegistry) ); @@ -1180,9 +1146,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); } @@ -1205,41 +1168,41 @@ public void testRetention() throws AssertionError { Status writeAspect2a = new Status().setRemoved(false); Status writeAspect2b = new Status().setRemoved(true); - List items = List.of( - AspectsBatchItem.builder() + List items = List.of( + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1) + .aspect(writeAspect1) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1a) + .aspect(writeAspect1a) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1b) + .aspect(writeAspect1b) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) - .value(writeAspect2) + .aspect(writeAspect2) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) - .value(writeAspect2a) + .aspect(writeAspect2a) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) - .value(writeAspect2b) + .aspect(writeAspect2b) .systemMetadata(metadata1) .build(_testEntityRegistry) ); @@ -1259,16 +1222,16 @@ public void testRetention() throws AssertionError { Status writeAspect2c = new Status().setRemoved(false); items = List.of( - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName) - .value(writeAspect1c) + .aspect(writeAspect1c) .systemMetadata(metadata1) .build(_testEntityRegistry), - AspectsBatchItem.builder() + UpsertBatchItem.builder() .urn(entityUrn) .aspectName(aspectName2) - .value(writeAspect2c) + .aspect(writeAspect2c) .systemMetadata(metadata1) .build(_testEntityRegistry) ); 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..45eb9e35226e1 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 @@ -86,7 +86,7 @@ private int getAndReIndexDataPlatforms(AuditStamp auditStamp, AspectSpec dataPla continue; } - _entityService.produceMetadataChangeLog( + _entityService.alwaysProduceMCLAsync( dpUrn, Constants.DATA_PLATFORM_ENTITY_NAME, Constants.DATA_PLATFORM_INFO_ASPECT_NAME, 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 35816e824d31f..d27324377163e 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 @@ -9,7 +9,7 @@ import com.linkedin.metadata.entity.AspectMigrationsDao; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; +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; @@ -65,16 +65,16 @@ public void execute() throws Exception { log.info("Reading urns {} to {} from the aspects table to generate dataplatform instance aspects", start, start + BATCH_SIZE); - List items = new LinkedList<>(); + 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()) { - items.add(AspectsBatchItem.builder() + items.add(UpsertBatchItem.builder() .urn(urn) .aspectName(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - .value(dataPlatformInstance.get()) + .aspect(dataPlatformInstance.get()) .build(_entityService.getEntityRegistry())); } } 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 800217c1d420a..c8efc7e309517 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 @@ -19,7 +19,7 @@ import java.util.stream.StreamSupport; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; +import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.core.io.ClassPathResource; @@ -57,7 +57,7 @@ public void execute() throws IOException, URISyntaxException { } // 2. For each JSON object, cast into a DataPlatformSnapshot object. - List dataPlatformAspects = StreamSupport.stream( + List dataPlatformAspects = StreamSupport.stream( Spliterators.spliteratorUnknownSize(dataPlatforms.iterator(), Spliterator.ORDERED), false) .map(dataPlatform -> { final String urnString; @@ -73,10 +73,10 @@ public void execute() throws IOException, URISyntaxException { final DataPlatformInfo info = RecordUtils.toRecordTemplate(DataPlatformInfo.class, dataPlatform.get("aspect").toString()); - return AspectsBatchItem.builder() + return UpsertBatchItem.builder() .urn(urn) .aspectName(PLATFORM_ASPECT_NAME) - .value(info) + .aspect(info) .build(_entityService.getEntityRegistry()); }).collect(Collectors.toList()); 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 0475e85157fa1..4b67c014b26e8 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 @@ -113,7 +113,7 @@ private void ingestRole(final Urn roleUrn, final DataHubRoleInfo dataHubRoleInfo 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/RestoreColumnLineageIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndices.java index 6e1522051bfab..654119287a6b1 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 @@ -99,7 +99,7 @@ private int getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStamp continue; } - _entityService.produceMetadataChangeLog( + _entityService.alwaysProduceMCLAsync( urn, Constants.DATASET_ENTITY_NAME, Constants.UPSTREAM_LINEAGE_ASPECT_NAME, @@ -150,7 +150,7 @@ private int getAndRestoreInputFieldsIndices(String entityName, int start, AuditS continue; } - _entityService.produceMetadataChangeLog( + _entityService.alwaysProduceMCLAsync( urn, entityName, Constants.INPUT_FIELDS_ASPECT_NAME, 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 617d8043c106b..ea71a41c7baa5 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 @@ -137,7 +137,7 @@ private void getAndRestoreUpstreamLineageIndices(int start, AuditStamp auditStam continue; } - _entityService.produceMetadataChangeLog( + _entityService.alwaysProduceMCLAsync( datasetUrn, DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, 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 486961c2c1f07..a5ea95a271b7f 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 @@ -99,7 +99,7 @@ null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) continue; } - _entityService.produceMetadataChangeLog( + _entityService.alwaysProduceMCLAsync( termUrn, Constants.GLOSSARY_TERM_ENTITY_NAME, Constants.GLOSSARY_TERM_INFO_ASPECT_NAME, @@ -142,7 +142,7 @@ null, null, start, BATCH_SIZE, new SearchFlags().setFulltext(false) continue; } - _entityService.produceMetadataChangeLog( + _entityService.alwaysProduceMCLAsync( nodeUrn, Constants.GLOSSARY_NODE_ENTITY_NAME, Constants.GLOSSARY_NODE_INFO_ASPECT_NAME, 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 f1950d3b798d9..30f8715c55cc1 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; @@ -101,7 +102,7 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw arg.getItems().stream() .allMatch(item -> item.getUrn().getEntityType().equals("chart") && item.getAspectName().equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - && item.getAspect() instanceof DataPlatformInstance) + && ((UpsertBatchItem) item).getAspect() instanceof DataPlatformInstance) ), any(), anyBoolean(), @@ -111,7 +112,7 @@ public void testExecuteWhenSomeEntitiesShouldReceiveDataPlatformInstance() throw !arg.getItems().stream() .allMatch(item -> item.getUrn().getEntityType().equals("chart") && item.getAspectName().equals(DATA_PLATFORM_INSTANCE_ASPECT_NAME) - && item.getAspect() instanceof DataPlatformInstance) + && ((UpsertBatchItem) item).getAspect() instanceof DataPlatformInstance) ), any(), anyBoolean(), anyBoolean()); } 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 52be8a5d364ee..29664436a28d9 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 @@ -64,7 +64,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 +76,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 +88,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 +126,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 +138,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 +150,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 +188,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 +200,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 +212,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), 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 1429da55940d5..2ebc16b3973d1 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 @@ -110,7 +110,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), @@ -122,7 +122,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), @@ -171,7 +171,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), @@ -183,7 +183,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), @@ -232,7 +232,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), @@ -244,7 +244,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/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 b89f9d75e7ee8..ea0d5ff2b24f0 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 @@ -21,7 +21,6 @@ import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RollbackRunResult; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; import com.linkedin.metadata.entity.validation.ValidationException; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.metadata.entity.AspectUtils; @@ -274,11 +273,11 @@ public static Pair ingestProposal(com.linkedin.mxe.MetadataChan AspectsBatch batch = AspectsBatch.builder().mcps(proposalStream.collect(Collectors.toList()), entityService.getEntityRegistry()).build(); - Set> proposalResult = + Set proposalResult = entityService.ingestProposal(batch, auditStamp, false); - Urn urn = proposalResult.stream().findFirst().get().getSecond().getUrn(); - return new Pair<>(urn.toString(), proposalResult.stream().anyMatch(resultPair -> resultPair.getSecond().isDidUpdate())); + Urn urn = proposalResult.stream().findFirst().get().getUrn(); + return new Pair<>(urn.toString(), proposalResult.stream().anyMatch(EntityService.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/mock/MockEntityService.java b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java index d9fd22b753340..0354761ecd96f 100644 --- a/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java +++ b/metadata-service/openapi-servlet/src/test/java/mock/MockEntityService.java @@ -170,13 +170,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 cf798b1461ade..846b098d1ab35 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 @@ -9,7 +9,6 @@ import com.google.common.collect.ImmutableList; import com.linkedin.aspect.GetTimeseriesAspectValuesResponse; import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchItem; import com.linkedin.metadata.resources.operations.Utils; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; @@ -39,7 +38,6 @@ import com.linkedin.restli.server.annotations.RestLiCollection; import com.linkedin.restli.server.annotations.RestMethod; import com.linkedin.restli.server.resources.CollectionResourceTaskTemplate; -import com.linkedin.util.Pair; import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; @@ -221,17 +219,16 @@ public Task ingestProposal( .build(); } - Set> results = + Set results = _entityService.ingestProposal(batch, auditStamp, asyncBool); - EntityService.IngestProposalResult one = results.stream() - .map(Pair::getSecond) + EntityService.IngestResult one = results.stream() .findFirst() .get(); - // Update runIds + // Update runIds, only works for existing documents, so ES document must exist Urn resultUrn = one.getUrn(); - if (!one.isQueued()) { + if (one.isProcessedMCL() || one.isUpdate()) { tryIndexRunId(resultUrn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); } return resultUrn.toString(); 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 12bd0c1628c13..5c877ac3fe758 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,6 +15,7 @@ import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.AspectDao; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -78,13 +80,34 @@ public void testAsyncDefaultAspects() throws URISyntaxException { reset(_producer, _aspectDao); + UpsertBatchItem req = UpsertBatchItem.builder() + .urn(urn) + .aspectName(mcp.getAspectName()) + .aspect(mcp.getAspect()) + .metadataChangeProposal(mcp) + .build(_entityRegistry); when(_aspectDao.runInTransactionWithRetry(any(), anyInt())) .thenReturn(List.of( - new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), - new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), - new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), - new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), - new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0))); + EntityService.UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name1")) + .auditStamp(new AuditStamp()) + .request(req).build(), + EntityService.UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name2")) + .auditStamp(new AuditStamp()) + .request(req).build(), + EntityService.UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name3")) + .auditStamp(new AuditStamp()) + .request(req).build(), + EntityService.UpdateAspectResult.builder().urn(urn) + .newValue(new DatasetProperties().setName("name4")) + .auditStamp(new AuditStamp()) + .request(req).build(), + EntityService.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/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js b/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js index b7d48703992e2..2fd5ad379252d 100644 --- a/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js +++ b/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js @@ -113,7 +113,7 @@ describe('siblings', () => { cy.clickOptionWithTestId('compress-lineage-toggle'); // check the subtypes - cy.get('text:contains(View)').should('have.length', 2); + cy.get('text:contains(View)').should('have.length', 3); cy.get('text:contains(Table)').should('have.length', 0); cy.get('text:contains(Seed)').should('have.length', 1); 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 f87b74e102011aa33010c6049b1ad54f9a40d89d Mon Sep 17 00:00:00 2001 From: David Leifker Date: Sun, 23 Jul 2023 17:07:11 -0500 Subject: [PATCH 07/27] Fix merge issues --- .../resolvers/embed/UpdateEmbedResolver.java | 2 +- .../resolvers/mutate/MutationUtils.java | 2 +- .../mutate/UpdateUserSettingResolver.java | 2 +- .../resolvers/mutate/util/DeleteUtils.java | 4 +- .../mutate/util/DeprecationUtils.java | 4 +- .../resolvers/mutate/util/DomainUtils.java | 4 +- .../resolvers/mutate/util/LabelUtils.java | 4 +- .../resolvers/mutate/util/OwnerUtils.java | 4 +- .../linkedin/datahub/graphql/TestUtils.java | 13 ++- .../BatchUpdateSoftDeletedResolverTest.java | 3 +- .../BatchUpdateDeprecationResolverTest.java | 3 +- .../domain/BatchSetDomainResolverTest.java | 3 +- .../embed/UpdateEmbedResolverTest.java | 5 +- .../owner/AddOwnersResolverTest.java | 3 +- .../owner/BatchAddOwnersResolverTest.java | 3 +- .../owner/BatchRemoveOwnersResolverTest.java | 3 +- .../resolvers/tag/AddTagsResolverTest.java | 3 +- .../tag/BatchAddTagsResolverTest.java | 9 +- .../tag/BatchRemoveTagsResolverTest.java | 7 +- .../resolvers/term/AddTermsResolverTest.java | 14 +-- .../term/BatchAddTermsResolverTest.java | 3 +- .../term/BatchRemoveTermsResolverTest.java | 3 +- .../metadata/client/JavaEntityClient.java | 8 +- .../metadata/entity/EntityServiceImpl.java | 59 ++++--------- .../linkedin/metadata/entity/EntityUtils.java | 2 +- .../cassandra/CassandraRetentionService.java | 10 +++ .../metadata/entity/ebean/EbeanAspectDao.java | 11 ++- .../entity/ebean/EbeanRetentionService.java | 10 +++ ...spectsBatch.java => AspectsBatchImpl.java} | 33 +++---- .../ebean/transactions/PatchBatchItem.java | 10 ++- .../ebean/transactions/UpsertBatchItem.java | 11 ++- .../entity/validation/ValidationUtils.java | 38 ++++++++ .../metadata/AspectIngestionUtils.java | 8 +- .../entity/EbeanEntityServiceTest.java | 6 +- .../metadata/entity/EntityServiceTest.java | 54 ++++++------ .../token/StatefulTokenService.java | 4 +- .../token/StatefulTokenServiceTest.java | 1 - .../linkedin/metadata/boot/BootstrapStep.java | 2 +- .../linkedin/metadata/boot/UpgradeStep.java | 4 +- .../boot/steps/BackfillBrowsePathsV2Step.java | 2 +- .../IngestDataPlatformInstancesStep.java | 4 +- .../boot/steps/IngestDataPlatformsStep.java | 4 +- .../IngestDefaultGlobalSettingsStep.java | 2 +- .../boot/steps/IngestOwnershipTypesStep.java | 4 +- .../boot/steps/IngestPoliciesStep.java | 4 +- .../metadata/boot/steps/IngestRolesStep.java | 4 +- .../boot/steps/RestoreDbtSiblingsIndices.java | 2 +- .../steps/UpgradeDefaultBrowsePathsStep.java | 2 +- .../steps/BackfillBrowsePathsV2StepTest.java | 4 +- .../IngestDefaultGlobalSettingsStepTest.java | 4 +- .../RestoreColumnLineageIndicesTest.java | 6 +- .../steps/RestoreGlossaryIndicesTest.java | 6 +- .../UpgradeDefaultBrowsePathsStepTest.java | 8 +- .../openapi/util/MappingUtil.java | 10 ++- .../java/entities/EntitiesControllerTest.java | 6 +- .../src/test/java/mock/MockEntityService.java | 5 -- .../resources/entity/AspectResource.java | 12 +-- .../entity/BatchIngestionRunResource.java | 2 +- .../resources/entity/AspectResourceTest.java | 12 +-- .../linkedin/metadata/entity/AspectUtils.java | 37 -------- .../metadata/entity/DeleteEntityService.java | 2 +- .../metadata/entity/EntityService.java | 86 +++++-------------- .../metadata/entity/IngestResult.java | 18 ++++ .../metadata/entity/RetentionService.java | 10 +-- .../metadata/entity/UpdateAspectResult.java | 8 ++ .../transactions/AbstractBatchItem.java | 12 ++- .../entity/transactions/AspectsBatch.java | 22 +++++ 67 files changed, 346 insertions(+), 324 deletions(-) rename metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/{AspectsBatch.java => AspectsBatchImpl.java} (59%) create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/IngestResult.java rename {metadata-io/src/main/java/com/linkedin/metadata/entity/ebean => metadata-service/services/src/main/java/com/linkedin/metadata/entity}/transactions/AbstractBatchItem.java (91%) create mode 100644 metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AspectsBatch.java diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java index 8fc931310570a..86b8eb5564152 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolver.java @@ -58,7 +58,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw updateEmbed(embed, input); final MetadataChangeProposal proposal = buildMetadataChangeProposalWithUrn(entityUrn, EMBED_ASPECT_NAME, embed); - _entityService.ingestSingleProposal( + _entityService.ingestProposal( proposal, new AuditStamp().setActor(UrnUtils.getUrn(context.getActorUrn())).setTime(System.currentTimeMillis()), false diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java index 524080bc94edb..0cf9acd62f736 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java @@ -29,7 +29,7 @@ private MutationUtils() { } public static void persistAspect(Urn urn, String aspectName, RecordTemplate aspect, Urn actor, EntityService entityService) { final MetadataChangeProposal proposal = buildMetadataChangeProposalWithUrn(urn, aspectName, aspect); - entityService.ingestSingleProposal(proposal, getAuditStamp(actor), false); + entityService.ingestProposal(proposal, getAuditStamp(actor), false); } /** diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java index e97c20b77434e..8c1d32c470f44 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java @@ -56,7 +56,7 @@ public CompletableFuture get(DataFetchingEnvironment environment) throw MetadataChangeProposal proposal = buildMetadataChangeProposalWithUrn(actor, CORP_USER_SETTINGS_ASPECT_NAME, newSettings); - _entityService.ingestSingleProposal(proposal, getAuditStamp(actor), false); + _entityService.ingestProposal(proposal, getAuditStamp(actor), false); return true; } catch (Exception e) { diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java index 9ec91b4784f08..94acb5a75918f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java @@ -12,7 +12,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -73,7 +73,7 @@ private static MetadataChangeProposal buildSoftDeleteProposal( } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - entityService.ingestProposal(AspectsBatch.builder() + entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } \ No newline at end of file diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java index 4d3b0770bd363..0a1712aaf9ab4 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java @@ -14,7 +14,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -88,7 +88,7 @@ private static MetadataChangeProposal buildUpdateDeprecationProposal( } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - entityService.ingestProposal(AspectsBatch.builder() + entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } \ No newline at end of file diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java index 4da6ee5c2db07..3809fcee28649 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java @@ -14,7 +14,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -87,7 +87,7 @@ public static void validateDomain(Urn domainUrn, EntityService entityService) { } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - entityService.ingestProposal(AspectsBatch.builder() + entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } \ No newline at end of file diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java index ffb499a3a833d..eb2d49ce987e3 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java @@ -20,7 +20,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.schema.EditableSchemaFieldInfo; import com.linkedin.schema.EditableSchemaMetadata; @@ -559,7 +559,7 @@ private static GlossaryTermAssociationArray removeTermsIfExists(GlossaryTerms te } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - entityService.ingestProposal(AspectsBatch.builder() + entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java index 9563e55f91f70..cf9c2a4ba6adc 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java @@ -21,7 +21,7 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; @@ -270,7 +270,7 @@ public static Boolean validateRemoveInput( } private static void ingestChangeProposals(List changes, EntityService entityService, Urn actor) { - entityService.ingestProposal(AspectsBatch.builder() + entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(changes, entityService.getEntityRegistry()).build(), getAuditStamp(actor), false); } 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 67afd4d3a0cc8..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,7 +9,7 @@ import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.UrnUtils; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +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; @@ -104,7 +104,7 @@ public static void verifyIngestProposal(EntityService mockService, int numberOfI } public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, List proposals) { - AspectsBatch batch = AspectsBatch.builder() + AspectsBatchImpl batch = AspectsBatchImpl.builder() .mcps(proposals, mockService.getEntityRegistry()) .build(); Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( @@ -115,7 +115,7 @@ public static void verifyIngestProposal(EntityService mockService, int numberOfI } public static void verifySingleIngestProposal(EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) { - Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( Mockito.eq(proposal), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -124,14 +124,14 @@ public static void verifySingleIngestProposal(EntityService mockService, int num public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations) { Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( - Mockito.any(AspectsBatch.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)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -140,8 +140,7 @@ public static void verifySingleIngestProposal(EntityService mockService, int num 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/delete/BatchUpdateSoftDeletedResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java index 6703d8726fa5d..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,6 +11,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -165,7 +166,7 @@ public void testGetEntityClientException() throws Exception { 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 444d7481a23b2..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,6 +12,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -183,7 +184,7 @@ public void testGetEntityClientException() throws Exception { 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/domain/BatchSetDomainResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java index 138fd1d9bf024..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,6 +14,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -277,7 +278,7 @@ public void testGetEntityClientException() throws Exception { 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/embed/UpdateEmbedResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java index 3cfb1c55b3275..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; @@ -136,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) );; @@ -156,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) ); 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 e1324edaf8a7f..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; @@ -198,7 +199,7 @@ public void testGetEntityClientException() throws Exception { 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 ff119451dd859..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; @@ -280,7 +281,7 @@ public void testGetEntityClientException() throws Exception { 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 c76f2ec4acb12..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; @@ -175,7 +176,7 @@ public void testGetEntityClientException() throws Exception { 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 0596924b54a25..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; @@ -206,7 +207,7 @@ public void testGetEntityClientException() throws Exception { 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 6009912f498ec..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,6 +15,7 @@ import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -184,7 +185,7 @@ 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()); } @@ -223,7 +224,7 @@ 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()); } @@ -247,7 +248,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.anyBoolean()); } @@ -256,7 +257,7 @@ public void testGetEntityClientException() throws Exception { 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 2e4314ef27e54..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,6 +16,7 @@ import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.utils.GenericRecordUtils; import com.linkedin.mxe.MetadataChangeProposal; import graphql.schema.DataFetchingEnvironment; @@ -179,7 +180,7 @@ 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()); } @@ -203,7 +204,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.anyBoolean()); } @@ -212,7 +213,7 @@ public void testGetEntityClientException() throws Exception { 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/term/AddTermsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/term/AddTermsResolverTest.java index 94584c3750d2c..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.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import graphql.schema.DataFetchingEnvironment; import java.util.concurrent.CompletionException; import org.mockito.Mockito; @@ -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(AspectsBatch.class), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -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(AspectsBatch.class), + Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false) ); @@ -143,7 +143,7 @@ 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()); } @@ -173,7 +173,7 @@ 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()); } @@ -194,7 +194,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.anyBoolean()); } @@ -203,7 +203,7 @@ public void testGetEntityClientException() throws Exception { 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 77200c8b72a22..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; @@ -220,7 +221,7 @@ public void testGetEntityClientException() throws Exception { 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 c0f1c06ba0556..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; @@ -182,7 +183,7 @@ public void testGetEntityClientException() throws Exception { 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/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java index 7cd9e57cde066..cae29f3bfee19 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,7 +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.ebean.transactions.AspectsBatch; +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; @@ -536,11 +538,11 @@ public String ingestProposal(@Nonnull final MetadataChangeProposal metadataChang Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), additionalChanges.stream()); - AspectsBatch batch = AspectsBatch.builder() + AspectsBatch batch = AspectsBatchImpl.builder() .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) .build(); - EntityService.IngestResult one = _entityService.ingestProposal(batch, auditStamp, async).stream() + IngestResult one = _entityService.ingestProposal(batch, auditStamp, async).stream() .findFirst().get(); Urn urn = one.getUrn(); 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 6430ab9afdd47..0d76ea1948214 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 @@ -34,14 +34,15 @@ 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.AbstractBatchItem; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +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.transactions.AspectsBatch; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.AspectSpec; import com.linkedin.metadata.models.EntitySpec; @@ -88,8 +89,6 @@ import javax.annotation.Nullable; import javax.persistence.EntityNotFoundException; -import lombok.Builder; -import lombok.Value; import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.*; @@ -133,34 +132,6 @@ public class EntityServiceImpl implements EntityService { * monotonically increasing version incrementing as usual once the latest version is replaced. */ - @Builder(toBuilder = true) - @Value - public static class UpdateAspectResult { - Urn urn; - UpsertBatchItem request; - RecordTemplate oldValue; - RecordTemplate newValue; - SystemMetadata oldSystemMetadata; - SystemMetadata newSystemMetadata; - MetadataAuditOperation operation; - AuditStamp auditStamp; - long maxVersion; - boolean processedMCL; - Future mclFuture; - } - - @Builder(toBuilder = true) - @Value - public static class IngestResult { - Urn urn; - AbstractBatchItem request; - boolean publishedMCL; - boolean processedMCL; - boolean publishedMCP; - boolean sqlCommitted; - boolean isUpdate; // update else insert - } - private static final int DEFAULT_MAX_TRANSACTION_RETRY = 3; protected final AspectDao _aspectDao; @@ -199,6 +170,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) { @@ -517,7 +489,8 @@ public ListResult listLatestAspects( * @param systemMetadata system metadata * @return update result */ - public List ingestAspects(Urn entityUrn, + @Override + public List ingestAspects(@Nonnull Urn entityUrn, List> pairList, @Nonnull final AuditStamp auditStamp, SystemMetadata systemMetadata) { @@ -529,7 +502,7 @@ public List ingestAspects(Urn entityUrn, .systemMetadata(systemMetadata) .build(_entityRegistry)) .collect(Collectors.toList()); - return ingestAspects(AspectsBatch.builder().items(items).build(), auditStamp, true, true); + return ingestAspects(AspectsBatchImpl.builder().items(items).build(), auditStamp, true, true); } /** @@ -541,6 +514,7 @@ public List ingestAspects(Urn entityUrn, * successful update * @return the {@link RecordTemplate} representation of the written aspect object */ + @Override public List ingestAspects(@Nonnull final AspectsBatch aspectsBatch, @Nonnull final AuditStamp auditStamp, boolean emitMCL, @@ -719,7 +693,7 @@ public RecordTemplate ingestAspectIfNotPresent(@Nonnull Urn urn, @Nonnull SystemMetadata systemMetadata) { log.debug("Invoked ingestAspectIfNotPresent with urn: {}, aspectName: {}, newValue: {}", urn, aspectName, newValue); - AspectsBatch aspectsBatch = AspectsBatch.builder() + AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() .one(UpsertBatchItem.builder() .urn(urn) .aspectName(aspectName) @@ -739,9 +713,10 @@ public RecordTemplate ingestAspectIfNotPresent(@Nonnull Urn urn, * @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 */ - public IngestResult ingestSingleProposal(MetadataChangeProposal proposal, AuditStamp auditStamp, final boolean async) { - return ingestProposal(AspectsBatch.builder().mcps(List.of(proposal), getEntityRegistry()).build(), auditStamp, async) - .stream().findFirst().get(); + @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(); } /** @@ -844,7 +819,7 @@ private Stream ingestProposalAsync(AspectsBatch aspectsBatch) { } private Stream ingestProposalSync(AspectsBatch aspectsBatch, AuditStamp auditStamp) { - AspectsBatch nonTimeseries = AspectsBatch.builder() + AspectsBatchImpl nonTimeseries = AspectsBatchImpl.builder() .items(aspectsBatch.getItems().stream() .filter(item -> !item.getAspectSpec().isTimeseries()) .collect(Collectors.toList())) @@ -1085,12 +1060,14 @@ public Map getEntities(@Nonnull final Set urns, @Nonnull Set toEntity(entry.getValue()))); } + @Override public Pair, 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)); } + @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, @@ -1276,7 +1253,7 @@ private void ingestSnapshotUnion(@Nonnull final Snapshot snapshotUnion, @Nonnull aspectRecordsToIngest.addAll(generateDefaultAspectsIfMissing(urn, aspectRecordsToIngest.stream().map(Pair::getFirst).collect(Collectors.toSet()))); - AspectsBatch aspectsBatch = AspectsBatch.builder() + AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() .items(aspectRecordsToIngest.stream().map(pair -> UpsertBatchItem.builder() .urn(urn) .aspectName(pair.getKey()) @@ -1594,7 +1571,7 @@ public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map mcps) { + return AspectsBatchImpl.builder() + .mcps(mcps, _entityService.getEntityRegistry()) + .build(); + } + @Override @WithSpan protected void applyRetention(List retentionContexts) { 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 2cfebaeab9839..93c40aab2a1fa 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 @@ -14,7 +14,16 @@ import com.linkedin.metadata.query.ExtraInfoArray; import com.linkedin.metadata.query.ListResultMetadata; import com.linkedin.metadata.search.utils.QueryUtils; -import io.ebean.*; +import io.ebean.DuplicateKeyException; +import io.ebean.EbeanServer; +import io.ebean.ExpressionList; +import io.ebean.Junction; +import io.ebean.PagedList; +import io.ebean.Query; +import io.ebean.RawSql; +import io.ebean.RawSqlBuilder; +import io.ebean.Transaction; +import io.ebean.TxScope; import io.ebean.annotation.Platform; import io.ebean.annotation.TxIsolation; import java.net.URISyntaxException; 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 21bbc38430436..33888594c4458 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,8 +4,11 @@ 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; @@ -48,6 +51,13 @@ public EntityService getEntityService() { return _entityService; } + @Override + protected AspectsBatch buildAspectsBatch(List mcps) { + return AspectsBatchImpl.builder() + .mcps(mcps, _entityService.getEntityRegistry()) + .build(); + } + @Override @WithSpan protected void applyRetention(List retentionContexts) { diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java similarity index 59% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java rename to metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java index 17037a4bced68..ca5e070bc5ca7 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatchImpl.java @@ -1,6 +1,8 @@ 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; @@ -8,40 +10,28 @@ import lombok.extern.slf4j.Slf4j; import java.util.List; -import java.util.Map; import java.util.Objects; -import java.util.Set; import java.util.stream.Collectors; + @Slf4j @Getter @Builder(toBuilder = true) -public class AspectsBatch { +public class AspectsBatchImpl implements AspectsBatch { private final List items; - public Map> getUrnAspectsMap() { - return items.stream() - .map(aspect -> Map.entry(aspect.getUrn().toString(), aspect.getAspectName())) - .collect(Collectors.groupingBy(Map.Entry::getKey, Collectors.mapping(Map.Entry::getValue, Collectors.toSet()))); - } - - public boolean containsDuplicateAspects() { - return items.stream().map(i -> String.format("%s_%s", i.getClass().getName(), i.hashCode())) - .distinct().count() != items.size(); - } - - public static class AspectsBatchBuilder { + public static class AspectsBatchImplBuilder { /** * Just one aspect record template * @param data aspect data * @return builder */ - public AspectsBatchBuilder one(AbstractBatchItem data) { + public AspectsBatchImplBuilder one(AbstractBatchItem data) { this.items = List.of(data); return this; } - public AspectsBatchBuilder mcps(List mcps, EntityRegistry entityRegistry) { + 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); @@ -61,8 +51,8 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) { return false; } - AspectsBatch that = (AspectsBatch) o; - return items.equals(that.items); + AspectsBatchImpl that = (AspectsBatchImpl) o; + return Objects.equals(items, that.items); } @Override @@ -72,9 +62,6 @@ public int hashCode() { @Override public String toString() { - return "AspectsBatch{" - + "items=" - + items - + '}'; + 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 index b5edccbc57fd8..cc0b3d915b407 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/PatchBatchItem.java @@ -10,8 +10,9 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.entity.AspectUtils; 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; @@ -58,6 +59,11 @@ 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()) @@ -94,7 +100,7 @@ public PatchBatchItem build(EntityRegistry entityRegistry) { entitySpec(entityRegistry.getEntitySpec(this.urn.getEntityType())); log.debug("entity spec = {}", this.entitySpec); - aspectSpec(AspectUtils.validate(this.entitySpec, this.aspectName)); + aspectSpec(ValidationUtils.validate(this.entitySpec, this.aspectName)); log.debug("aspect spec = {}", this.aspectSpec); if (this.patch == null) { 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 index 9c89858f33c31..bd58d267a8308 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/UpsertBatchItem.java @@ -5,9 +5,9 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.metadata.entity.AspectUtils; 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; @@ -50,6 +50,11 @@ 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()); @@ -70,10 +75,10 @@ public UpsertBatchItem build(EntityRegistry entityRegistry) { entitySpec(entityRegistry.getEntitySpec(this.urn.getEntityType())); log.debug("entity spec = {}", this.entitySpec); - aspectSpec(AspectUtils.validate(this.entitySpec, this.aspectName)); + aspectSpec(ValidationUtils.validate(this.entitySpec, this.aspectName)); log.debug("aspect spec = {}", this.aspectSpec); - AspectUtils.validateRecordTemplate(entityRegistry, this.entitySpec, this.urn, this.aspect); + 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); 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/test/java/com/linkedin/metadata/AspectIngestionUtils.java b/metadata-io/src/test/java/com/linkedin/metadata/AspectIngestionUtils.java index cf1e262a04343..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,7 +5,7 @@ import com.linkedin.common.urn.UrnUtils; import com.linkedin.identity.CorpUserInfo; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +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; @@ -41,7 +41,7 @@ public static Map ingestCorpUserKeyAspects(EntityService entit .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService.getEntityRegistry())); } - entityService.ingestAspects(AspectsBatch.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } @@ -67,7 +67,7 @@ public static Map ingestCorpUserInfoAspects(@Nonnull final En .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService.getEntityRegistry())); } - entityService.ingestAspects(AspectsBatch.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } @@ -94,7 +94,7 @@ public static Map ingestChartInfoAspects(@Nonnull final EntitySe .systemMetadata(AspectGenerationUtils.createSystemMetadata()) .build(entityService.getEntityRegistry())); } - entityService.ingestAspects(AspectsBatch.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); + entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), AspectGenerationUtils.createAuditStamp(), true, true); return aspects; } } 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 87c8f0ae30498..bf720166fdd9a 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 @@ -10,7 +10,7 @@ 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.AspectsBatch; +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; @@ -117,7 +117,7 @@ public void testIngestListLatestAspects() throws AssertionError { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // List aspects ListResult batch1 = _entityServiceImpl.listLatestAspects(entityUrn1.getEntityType(), aspectName, 0, 2); @@ -181,7 +181,7 @@ public void testIngestListUrns() throws AssertionError { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // List aspects urns ListUrnsResult batch1 = _entityServiceImpl.listUrns(entityUrn1.getEntityType(), 0, 2); 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 1c116dee79f8f..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,7 +32,7 @@ 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.AspectsBatch; +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; @@ -48,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; @@ -544,7 +543,7 @@ public void testReingestLineageProposal() throws Exception { mcp1.setSystemMetadata(metadata1); mcp1.setAspectName(UPSTREAM_LINEAGE_ASPECT_NAME); - _entityServiceImpl.ingestSingleProposal(mcp1, TEST_AUDIT_STAMP, false); + _entityServiceImpl.ingestProposal(mcp1, TEST_AUDIT_STAMP, false); final MetadataChangeLog initialChangeLog = new MetadataChangeLog(); initialChangeLog.setEntityType(entityUrn.getEntityType()); @@ -579,7 +578,7 @@ public void testReingestLineageProposal() throws Exception { // Mockito detects the previous invocation and throws an error in verifying the second call unless invocations are cleared clearInvocations(_mockProducer); - _entityServiceImpl.ingestSingleProposal(mcp1, TEST_AUDIT_STAMP, false); + _entityServiceImpl.ingestProposal(mcp1, TEST_AUDIT_STAMP, false); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.eq(restateChangeLog)); @@ -606,7 +605,7 @@ public void testIngestTimeseriesAspect() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetProfileSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityServiceImpl.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, false); + _entityServiceImpl.ingestProposal(gmce, TEST_AUDIT_STAMP, false); } @Test @@ -625,7 +624,7 @@ public void testAsyncProposalVersioned() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetPropertiesSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityServiceImpl.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, true); + _entityServiceImpl.ingestProposal(gmce, TEST_AUDIT_STAMP, true); verify(_mockProducer, times(0)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.any()); verify(_mockProducer, times(1)).produceMetadataChangeProposal(Mockito.eq(entityUrn), @@ -651,7 +650,7 @@ public void testAsyncProposalTimeseries() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetProfileSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityServiceImpl.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, true); + _entityServiceImpl.ingestProposal(gmce, TEST_AUDIT_STAMP, true); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), Mockito.any()); verify(_mockProducer, times(0)).produceMetadataChangeProposal(Mockito.eq(entityUrn), @@ -719,7 +718,7 @@ public void testGetAspectAtVersion() throws AssertionError { assertTrue(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); // Validate retrieval of CorpUserInfo Aspect #2 - _entityService.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect2)), TEST_AUDIT_STAMP, null); + _entityServiceImpl.ingestAspects(entityUrn, List.of(Pair.of(aspectName, writeAspect2)), TEST_AUDIT_STAMP, null); VersionedAspect writtenVersionedAspect2 = new VersionedAspect(); writtenVersionedAspect2.setAspect(Aspect.create(writeAspect2)); @@ -731,7 +730,7 @@ public void testGetAspectAtVersion() throws AssertionError { verify(_mockProducer, times(2)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.eq(corpUserInfoSpec), Mockito.any()); - readAspect1 = _entityService.getVersionedAspect(entityUrn, aspectName, -1); + readAspect1 = _entityServiceImpl.getVersionedAspect(entityUrn, aspectName, -1); assertFalse(DataTemplateUtil.areEqual(writtenVersionedAspect1, readAspect1)); verifyNoMoreInteractions(_mockProducer); @@ -786,7 +785,7 @@ public void testRollbackAspect() throws AssertionError { .systemMetadata(metadata2) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _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(); @@ -854,7 +853,7 @@ public void testRollbackKey() throws AssertionError { .systemMetadata(metadata2) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _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(); @@ -898,7 +897,6 @@ 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 = EntityUtils.buildKeyAspect(_testEntityRegistry, entityUrn1); @@ -943,7 +941,7 @@ public void testRollbackUrn() throws AssertionError { .systemMetadata(metadata2) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _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(); @@ -981,10 +979,10 @@ public void testIngestGetLatestAspect() throws AssertionError { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // Validate retrieval of CorpUserInfo Aspect #1 - RecordTemplate readAspect1 = _entityService.getLatestAspect(entityUrn, aspectName); + RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, readAspect1)); ArgumentCaptor mclCaptor = ArgumentCaptor.forClass(MetadataChangeLog.class); @@ -1010,10 +1008,10 @@ public void testIngestGetLatestAspect() throws AssertionError { .systemMetadata(metadata2) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // Validate retrieval of CorpUserInfo Aspect #2 - RecordTemplate readAspect2 = _entityService.getLatestAspect(entityUrn, aspectName); + RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); EntityAspect readAspectDao1 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 1); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 0); @@ -1050,10 +1048,10 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // Validate retrieval of CorpUserInfo Aspect #1 - EnvelopedAspect readAspect1 = _entityService.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); + EnvelopedAspect readAspect1 = _entityServiceImpl.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, new CorpUserInfo(readAspect1.getValue().data()))); // Ingest CorpUserInfo Aspect #2 @@ -1067,10 +1065,10 @@ public void testIngestGetLatestEnvelopedAspect() throws Exception { .systemMetadata(metadata2) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // Validate retrieval of CorpUserInfo Aspect #2 - EnvelopedAspect readAspect2 = _entityService.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); + EnvelopedAspect readAspect2 = _entityServiceImpl.getLatestEnvelopedAspect("corpuser", entityUrn, aspectName); EntityAspect readAspectDao1 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 1); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, 0); @@ -1104,10 +1102,10 @@ public void testIngestSameAspect() throws AssertionError { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // Validate retrieval of CorpUserInfo Aspect #1 - RecordTemplate readAspect1 = _entityService.getLatestAspect(entityUrn, aspectName); + RecordTemplate readAspect1 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); assertTrue(DataTemplateUtil.areEqual(writeAspect1, readAspect1)); ArgumentCaptor mclCaptor = ArgumentCaptor.forClass(MetadataChangeLog.class); @@ -1133,10 +1131,10 @@ public void testIngestSameAspect() throws AssertionError { .systemMetadata(metadata2) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _entityServiceImpl.ingestAspects(AspectsBatchImpl.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); // Validate retrieval of CorpUserInfo Aspect #2 - RecordTemplate readAspect2 = _entityService.getLatestAspect(entityUrn, aspectName); + RecordTemplate readAspect2 = _entityServiceImpl.getLatestAspect(entityUrn, aspectName); EntityAspect readAspectDao2 = _aspectDao.getAspect(entityUrn.toString(), aspectName, ASPECT_LATEST_VERSION); assertTrue(DataTemplateUtil.areEqual(writeAspect2, readAspect2)); @@ -1207,7 +1205,7 @@ public void testRetention() throws AssertionError { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _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); @@ -1236,7 +1234,7 @@ public void testRetention() throws AssertionError { .systemMetadata(metadata1) .build(_testEntityRegistry) ); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), TEST_AUDIT_STAMP, true, true); + _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); @@ -1451,7 +1449,7 @@ public void testUIPreProcessedProposal() throws Exception { genericAspect.setValue(ByteString.unsafeWrap(datasetPropertiesSerialized)); genericAspect.setContentType("application/json"); gmce.setAspect(genericAspect); - _entityServiceImpl.ingestSingleProposal(gmce, TEST_AUDIT_STAMP, false); + _entityServiceImpl.ingestProposal(gmce, TEST_AUDIT_STAMP, false); ArgumentCaptor captor = ArgumentCaptor.forClass(MetadataChangeLog.class); verify(_mockProducer, times(1)).produceMetadataChangeLog(Mockito.eq(entityUrn), Mockito.any(), captor.capture()); 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 40cf3ae1d5b9a..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,7 +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.AspectsBatch; +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; @@ -129,7 +129,7 @@ public String generateAccessToken(@Nonnull final TokenType type, @Nonnull final Stream proposalStream = Stream.concat(Stream.of(proposal), AspectUtils.getAdditionalChanges(proposal, _entityService).stream()); - _entityService.ingestProposal(AspectsBatch.builder() + _entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) .build(), auditStamp, false); 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 12c1307f996d6..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 @@ -14,7 +14,6 @@ import java.util.Date; import java.util.Map; -import com.linkedin.metadata.models.registry.EntityRegistry; import org.mockito.Mockito; import org.testng.annotations.Test; diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java index b7477fdcc0af9..876a0871fa4cb 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/BootstrapStep.java @@ -63,6 +63,6 @@ static void setUpgradeResult(Urn urn, EntityService entityService) throws URISyn upgradeProposal.setAspectName(Constants.DATA_HUB_UPGRADE_RESULT_ASPECT_NAME); upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(upgradeResult)); upgradeProposal.setChangeType(ChangeType.UPSERT); - entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); + entityService.ingestProposal(upgradeProposal, auditStamp, false); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java index b408af8b66720..dbbcf3a139bf1 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/UpgradeStep.java @@ -92,7 +92,7 @@ private void ingestUpgradeRequestAspect() throws URISyntaxException { upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(upgradeRequest)); upgradeProposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); + _entityService.ingestProposal(upgradeProposal, auditStamp, false); } private void ingestUpgradeResultAspect() throws URISyntaxException { @@ -107,7 +107,7 @@ private void ingestUpgradeResultAspect() throws URISyntaxException { upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(upgradeResult)); upgradeProposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); + _entityService.ingestProposal(upgradeProposal, auditStamp, false); } private void cleanUpgradeAfterError(Exception e, String errorMessage) { diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java index 51b1da92756a7..ea9ac57778550 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2Step.java @@ -138,7 +138,7 @@ private void ingestBrowsePathsV2(Urn urn, AuditStamp auditStamp) throws Exceptio proposal.setChangeType(ChangeType.UPSERT); proposal.setSystemMetadata(new SystemMetadata().setRunId(DEFAULT_RUN_ID).setLastObserved(System.currentTimeMillis())); proposal.setAspect(GenericRecordUtils.serializeAspect(browsePathsV2)); - _entityService.ingestSingleProposal( + _entityService.ingestProposal( proposal, auditStamp, false 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 d27324377163e..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,7 +8,7 @@ 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.AspectsBatch; +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; @@ -81,7 +81,7 @@ public void execute() throws Exception { final AuditStamp aspectAuditStamp = new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); - _entityService.ingestAspects(AspectsBatch.builder().items(items).build(), aspectAuditStamp, true, true); + _entityService.ingestAspects(AspectsBatchImpl.builder().items(items).build(), aspectAuditStamp, true, true); 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 c8efc7e309517..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 @@ -18,7 +18,7 @@ import java.util.stream.Collectors; import java.util.stream.StreamSupport; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -80,7 +80,7 @@ public void execute() throws IOException, URISyntaxException { .build(_entityService.getEntityRegistry()); }).collect(Collectors.toList()); - _entityService.ingestAspects(AspectsBatch.builder().items(dataPlatformAspects).build(), + _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/IngestDefaultGlobalSettingsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java index 1541ca40d8911..5bc80f46e6478 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStep.java @@ -116,7 +116,7 @@ public void execute() throws IOException, URISyntaxException { proposal.setAspect(GenericRecordUtils.serializeAspect(newSettings)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestSingleProposal( + _entityService.ingestProposal( proposal, new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestOwnershipTypesStep.java index 596af9cb857f6..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,7 +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.AspectsBatch; +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; @@ -108,7 +108,7 @@ private void ingestOwnershipType(final Urn ownershipTypeUrn, final OwnershipType proposal.setAspect(GenericRecordUtils.serializeAspect(info)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(AspectsBatch.builder() + _entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(List.of(keyAspectProposal, proposal), _entityService.getEntityRegistry()).build(), auditStamp, false); } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java index 67291946f83dd..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,7 +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.AspectsBatch; +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; @@ -182,7 +182,7 @@ private void ingestPolicy(final Urn urn, final DataHubPolicyInfo info) throws UR proposal.setAspect(GenericRecordUtils.serializeAspect(info)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(AspectsBatch.builder() + _entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(List.of(keyAspectProposal, proposal), _entityRegistry) .build(), new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), 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 4b67c014b26e8..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,7 +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.AspectsBatch; +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; @@ -108,7 +108,7 @@ private void ingestRole(final Urn roleUrn, final DataHubRoleInfo dataHubRoleInfo proposal.setAspect(GenericRecordUtils.serializeAspect(dataHubRoleInfo)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestProposal(AspectsBatch.builder() + _entityService.ingestProposal(AspectsBatchImpl.builder() .mcps(List.of(keyAspectProposal, proposal), _entityRegistry).build(), new AuditStamp().setActor(Urn.createFromString(SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); 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 ea71a41c7baa5..4828e3b2b2b28 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 @@ -168,6 +168,6 @@ private void ingestUpgradeAspect(String aspectName, RecordTemplate aspect, Audit upgradeProposal.setAspect(GenericRecordUtils.serializeAspect(aspect)); upgradeProposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestSingleProposal(upgradeProposal, auditStamp, false); + _entityService.ingestProposal(upgradeProposal, auditStamp, false); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java index 38cf773e577b4..7fcafa24d7b45 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStep.java @@ -128,7 +128,7 @@ private void migrateBrowsePath(Urn urn, AuditStamp auditStamp) throws Exception proposal.setChangeType(ChangeType.UPSERT); proposal.setSystemMetadata(new SystemMetadata().setRunId(DEFAULT_RUN_ID).setLastObserved(System.currentTimeMillis())); proposal.setAspect(GenericRecordUtils.serializeAspect(newPaths)); - _entityService.ingestSingleProposal( + _entityService.ingestProposal( proposal, auditStamp, false diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java index 1fb979cba2d3c..49fce75ab7c61 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/BackfillBrowsePathsV2StepTest.java @@ -97,7 +97,7 @@ public void testExecuteNoExistingBrowsePaths() throws Exception { Mockito.eq(null) ); // Verify that 11 aspects are ingested, 2 for the upgrade request / result, 9 for ingesting 1 of each entity type - Mockito.verify(mockService, Mockito.times(11)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(11)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -124,7 +124,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { BackfillBrowsePathsV2Step backfillBrowsePathsV2Step = new BackfillBrowsePathsV2Step(mockService, mockSearchService); backfillBrowsePathsV2Step.execute(); - Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean() diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java index 50ab1f53455f1..24bdd193a39c8 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDefaultGlobalSettingsStepTest.java @@ -38,7 +38,7 @@ public void testExecuteValidSettingsNoExistingSettings() throws Exception { GlobalSettingsInfo expectedResult = new GlobalSettingsInfo(); expectedResult.setViews(new GlobalViewsSettings().setDefaultView(UrnUtils.getUrn("urn:li:dataHubView:test"))); - Mockito.verify(entityService, times(1)).ingestSingleProposal( + Mockito.verify(entityService, times(1)).ingestProposal( Mockito.eq(buildUpdateSettingsProposal(expectedResult)), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -65,7 +65,7 @@ public void testExecuteValidSettingsExistingSettings() throws Exception { GlobalSettingsInfo expectedResult = new GlobalSettingsInfo(); expectedResult.setViews(new GlobalViewsSettings().setDefaultView(UrnUtils.getUrn("urn:li:dataHubView:custom"))); - Mockito.verify(entityService, times(1)).ingestSingleProposal( + Mockito.verify(entityService, times(1)).ingestProposal( Mockito.eq(buildUpdateSettingsProposal(expectedResult)), Mockito.any(AuditStamp.class), Mockito.eq(false) 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 29664436a28d9..d90c1947b3e5e 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 @@ -59,7 +59,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.CHART_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.DASHBOARD_ENTITY_NAME); // creates upgradeRequest and upgradeResult aspects - Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -121,7 +121,7 @@ public void testExecuteWithNewVersion() throws Exception { Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.CHART_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(1)).getEntitySpec(Constants.DASHBOARD_ENTITY_NAME); // creates upgradeRequest and upgradeResult aspects - Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -183,7 +183,7 @@ public void testDoesNotExecuteWithSameVersion() throws Exception { Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.CHART_ENTITY_NAME); Mockito.verify(mockRegistry, Mockito.times(0)).getEntitySpec(Constants.DASHBOARD_ENTITY_NAME); // creates upgradeRequest and upgradeResult aspects - Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) 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 2ebc16b3973d1..e6104c9c59063 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 @@ -105,7 +105,7 @@ public void testExecuteFirstTime() throws Exception { 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)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -166,7 +166,7 @@ public void testExecutesWithNewVersion() throws Exception { 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)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.eq(false) @@ -227,7 +227,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { Mockito.verify(mockSearchService, Mockito.times(0)).search(Constants.GLOSSARY_NODE_ENTITY_NAME, "", null, null, 0, 1000, new SearchFlags().setFulltext(false) .setSkipAggregates(true).setSkipHighlighting(true)); - Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean() diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java index a9410fb2522f2..5e4ad6e7fe880 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/UpgradeDefaultBrowsePathsStepTest.java @@ -87,7 +87,7 @@ public void testExecuteNoExistingBrowsePaths() throws Exception { Mockito.eq(5000) ); // Verify that 4 aspects are ingested, 2 for the upgrade request / result, but none for ingesting - Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -155,7 +155,7 @@ public void testExecuteFirstTime() throws Exception { Mockito.eq(5000) ); // Verify that 4 aspects are ingested, 2 for the upgrade request / result and 2 for the browse pahts - Mockito.verify(mockService, Mockito.times(4)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(4)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -223,7 +223,7 @@ public void testDoesNotRunWhenBrowsePathIsNotQualified() throws Exception { Mockito.eq(5000) ); // Verify that 2 aspects are ingested, only those for the upgrade step - Mockito.verify(mockService, Mockito.times(2)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(2)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(), Mockito.eq(false) @@ -249,7 +249,7 @@ public void testDoesNotRunWhenAlreadyExecuted() throws Exception { UpgradeDefaultBrowsePathsStep step = new UpgradeDefaultBrowsePathsStep(mockService); step.execute(); - Mockito.verify(mockService, Mockito.times(0)).ingestSingleProposal( + Mockito.verify(mockService, Mockito.times(0)).ingestProposal( Mockito.any(MetadataChangeProposal.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean() 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 ea0d5ff2b24f0..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.IngestResult; import com.linkedin.metadata.entity.RollbackRunResult; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +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; @@ -270,14 +272,14 @@ public static Pair ingestProposal(com.linkedin.mxe.MetadataChan Stream proposalStream = Stream.concat(Stream.of(serviceProposal), AspectUtils.getAdditionalChanges(serviceProposal, entityService).stream()); - AspectsBatch batch = AspectsBatch.builder().mcps(proposalStream.collect(Collectors.toList()), + AspectsBatch batch = AspectsBatchImpl.builder().mcps(proposalStream.collect(Collectors.toList()), entityService.getEntityRegistry()).build(); - Set proposalResult = + Set proposalResult = entityService.ingestProposal(batch, auditStamp, false); Urn urn = proposalResult.stream().findFirst().get().getUrn(); - return new Pair<>(urn.toString(), proposalResult.stream().anyMatch(EntityService.IngestResult::isSqlCommitted)); + 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 db155e192e653..6a0c336e42bfc 100644 --- a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java +++ b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java @@ -9,7 +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.EntityService; +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; @@ -66,8 +66,8 @@ public void setup() EntityRegistry mockEntityRegistry = new MockEntityRegistry(); AspectDao aspectDao = Mockito.mock(AspectDao.class); Mockito.when(aspectDao.runInTransactionWithRetry( - ArgumentMatchers.>any(), anyInt())).thenAnswer(i -> - ((Function) i.getArgument(0)).apply(Mockito.mock(Transaction.class)) + ArgumentMatchers.>any(), anyInt())).thenAnswer(i -> + ((Function) i.getArgument(0)).apply(Mockito.mock(Transaction.class)) ); EventProducer mockEntityEventProducer = Mockito.mock(EventProducer.class); 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 456c8c293ef1f..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; 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 846b098d1ab35..1ce57b4b95d9e 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,7 +8,9 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.linkedin.aspect.GetTimeseriesAspectValuesResponse; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; +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; @@ -207,22 +209,22 @@ public Task ingestProposal( final AspectsBatch batch; if (asyncBool) { // if async we'll expand the additional changes later, no need to do this early - batch = AspectsBatch.builder() + batch = AspectsBatchImpl.builder() .mcps(List.of(metadataChangeProposal), _entityService.getEntityRegistry()) .build(); } else { Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService).stream()); - batch = AspectsBatch.builder() + batch = AspectsBatchImpl.builder() .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) .build(); } - Set results = + Set results = _entityService.ingestProposal(batch, auditStamp, asyncBool); - EntityService.IngestResult one = results.stream() + IngestResult one = results.stream() .findFirst() .get(); diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java index f5a009b701888..3ff22fb767676 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/BatchIngestionRunResource.java @@ -294,7 +294,7 @@ private void updateExecutionRequestStatus(String runId, String status) { proposal.setAspect(GenericRecordUtils.serializeAspect(requestResult)); proposal.setChangeType(ChangeType.UPSERT); - _entityService.ingestSingleProposal(proposal, + _entityService.ingestProposal(proposal, new AuditStamp().setActor(UrnUtils.getUrn(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()), false); } } catch (Exception e) { 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 0c7a0bb6d47e2..59201d875e5ce 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 @@ -15,9 +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.UpdateAspectResult; import com.linkedin.metadata.entity.ebean.transactions.UpsertBatchItem; import com.linkedin.metadata.entity.EntityServiceImpl; -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; @@ -90,23 +90,23 @@ public void testAsyncDefaultAspects() throws URISyntaxException { .build(_entityRegistry); when(_aspectDao.runInTransactionWithRetry(any(), anyInt())) .thenReturn(List.of( - EntityService.UpdateAspectResult.builder().urn(urn) + UpdateAspectResult.builder().urn(urn) .newValue(new DatasetProperties().setName("name1")) .auditStamp(new AuditStamp()) .request(req).build(), - EntityService.UpdateAspectResult.builder().urn(urn) + UpdateAspectResult.builder().urn(urn) .newValue(new DatasetProperties().setName("name2")) .auditStamp(new AuditStamp()) .request(req).build(), - EntityService.UpdateAspectResult.builder().urn(urn) + UpdateAspectResult.builder().urn(urn) .newValue(new DatasetProperties().setName("name3")) .auditStamp(new AuditStamp()) .request(req).build(), - EntityService.UpdateAspectResult.builder().urn(urn) + UpdateAspectResult.builder().urn(urn) .newValue(new DatasetProperties().setName("name4")) .auditStamp(new AuditStamp()) .request(req).build(), - EntityService.UpdateAspectResult.builder().urn(urn) + UpdateAspectResult.builder().urn(urn) .newValue(new DatasetProperties().setName("name5")) .auditStamp(new AuditStamp()) .request(req).build())); 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 58e24825d1f2b..e062d55254f90 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 @@ -4,17 +4,11 @@ import com.google.common.collect.ImmutableSet; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; -import com.linkedin.data.schema.validation.ValidationResult; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.events.metadata.ChangeType; -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.GenericRecordUtils; import com.linkedin.mxe.GenericAspect; @@ -25,7 +19,6 @@ import java.util.Map; import java.util.Objects; import java.util.Set; -import java.util.function.Consumer; import java.util.stream.Collectors; import javax.annotation.Nonnull; import lombok.extern.slf4j.Slf4j; @@ -138,34 +131,4 @@ public static AuditStamp getAuditStamp(Urn actor) { auditStamp.setActor(actor); return auditStamp; } - - 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); - } } \ No newline at end of file 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 f95227373f556..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,7 +270,7 @@ 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 EntityService.IngestResult ingestProposalResult = _entityService.ingestSingleProposal(proposal, auditStamp, false); + final IngestResult ingestProposalResult = _entityService.ingestProposal(proposal, auditStamp, false); if (!ingestProposalResult.isSqlCommitted()) { log.error("Failed to ingest aspect with references removed. Before {}, after: {}, please check MCP processor" 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..a84b8aabee9e6 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); @@ -289,8 +237,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 b1d654864d674..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 @@ -7,9 +7,9 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.Constants; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatch; 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; @@ -111,14 +111,14 @@ public boolean setRetention(@Nullable String entityName, @Nullable String aspect AuditStamp auditStamp = new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()); - AspectsBatch batch = AspectsBatch.builder() - .mcps(List.of(keyProposal, aspectProposal), getEntityService().getEntityRegistry()) - .build(); + AspectsBatch batch = buildAspectsBatch(List.of(keyProposal, aspectProposal)); return getEntityService().ingestProposal(batch, auditStamp, false).stream() - .anyMatch(EntityService.IngestResult::isSqlCommitted); + .anyMatch(IngestResult::isSqlCommitted); } + protected abstract AspectsBatch buildAspectsBatch(List mcps); + /** * Delete the retention policy set for given entity and aspect. * diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/UpdateAspectResult.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/UpdateAspectResult.java index 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-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java similarity index 91% rename from metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java rename to metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java index a51bf0c943687..03a2b4e2a7f73 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AbstractBatchItem.java +++ b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/transactions/AbstractBatchItem.java @@ -1,11 +1,10 @@ -package com.linkedin.metadata.entity.ebean.transactions; - +package com.linkedin.metadata.entity.transactions; import com.linkedin.common.urn.Urn; import com.linkedin.events.metadata.ChangeType; -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.models.registry.template.AspectTemplateEngine; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; @@ -13,6 +12,9 @@ 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(); @@ -30,11 +32,13 @@ public abstract class AbstractBatchItem { 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(EntityService.DEFAULT_RUN_ID); + systemMetadata.setRunId(DEFAULT_RUN_ID); systemMetadata.setLastObserved(System.currentTimeMillis()); } return systemMetadata; 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()))); + } +} From fc4e48bcbf2d5164cadef5b72daaa68136f23674 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Mon, 24 Jul 2023 00:34:16 -0500 Subject: [PATCH 08/27] Complete futures for async produce --- .../metadata/entity/EntityServiceImpl.java | 41 +++++++++++++------ .../com/linkedin/metadata/ESTestUtils.java | 5 ++- .../boot/steps/IndexDataPlatformsStep.java | 18 +++++++- .../steps/RestoreColumnLineageIndices.java | 30 ++++++++++++-- .../boot/steps/RestoreDbtSiblingsIndices.java | 17 +++++++- .../boot/steps/RestoreGlossaryIndices.java | 31 ++++++++++++-- .../RestoreColumnLineageIndicesTest.java | 8 ++++ .../steps/RestoreGlossaryIndicesTest.java | 13 +++++- 8 files changed, 136 insertions(+), 27 deletions(-) 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 0d76ea1948214..6fbd027f34daf 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 @@ -1062,17 +1062,17 @@ public Map getEntities(@Nonnull final Set urns, @Nonnull Set, Boolean> alwaysProduceMCLAsync(@Nonnull final Urn urn, @Nonnull final AspectSpec aspectSpec, - @Nonnull final MetadataChangeLog metadataChangeLog) { + @Nonnull final MetadataChangeLog metadataChangeLog) { Future future = _producer.produceMetadataChangeLog(urn, aspectSpec, metadataChangeLog); return Pair.of(future, preprocessEvent(metadataChangeLog)); } @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) { + @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); @@ -1382,24 +1382,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); - alwaysProduceMCLAsync(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); } }); @@ -1438,11 +1447,19 @@ public RollbackRunResult deleteUrn(Urn urn) { rowsDeletedFromEntityDeletion = result.additionalRowsAffected; removedAspects.add(summary); - alwaysProduceMCLAsync(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); 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 c1297866edcc4..bf0d0a214135d 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-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 45eb9e35226e1..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.alwaysProduceMCLAsync( + 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/RestoreColumnLineageIndices.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndices.java index 654119287a6b1..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.alwaysProduceMCLAsync( + 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.alwaysProduceMCLAsync( + 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 4828e3b2b2b28..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.alwaysProduceMCLAsync( + 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 a5ea95a271b7f..23dd2284cb3a6 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.alwaysProduceMCLAsync( + 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.alwaysProduceMCLAsync( + 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/RestoreColumnLineageIndicesTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/RestoreColumnLineageIndicesTest.java index d90c1947b3e5e..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 { @@ -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 e6104c9c59063..cf23d001ae2cc 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 org.mockito.Mockito; import org.testng.annotations.Test; @@ -28,6 +29,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.Map; +import java.util.concurrent.Future; public class RestoreGlossaryIndicesTest { @@ -93,6 +95,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); @@ -154,6 +161,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); @@ -163,7 +175,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( From 85b51e450a46dbabd9b51dce0b79d05c70779889 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Mon, 24 Jul 2023 14:09:01 -0500 Subject: [PATCH 09/27] Add metrics around exceptions, skip retry on key aspects --- .../linkedin/metadata/entity/AspectDao.java | 7 ++++ .../metadata/entity/EntityServiceImpl.java | 2 +- .../metadata/entity/ebean/EbeanAspectDao.java | 34 +++++++++++++++++-- .../java/entities/EntitiesControllerTest.java | 2 +- .../resources/entity/AspectResourceTest.java | 2 +- 5 files changed, 42 insertions(+), 5 deletions(-) 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 4bd8547633dcb..bb1a35f5e130e 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,6 +3,7 @@ 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 io.ebean.PagedList; import io.ebean.Transaction; @@ -130,4 +131,10 @@ default Map getNextVersions(@Nonnull final String urn, @Nonnull fi @Nonnull 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); + } } 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 6fbd027f34daf..922e2b09113fa 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 @@ -635,7 +635,7 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa } return upsertResults; - }, DEFAULT_MAX_TRANSACTION_RETRY); + }, aspectsBatch, DEFAULT_MAX_TRANSACTION_RETRY); } @Nonnull 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 93c40aab2a1fa..dcd145893bf63 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java @@ -10,10 +10,15 @@ 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.elasticsearch.update.BulkListener; 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.ExpressionList; @@ -36,6 +41,7 @@ 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 java.util.stream.Collectors; @@ -489,9 +495,16 @@ public ListResult listLatestAspectMetadata( @Override @Nonnull 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 { @@ -501,8 +514,14 @@ public T runInTransactionWithRetry(@Nonnull final Function b transaction.commit(); lastException = null; break; - } catch (RollbackException | DuplicateKeyException exception) { + } catch (RollbackException exception) { lastException = exception; + } catch (DuplicateKeyException exception) { + 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."); + } else { + lastException = exception; + } } catch (PersistenceException exception) { // TODO: replace this logic by catching SerializableConflictException above once the exception is available SpiServer pluginApi = _server.getPluginApi(); @@ -533,6 +552,7 @@ public T runInTransactionWithRetry(@Nonnull final Function b } while (++retryCount <= maxTransactionRetry); if (lastException != null) { + incrementExceptionMetrics(batch, lastException); throw new RetryLimitReached("Failed to add after " + maxTransactionRetry + " retries", lastException); } @@ -684,4 +704,14 @@ private static Map> toUrnAspectMap(Collection< .map(e -> Map.entry(e.getKey(), toAspectMap(e.getValue()))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } + + private static void incrementExceptionMetrics(AspectsBatch batchOpt, Throwable failure) { + Optional.ofNullable(batchOpt).ifPresent(batch -> batch.getItems().stream() + .map(item -> buildMetricName(item.getEntitySpec(), item.getAspectSpec(), "retryLimitException")) + .forEach(metricName -> MetricUtils.exceptionCounter(BulkListener.class, metricName, failure))); + } + + 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-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java index 6a0c336e42bfc..229e71168557d 100644 --- a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java +++ b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java @@ -66,7 +66,7 @@ public void setup() EntityRegistry mockEntityRegistry = new MockEntityRegistry(); AspectDao aspectDao = Mockito.mock(AspectDao.class); Mockito.when(aspectDao.runInTransactionWithRetry( - ArgumentMatchers.>any(), anyInt())).thenAnswer(i -> + ArgumentMatchers.>any(), any(), anyInt())).thenAnswer(i -> ((Function) i.getArgument(0)).apply(Mockito.mock(Transaction.class)) ); 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 59201d875e5ce..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 @@ -88,7 +88,7 @@ public void testAsyncDefaultAspects() throws URISyntaxException { .aspect(mcp.getAspect()) .metadataChangeProposal(mcp) .build(_entityRegistry); - when(_aspectDao.runInTransactionWithRetry(any(), anyInt())) + when(_aspectDao.runInTransactionWithRetry(any(), any(), anyInt())) .thenReturn(List.of( UpdateAspectResult.builder().urn(urn) .newValue(new DatasetProperties().setName("name1")) From 843cf0675908c65e7c98daa35c738678b8e3c925 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 25 Jul 2023 11:06:36 -0500 Subject: [PATCH 10/27] Use fully qualified path for docker task lookup --- datahub-frontend/build.gradle | 2 +- datahub-upgrade/build.gradle | 2 +- docker/datahub-ingestion-base/build.gradle | 2 +- docker/datahub-ingestion-slim/build.gradle | 2 +- docker/datahub-ingestion/build.gradle | 2 +- docker/elasticsearch-setup/build.gradle | 2 +- docker/kafka-setup/build.gradle | 2 +- docker/mysql-setup/build.gradle | 2 +- docker/postgres-setup/build.gradle | 2 +- metadata-jobs/mae-consumer-job/build.gradle | 2 +- metadata-jobs/mce-consumer-job/build.gradle | 2 +- metadata-service/war/build.gradle | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) diff --git a/datahub-frontend/build.gradle b/datahub-frontend/build.gradle index f21d10d8f3842..22d57100cd924 100644 --- a/datahub-frontend/build.gradle +++ b/datahub-frontend/build.gradle @@ -94,7 +94,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-upgrade/build.gradle b/datahub-upgrade/build.gradle index ad2bf02bfdcc7..b246d7bee1f9c 100644 --- a/datahub-upgrade/build.gradle +++ b/datahub-upgrade/build.gradle @@ -97,7 +97,7 @@ docker { load(true) push(false) } -tasks.getByName("docker").dependsOn([bootJar]) +tasks.getByPath(":datahub-upgrade:docker").dependsOn([bootJar]) task cleanLocalDockerImages { doLast { diff --git a/docker/datahub-ingestion-base/build.gradle b/docker/datahub-ingestion-base/build.gradle index fe3c12a59886f..18d1ad3d193e9 100644 --- a/docker/datahub-ingestion-base/build.gradle +++ b/docker/datahub-ingestion-base/build.gradle @@ -19,7 +19,7 @@ docker { include "docker/${docker_dir}/*" } } -tasks.getByPath('docker').dependsOn('build') +tasks.getByPath(':docker:datahub-ingestion-base:docker').dependsOn('build') task mkdirBuildDocker { doFirst { diff --git a/docker/datahub-ingestion-slim/build.gradle b/docker/datahub-ingestion-slim/build.gradle index f21b66b576a0c..7d0541fd5102c 100644 --- a/docker/datahub-ingestion-slim/build.gradle +++ b/docker/datahub-ingestion-slim/build.gradle @@ -22,7 +22,7 @@ docker { buildx(false) } -tasks.getByPath('docker').dependsOn(['build', ':docker:datahub-ingestion:docker']) +tasks.getByPath(':docker:datahub-ingestion-slim:docker').dependsOn(['build', ':docker:datahub-ingestion:docker']) task mkdirBuildDocker { doFirst { diff --git a/docker/datahub-ingestion/build.gradle b/docker/datahub-ingestion/build.gradle index 7a24d87794c0e..d7999556e1c2c 100644 --- a/docker/datahub-ingestion/build.gradle +++ b/docker/datahub-ingestion/build.gradle @@ -28,7 +28,7 @@ docker { } buildArgs([DOCKER_VERSION: version]) } -tasks.getByPath('docker').dependsOn(['build', ':docker:datahub-ingestion-base:docker']) +tasks.getByPath(':docker:datahub-ingestion:docker').dependsOn(['build', ':docker:datahub-ingestion-base:docker']) task mkdirBuildDocker { doFirst { diff --git a/docker/elasticsearch-setup/build.gradle b/docker/elasticsearch-setup/build.gradle index cc2fe1ec5c4db..73c0b85abdfcc 100644 --- a/docker/elasticsearch-setup/build.gradle +++ b/docker/elasticsearch-setup/build.gradle @@ -25,7 +25,7 @@ docker { load(true) push(false) } -tasks.getByPath('docker').dependsOn('build') +tasks.getByPath(':docker:elasticsearch-setup:docker').dependsOn('build') task mkdirBuildDocker { doFirst { diff --git a/docker/kafka-setup/build.gradle b/docker/kafka-setup/build.gradle index a5d33457e45f7..7f43a79b9d9a6 100644 --- a/docker/kafka-setup/build.gradle +++ b/docker/kafka-setup/build.gradle @@ -24,7 +24,7 @@ docker { load(true) push(false) } -tasks.getByPath('docker').dependsOn('build') +tasks.getByPath(':docker:kafka-setup:docker').dependsOn('build') task mkdirBuildDocker { doFirst { diff --git a/docker/mysql-setup/build.gradle b/docker/mysql-setup/build.gradle index 48a28f15a581d..40a25c13de0bf 100644 --- a/docker/mysql-setup/build.gradle +++ b/docker/mysql-setup/build.gradle @@ -25,7 +25,7 @@ docker { load(true) push(false) } -tasks.getByPath('docker').dependsOn('build') +tasks.getByPath(':docker:mysql-setup:docker').dependsOn('build') task mkdirBuildDocker { doFirst { diff --git a/docker/postgres-setup/build.gradle b/docker/postgres-setup/build.gradle index a5b0413ec4be8..47825ad431619 100644 --- a/docker/postgres-setup/build.gradle +++ b/docker/postgres-setup/build.gradle @@ -25,7 +25,7 @@ docker { load(true) push(false) } -tasks.getByPath('docker').dependsOn('build') +tasks.getByPath(':docker:postgres-setup:docker').dependsOn('build') task mkdirBuildDocker { doFirst { diff --git a/metadata-jobs/mae-consumer-job/build.gradle b/metadata-jobs/mae-consumer-job/build.gradle index e7941a04224e3..51e37edc9fb92 100644 --- a/metadata-jobs/mae-consumer-job/build.gradle +++ b/metadata-jobs/mae-consumer-job/build.gradle @@ -51,7 +51,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/mce-consumer-job/build.gradle b/metadata-jobs/mce-consumer-job/build.gradle index 5981284e9da3f..587b744923276 100644 --- a/metadata-jobs/mce-consumer-job/build.gradle +++ b/metadata-jobs/mce-consumer-job/build.gradle @@ -64,7 +64,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-service/war/build.gradle b/metadata-service/war/build.gradle index 7e9aa90664611..5aeb32097d210 100644 --- a/metadata-service/war/build.gradle +++ b/metadata-service/war/build.gradle @@ -80,7 +80,7 @@ docker { load(true) push(false) } -tasks.getByName("docker").dependsOn([build, war]) +tasks.getByPath(":metadata-service:war:docker").dependsOn([build, war]) task cleanLocalDockerImages { doLast { From d19c533a5e22dc63bd8587768b9bc688d292e77d Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 25 Jul 2023 11:06:54 -0500 Subject: [PATCH 11/27] Avoid pulling mae/mce consumers in smoke-tests (not used) --- .github/workflows/docker-unified.yml | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index a23e5a134d294..b8b4dbee88037 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -429,8 +429,8 @@ jobs: kafka_setup_build, mysql_setup_build, elasticsearch_setup_build, - mae_consumer_build, - mce_consumer_build, +# mae_consumer_build, +# mce_consumer_build, datahub_upgrade_build, ] steps: @@ -475,16 +475,16 @@ jobs: if: ${{ needs.setup.outputs.publish != 'true' }} with: image: ${{ env.DATAHUB_ELASTIC_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - - name: Download MCE Consumer image - uses: ishworkh/docker-image-artifact-download@v1 - if: ${{ needs.setup.outputs.publish != 'true' }} - with: - image: ${{ env.DATAHUB_MCE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - - name: Download MAE Consumer image - uses: ishworkh/docker-image-artifact-download@v1 - if: ${{ needs.setup.outputs.publish != 'true' }} - with: - image: ${{ env.DATAHUB_MAE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} +# - name: Download MCE Consumer image +# uses: ishworkh/docker-image-artifact-download@v1 +# if: ${{ needs.setup.outputs.publish != 'true' }} +# with: +# image: ${{ env.DATAHUB_MCE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} +# - name: Download MAE Consumer image +# uses: ishworkh/docker-image-artifact-download@v1 +# if: ${{ needs.setup.outputs.publish != 'true' }} +# with: +# image: ${{ env.DATAHUB_MAE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - name: Download upgrade image uses: ishworkh/docker-image-artifact-download@v1 if: ${{ needs.setup.outputs.publish != 'true' }} From f448e0e5cccab39c6a462894a336665b28529d2f Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 25 Jul 2023 14:27:20 -0500 Subject: [PATCH 12/27] try non-standalone --- smoke-test/run-quickstart.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/smoke-test/run-quickstart.sh b/smoke-test/run-quickstart.sh index 050b5d2db95c9..da237f4c69495 100755 --- a/smoke-test/run-quickstart.sh +++ b/smoke-test/run-quickstart.sh @@ -15,4 +15,5 @@ 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 +# --standalone_consumers +datahub docker quickstart --version ${DATAHUB_VERSION} --dump-logs-on-failure --kafka-setup From 38f2ba82d8b571f085e7af6457110a940cbce8d2 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 25 Jul 2023 14:28:29 -0500 Subject: [PATCH 13/27] drop kafka-setup --- .github/workflows/docker-unified.yml | 12 ++++++------ smoke-test/run-quickstart.sh | 4 +++- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index b8b4dbee88037..2b3a4d4da9dfc 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -426,7 +426,7 @@ jobs: setup, gms_build, frontend_build, - kafka_setup_build, +# kafka_setup_build, mysql_setup_build, elasticsearch_setup_build, # mae_consumer_build, @@ -460,11 +460,11 @@ jobs: if: ${{ needs.setup.outputs.publish != 'true' }} with: image: ${{ env.DATAHUB_FRONTEND_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - - name: Download Kafka Setup image - uses: ishworkh/docker-image-artifact-download@v1 - if: ${{ needs.setup.outputs.publish != 'true' }} - with: - image: ${{ env.DATAHUB_KAFKA_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }} +# - name: Download Kafka Setup image +# uses: ishworkh/docker-image-artifact-download@v1 +# if: ${{ needs.setup.outputs.publish != 'true' }} +# with: +# image: ${{ env.DATAHUB_KAFKA_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - name: Download Mysql Setup image uses: ishworkh/docker-image-artifact-download@v1 if: ${{ needs.setup.outputs.publish != 'true' }} diff --git a/smoke-test/run-quickstart.sh b/smoke-test/run-quickstart.sh index da237f4c69495..6998267981b3c 100755 --- a/smoke-test/run-quickstart.sh +++ b/smoke-test/run-quickstart.sh @@ -15,5 +15,7 @@ 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} --dump-logs-on-failure + # --standalone_consumers -datahub docker quickstart --version ${DATAHUB_VERSION} --dump-logs-on-failure --kafka-setup +# --kafka-setup From 3122414326f541631e3ce3f3b075a5d9a63e18a4 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 25 Jul 2023 20:06:37 -0500 Subject: [PATCH 14/27] kafka is required --- .github/workflows/docker-unified.yml | 12 ++++++------ metadata-ingestion/src/datahub/cli/docker_cli.py | 1 + smoke-test/run-quickstart.sh | 3 +-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index 2b3a4d4da9dfc..b8b4dbee88037 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -426,7 +426,7 @@ jobs: setup, gms_build, frontend_build, -# kafka_setup_build, + kafka_setup_build, mysql_setup_build, elasticsearch_setup_build, # mae_consumer_build, @@ -460,11 +460,11 @@ jobs: if: ${{ needs.setup.outputs.publish != 'true' }} with: image: ${{ env.DATAHUB_FRONTEND_IMAGE }}:${{ needs.setup.outputs.unique_tag }} -# - name: Download Kafka Setup image -# uses: ishworkh/docker-image-artifact-download@v1 -# if: ${{ needs.setup.outputs.publish != 'true' }} -# with: -# image: ${{ env.DATAHUB_KAFKA_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }} + - name: Download Kafka Setup image + uses: ishworkh/docker-image-artifact-download@v1 + if: ${{ needs.setup.outputs.publish != 'true' }} + with: + image: ${{ env.DATAHUB_KAFKA_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - name: Download Mysql Setup image uses: ishworkh/docker-image-artifact-download@v1 if: ${{ needs.setup.outputs.publish != 'true' }} diff --git a/metadata-ingestion/src/datahub/cli/docker_cli.py b/metadata-ingestion/src/datahub/cli/docker_cli.py index e2b7b2a2e1ff4..c516f53cb5755 100644 --- a/metadata-ingestion/src/datahub/cli/docker_cli.py +++ b/metadata-ingestion/src/datahub/cli/docker_cli.py @@ -893,6 +893,7 @@ def download_compose_files( tmp_file.write(quickstart_download_response.content) logger.debug(f"Copied to {path}") if kafka_setup: + base_url = get_docker_compose_base_url(compose_git_ref) kafka_setup_github_file = f"{base_url}/{KAFKA_SETUP_QUICKSTART_COMPOSE_FILE}" default_kafka_compose_file = ( diff --git a/smoke-test/run-quickstart.sh b/smoke-test/run-quickstart.sh index 6998267981b3c..74ba6c9d24cc9 100755 --- a/smoke-test/run-quickstart.sh +++ b/smoke-test/run-quickstart.sh @@ -15,7 +15,6 @@ 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} --dump-logs-on-failure +datahub docker quickstart --version ${DATAHUB_VERSION} --kafka-setup --dump-logs-on-failure # --standalone_consumers -# --kafka-setup From da47ea14b0002160be0d0387f6d7eb68a508d402 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Thu, 27 Jul 2023 08:09:43 -0500 Subject: [PATCH 15/27] disable ES threshold --- .github/workflows/docker-unified.yml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index b8b4dbee88037..53a3123378142 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -501,6 +501,18 @@ jobs: # we are doing this because gms takes time to get ready # and we don't have a better readiness check when bootstrap is done sleep 60s + - name: Disable ES Disk Threshold + run: | + curl -XPUT "http://localhost:9200/_cluster/settings" \ + -H 'Content-Type: application/json' -d'{ + "persistent": { + "cluster": { + "routing": { + "allocation.disk.threshold_enabled": false + } + } + } + }' - name: Smoke test env: RUN_QUICKSTART: false From 846d95b110c485f2bbc455866ee239416ba92471 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 1 Aug 2023 16:12:47 -0500 Subject: [PATCH 16/27] fix(test): increase siblings.js test stability --- datahub-web-react/src/app/lineage/LineageEntityNode.tsx | 9 ++++++--- .../tests/cypress/cypress/e2e/siblings/siblings.js | 8 +++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx index 5b39d6aaf19c3..4526e3a225ce2 100644 --- a/datahub-web-react/src/app/lineage/LineageEntityNode.tsx +++ b/datahub-web-react/src/app/lineage/LineageEntityNode.tsx @@ -132,6 +132,10 @@ export default function LineageEntityNode({ areColumnsCollapsed, ); + const entityName = + capitalizeFirstLetterOnly(node.data.subtype) || + (node.data.type && entityRegistry.getEntityName(node.data.type)); + return ( {unexploredHiddenChildren && (isHovered || isSelected) ? ( @@ -335,9 +339,8 @@ export default function LineageEntityNode({ {' '} |{' '} - - {capitalizeFirstLetterOnly(node.data.subtype) || - (node.data.type && entityRegistry.getEntityName(node.data.type))} + + {entityName} {expandTitles ? ( diff --git a/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js b/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js index b7d48703992e2..00de08e77a185 100644 --- a/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js +++ b/smoke-test/tests/cypress/cypress/e2e/siblings/siblings.js @@ -113,9 +113,11 @@ describe('siblings', () => { cy.clickOptionWithTestId('compress-lineage-toggle'); // check the subtypes - cy.get('text:contains(View)').should('have.length', 2); - cy.get('text:contains(Table)').should('have.length', 0); - cy.get('text:contains(Seed)').should('have.length', 1); + cy.get('[data-testid="Seed"]').should('have.length', 1); + // center counts twice since we secretely render two center nodes, plus the downstream bigquery + cy.get('[data-testid="View"]').should('have.length', 3); + cy.get('[data-testid="Table"]').should('have.length', 0); + // check the names cy.get('text:contains(raw_orders)').should('have.length', 1); From 9ef881b3f7c889af66c95d8d3cbd8c9a80f7a6e5 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 1 Aug 2023 18:48:20 -0500 Subject: [PATCH 17/27] lint --- .../java/com/linkedin/metadata/entity/EntityServiceImpl.java | 1 + 1 file changed, 1 insertion(+) 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 e87c80318331b..121e14b8288de 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 @@ -66,6 +66,7 @@ import io.ebean.PagedList; import java.net.URISyntaxException; +import java.nio.charset.StandardCharsets; import java.sql.Timestamp; import java.util.ArrayList; import java.util.Arrays; From 1fe33b6689743d1c601b462f4fd30d875960020a Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 1 Aug 2023 20:14:17 -0500 Subject: [PATCH 18/27] drop mce/mae consumers, space issues --- .github/workflows/docker-unified.yml | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index 8389b2768d5be..b4df9ae52aad6 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -429,8 +429,8 @@ jobs: kafka_setup_build, mysql_setup_build, elasticsearch_setup_build, - mae_consumer_build, - mce_consumer_build, +# mae_consumer_build, +# mce_consumer_build, datahub_upgrade_build, ] steps: @@ -475,16 +475,16 @@ jobs: if: ${{ needs.setup.outputs.publish != 'true' }} with: image: ${{ env.DATAHUB_ELASTIC_SETUP_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - - name: Download MCE Consumer image - uses: ishworkh/docker-image-artifact-download@v1 - if: ${{ needs.setup.outputs.publish != 'true' }} - with: - image: ${{ env.DATAHUB_MCE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - - name: Download MAE Consumer image - uses: ishworkh/docker-image-artifact-download@v1 - if: ${{ needs.setup.outputs.publish != 'true' }} - with: - image: ${{ env.DATAHUB_MAE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} +# - name: Download MCE Consumer image +# uses: ishworkh/docker-image-artifact-download@v1 +# if: ${{ needs.setup.outputs.publish != 'true' }} +# with: +# image: ${{ env.DATAHUB_MCE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} +# - name: Download MAE Consumer image +# uses: ishworkh/docker-image-artifact-download@v1 +# if: ${{ needs.setup.outputs.publish != 'true' }} +# with: +# image: ${{ env.DATAHUB_MAE_CONSUMER_IMAGE }}:${{ needs.setup.outputs.unique_tag }} - name: Download upgrade image uses: ishworkh/docker-image-artifact-download@v1 if: ${{ needs.setup.outputs.publish != 'true' }} From bd752be640b1a3f0c9d4ee021e3e243735a9d71b Mon Sep 17 00:00:00 2001 From: David Leifker Date: Thu, 10 Aug 2023 21:19:52 -0500 Subject: [PATCH 19/27] add multi-threading test --- metadata-io/build.gradle | 1 + .../metadata/entity/EntityServiceImpl.java | 2 +- .../ebean/AspectStorageValidationUtil.java | 4 +- .../com/linkedin/metadata/EbeanTestUtils.java | 2 +- .../entity/EbeanEntityServiceTest.java | 110 ++++++- .../java/io/datahub/test/DataGenerator.java | 303 ++++++++++++++++++ 6 files changed, 417 insertions(+), 5 deletions(-) create mode 100644 metadata-io/src/test/java/io/datahub/test/DataGenerator.java diff --git a/metadata-io/build.gradle b/metadata-io/build.gradle index 98b741b6f51a6..405603eebca9e 100644 --- a/metadata-io/build.gradle +++ b/metadata-io/build.gradle @@ -66,6 +66,7 @@ dependencies { testImplementation project(':datahub-graphql-core') // 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/entity/EntityServiceImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java index 121e14b8288de..1cf622277894b 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 @@ -1156,7 +1156,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); } 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..1474444c7ef45 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 @@ -30,7 +30,7 @@ public static long getV2NonSystemRowCount(EbeanServer server) { public static boolean checkV2TableExists(EbeanServer 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 List rows = query.findList(); @@ -40,7 +40,7 @@ public static boolean checkV2TableExists(EbeanServer server) { public static boolean checkV1TableExists(EbeanServer 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 List rows = query.findList(); 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..16d5888bd122d 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java @@ -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/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java index bf720166fdd9a..9bc430267f13e 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; @@ -18,16 +20,25 @@ 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.datahub.test.DataGenerator; import io.ebean.EbeanServer; 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.List; +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; @@ -51,7 +62,7 @@ public void setupTest() { EbeanServer 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); @@ -225,4 +236,101 @@ public void testNestedTransactions() throws AssertionError { } System.out.println("done"); } + + /** + * 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.getEntityRegistry()); + EbeanServer server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); + final int testEntityCount = 25; + + int count = Objects.requireNonNull(server.createSqlQuery( + "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 int testThreads = 15; + final LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(testThreads * 2); + + // Spin up workers + List writeThreads = IntStream.range(0, testThreads) + .mapToObj(threadId -> new Thread(new MultiThreadTestWorker(queue, _entityServiceImpl))) + .collect(Collectors.toList()); + writeThreads.forEach(Thread::start); + + // Add data + List aspects = List.of("status", "globalTags", "glossaryTerms"); + + List generatedMCPs = dataGenerator.generateMCPs("dataset", testEntityCount, aspects) + .collect(Collectors.toList()); + + generatedMCPs.forEach(mcp -> { + try { + queue.put(List.of(mcp)); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + + // Terminate workers with empty mcp + IntStream.range(0, testThreads).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(5000); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + + // Expected aspects + Set> generatedAspectIds = generatedMCPs.stream() + .map(mcp -> Triple.of(mcp.getEntityUrn().toString(), mcp.getAspectName(), 0L)) + .collect(Collectors.toSet()); + + // Actual inserts + Set> actualAspectIds = server.createSqlQuery( + "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()); + assertEquals(actualAspectIds, generatedAspectIds); + } + + public 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 mcp = queue.take(); + if (mcp.isEmpty()) { + break; + } + final AuditStamp auditStamp = new AuditStamp(); + auditStamp.setActor(Urn.createFromString(Constants.DATAHUB_ACTOR)); + auditStamp.setTime(System.currentTimeMillis()); + entityService.ingestProposal(mcp.get(0), auditStamp, false); + } + } catch (InterruptedException | URISyntaxException ie) { + throw new RuntimeException(ie); + } + } + } } 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..e26c26b467c23 --- /dev/null +++ b/metadata-io/src/test/java/io/datahub/test/DataGenerator.java @@ -0,0 +1,303 @@ +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.metadata.Constants; +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.LinkedList; +import java.util.List; +import java.util.Map; +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; + + public DataGenerator(EntityRegistry entityRegistry) { + this.entityRegistry = entityRegistry; + } + + public Stream generateDatasets() { + return generateMCPs("dataset", 10, List.of()); + } + + public Stream generateMCPs(String entityName, long count, List aspects) { + EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName); + + 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 -> { + List additionalMCPs = new LinkedList<>(); + + // Expand with aspects + for (String aspectName : aspects) { + AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); + if (aspectSpec == null) { + throw new IllegalStateException("Aspect " + aspectName + " not found for entity " + entityName); + } + + RecordTemplate aspect = overrideRandomAspectSuppliers.getOrDefault(aspectName, + DataGenerator::defaultRandomAspect).apply(entitySpec, aspectSpec); + + // Maybe nested, like globalTags/glossaryTerms + additionalMCPs.addAll(generateRandomNestedArray(aspect, aspectSpec, 5)); + + 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()); + }); + } + + public static Map> overrideRandomAspectSuppliers = Map.of( + ); + + 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 = FAKER.lorem().characters(8, 16, false); + + switch (aspectSpec.getName()) { + case "glossaryTermInfo": + if (stringMethod.getName().equals("setName")) { + value = normalize(FAKER.company().buzzword()); + } + break; + default: + break; + } + + // global + if (stringMethod.getName().toLowerCase().contains("description") + || stringMethod.getName().toLowerCase().contains("definition")) { + value = FAKER.lorem().paragraph(); + } + + stringMethod.invoke(aspect, value); + } + + 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 List generateRandomNestedArray(RecordTemplate aspect, AspectSpec nestedAspect, int count) { + try { + switch (nestedAspect.getName()) { + case "globalTags": + List tags = generateMCPs("tag", count, List.of()) + .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; + case "glossaryTerms": + List terms = generateMCPs("glossaryTerm", count, + List.of("glossaryTermInfo")).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; + default: + return List.of(); + } + } 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; + 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; + } + } + + 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(FAKER.device().platform()))); + 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+", "_"); + } +} From 8b666429858f7a474d7bb215e58879c73a9cf407 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Thu, 10 Aug 2023 22:32:59 -0500 Subject: [PATCH 20/27] fix merge, ebean upgrade --- build.gradle | 6 ++++-- .../resolvers/mutate/util/DeleteUtils.java | 1 - .../resolvers/mutate/util/DeprecationUtils.java | 1 - .../resolvers/mutate/util/DomainUtils.java | 1 - .../resolvers/mutate/util/LabelUtils.java | 1 - .../resolvers/mutate/util/OwnerUtils.java | 1 - .../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 | 6 +++--- .../datahub/upgrade/nocode/NoCodeUpgrade.java | 8 ++++---- .../upgrade/nocode/RemoveAspectV2TableStep.java | 8 ++++---- .../upgrade/nocode/UpgradeQualificationStep.java | 8 ++++---- .../nocodecleanup/DeleteAspectTableStep.java | 8 ++++---- .../nocodecleanup/NoCodeCleanupUpgrade.java | 8 ++++---- .../NoCodeUpgradeQualificationStep.java | 6 +++--- .../restorebackup/ClearAspectV2TableStep.java | 6 +++--- .../upgrade/restorebackup/RestoreBackup.java | 6 +++--- .../upgrade/restoreindices/RestoreIndices.java | 6 +++--- .../upgrade/restoreindices/SendMAEStep.java | 6 +++--- .../UpgradeCliApplicationTestConfiguration.java | 4 ++-- metadata-io/build.gradle | 1 + .../metadata/entity/EntityServiceImpl.java | 2 +- .../linkedin/metadata/entity/EntityUtils.java | 10 ++++------ .../ebean/AspectStorageValidationUtil.java | 14 +++++++------- .../metadata/entity/ebean/EbeanAspectDao.java | 16 ++++++++-------- .../entity/ebean/EbeanRetentionService.java | 4 ++-- .../com/linkedin/metadata/AspectUtilsTest.java | 4 ++-- .../com/linkedin/metadata/EbeanTestUtils.java | 8 ++++---- .../entity/EbeanAspectMigrationsDaoTest.java | 4 ++-- .../metadata/entity/EbeanEntityServiceTest.java | 12 ++++++------ .../timeline/EbeanTimelineServiceTest.java | 4 ++-- .../MaeConsumerApplicationTestConfiguration.java | 4 ++-- .../MceConsumerApplicationTestConfiguration.java | 4 ++-- .../gms/factory/entity/EbeanServerFactory.java | 6 +++--- .../factory/entity/EntityAspectDaoFactory.java | 4 ++-- .../entity/EntityAspectMigrationsDaoFactory.java | 4 ++-- .../factory/entity/RetentionServiceFactory.java | 4 ++-- 40 files changed, 107 insertions(+), 111 deletions(-) diff --git a/build.gradle b/build.gradle index 605b4fcc050e7..e747409b187e5 100644 --- a/build.gradle +++ b/build.gradle @@ -18,6 +18,7 @@ buildscript { ext.logbackClassic = '1.2.12' ext.hadoop3Version = '3.3.5' ext.kafkaVersion = '2.3.0' + 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-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java index 5556c818359ca..7d4c5bee61e19 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java @@ -12,7 +12,6 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java index 78862ecf3057d..bd82bbb8e514f 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java @@ -14,7 +14,6 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java index eca9b19e207ee..b57160be09d32 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java @@ -14,7 +14,6 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java index 96ed8b38fd835..a93c7d5b333da 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java @@ -20,7 +20,6 @@ import com.linkedin.metadata.Constants; import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.EntityUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.schema.EditableSchemaFieldInfo; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java index c006f9c93972f..d2f7f896e5953 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java @@ -21,7 +21,6 @@ import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.EntityUtils; -import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.mxe.MetadataChangeProposal; import java.util.ArrayList; import java.util.List; 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 c321a2700091a..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 @@ -18,7 +18,7 @@ import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import com.linkedin.metadata.models.EntitySpec; import com.linkedin.util.Pair; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.PagedList; import java.net.URISyntaxException; import java.util.HashSet; @@ -37,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; 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/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/metadata-io/build.gradle b/metadata-io/build.gradle index 405603eebca9e..446b9decee3be 100644 --- a/metadata-io/build.gradle +++ b/metadata-io/build.gradle @@ -37,6 +37,7 @@ dependencies { compile externalDependency.kafkaClients compile externalDependency.ebean enhance externalDependency.ebeanAgent + implementation externalDependency.ebeanDdl compile externalDependency.opentelemetryAnnotations compile externalDependency.resilience4j compile externalDependency.springContext 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 1cf622277894b..ddccf0bacd80d 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 @@ -133,7 +133,7 @@ public class EntityServiceImpl implements EntityService { * monotonically increasing version incrementing as usual once the latest version is replaced. */ - 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; 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 71df144970e2c..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,6 +8,7 @@ import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.template.RecordTemplate; import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl; import com.linkedin.metadata.entity.validation.EntityRegistryUrnValidator; import com.linkedin.metadata.entity.validation.RecordTemplateValidator; import com.linkedin.metadata.models.AspectSpec; @@ -22,10 +23,9 @@ import lombok.extern.slf4j.Slf4j; -import java.net.URISyntaxException; -import java.util.List; 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; @@ -68,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); } /** 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 1474444c7ef45..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 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 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 052d51eb45f61..7164339cba3a9 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 @@ -21,7 +21,7 @@ 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; @@ -48,7 +48,7 @@ import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import javax.persistence.PersistenceException; +import io.ebean.SerializableConflictException; import javax.persistence.RollbackException; import javax.persistence.Table; @@ -61,7 +61,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(); @@ -74,7 +74,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; } @@ -84,10 +84,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; } @@ -525,9 +525,9 @@ public T runInTransactionWithRetry(@Nonnull final Function b MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); lastException = exception; } - } catch (PersistenceException exception) { + } catch (SerializableConflictException 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(); 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 33888594c4458..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 @@ -14,7 +14,7 @@ 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; @@ -41,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(); 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/EbeanTestUtils.java b/metadata-io/src/test/java/com/linkedin/metadata/EbeanTestUtils.java index 16d5888bd122d..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 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 9bc430267f13e..e6c94050edec2 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 @@ -23,7 +23,7 @@ import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; import io.datahub.test.DataGenerator; -import io.ebean.EbeanServer; +import io.ebean.Database; import io.ebean.Transaction; import io.ebean.TxScope; import io.ebean.annotation.TxIsolation; @@ -59,7 +59,7 @@ public EbeanEntityServiceTest() throws EntityRegistryException { @BeforeMethod public void setupTest() { - EbeanServer server = EbeanTestUtils.createTestServer(); + Database server = EbeanTestUtils.createTestServer(); _mockProducer = mock(EventProducer.class); _aspectDao = new EbeanAspectDao(server); @@ -216,7 +216,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))) { @@ -244,10 +244,10 @@ public void testNestedTransactions() throws AssertionError { @Test public void multiThreadingTest() { DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl.getEntityRegistry()); - EbeanServer server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); + Database server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); final int testEntityCount = 25; - int count = Objects.requireNonNull(server.createSqlQuery( + 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."); @@ -300,7 +300,7 @@ public void multiThreadingTest() { .collect(Collectors.toSet()); // Actual inserts - Set> actualAspectIds = server.createSqlQuery( + 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()); 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-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/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/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; From 2f6da4c33b9a0991e0393a248dbca374aeb69ac7 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Fri, 11 Aug 2023 17:56:58 -0500 Subject: [PATCH 21/27] Improve tests Improve tx handling Prevent deadlock dropped writes --- .../linkedin/metadata/entity/AspectDao.java | 54 +++--- .../metadata/entity/EntityServiceImpl.java | 152 +++++++++------ .../entity/cassandra/CassandraAspectDao.java | 54 +++--- .../metadata/entity/ebean/EbeanAspectDao.java | 128 +++++-------- .../entity/EbeanEntityServiceTest.java | 149 +++++++++++---- .../java/io/datahub/test/DataGenerator.java | 180 ++++++++++++------ .../linkedin/metadata/entity/AspectUtils.java | 42 +++- .../metadata/entity/EntityService.java | 27 ++- .../src/main/java/mock/MockEntitySpec.java | 2 +- 9 files changed, 485 insertions(+), 303 deletions(-) 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 02bad2fd8c9c0..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 @@ -57,34 +57,36 @@ default EntityAspect getLatestAspect(@Nonnull final String urn, @Nonnull final S 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( @@ -101,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( 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 ddccf0bacd80d..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 @@ -90,6 +90,7 @@ import javax.annotation.Nullable; import javax.persistence.EntityNotFoundException; +import io.ebean.Transaction; import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.*; @@ -595,7 +596,7 @@ private List ingestAspectsToLocalDB(@Nonnull final AspectsBa final UpdateAspectResult result; if (overwrite || latest == null) { - result = ingestAspectToLocalDBNoTransaction(item.getUrn(), item.getAspectName(), item.getAspect(), + result = ingestAspectToLocalDB(tx, item.getUrn(), item.getAspectName(), item.getAspect(), auditStamp, item.getSystemMetadata(), latest, nextVersion).toBuilder().request(item).build(); // support inner-batch upserts @@ -1198,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 = EntityUtils.buildKeyAspect(_entityRegistry, 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, @@ -1266,7 +1297,7 @@ 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::getFirst).collect(Collectors.toSet()))); + aspectRecordsToIngest.stream().collect(Collectors.toMap(Pair::getKey, Pair::getValue)))); AspectsBatchImpl aspectsBatch = AspectsBatchImpl.builder() .items(aspectRecordsToIngest.stream().map(pair -> UpsertBatchItem.builder() @@ -1572,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 @@ -1583,15 +1614,15 @@ public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull 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); @@ -333,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())) @@ -346,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)) @@ -489,19 +489,20 @@ public Map> getNextVersions(Map> u @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(); @@ -587,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( @@ -609,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/ebean/EbeanAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/EbeanAspectDao.java index 7164339cba3a9..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 @@ -17,7 +17,6 @@ import com.linkedin.metadata.query.ExtraInfo; import com.linkedin.metadata.query.ExtraInfoArray; import com.linkedin.metadata.query.ListResultMetadata; -import com.linkedin.metadata.search.elasticsearch.update.BulkListener; import com.linkedin.metadata.search.utils.QueryUtils; import com.linkedin.metadata.utils.metrics.MetricUtils; import io.ebean.DuplicateKeyException; @@ -30,10 +29,8 @@ import io.ebean.RawSqlBuilder; import io.ebean.Transaction; import io.ebean.TxScope; -import io.ebean.annotation.Platform; import io.ebean.annotation.TxIsolation; import java.net.URISyntaxException; -import java.sql.SQLException; import java.sql.Timestamp; import java.time.Clock; import java.util.ArrayList; @@ -42,18 +39,15 @@ 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 java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import io.ebean.SerializableConflictException; -import javax.persistence.RollbackException; + +import javax.persistence.PersistenceException; import javax.persistence.Table; -import io.ebean.config.dbplatform.DatabasePlatform; -import io.ebean.plugin.SpiServer; import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION; @@ -114,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(); @@ -137,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(); @@ -170,21 +166,21 @@ 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); } } @@ -238,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 @@ -515,43 +511,17 @@ public T runInTransactionWithRetry(@Nonnull final Function b transaction.commit(); lastException = null; break; - } catch (RollbackException exception) { - MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); - lastException = exception; - } catch (DuplicateKeyException exception) { - 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."); - } else { - MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); - lastException = exception; - } - } catch (SerializableConflictException exception) { - MetricUtils.counter(MetricRegistry.name(this.getClass(), "txFailed")).inc(); - - 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; - } + } catch (PersistenceException exception) { + 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); @@ -709,12 +679,6 @@ private static Map> toUrnAspectMap(Collection< .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } - private static void incrementExceptionMetrics(AspectsBatch batchOpt, Throwable failure) { - Optional.ofNullable(batchOpt).ifPresent(batch -> batch.getItems().stream() - .map(item -> buildMetricName(item.getEntitySpec(), item.getAspectSpec(), "retryLimitException")) - .forEach(metricName -> MetricUtils.exceptionCounter(BulkListener.class, metricName, failure))); - } - 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/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java index e6c94050edec2..75db5b08cec6e 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 @@ -33,9 +33,7 @@ import org.testng.annotations.Test; import java.net.URISyntaxException; -import java.util.List; -import java.util.Objects; -import java.util.Set; +import java.util.*; import java.util.concurrent.LinkedBlockingQueue; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -237,47 +235,133 @@ 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.getEntityRegistry()); + DataGenerator dataGenerator = new DataGenerator(_entityServiceImpl); Database server = ((EbeanAspectDao) _entityServiceImpl._aspectDao).getServer(); - final int testEntityCount = 25; + + // 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"); + "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 int testThreads = 15; - final LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(testThreads * 2); + final LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(threadCount * 2); // Spin up workers - List writeThreads = IntStream.range(0, testThreads) - .mapToObj(threadId -> new Thread(new MultiThreadTestWorker(queue, _entityServiceImpl))) + List writeThreads = IntStream.range(0, threadCount) + .mapToObj(threadId -> new Thread(new MultiThreadTestWorker(queue, entityService))) .collect(Collectors.toList()); writeThreads.forEach(Thread::start); - // Add data - List aspects = List.of("status", "globalTags", "glossaryTerms"); - - List generatedMCPs = dataGenerator.generateMCPs("dataset", testEntityCount, aspects) - .collect(Collectors.toList()); - - generatedMCPs.forEach(mcp -> { + testData.forEach(mcps -> { try { - queue.put(List.of(mcp)); + queue.put(mcps); } catch (InterruptedException e) { throw new RuntimeException(e); } }); // Terminate workers with empty mcp - IntStream.range(0, testThreads).forEach(threadId -> { + IntStream.range(0, threadCount).forEach(threadId -> { try { queue.put(List.of()); } catch (InterruptedException e) { @@ -288,26 +372,14 @@ public void multiThreadingTest() { // Wait for threads to finish writeThreads.forEach(thread -> { try { - thread.join(5000); + thread.join(10000); } catch (InterruptedException e) { throw new RuntimeException(e); } }); - - // Expected aspects - Set> generatedAspectIds = generatedMCPs.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()); - assertEquals(actualAspectIds, generatedAspectIds); } - public static class MultiThreadTestWorker implements Runnable { + private static class MultiThreadTestWorker implements Runnable { private final EntityServiceImpl entityService; private final LinkedBlockingQueue> queue; @@ -319,14 +391,17 @@ public MultiThreadTestWorker(LinkedBlockingQueue> q public void run() { try { while (true) { - List mcp = queue.take(); - if (mcp.isEmpty()) { + List mcps = queue.take(); + if (mcps.isEmpty()) { break; } final AuditStamp auditStamp = new AuditStamp(); auditStamp.setActor(Urn.createFromString(Constants.DATAHUB_ACTOR)); auditStamp.setTime(System.currentTimeMillis()); - entityService.ingestProposal(mcp.get(0), auditStamp, false); + 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/io/datahub/test/DataGenerator.java b/metadata-io/src/test/java/io/datahub/test/DataGenerator.java index e26c26b467c23..3b374993cde16 100644 --- a/metadata-io/src/test/java/io/datahub/test/DataGenerator.java +++ b/metadata-io/src/test/java/io/datahub/test/DataGenerator.java @@ -10,7 +10,10 @@ 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; @@ -27,9 +30,11 @@ 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; @@ -40,18 +45,23 @@ public class DataGenerator { private final static Faker FAKER = new Faker(); private final EntityRegistry entityRegistry; + private final EntityService entityService; - public DataGenerator(EntityRegistry entityRegistry) { - this.entityRegistry = entityRegistry; + public DataGenerator(EntityService entityService) { + this.entityService = entityService; + this.entityRegistry = entityService.getEntityRegistry(); } - public Stream generateDatasets() { + public Stream> generateDatasets() { return generateMCPs("dataset", 10, List.of()); } - public Stream generateMCPs(String entityName, long count, List aspects) { + 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(); @@ -62,20 +72,30 @@ public Stream generateMCPs(String entityName, long count mcp.setChangeType(ChangeType.UPSERT); return mcp; }).flatMap(mcp -> { + // Expand with additional random aspects List additionalMCPs = new LinkedList<>(); - // Expand with aspects for (String aspectName : aspects) { AspectSpec aspectSpec = entitySpec.getAspectSpec(aspectName); if (aspectSpec == null) { throw new IllegalStateException("Aspect " + aspectName + " not found for entity " + entityName); } - RecordTemplate aspect = overrideRandomAspectSuppliers.getOrDefault(aspectName, + RecordTemplate aspect = randomAspectGenerators.getOrDefault(aspectName, DataGenerator::defaultRandomAspect).apply(entitySpec, aspectSpec); - // Maybe nested, like globalTags/glossaryTerms - additionalMCPs.addAll(generateRandomNestedArray(aspect, aspectSpec, 5)); + // 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()); @@ -88,10 +108,68 @@ public Stream generateMCPs(String entityName, long count } 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> overrideRandomAspectSuppliers = Map.of( + 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) { @@ -116,15 +194,10 @@ private static RecordTemplate defaultRandomAspect(@Nonnull EntitySpec entitySpec .collect(Collectors.toList()); for (Method stringMethod : stringMethods) { - String value = FAKER.lorem().characters(8, 16, false); - - switch (aspectSpec.getName()) { - case "glossaryTermInfo": - if (stringMethod.getName().equals("setName")) { - value = normalize(FAKER.company().buzzword()); - } - break; + String value; + switch (aspectSpec.getName() + "_" + stringMethod.getName()) { default: + value = FAKER.lorem().characters(8, 16, false); break; } @@ -137,6 +210,19 @@ private static RecordTemplate defaultRandomAspect(@Nonnull EntitySpec entitySpec 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 @@ -158,50 +244,6 @@ private static RecordTemplate defaultRandomAspect(@Nonnull EntitySpec entitySpec } } - private List generateRandomNestedArray(RecordTemplate aspect, AspectSpec nestedAspect, int count) { - try { - switch (nestedAspect.getName()) { - case "globalTags": - List tags = generateMCPs("tag", count, List.of()) - .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; - case "glossaryTerms": - List terms = generateMCPs("glossaryTerm", count, - List.of("glossaryTermInfo")).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; - default: - return List.of(); - } - } catch (Exception e) { - throw new RuntimeException(e); - } - } - private static RecordTemplate randomKeyAspect(EntitySpec entitySpec) { Class keyClass = entitySpec.getKeyAspectSpec().getDataTemplateClass(); try { @@ -220,6 +262,9 @@ private static RecordTemplate randomKeyAspect(EntitySpec entitySpec) { 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: @@ -237,6 +282,7 @@ private static RecordTemplate randomKeyAspect(EntitySpec entitySpec) { stringMethods.get(2).invoke(key, animal.name().toLowerCase()); break; } + break; } List urnMethods = Arrays.stream(keyClass.getMethods()) @@ -249,7 +295,7 @@ private static RecordTemplate randomKeyAspect(EntitySpec entitySpec) { switch (entitySpec.getName()) { case "dataset": urnMethod.invoke(key, randomUrnLowerCase("dataPlatform", - List.of(FAKER.device().platform()))); + List.of(randomDataPlatform()))); break; default: throw new NotImplementedException(entitySpec.getName()); @@ -300,4 +346,14 @@ private static Urn randomUrnLowerCase(String entityType, List tuple) { 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-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..e22fffd3fef50 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; @@ -24,6 +25,8 @@ import lombok.extern.slf4j.Slf4j; import org.joda.time.DateTimeUtils; +import static com.linkedin.metadata.utils.GenericRecordUtils.JSON; + @Slf4j public class AspectUtils { @@ -32,21 +35,44 @@ private AspectUtils() { } public static List getAdditionalChanges( - @Nonnull MetadataChangeProposal metadataChangeProposal, - @Nonnull EntityService entityService) { + @Nonnull MetadataChangeProposal metadataChangeProposal, + @Nonnull EntityService entityService, + boolean onPrimaryKeyInsertOnly) { + // No additional changes for delete operation if (metadataChangeProposal.getChangeType() == ChangeType.DELETE) { return Collections.emptyList(); } final Urn urn = EntityKeyUtils.getUrnFromProposal(metadataChangeProposal, - entityService.getKeyAspectSpec(metadataChangeProposal.getEntityType())); + entityService.getKeyAspectSpec(metadataChangeProposal.getEntityType())); + + RecordTemplate aspectRecord = GenericRecordUtils.deserializeAspect( + metadataChangeProposal.getAspect().getValue(), + JSON, + entityService.getEntityRegistry().getEntitySpec(urn.getEntityType()).getAspectSpec(metadataChangeProposal.getAspectName())); + + if (onPrimaryKeyInsertOnly) { + return entityService.generateDefaultAspectsOnFirstWrite(urn, ImmutableMap.of(metadataChangeProposal.getAspectName(), aspectRecord)) + .getValue() + .stream() + .map(entry -> getProposalFromAspect(entry.getKey(), entry.getValue(), metadataChangeProposal)) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } else { + return entityService.generateDefaultAspectsIfMissing(urn, ImmutableMap.of(metadataChangeProposal.getAspectName(), aspectRecord)) + .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/EntityService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java index a84b8aabee9e6..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 @@ -220,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); 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 From 12eb77ebdf96f934d3dbc578d6d72b3bc493bcd2 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Fri, 11 Aug 2023 21:37:58 -0500 Subject: [PATCH 22/27] Fix deserialization Fix lint --- .../linkedin/metadata/entity/EbeanEntityServiceTest.java | 6 +++++- .../java/com/linkedin/metadata/entity/AspectUtils.java | 9 +++------ 2 files changed, 8 insertions(+), 7 deletions(-) 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 75db5b08cec6e..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 @@ -33,7 +33,11 @@ import org.testng.annotations.Test; import java.net.URISyntaxException; -import java.util.*; +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; 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 e22fffd3fef50..2cb2ae790e32a 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 @@ -25,8 +25,6 @@ import lombok.extern.slf4j.Slf4j; import org.joda.time.DateTimeUtils; -import static com.linkedin.metadata.utils.GenericRecordUtils.JSON; - @Slf4j public class AspectUtils { @@ -47,10 +45,9 @@ public static List getAdditionalChanges( final Urn urn = EntityKeyUtils.getUrnFromProposal(metadataChangeProposal, entityService.getKeyAspectSpec(metadataChangeProposal.getEntityType())); - RecordTemplate aspectRecord = GenericRecordUtils.deserializeAspect( - metadataChangeProposal.getAspect().getValue(), - JSON, - entityService.getEntityRegistry().getEntitySpec(urn.getEntityType()).getAspectSpec(metadataChangeProposal.getAspectName())); + RecordTemplate aspectRecord = GenericRecordUtils.deserializeAspect(metadataChangeProposal.getAspect().getValue(), + metadataChangeProposal.getAspect().getContentType(), entityService.getEntityRegistry() + .getEntitySpec(urn.getEntityType()).getAspectSpec(metadataChangeProposal.getAspectName())); if (onPrimaryKeyInsertOnly) { return entityService.generateDefaultAspectsOnFirstWrite(urn, ImmutableMap.of(metadataChangeProposal.getAspectName(), aspectRecord)) From 01e489bc39a6083426b774a3e0a7eb60965ba54b Mon Sep 17 00:00:00 2001 From: David Leifker Date: Fri, 11 Aug 2023 22:29:14 -0500 Subject: [PATCH 23/27] only add additional mcp for upsert/create ops --- .../main/java/com/linkedin/metadata/entity/AspectUtils.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 2cb2ae790e32a..ef6a0bfb1e6a4 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 @@ -32,13 +32,15 @@ public class AspectUtils { private AspectUtils() { } + public static final Set SUPPORTED_TYPES = Set.of(ChangeType.UPSERT, ChangeType.CREATE); + public static List getAdditionalChanges( @Nonnull MetadataChangeProposal metadataChangeProposal, @Nonnull EntityService entityService, boolean onPrimaryKeyInsertOnly) { - // No additional changes for delete operation - if (metadataChangeProposal.getChangeType() == ChangeType.DELETE) { + // No additional changes for unsupported operations + if (!SUPPORTED_TYPES.contains(metadataChangeProposal.getChangeType())) { return Collections.emptyList(); } From 028f3aa8bc263b2c9f971a328f2f6db1445eb0fd Mon Sep 17 00:00:00 2001 From: David Leifker Date: Sat, 12 Aug 2023 07:37:40 -0500 Subject: [PATCH 24/27] fix AspectUtils test from patch to upsert --- .../src/test/java/com/linkedin/metadata/AspectUtilsTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 46d08bc8887b9..b09f8c6ea53e2 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java @@ -54,7 +54,7 @@ public void testAdditionalChanges() { DatasetProperties datasetProperties = new DatasetProperties().setName("name"); proposal1.setAspect(GenericRecordUtils.serializeAspect(datasetProperties)); proposal1.setEntityType("dataset"); - proposal1.setChangeType(ChangeType.PATCH); + proposal1.setChangeType(ChangeType.UPSERT); List proposalList = AspectUtils.getAdditionalChanges(proposal1, entityServiceImpl); // proposals for key aspect, browsePath, browsePathV2, dataPlatformInstance From cb98802bde4b20664c5e810a460f24e39860250b Mon Sep 17 00:00:00 2001 From: David Leifker Date: Fri, 18 Aug 2023 22:17:13 -0500 Subject: [PATCH 25/27] temp fix incorrect head tag --- .github/workflows/docker-unified.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/docker-unified.yml b/.github/workflows/docker-unified.yml index c268a66938945..232d74f0e53d0 100644 --- a/.github/workflows/docker-unified.yml +++ b/.github/workflows/docker-unified.yml @@ -595,7 +595,8 @@ jobs: platforms: linux/amd64,linux/arm64/v8 - name: Compute Tag id: tag - run: echo "tag=${{ (steps.filter.outputs.datahub-ingestion-base == 'true' || steps.filter.outputs.datahub-ingestion == 'true') && needs.setup.outputs.slim_tag || 'head' }}" >> $GITHUB_OUTPUT + # TODO: Replace with `head` once publishing is fixed + run: echo "tag=${{ (steps.filter.outputs.datahub-ingestion-base == 'true' || steps.filter.outputs.datahub-ingestion == 'true') && needs.setup.outputs.slim_tag || 'pr8515-full' }}" >> $GITHUB_OUTPUT datahub_ingestion_slim_scan: permissions: contents: read # for actions/checkout to fetch code From 7ff7abb0682d65303a1688069dfd274bc091ae98 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Tue, 29 Aug 2023 10:52:08 -0500 Subject: [PATCH 26/27] generate default aspects even on patch --- .../src/test/java/com/linkedin/metadata/AspectUtilsTest.java | 2 +- .../src/main/java/com/linkedin/metadata/entity/AspectUtils.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 b09f8c6ea53e2..46d08bc8887b9 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/AspectUtilsTest.java @@ -54,7 +54,7 @@ public void testAdditionalChanges() { DatasetProperties datasetProperties = new DatasetProperties().setName("name"); proposal1.setAspect(GenericRecordUtils.serializeAspect(datasetProperties)); proposal1.setEntityType("dataset"); - proposal1.setChangeType(ChangeType.UPSERT); + proposal1.setChangeType(ChangeType.PATCH); List proposalList = AspectUtils.getAdditionalChanges(proposal1, entityServiceImpl); // proposals for key aspect, browsePath, browsePathV2, dataPlatformInstance 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 ef6a0bfb1e6a4..9f317d22db5c9 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 @@ -32,7 +32,7 @@ public class AspectUtils { private AspectUtils() { } - public static final Set SUPPORTED_TYPES = Set.of(ChangeType.UPSERT, ChangeType.CREATE); + public static final Set SUPPORTED_TYPES = Set.of(ChangeType.UPSERT, ChangeType.CREATE, ChangeType.PATCH); public static List getAdditionalChanges( @Nonnull MetadataChangeProposal metadataChangeProposal, From 7f66861c76039edb448cb6f38cf6a02c13fe6e40 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Sat, 2 Sep 2023 18:03:49 -0500 Subject: [PATCH 27/27] fix patch with default aspects --- .../resources/entity/AspectResource.java | 2 +- .../linkedin/metadata/entity/AspectUtils.java | 16 +++++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) 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 1ce57b4b95d9e..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 @@ -208,7 +208,7 @@ public Task ingestProposal( try { final AspectsBatch batch; if (asyncBool) { - // if async we'll expand the additional changes later, no need to do this early + // if async we'll expand the getAdditionalChanges later, no need to do this early batch = AspectsBatchImpl.builder() .mcps(List.of(metadataChangeProposal), _entityService.getEntityRegistry()) .build(); 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 9f317d22db5c9..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 @@ -47,19 +47,25 @@ public static List getAdditionalChanges( final Urn urn = EntityKeyUtils.getUrnFromProposal(metadataChangeProposal, entityService.getKeyAspectSpec(metadataChangeProposal.getEntityType())); - RecordTemplate aspectRecord = GenericRecordUtils.deserializeAspect(metadataChangeProposal.getAspect().getValue(), - metadataChangeProposal.getAspect().getContentType(), entityService.getEntityRegistry() - .getEntitySpec(urn.getEntityType()).getAspectSpec(metadataChangeProposal.getAspectName())); + 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, ImmutableMap.of(metadataChangeProposal.getAspectName(), aspectRecord)) + 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, ImmutableMap.of(metadataChangeProposal.getAspectName(), aspectRecord)) + return entityService.generateDefaultAspectsIfMissing(urn, includedAspects) .stream() .map(entry -> getProposalFromAspect(entry.getKey(), entry.getValue(), metadataChangeProposal)) .filter(Objects::nonNull)