feat(ingest): add async option to ingest proposal endpoint (#6097)

* feat(ingest): add async option to ingest proposal endpoint * small tweak to validate before write to K, also keep existing path for timeseries aspects * avoid double convert Co-authored-by: Shirshanka Das <[email protected]>
datahub-project · Oct 4, 2022 · bfb903c · bfb903c
1 parent 5fb875a
commit bfb903c
Show file tree

Hide file tree

Showing 59 changed files with 396 additions and 405 deletions.
diff --git a/...aphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java b/...aphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutationUtils.java
@@ -30,7 +30,7 @@ public static void persistAspect(Urn urn, String aspectName, RecordTemplate aspe
     proposal.setAspectName(aspectName);
     proposal.setAspect(GenericRecordUtils.serializeAspect(aspect));
     proposal.setChangeType(ChangeType.UPSERT);
-    entityService.ingestProposal(proposal, getAuditStamp(actor));
+    entityService.ingestProposal(proposal, getAuditStamp(actor), false);
   }
 
   public static MetadataChangeProposal buildMetadataChangeProposal(Urn urn, String aspectName, RecordTemplate aspect, Urn actor, EntityService entityService) {

diff --git a/...rc/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java b/...rc/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateUserSettingResolver.java
@@ -51,7 +51,7 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
         MetadataChangeProposal proposal =
             buildMetadataChangeProposal(actor, CORP_USER_SETTINGS_ASPECT_NAME, newSettings, actor, _entityService);
 
-        _entityService.ingestProposal(proposal, getAuditStamp(actor));
+        _entityService.ingestProposal(proposal, getAuditStamp(actor), false);
 
         return true;
       } catch (Exception e) {

diff --git a/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java b/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeleteUtils.java
@@ -74,7 +74,7 @@ private static MetadataChangeProposal buildSoftDeleteProposal(
   private static void ingestChangeProposals(List<MetadataChangeProposal> changes, EntityService entityService, Urn actor) {
     // TODO: Replace this with a batch ingest proposals endpoint.
     for (MetadataChangeProposal change : changes) {
-      entityService.ingestProposal(change, getAuditStamp(actor));
+      entityService.ingestProposal(change, getAuditStamp(actor), false);
     }
   }
 }
diff --git a/...re/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java b/...re/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DeprecationUtils.java
@@ -89,7 +89,7 @@ private static MetadataChangeProposal buildUpdateDeprecationProposal(
   private static void ingestChangeProposals(List<MetadataChangeProposal> changes, EntityService entityService, Urn actor) {
     // TODO: Replace this with a batch ingest proposals endpoint.
     for (MetadataChangeProposal change : changes) {
-      entityService.ingestProposal(change, getAuditStamp(actor));
+      entityService.ingestProposal(change, getAuditStamp(actor), false);
     }
   }
 }
diff --git a/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java b/...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/DomainUtils.java
@@ -88,7 +88,7 @@ public static void validateDomain(Urn domainUrn, EntityService entityService) {
   private static void ingestChangeProposals(List<MetadataChangeProposal> changes, EntityService entityService, Urn actor) {
     // TODO: Replace this with a batch ingest proposals endpoint.
     for (MetadataChangeProposal change : changes) {
-      entityService.ingestProposal(change, getAuditStamp(actor));
+      entityService.ingestProposal(change, getAuditStamp(actor), false);
     }
   }
 }
diff --git a/...hql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java b/...hql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java
@@ -556,7 +556,7 @@ private static GlossaryTermAssociationArray removeTermsIfExists(GlossaryTerms te
   private static void ingestChangeProposals(List<MetadataChangeProposal> changes, EntityService entityService, Urn actor) {
     // TODO: Replace this with a batch ingest proposals endpoint.
     for (MetadataChangeProposal change : changes) {
-      entityService.ingestProposal(change, getAuditStamp(actor));
+      entityService.ingestProposal(change, getAuditStamp(actor), false);
     }
   }
 }
diff --git a/...hql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java b/...hql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/OwnerUtils.java
@@ -215,7 +215,7 @@ public static Boolean validateRemoveInput(
   private static void ingestChangeProposals(List<MetadataChangeProposal> changes, EntityService entityService, Urn actor) {
     // TODO: Replace this with a batch ingest proposals endpoint.
     for (MetadataChangeProposal change : changes) {
-      entityService.ingestProposal(change, getAuditStamp(actor));
+      entityService.ingestProposal(change, getAuditStamp(actor), false);
     }
   }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java
@@ -203,7 +203,7 @@ public Chart update(@Nonnull String urn, @Nonnull ChartUpdateInput input, @Nonnu
             proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
 
             try {
-                _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+                _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
             } catch (RemoteInvocationException e) {
                 throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
             }

diff --git a/...raphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java b/...raphql-core/src/main/java/com/linkedin/datahub/graphql/types/dashboard/DashboardType.java
@@ -192,7 +192,7 @@ public Dashboard update(@Nonnull String urn, @Nonnull DashboardUpdateInput input
             proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
 
             try {
-                _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+                _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
             } catch (RemoteInvocationException e) {
                 throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
             }

diff --git a/...-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java b/...-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/DataFlowType.java
@@ -179,7 +179,7 @@ public DataFlow update(@Nonnull String urn, @Nonnull DataFlowUpdateInput input,
             proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
 
             try {
-                _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+                _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
             } catch (RemoteInvocationException e) {
                 throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
             }

diff --git a/...ub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java b/...ub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/DataJobType.java
@@ -179,7 +179,7 @@ public DataJob update(@Nonnull String urn, @Nonnull DataJobUpdateInput input, @N
             proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
 
             try {
-                _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+                _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
             } catch (RemoteInvocationException e) {
                 throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
             }

diff --git a/...ub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java b/...ub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java
@@ -208,7 +208,7 @@ public List<Dataset> batchUpdate(@Nonnull BatchDatasetUpdateInput[] input, @Nonn
         final List<String> urns = Arrays.stream(input).map(BatchDatasetUpdateInput::getUrn).collect(Collectors.toList());
 
         try {
-            _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+            _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
         } catch (RemoteInvocationException e) {
             throw new RuntimeException(String.format("Failed to write entity with urn %s", urns), e);
         }
@@ -224,7 +224,7 @@ public Dataset update(@Nonnull String urn, @Nonnull DatasetUpdateInput input, @N
             proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
 
             try {
-                _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+                _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
             } catch (RemoteInvocationException e) {
                 throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
             }

diff --git a/...-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/notebook/NotebookType.java b/...-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/notebook/NotebookType.java
@@ -175,7 +175,7 @@ public Notebook update(@Nonnull String urn, @Nonnull NotebookUpdateInput input,
     proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
 
     try {
-      _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+      _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
     } catch (RemoteInvocationException e) {
       throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
     }

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/TagType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/tag/TagType.java
@@ -132,7 +132,7 @@ public Tag update(@Nonnull String urn, @Nonnull TagUpdateInput input, @Nonnull Q
             final Collection<MetadataChangeProposal> proposals = TagUpdateInputMapper.map(input, actor);
             proposals.forEach(proposal -> proposal.setEntityUrn(UrnUtils.getUrn(urn)));
             try {
-                _entityClient.batchIngestProposals(proposals, context.getAuthentication());
+                _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
             } catch (RemoteInvocationException e) {
                 throw new RuntimeException(String.format("Failed to write entity with urn %s", urn), e);
             }

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
@@ -3,6 +3,9 @@
 import com.datahub.authentication.Authentication;
 import com.datahub.authorization.AuthorizationResult;
 import com.datahub.authorization.Authorizer;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.mxe.MetadataChangeProposal;
 import org.mockito.Mockito;
 
 
@@ -36,5 +39,27 @@ public static QueryContext getMockDenyContext() {
     return mockContext;
   }
 
+  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)
+    );
+  }
+
+  public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations) {
+    Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal(
+        Mockito.any(MetadataChangeProposal.class),
+        Mockito.any(AuditStamp.class),
+        Mockito.eq(false)
+    );
+  }
+
+  public static void verifyNoIngestProposal(EntityService mockService) {
+    Mockito.verify(mockService, Mockito.times(0)).ingestProposal(
+        Mockito.any(),
+        Mockito.any(AuditStamp.class), Mockito.anyBoolean());
+  }
+
   private TestUtils() { }
 }
diff --git a/...ava/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java b/...ava/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java
@@ -66,10 +66,7 @@ public void testGetSuccessNoExistingStatus() throws Exception {
     proposal1.setAspect(GenericRecordUtils.serializeAspect(newStatus));
     proposal1.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal1),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal1);
 
     final MetadataChangeProposal proposal2 = new MetadataChangeProposal();
     proposal2.setEntityUrn(Urn.createFromString(TEST_ENTITY_URN_2));
@@ -78,10 +75,7 @@ public void testGetSuccessNoExistingStatus() throws Exception {
     proposal2.setAspect(GenericRecordUtils.serializeAspect(newStatus));
     proposal2.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal2),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal2);
   }
 
   @Test
@@ -124,10 +118,7 @@ public void testGetSuccessExistingStatus() throws Exception {
     proposal1.setAspect(GenericRecordUtils.serializeAspect(newStatus));
     proposal1.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal1),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal1);
 
     final MetadataChangeProposal proposal2 = new MetadataChangeProposal();
     proposal2.setEntityUrn(Urn.createFromString(TEST_ENTITY_URN_2));
@@ -136,10 +127,7 @@ public void testGetSuccessExistingStatus() throws Exception {
     proposal2.setAspect(GenericRecordUtils.serializeAspect(newStatus));
     proposal2.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal2),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal2);
   }
 
   @Test
