feat(EntityService): batched transactions and ebean updates (#8456)

build.gradle

@@ -18,6 +18,7 @@ buildscript {
   ext.hadoop3Version = '3.3.5'
   ext.kafkaVersion = '2.3.0'
   ext.hazelcastVersion = '5.3.1'
+  ext.ebeanVersion = '12.16.1'
 
   ext.docker_registry = 'linkedin'
 
@@ -86,8 +87,9 @@ project.ext.externalDependency = [
     'dgraph4j' : 'io.dgraph:dgraph4j:21.03.1',
     'dropwizardMetricsCore': 'io.dropwizard.metrics:metrics-core:4.2.3',
     'dropwizardMetricsJmx': 'io.dropwizard.metrics:metrics-jmx:4.2.3',
-    'ebean': 'io.ebean:ebean:11.33.3',
-    'ebeanAgent': 'io.ebean:ebean-agent:11.27.1',
+    'ebean': 'io.ebean:ebean:' + ebeanVersion,
+    'ebeanAgent': 'io.ebean:ebean-agent:' + ebeanVersion,
+    'ebeanDdl': 'io.ebean:ebean-ddl-generator:' + ebeanVersion,
     'elasticSearchRest': 'org.elasticsearch.client:elasticsearch-rest-high-level-client:' + elasticsearchVersion,
     'elasticSearchTransport': 'org.elasticsearch.client:transport:' + elasticsearchVersion,
     'findbugsAnnotations': 'com.google.code.findbugs:annotations:3.0.1',

datahub-frontend/build.gradle

@@ -96,7 +96,7 @@ task unversionZip(type: Copy, dependsOn: [':datahub-web-react:build', dist]) {
   into "${buildDir}/docker/"
   rename "datahub-frontend-${version}.zip", "datahub-frontend.zip"
 }
-tasks.getByName("docker").dependsOn(unversionZip)
+tasks.getByPath(":datahub-frontend:docker").dependsOn(unversionZip)
 
 task cleanLocalDockerImages {
   doLast {

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java

@@ -9,12 +9,24 @@
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
+import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.mxe.MetadataChangeProposal;
 import org.mockito.Mockito;
 
+import java.util.List;
+
 
 public class TestUtils {
 
+  public static EntityService getMockEntityService() {
+    EntityRegistry registry = new ConfigEntityRegistry(TestUtils.class.getResourceAsStream("/test-entity-registry.yaml"));
+    EntityService mockEntityService = Mockito.mock(EntityService.class);
+    Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry);
+    return mockEntityService;
+  }
+
   public static QueryContext getMockAllowContext() {
     return getMockAllowContext("urn:li:corpuser:test");
   }
@@ -88,25 +100,47 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ
   }
 
   public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) {
+    verifyIngestProposal(mockService, numberOfInvocations, List.of(proposal));
+  }
+
+  public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations, List<MetadataChangeProposal> proposals) {
+    AspectsBatchImpl batch = AspectsBatchImpl.builder()
+              .mcps(proposals, mockService.getEntityRegistry())
+              .build();
+      Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal(
+              Mockito.eq(batch),
+              Mockito.any(AuditStamp.class),
+              Mockito.eq(false)
+      );
+  }
+
+  public static void verifySingleIngestProposal(EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) {
     Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal(
-        Mockito.eq(proposal),
-        Mockito.any(AuditStamp.class),
-        Mockito.eq(false)
+            Mockito.eq(proposal),
+            Mockito.any(AuditStamp.class),
+            Mockito.eq(false)
     );
   }
 
   public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations) {
     Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal(
-        Mockito.any(MetadataChangeProposal.class),
+        Mockito.any(AspectsBatchImpl.class),
         Mockito.any(AuditStamp.class),
         Mockito.eq(false)
     );
   }
 
+  public static void verifySingleIngestProposal(EntityService mockService, int numberOfInvocations) {
+    Mockito.verify(mockService, Mockito.times(numberOfInvocations)).ingestProposal(
+            Mockito.any(MetadataChangeProposal.class),
+            Mockito.any(AuditStamp.class),
+            Mockito.eq(false)
+    );
+  }
+
   public static void verifyNoIngestProposal(EntityService mockService) {
     Mockito.verify(mockService, Mockito.times(0)).ingestProposal(
-        Mockito.any(),
-        Mockito.any(AuditStamp.class), Mockito.anyBoolean());
+        Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.anyBoolean());
   }
 
   private TestUtils() { }

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);

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java

@@ -11,8 +11,11 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.schema.DataFetchingEnvironment;
+
+import java.util.List;
 import java.util.concurrent.CompletionException;
 import org.mockito.Mockito;
 import org.testng.annotations.Test;
