feat(entity-registry): entity registry plugins (datahub-project#9538)

sleeperdeep · Jan 8, 2024 · 8415fc2 · 8415fc2
1 parent ade7b61
commit 8415fc2
Show file tree

Hide file tree

Showing 120 changed files with 3,877 additions and 1,000 deletions.
diff --git a/build.gradle b/build.gradle
@@ -231,7 +231,8 @@ project.ext.externalDependency = [
     'common': 'commons-io:commons-io:2.7',
     'jline':'jline:jline:1.4.1',
     'jetbrains':' org.jetbrains.kotlin:kotlin-stdlib:1.6.0',
-    'annotationApi': 'javax.annotation:javax.annotation-api:1.3.2'
+    'annotationApi': 'javax.annotation:javax.annotation-api:1.3.2',
+    'classGraph': 'io.github.classgraph:classgraph:4.8.165',
 ]
 
 allprojects {

diff --git a/buildSrc/src/main/java/com/linkedin/metadata/aspect/plugins/config b/buildSrc/src/main/java/com/linkedin/metadata/aspect/plugins/config
@@ -0,0 +1 @@
+../../../../../../../../../entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config
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
@@ -1,5 +1,7 @@
 package com.linkedin.datahub.graphql;
 
+import static org.mockito.Mockito.mock;
+
 import com.datahub.authentication.Actor;
 import com.datahub.authentication.ActorType;
 import com.datahub.authentication.Authentication;
@@ -10,7 +12,8 @@
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.MCPUpsertBatchItem;
 import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.mxe.MetadataChangeProposal;
@@ -19,13 +22,14 @@
 
 public class TestUtils {
 
-  public static EntityService getMockEntityService() {
+  public static EntityService<MCPUpsertBatchItem> getMockEntityService() {
     PathSpecBasedSchemaAnnotationVisitor.class
         .getClassLoader()
         .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false);
     EntityRegistry registry =
         new ConfigEntityRegistry(TestUtils.class.getResourceAsStream("/test-entity-registry.yaml"));
-    EntityService mockEntityService = Mockito.mock(EntityService.class);
+    EntityService<MCPUpsertBatchItem> mockEntityService =
+        (EntityService<MCPUpsertBatchItem>) Mockito.mock(EntityService.class);
     Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry);
     return mockEntityService;
   }
@@ -35,11 +39,11 @@ public static QueryContext getMockAllowContext() {
   }
 
   public static QueryContext getMockAllowContext(String actorUrn) {
-    QueryContext mockContext = Mockito.mock(QueryContext.class);
+    QueryContext mockContext = mock(QueryContext.class);
     Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn);
 
-    Authorizer mockAuthorizer = Mockito.mock(Authorizer.class);
-    AuthorizationResult result = Mockito.mock(AuthorizationResult.class);
+    Authorizer mockAuthorizer = mock(Authorizer.class);
+    AuthorizationResult result = mock(AuthorizationResult.class);
     Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.ALLOW);
     Mockito.when(mockAuthorizer.authorize(Mockito.any())).thenReturn(result);
 
@@ -52,11 +56,11 @@ public static QueryContext getMockAllowContext(String actorUrn) {
   }
 
   public static QueryContext getMockAllowContext(String actorUrn, AuthorizationRequest request) {
-    QueryContext mockContext = Mockito.mock(QueryContext.class);
+    QueryContext mockContext = mock(QueryContext.class);
     Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn);
 
-    Authorizer mockAuthorizer = Mockito.mock(Authorizer.class);
-    AuthorizationResult result = Mockito.mock(AuthorizationResult.class);
+    Authorizer mockAuthorizer = mock(Authorizer.class);
+    AuthorizationResult result = mock(AuthorizationResult.class);
     Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.ALLOW);
     Mockito.when(mockAuthorizer.authorize(Mockito.eq(request))).thenReturn(result);
 
@@ -73,11 +77,11 @@ public static QueryContext getMockDenyContext() {
   }
 
   public static QueryContext getMockDenyContext(String actorUrn) {
-    QueryContext mockContext = Mockito.mock(QueryContext.class);
+    QueryContext mockContext = mock(QueryContext.class);
     Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn);
 
-    Authorizer mockAuthorizer = Mockito.mock(Authorizer.class);
-    AuthorizationResult result = Mockito.mock(AuthorizationResult.class);
+    Authorizer mockAuthorizer = mock(Authorizer.class);
+    AuthorizationResult result = mock(AuthorizationResult.class);
     Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.DENY);
     Mockito.when(mockAuthorizer.authorize(Mockito.any())).thenReturn(result);
 
@@ -90,11 +94,11 @@ public static QueryContext getMockDenyContext(String actorUrn) {
   }
 
   public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequest request) {
-    QueryContext mockContext = Mockito.mock(QueryContext.class);
+    QueryContext mockContext = mock(QueryContext.class);
     Mockito.when(mockContext.getActorUrn()).thenReturn(actorUrn);
 
-    Authorizer mockAuthorizer = Mockito.mock(Authorizer.class);
-    AuthorizationResult result = Mockito.mock(AuthorizationResult.class);
+    Authorizer mockAuthorizer = mock(Authorizer.class);
+    AuthorizationResult result = mock(AuthorizationResult.class);
     Mockito.when(result.getType()).thenReturn(AuthorizationResult.Type.DENY);
     Mockito.when(mockAuthorizer.authorize(Mockito.eq(request))).thenReturn(result);
 
@@ -107,45 +111,54 @@ public static QueryContext getMockDenyContext(String actorUrn, AuthorizationRequ
   }
 
   public static void verifyIngestProposal(
-      EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) {
+      EntityService<MCPUpsertBatchItem> mockService,
+      int numberOfInvocations,
+      MetadataChangeProposal proposal) {
     verifyIngestProposal(mockService, numberOfInvocations, List.of(proposal));
   }
 
   public static void verifyIngestProposal(
-      EntityService mockService, int numberOfInvocations, List<MetadataChangeProposal> proposals) {
+      EntityService<MCPUpsertBatchItem> mockService,
+      int numberOfInvocations,
+      List<MetadataChangeProposal> proposals) {
     AspectsBatchImpl batch =
-        AspectsBatchImpl.builder().mcps(proposals, mockService.getEntityRegistry()).build();
+        AspectsBatchImpl.builder()
+            .mcps(
+                proposals,
+                mock(AuditStamp.class),
+                mockService.getEntityRegistry(),
+                mockService.getSystemEntityClient())
+            .build();
     Mockito.verify(mockService, Mockito.times(numberOfInvocations))
-        .ingestProposal(Mockito.eq(batch), Mockito.any(AuditStamp.class), Mockito.eq(false));
+        .ingestProposal(Mockito.eq(batch), Mockito.eq(false));
   }
 
   public static void verifySingleIngestProposal(
-      EntityService mockService, int numberOfInvocations, MetadataChangeProposal proposal) {
+      EntityService<MCPUpsertBatchItem> mockService,
+      int numberOfInvocations,
+      MetadataChangeProposal proposal) {
     Mockito.verify(mockService, Mockito.times(numberOfInvocations))
         .ingestProposal(Mockito.eq(proposal), Mockito.any(AuditStamp.class), Mockito.eq(false));
   }
 
-  public static void verifyIngestProposal(EntityService mockService, int numberOfInvocations) {
+  public static void verifyIngestProposal(
+      EntityService<MCPUpsertBatchItem> mockService, int numberOfInvocations) {
     Mockito.verify(mockService, Mockito.times(numberOfInvocations))
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false));
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false));
   }
 
   public static void verifySingleIngestProposal(
-      EntityService mockService, int numberOfInvocations) {
+      EntityService<MCPUpsertBatchItem> 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) {
+  public static void verifyNoIngestProposal(EntityService<MCPUpsertBatchItem> mockService) {
     Mockito.verify(mockService, Mockito.times(0))
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
   }
 
   private TestUtils() {}

diff --git a/...ava/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java b/...ava/com/linkedin/datahub/graphql/resolvers/delete/BatchUpdateSoftDeletedResolverTest.java
@@ -5,7 +5,6 @@
 import static org.testng.Assert.*;
 
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.Status;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
@@ -15,7 +14,7 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.List;
@@ -184,10 +183,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class)
         .when(mockService)
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
 
     BatchUpdateSoftDeletedResolver resolver = new BatchUpdateSoftDeletedResolver(mockService);
 

diff --git a/...om/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java b/...om/linkedin/datahub/graphql/resolvers/deprecation/BatchUpdateDeprecationResolverTest.java
@@ -5,7 +5,6 @@
 import static org.testng.Assert.*;
 
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.Deprecation;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
@@ -16,7 +15,7 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.List;
@@ -217,10 +216,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class)
         .when(mockService)
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
 
     BatchUpdateDeprecationResolver resolver = new BatchUpdateDeprecationResolver(mockService);
 