@@ -171,9 +159,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception {
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
-    Mockito.verify(mockService, Mockito.times(0)).ingestProposal(
-        Mockito.any(),
-        Mockito.any(AuditStamp.class));
+    verifyNoIngestProposal(mockService);
   }
 
   @Test
@@ -191,9 +177,7 @@ public void testGetUnauthorized() throws Exception {
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
-    Mockito.verify(mockService, Mockito.times(0)).ingestProposal(
-        Mockito.any(),
-        Mockito.any(AuditStamp.class));
+    verifyNoIngestProposal(mockService);
   }
 
   @Test
@@ -202,7 +186,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal(
         Mockito.any(),
-        Mockito.any(AuditStamp.class));
+        Mockito.any(AuditStamp.class), Mockito.anyBoolean());
 
     BatchUpdateSoftDeletedResolver resolver = new BatchUpdateSoftDeletedResolver(mockService);
 

diff --git a/...om/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java b/...om/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java
@@ -73,10 +73,7 @@ public void testGetSuccessNoExistingDeprecation() throws Exception {
     proposal1.setAspect(GenericRecordUtils.serializeAspect(newDeprecation));
     proposal1.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal1),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal1);
 
     final MetadataChangeProposal proposal2 = new MetadataChangeProposal();
     proposal2.setEntityUrn(Urn.createFromString(TEST_ENTITY_URN_2));
