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

Author: asikowitz

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

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;

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

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
         }
     }
-}
+}

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 {

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
+
 }

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