feat(models,ui): Support query nodes in lineage (#9850)

asikowitz · web-flow · commit 37952036ad16 · 2024-02-15T08:58:25.000-05:00
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -90,6 +90,7 @@
 import com.linkedin.datahub.graphql.generated.ParentDomainsResult;
 import com.linkedin.datahub.graphql.generated.PolicyMatchCriterionValue;
 import com.linkedin.datahub.graphql.generated.QueryEntity;
+import com.linkedin.datahub.graphql.generated.QueryProperties;
 import com.linkedin.datahub.graphql.generated.QuerySubject;
 import com.linkedin.datahub.graphql.generated.QuickFilter;
 import com.linkedin.datahub.graphql.generated.RecommendationContent;
@@ -2583,9 +2584,27 @@ private void configureQueryEntityResolvers(final RuntimeWiring.Builder builder)
     builder
         .type(
             "QueryEntity",
+            typeWiring ->
+                typeWiring
+                    .dataFetcher(
+                        "relationships", new EntityRelationshipsResultResolver(graphClient))
+                    .dataFetcher(
+                        "platform",
+                        new LoadableTypeResolver<>(
+                            dataPlatformType,
+                            (env) -> {
+                              final QueryEntity query = env.getSource();
+                              return query.getPlatform() != null
+                                  ? query.getPlatform().getUrn()
+                                  : null;
+                            })))
+        .type(
+            "QueryProperties",
             typeWiring ->
                 typeWiring.dataFetcher(
-                    "relationships", new EntityRelationshipsResultResolver(graphClient)))
+                    "origin",
+                    new EntityTypeResolver(
+                        entityTypes, (env) -> ((QueryProperties) env.getSource()).getOrigin())))
         .type(
             "ListQueriesResult",
             typeWiring ->
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/FineGrainedLineagesMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/FineGrainedLineagesMapper.java
@@ -45,6 +45,12 @@ public List<com.linkedin.datahub.graphql.generated.FineGrainedLineage> apply(
                 .map(FineGrainedLineagesMapper::mapDatasetSchemaField)
                 .collect(Collectors.toList()));
       }
+      if (fineGrainedLineage.hasQuery()) {
+        resultEntry.setQuery(fineGrainedLineage.getQuery().toString());
+      }
+      if (fineGrainedLineage.hasTransformOperation()) {
+        resultEntry.setTransformOperation(fineGrainedLineage.getTransformOperation());
+      }
       result.add(resultEntry);
     }
     return result;
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/query/QueryMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/query/QueryMapper.java
@@ -2,17 +2,20 @@
 
 import static com.linkedin.metadata.Constants.*;
 
+import com.linkedin.common.DataPlatformInstance;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.data.template.GetMode;
 import com.linkedin.datahub.graphql.generated.AuditStamp;
+import com.linkedin.datahub.graphql.generated.DataPlatform;
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.QueryEntity;
 import com.linkedin.datahub.graphql.generated.QueryLanguage;
 import com.linkedin.datahub.graphql.generated.QuerySource;
 import com.linkedin.datahub.graphql.generated.QueryStatement;
 import com.linkedin.datahub.graphql.generated.QuerySubject;
+import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
 import com.linkedin.datahub.graphql.types.common.mappers.util.MappingHelper;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
 import com.linkedin.entity.EntityResponse;
@@ -22,7 +25,9 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import lombok.extern.slf4j.Slf4j;
 