@@ -85,10 +82,7 @@ public void testGetSuccessNoExistingDeprecation() throws Exception {
     proposal2.setAspect(GenericRecordUtils.serializeAspect(newDeprecation));
     proposal2.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal2),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal2);
   }
 
   @Test
@@ -140,10 +134,7 @@ public void testGetSuccessExistingDeprecation() throws Exception {
     proposal1.setAspect(GenericRecordUtils.serializeAspect(newDeprecation));
     proposal1.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal1),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal1);
 
     final MetadataChangeProposal proposal2 = new MetadataChangeProposal();
     proposal2.setEntityUrn(Urn.createFromString(TEST_ENTITY_URN_2));
@@ -152,10 +143,7 @@ public void testGetSuccessExistingDeprecation() throws Exception {
     proposal2.setAspect(GenericRecordUtils.serializeAspect(newDeprecation));
     proposal2.setChangeType(ChangeType.UPSERT);
 
-    Mockito.verify(mockService, Mockito.times(1)).ingestProposal(
-        Mockito.eq(proposal2),
-        Mockito.any(AuditStamp.class)
-    );
+    verifyIngestProposal(mockService, 1, proposal2);
   }
 
   @Test
@@ -188,9 +176,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception {
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
-    Mockito.verify(mockService, Mockito.times(0)).ingestProposal(
-        Mockito.any(),
-        Mockito.any(AuditStamp.class));
+    verifyNoIngestProposal(mockService);
   }
 
   @Test
@@ -209,9 +195,7 @@ public void testGetUnauthorized() throws Exception {
     Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
-    Mockito.verify(mockService, Mockito.times(0)).ingestProposal(
-        Mockito.any(),
-        Mockito.any(AuditStamp.class));
+    verifyNoIngestProposal(mockService);
   }
 
   @Test
@@ -220,7 +204,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal(
         Mockito.any(),
-        Mockito.any(AuditStamp.class));
+        Mockito.any(AuditStamp.class), Mockito.anyBoolean());
 
     BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService);