feat(models) : Joins (Datasets) schema, resolvers and UI

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -53,6 +53,7 @@
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.DatasetStatsSummary;
 import com.linkedin.datahub.graphql.generated.Domain;
+import com.linkedin.datahub.graphql.generated.ERModelRelationshipProperties;
 import com.linkedin.datahub.graphql.generated.EntityPath;
 import com.linkedin.datahub.graphql.generated.EntityRelationship;
 import com.linkedin.datahub.graphql.generated.EntityRelationshipLegacy;
@@ -308,6 +309,9 @@
 import com.linkedin.datahub.graphql.types.datatype.DataTypeType;
 import com.linkedin.datahub.graphql.types.domain.DomainType;
 import com.linkedin.datahub.graphql.types.entitytype.EntityTypeType;
+import com.linkedin.datahub.graphql.types.ermodelrelationship.CreateERModelRelationshipResolver;
+import com.linkedin.datahub.graphql.types.ermodelrelationship.ERModelRelationshipType;
+import com.linkedin.datahub.graphql.types.ermodelrelationship.UpdateERModelRelationshipResolver;
 import com.linkedin.datahub.graphql.types.form.FormType;
 import com.linkedin.datahub.graphql.types.glossary.GlossaryNodeType;
 import com.linkedin.datahub.graphql.types.glossary.GlossaryTermType;
@@ -346,6 +350,7 @@
 import com.linkedin.metadata.recommendation.RecommendationsService;
 import com.linkedin.metadata.secret.SecretService;
 import com.linkedin.metadata.service.DataProductService;
+import com.linkedin.metadata.service.ERModelRelationshipService;
 import com.linkedin.metadata.service.FormService;
 import com.linkedin.metadata.service.LineageService;
 import com.linkedin.metadata.service.OwnershipTypeService;