diff --git a/...c/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java b/...c/test/java/com/linkedin/datahub/graphql/resolvers/domain/BatchSetDomainResolverTest.java
@@ -5,7 +5,6 @@
 import static org.testng.Assert.*;
 
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.UrnArray;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
@@ -18,7 +17,7 @@
 import com.linkedin.events.metadata.ChangeType;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import com.linkedin.metadata.utils.GenericRecordUtils;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.schema.DataFetchingEnvironment;
@@ -311,10 +310,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class)
         .when(mockService)
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
 
     BatchSetDomainResolver resolver = new BatchSetDomainResolver(mockService);
 

diff --git a/...e/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java b/...e/src/test/java/com/linkedin/datahub/graphql/resolvers/embed/UpdateEmbedResolverTest.java
@@ -7,7 +7,6 @@
 import com.datahub.authentication.Authentication;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.Embed;
 import com.linkedin.common.urn.CorpuserUrn;
 import com.linkedin.common.urn.Urn;
@@ -19,7 +18,7 @@
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import com.linkedin.mxe.MetadataChangeProposal;
 import com.linkedin.r2.RemoteInvocationException;
 import graphql.schema.DataFetchingEnvironment;
@@ -142,8 +141,7 @@ public void testGetFailureEntityDoesNotExist() throws Exception {
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
     Mockito.verify(mockService, Mockito.times(0))
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false));
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false));
     ;
   }
 
@@ -161,8 +159,7 @@ public void testGetUnauthorized() throws Exception {
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
     Mockito.verify(mockService, Mockito.times(0))
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class), Mockito.any(AuditStamp.class), Mockito.eq(false));
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.eq(false));
   }
 
   @Test

diff --git a/...ore/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java b/...ore/src/test/java/com/linkedin/datahub/graphql/resolvers/owner/AddOwnersResolverTest.java
@@ -4,7 +4,6 @@
 import static org.testng.Assert.*;
 
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.Owner;
 import com.linkedin.common.OwnerArray;
 import com.linkedin.common.Ownership;
@@ -21,7 +20,7 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletionException;
 import org.mockito.Mockito;
@@ -399,10 +398,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class)
         .when(mockService)
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
 
     AddOwnersResolver resolver = new AddOwnersResolver(Mockito.mock(EntityService.class));
 

diff --git a/...rc/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchAddOwnersResolverTest.java b/...rc/test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchAddOwnersResolverTest.java
@@ -4,7 +4,6 @@
 import static org.testng.Assert.*;
 
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.Owner;
 import com.linkedin.common.OwnerArray;
 import com.linkedin.common.Ownership;
@@ -20,7 +19,7 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.util.OwnerUtils;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletionException;
 import org.mockito.Mockito;
@@ -337,10 +336,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class)
         .when(mockService)
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
 
     BatchAddOwnersResolver resolver = new BatchAddOwnersResolver(mockService);
 

diff --git a/...test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchRemoveOwnersResolverTest.java b/...test/java/com/linkedin/datahub/graphql/resolvers/owner/BatchRemoveOwnersResolverTest.java
@@ -4,7 +4,6 @@
 import static org.testng.Assert.*;
 
 import com.google.common.collect.ImmutableList;
-import com.linkedin.common.AuditStamp;
 import com.linkedin.common.Owner;
 import com.linkedin.common.OwnerArray;
 import com.linkedin.common.Ownership;
@@ -17,7 +16,7 @@
 import com.linkedin.datahub.graphql.resolvers.mutate.BatchRemoveOwnersResolver;
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.entity.EntityService;
-import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
+import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletionException;
 import org.mockito.Mockito;
@@ -204,10 +203,7 @@ public void testGetEntityClientException() throws Exception {
 
     Mockito.doThrow(RuntimeException.class)
         .when(mockService)
-        .ingestProposal(
-            Mockito.any(AspectsBatchImpl.class),
-            Mockito.any(AuditStamp.class),
-            Mockito.anyBoolean());
+        .ingestProposal(Mockito.any(AspectsBatchImpl.class), Mockito.anyBoolean());
 
     BatchRemoveOwnersResolver resolver = new BatchRemoveOwnersResolver(mockService);
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		../../../../../../../../../entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config