+@Slf4j
 public class QueryMapper implements ModelMapper<EntityResponse, QueryEntity> {
 
   public static final QueryMapper INSTANCE = new QueryMapper();
@@ -41,9 +46,20 @@ public QueryEntity apply(@Nonnull final EntityResponse entityResponse) {
     MappingHelper<QueryEntity> mappingHelper = new MappingHelper<>(aspectMap, result);
     mappingHelper.mapToResult(QUERY_PROPERTIES_ASPECT_NAME, this::mapQueryProperties);
     mappingHelper.mapToResult(QUERY_SUBJECTS_ASPECT_NAME, this::mapQuerySubjects);
+    mappingHelper.mapToResult(DATA_PLATFORM_INSTANCE_ASPECT_NAME, this::mapPlatform);
     return mappingHelper.getResult();
   }
 
+  private void mapPlatform(@Nonnull QueryEntity query, @Nonnull DataMap dataMap) {
+    DataPlatformInstance aspect = new DataPlatformInstance(dataMap);
+    if (aspect.hasPlatform()) {
+      final DataPlatform platform = new DataPlatform();
+      platform.setUrn(aspect.getPlatform().toString());
+      platform.setType(EntityType.DATA_PLATFORM);
+      query.setPlatform(platform);
+    }
+  }
+
   private void mapQueryProperties(@Nonnull QueryEntity query, @Nonnull DataMap dataMap) {
     QueryProperties queryProperties = new QueryProperties(dataMap);
     com.linkedin.datahub.graphql.generated.QueryProperties res =
@@ -57,6 +73,9 @@ private void mapQueryProperties(@Nonnull QueryEntity query, @Nonnull DataMap dat
             QueryLanguage.valueOf(queryProperties.getStatement().getLanguage().toString())));
     res.setName(queryProperties.getName(GetMode.NULL));
     res.setDescription(queryProperties.getDescription(GetMode.NULL));
+    if (queryProperties.hasOrigin() && queryProperties.getOrigin() != null) {
+      res.setOrigin(UrnToEntityMapper.map(queryProperties.getOrigin()));
+    }
 
     AuditStamp created = new AuditStamp();
     created.setTime(queryProperties.getCreated().getTime());
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -1561,7 +1561,7 @@ type RoleUser {
 
 type RoleProperties {
     """
-    Name of the Role in an organisation 
+    Name of the Role in an organisation
     """
     name: String!
 
@@ -1584,6 +1584,8 @@ type RoleProperties {
 type FineGrainedLineage {
     upstreams: [SchemaFieldRef!]
     downstreams: [SchemaFieldRef!]
+    query: String
+    transformOperation: String
 }
 
 """
@@ -11200,6 +11202,11 @@ type QueryProperties {
   An Audit Stamp corresponding to the update of this resource
   """
   lastModified: AuditStamp!
+
+  """
+  The asset that this query originated from, e.g. a View, a dbt Model, etc.
+  """
+  origin: Entity
 }
 
 """
@@ -11240,6 +11247,11 @@ type QueryEntity implements Entity {
   Granular API for querying edges extending from this entity
   """
   relationships(input: RelationshipsInput!): EntityRelationshipsResult
+
+  """
+  Platform from which the Query was detected
+  """
+  platform: DataPlatform
 }
 
 """
diff --git a/datahub-web-react/src/graphql/lineage.graphql b/datahub-web-react/src/graphql/lineage.graphql
@@ -55,6 +55,8 @@ fragment lineageNodeProperties on EntityWithRelationships {
                     urn
                     path
                 }
+                query
+                transformOperation
             }
         }
         health {
@@ -562,4 +564,4 @@ query getSearchAcrossLineageCounts(
             ...facetFields
         }
     }
-}
+}
diff --git a/datahub-web-react/src/graphql/query.graphql b/datahub-web-react/src/graphql/query.graphql
@@ -11,8 +11,8 @@ query getQuery($urn: String!) {
 fragment query on QueryEntity {
     urn
     properties {
-        description
         name
+        description
         source
         statement {
             value
@@ -26,6 +26,12 @@ fragment query on QueryEntity {
             time
             actor
         }
+        origin {
+            ...searchResultFields
+        }
+    }
+    platform {
+        ...platformFields
     }
     subjects {
         dataset {
diff --git a/metadata-models/src/main/pegasus/com/linkedin/query/QueryProperties.pdl b/metadata-models/src/main/pegasus/com/linkedin/query/QueryProperties.pdl
@@ -76,7 +76,9 @@ record QueryProperties {
   lastModified: AuditStamp
 
   /**
-   * The urn of the DataPlatform where the Query was executed.
+   * The origin of the Query.
+   * This is the source of the Query (e.g. a View, Stored Procedure, dbt Model, etc.) that the Query was created from.
    */
-  dataPlatform: optional Urn
+   origin: optional Urn
+
 }
diff --git a/metadata-models/src/main/resources/entity-registry.yml b/metadata-models/src/main/resources/entity-registry.yml
@@ -488,6 +488,8 @@ entities:
       - queryProperties
       - querySubjects
       - status
+      - dataPlatformInstance
+      - subTypes
   - name: dataProduct
     category: core
     keyAspect: dataProductKey
@@ -562,4 +564,4 @@ plugins:
         - UPSERT
       supportedEntityAspectNames:
         - entityName: '*'
-          aspectName: structuredProperties
+          aspectName: structuredProperties

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,12 @@ public List<com.linkedin.datahub.graphql.generated.FineGrainedLineage> apply(`
`45`	`45`	`.map(FineGrainedLineagesMapper::mapDatasetSchemaField)`
`46`	`46`	`.collect(Collectors.toList()));`
`47`	`47`	`}`
	`48`	`+ if (fineGrainedLineage.hasQuery()) {`
	`49`	`+ resultEntry.setQuery(fineGrainedLineage.getQuery().toString());`
	`50`	`+ }`
	`51`	`+ if (fineGrainedLineage.hasTransformOperation()) {`
	`52`	`+ resultEntry.setTransformOperation(fineGrainedLineage.getTransformOperation());`
	`53`	`+ }`
`48`	`54`	`result.add(resultEntry);`
`49`	`55`	`}`
`50`	`56`	`return result;`
Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,8 @@ fragment lineageNodeProperties on EntityWithRelationships {`
`55`	`55`	`urn`
`56`	`56`	`path`
`57`	`57`	`}`
	`58`	`+ query`
	`59`	`+ transformOperation`
`58`	`60`	`}`
`59`	`61`	`}`
`60`	`62`	`health {`
`@@ -562,4 +564,4 @@ query getSearchAcrossLineageCounts(`
`562`	`564`	`...facetFields`
`563`	`565`	`}`
`564`	`566`	`}`
`565`		`-}`
	`567`	`+}`