@@ -417,6 +422,7 @@ public class GmsGraphQLEngine {
   private final LineageService lineageService;
   private final QueryService queryService;
   private final DataProductService dataProductService;
+  private final ERModelRelationshipService erModelRelationshipService;
   private final FormService formService;
   private final RestrictedService restrictedService;
 
@@ -462,6 +468,7 @@ public class GmsGraphQLEngine {
   private final DataHubPolicyType dataHubPolicyType;
   private final DataHubRoleType dataHubRoleType;
   private final SchemaFieldType schemaFieldType;
+  private final ERModelRelationshipType erModelRelationshipType;
   private final DataHubViewType dataHubViewType;
   private final QueryType queryType;
   private final DataProductType dataProductType;
@@ -529,6 +536,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
     this.settingsService = args.settingsService;
     this.lineageService = args.lineageService;
     this.queryService = args.queryService;
+    this.erModelRelationshipService = args.erModelRelationshipService;
     this.dataProductService = args.dataProductService;
     this.formService = args.formService;
     this.restrictedService = args.restrictedService;
@@ -572,6 +580,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
     this.dataHubPolicyType = new DataHubPolicyType(entityClient);
     this.dataHubRoleType = new DataHubRoleType(entityClient);
     this.schemaFieldType = new SchemaFieldType(entityClient, featureFlags);
+    this.erModelRelationshipType = new ERModelRelationshipType(entityClient, featureFlags);
     this.dataHubViewType = new DataHubViewType(entityClient);
     this.queryType = new QueryType(entityClient);
     this.dataProductType = new DataProductType(entityClient);
@@ -617,6 +626,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
             dataHubPolicyType,
             dataHubRoleType,
             schemaFieldType,
+            erModelRelationshipType,
             dataHubViewType,
             queryType,
             dataProductType,
@@ -707,6 +717,7 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
     configureTestResultResolvers(builder);
     configureRoleResolvers(builder);
     configureSchemaFieldResolvers(builder);
+    configureERModelRelationshipResolvers(builder);
     configureEntityPathResolvers(builder);
     configureResolvedAuditStampResolvers(builder);
     configureViewResolvers(builder);
@@ -939,6 +950,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher("glossaryTerm", getResolver(glossaryTermType))
                 .dataFetcher("glossaryNode", getResolver(glossaryNodeType))
                 .dataFetcher("domain", getResolver((domainType)))
+                .dataFetcher("erModelRelationship", getResolver(erModelRelationshipType))
                 .dataFetcher("dataPlatform", getResolver(dataPlatformType))
                 .dataFetcher("dataPlatformInstance", getResolver(dataPlatformInstanceType))
                 .dataFetcher("mlFeatureTable", getResolver(mlFeatureTableType))
@@ -1069,6 +1081,13 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher("updateDataFlow", new MutableTypeResolver<>(dataFlowType))
                 .dataFetcher("updateCorpUserProperties", new MutableTypeResolver<>(corpUserType))
                 .dataFetcher("updateCorpGroupProperties", new MutableTypeResolver<>(corpGroupType))
+                .dataFetcher(
+                    "updateERModelRelationship",
+                    new UpdateERModelRelationshipResolver(this.entityClient))
+                .dataFetcher(
+                    "createERModelRelationship",
+                    new CreateERModelRelationshipResolver(
+                        this.entityClient, this.erModelRelationshipService))
                 .dataFetcher("addTag", new AddTagResolver(entityService))
                 .dataFetcher("addTags", new AddTagsResolver(entityService))
                 .dataFetcher("batchAddTags", new BatchAddTagsResolver(entityService))
@@ -2078,6 +2097,60 @@ private void configureTypeExtensions(final RuntimeWiring.Builder builder) {
     builder.scalar(GraphQLLong);
   }
 
+  /** Configures resolvers responsible for resolving the {@link ERModelRelationship} type. */
+  private void configureERModelRelationshipResolvers(final RuntimeWiring.Builder builder) {
+    builder
+        .type(
+            "ERModelRelationship",
+            typeWiring ->
+                typeWiring
+                    .dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
+                    .dataFetcher(
+                        "relationships", new EntityRelationshipsResultResolver(graphClient)))
+        .type(
+            "ERModelRelationshipProperties",
+            typeWiring ->
+                typeWiring
+                    .dataFetcher(
+                        "source",
+                        new LoadableTypeResolver<>(
+                            datasetType,
+                            (env) -> {
+                              final ERModelRelationshipProperties erModelRelationshipProperties =
+                                  env.getSource();
+                              return erModelRelationshipProperties.getSource() != null
+                                  ? erModelRelationshipProperties.getSource().getUrn()
+                                  : null;
+                            }))
+                    .dataFetcher(
+                        "destination",
+                        new LoadableTypeResolver<>(
+                            datasetType,
+                            (env) -> {
+                              final ERModelRelationshipProperties erModelRelationshipProperties =
+                                  env.getSource();
+                              return erModelRelationshipProperties.getDestination() != null
+                                  ? erModelRelationshipProperties.getDestination().getUrn()
+                                  : null;
+                            })))
+        .type(
+            "Owner",
+            typeWiring ->
+                typeWiring.dataFetcher(
+                    "owner",
+                    new OwnerTypeResolver<>(
+                        ownerTypes, (env) -> ((Owner) env.getSource()).getOwner())))
+        .type(
+            "InstitutionalMemoryMetadata",
+            typeWiring ->
+                typeWiring.dataFetcher(
+                    "author",
+                    new LoadableTypeResolver<>(
+                        corpUserType,
+                        (env) ->
+                            ((InstitutionalMemoryMetadata) env.getSource()).getAuthor().getUrn())));
+  }
+
   /**
    * Configures resolvers responsible for resolving the {@link
    * com.linkedin.datahub.graphql.generated.DataJob} type.

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java

@@ -25,6 +25,7 @@
 import com.linkedin.metadata.recommendation.RecommendationsService;
 import com.linkedin.metadata.secret.SecretService;
 import com.linkedin.metadata.service.DataProductService;
+import com.linkedin.metadata.service.ERModelRelationshipService;
 import com.linkedin.metadata.service.FormService;
 import com.linkedin.metadata.service.LineageService;
 import com.linkedin.metadata.service.OwnershipTypeService;
@@ -75,6 +76,7 @@ public class GmsGraphQLEngineArgs {
   QueryService queryService;
   FeatureFlags featureFlags;
   DataProductService dataProductService;
+  ERModelRelationshipService erModelRelationshipService;
   FormService formService;
   RestrictedService restrictedService;
   int graphQLQueryComplexityLimit;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/featureflags/FeatureFlags.java

@@ -15,6 +15,7 @@ public class FeatureFlags {
   private boolean platformBrowseV2 = false;
   private PreProcessHooks preProcessHooks;
   private boolean showAcrylInfo = false;
+  private boolean erModelRelationshipFeatureEnabled = false;
   private boolean showAccessManagement = false;
   private boolean nestedDomainsEnabled = false;
   private boolean schemaFieldEntityFetchEnabled = false;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java

@@ -179,6 +179,8 @@ public CompletableFuture<AppConfig> get(final DataFetchingEnvironment environmen
             .setReadOnlyModeEnabled(_featureFlags.isReadOnlyModeEnabled())
             .setShowBrowseV2(_featureFlags.isShowBrowseV2())
             .setShowAcrylInfo(_featureFlags.isShowAcrylInfo())
+            .setErModelRelationshipFeatureEnabled(
+                _featureFlags.isErModelRelationshipFeatureEnabled())
             .setShowAccessManagement(_featureFlags.isShowAccessManagement())
             .setNestedDomainsEnabled(_featureFlags.isNestedDomainsEnabled())
             .setPlatformBrowseV2(_featureFlags.isPlatformBrowseV2())
@@ -262,6 +264,10 @@ private EntityType mapResourceTypeToEntityType(final String resourceType) {
         .getResourceType()
         .equals(resourceType)) {
       return EntityType.CORP_USER;
+    } else if (com.linkedin.metadata.authorization.PoliciesConfig.ER_MODEL_RELATIONSHIP_PRIVILEGES
+        .getResourceType()
+        .equals(resourceType)) {
+      return EntityType.ER_MODEL_RELATIONSHIP;
     } else {
       return null;
     }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java

@@ -19,6 +19,7 @@
 import com.linkedin.datahub.graphql.generated.DataProduct;
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.Domain;
+import com.linkedin.datahub.graphql.generated.ERModelRelationship;
 import com.linkedin.datahub.graphql.generated.Entity;
 import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.GlossaryNode;
@@ -155,6 +156,11 @@ public Entity apply(Urn input) {
       ((Domain) partialEntity).setUrn(input.toString());
       ((Domain) partialEntity).setType(EntityType.DOMAIN);
     }
+    if (input.getEntityType().equals("erModelRelationship")) {
+      partialEntity = new ERModelRelationship();
+      ((ERModelRelationship) partialEntity).setUrn(input.toString());
+      ((ERModelRelationship) partialEntity).setType(EntityType.ER_MODEL_RELATIONSHIP);
+    }
     if (input.getEntityType().equals("assertion")) {
       partialEntity = new Assertion();
       ((Assertion) partialEntity).setUrn(input.toString());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeMapper.java

@@ -38,6 +38,7 @@ public class EntityTypeMapper {
           .put(EntityType.NOTEBOOK, "notebook")
           .put(EntityType.DATA_PLATFORM_INSTANCE, "dataPlatformInstance")
           .put(EntityType.TEST, "test")
+          .put(EntityType.ER_MODEL_RELATIONSHIP, Constants.ER_MODEL_RELATIONSHIP_ENTITY_NAME)
           .put(EntityType.DATAHUB_VIEW, Constants.DATAHUB_VIEW_ENTITY_NAME)
           .put(EntityType.DATA_PRODUCT, Constants.DATA_PRODUCT_ENTITY_NAME)
           .put(EntityType.SCHEMA_FIELD, "schemaField")

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/ermodelrelationship/CreateERModelRelationshipResolver.java

@@ -0,0 +1,113 @@
+package com.linkedin.datahub.graphql.types.ermodelrelationship;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+
+import com.datahub.authentication.Authentication;
+import com.linkedin.common.urn.CorpuserUrn;
+import com.linkedin.common.urn.ERModelRelationshipUrn;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.exception.AuthorizationException;
+import com.linkedin.datahub.graphql.generated.ERModelRelationship;
+import com.linkedin.datahub.graphql.generated.ERModelRelationshipPropertiesInput;
+import com.linkedin.datahub.graphql.generated.ERModelRelationshipUpdateInput;
+import com.linkedin.datahub.graphql.types.ermodelrelationship.mappers.ERModelRelationMapper;
+import com.linkedin.datahub.graphql.types.ermodelrelationship.mappers.ERModelRelationshipUpdateInputMapper;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.service.ERModelRelationshipService;
+import com.linkedin.mxe.MetadataChangeProposal;
+import com.linkedin.r2.RemoteInvocationException;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.nio.charset.StandardCharsets;
+import java.util.Collection;
+import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.codec.digest.DigestUtils;
+
+@Slf4j
+@RequiredArgsConstructor
+public class CreateERModelRelationshipResolver
+    implements DataFetcher<CompletableFuture<ERModelRelationship>> {
+
+  private final EntityClient _entityClient;
+  private final ERModelRelationshipService _erModelRelationshipService;
+
+  @Override
+  public CompletableFuture<ERModelRelationship> get(DataFetchingEnvironment environment)
+      throws Exception {
+    final ERModelRelationshipUpdateInput input =
+        bindArgument(environment.getArgument("input"), ERModelRelationshipUpdateInput.class);
+
+    final ERModelRelationshipPropertiesInput erModelRelationshipPropertiesInput =
+        input.getProperties();
+    String ermodelrelationName = erModelRelationshipPropertiesInput.getName();
+    String source = erModelRelationshipPropertiesInput.getSource();
+    String destination = erModelRelationshipPropertiesInput.getDestination();
+
+    String lowDataset = source;
+    String highDataset = destination;
+    if (source.compareTo(destination) > 0) {
+      lowDataset = destination;
+      highDataset = source;
+    }
+    // The following sequence mimics datahub.emitter.mce_builder.datahub_guid
+
+    String ermodelrelationKey =
+        "{\"Source\":\""
+            + lowDataset
+            + "\",\"Destination\":\""
+            + highDataset
+            + "\",\"ERModelRelationName\":\""
+            + ermodelrelationName
+            + "\"}";
+
+    byte[] mybytes = ermodelrelationKey.getBytes(StandardCharsets.UTF_8);
+
+    String ermodelrelationKeyEncoded = new String(mybytes, StandardCharsets.UTF_8);
+    String ermodelrelationGuid = DigestUtils.md5Hex(ermodelrelationKeyEncoded);
+    log.info(
+        "ermodelrelationkey {}, ermodelrelationGuid {}",
+        ermodelrelationKeyEncoded,
+        ermodelrelationGuid);
+
+    ERModelRelationshipUrn inputUrn = new ERModelRelationshipUrn(ermodelrelationGuid);
+    QueryContext context = environment.getContext();
+    final Authentication authentication = context.getAuthentication();
+    final CorpuserUrn actor = CorpuserUrn.createFromString(context.getActorUrn());
+    if (!ERModelRelationshipType.canCreateERModelRelation(
+        context,
+        Urn.createFromString(input.getProperties().getSource()),
+        Urn.createFromString(input.getProperties().getDestination()))) {
+      throw new AuthorizationException(
+          "Unauthorized to create erModelRelationship. Please contact your DataHub administrator.");
+    }
+    return CompletableFuture.supplyAsync(
+        () -> {
+          try {
+            log.debug("Create ERModelRelation input: {}", input);
+            final Collection<MetadataChangeProposal> proposals =
+                ERModelRelationshipUpdateInputMapper.map(input, actor);
+            proposals.forEach(proposal -> proposal.setEntityUrn(inputUrn));
+            try {
+              _entityClient.batchIngestProposals(proposals, context.getAuthentication(), false);
+            } catch (RemoteInvocationException e) {
+              throw new RuntimeException("Failed to create erModelRelationship entity", e);
+            }
+            return ERModelRelationMapper.map(
+                _erModelRelationshipService.getERModelRelationshipResponse(
+                    Urn.createFromString(inputUrn.toString()), authentication));
+          } catch (Exception e) {
+            log.error(
+                "Failed to create ERModelRelation to resource with input {}, {}",
+                input,
+                e.getMessage());
+            throw new RuntimeException(
+                String.format(
+                    "Failed to create erModelRelationship to resource with input %s", input),
+                e);
+          }
+        });
+  }
+}