@@ -29,7 +32,7 @@ public class BatchUpdateSoftDeletedResolverTest {
 
   @Test
   public void testGetSuccessNoExistingStatus() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.when(mockService.getAspect(
         Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)),
@@ -61,20 +64,17 @@ public void testGetSuccessNoExistingStatus() throws Exception {
 
     final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1),
         STATUS_ASPECT_NAME, newStatus);
-
-    verifyIngestProposal(mockService, 1, proposal1);
-
     final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2),
         STATUS_ASPECT_NAME, newStatus);
 
-    verifyIngestProposal(mockService, 1, proposal2);
+    verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2));
   }
 
   @Test
   public void testGetSuccessExistingStatus() throws Exception {
     final Status originalStatus = new Status().setRemoved(true);
 
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.when(mockService.getAspect(
         Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)),
@@ -105,18 +105,15 @@ public void testGetSuccessExistingStatus() throws Exception {
 
     final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1),
         STATUS_ASPECT_NAME, newStatus);
-
-    verifyIngestProposal(mockService, 1, proposal1);
-
     final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2),
         STATUS_ASPECT_NAME, newStatus);
 
-    verifyIngestProposal(mockService, 1, proposal2);
+    verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2));
   }
 
   @Test
   public void testGetFailureResourceDoesNotExist() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.when(mockService.getAspect(
         Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)),
@@ -148,7 +145,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception {
 
   @Test
   public void testGetUnauthorized() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     BatchUpdateSoftDeletedResolver resolver = new BatchUpdateSoftDeletedResolver(mockService);
 
@@ -166,10 +163,10 @@ public void testGetUnauthorized() throws Exception {
 
   @Test
   public void testGetEntityClientException() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal(
-        Mockito.any(),
+        Mockito.any(AspectsBatchImpl.class),
         Mockito.any(AuditStamp.class), Mockito.anyBoolean());
 
     BatchUpdateSoftDeletedResolver resolver = new BatchUpdateSoftDeletedResolver(mockService);

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java

@@ -12,8 +12,11 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.schema.DataFetchingEnvironment;
+
+import java.util.List;
 import java.util.concurrent.CompletionException;
 import org.mockito.Mockito;
 import org.testng.annotations.Test;
@@ -30,7 +33,7 @@ public class BatchUpdateDeprecationResolverTest {
 
   @Test
   public void testGetSuccessNoExistingDeprecation() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.when(mockService.getAspect(
         Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)),
@@ -68,12 +71,10 @@ public void testGetSuccessNoExistingDeprecation() throws Exception {
 
     final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1),
         DEPRECATION_ASPECT_NAME, newDeprecation);
-
-    verifyIngestProposal(mockService, 1, proposal1);
-
     final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2),
         DEPRECATION_ASPECT_NAME, newDeprecation);
-    verifyIngestProposal(mockService, 1, proposal2);
+
+    verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2));
   }
 
   @Test
@@ -83,7 +84,7 @@ public void testGetSuccessExistingDeprecation() throws Exception {
         .setNote("")
         .setActor(UrnUtils.getUrn("urn:li:corpuser:test"));
 
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.when(mockService.getAspect(
         Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)),
@@ -120,18 +121,15 @@ public void testGetSuccessExistingDeprecation() throws Exception {
 
     final MetadataChangeProposal proposal1 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_1),
         DEPRECATION_ASPECT_NAME, newDeprecation);
-
-    verifyIngestProposal(mockService, 1, proposal1);
-
     final MetadataChangeProposal proposal2 = MutationUtils.buildMetadataChangeProposalWithUrn(Urn.createFromString(TEST_ENTITY_URN_2),
         DEPRECATION_ASPECT_NAME, newDeprecation);
 
-    verifyIngestProposal(mockService, 1, proposal2);
+    verifyIngestProposal(mockService, 1, List.of(proposal1, proposal2));
   }
 
   @Test
   public void testGetFailureResourceDoesNotExist() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.when(mockService.getAspect(
         Mockito.eq(UrnUtils.getUrn(TEST_ENTITY_URN_1)),
@@ -164,7 +162,7 @@ public void testGetFailureResourceDoesNotExist() throws Exception {
 
   @Test
   public void testGetUnauthorized() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService);
 
@@ -183,10 +181,10 @@ public void testGetUnauthorized() throws Exception {
 
   @Test
   public void testGetEntityClientException() throws Exception {
-    EntityService mockService = Mockito.mock(EntityService.class);
+    EntityService mockService = getMockEntityService();
 
     Mockito.doThrow(RuntimeException.class).when(mockService).ingestProposal(
-        Mockito.any(),
+        Mockito.any(AspectsBatchImpl.class),
         Mockito.any(AuditStamp.class), Mockito.anyBoolean());
 
     BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService);

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);