feat(lineage) Add column-level impact analysis feature (datahub-proje…

…ct#6272)
CybercentreCanada · Nov 18, 2022 · 1627155 · 1627155
1 parent e82b3c5
commit 1627155
Show file tree

Hide file tree

Showing 41 changed files with 989 additions and 109 deletions.
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
@@ -44,6 +44,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.EntityPath;
 import com.linkedin.datahub.graphql.generated.EntityRelationship;
 import com.linkedin.datahub.graphql.generated.EntityRelationshipLegacy;
 import com.linkedin.datahub.graphql.generated.ForeignKeyConstraint;
@@ -71,6 +72,7 @@
 import com.linkedin.datahub.graphql.generated.Owner;
 import com.linkedin.datahub.graphql.generated.PolicyMatchCriterionValue;
 import com.linkedin.datahub.graphql.generated.RecommendationContent;
+import com.linkedin.datahub.graphql.generated.SchemaFieldEntity;
 import com.linkedin.datahub.graphql.generated.SearchAcrossLineageResult;
 import com.linkedin.datahub.graphql.generated.SearchResult;
 import com.linkedin.datahub.graphql.generated.SiblingProperties;
@@ -242,6 +244,7 @@
 import com.linkedin.datahub.graphql.types.notebook.NotebookType;
 import com.linkedin.datahub.graphql.types.policy.DataHubPolicyType;
 import com.linkedin.datahub.graphql.types.role.DataHubRoleType;
+import com.linkedin.datahub.graphql.types.schemafield.SchemaFieldType;
 import com.linkedin.datahub.graphql.types.tag.TagType;
 import com.linkedin.datahub.graphql.types.test.TestType;
 import com.linkedin.entity.client.EntityClient;
@@ -352,6 +355,7 @@ public class GmsGraphQLEngine {
     private final TestType testType;
     private final DataHubPolicyType dataHubPolicyType;
     private final DataHubRoleType dataHubRoleType;
+    private final SchemaFieldType schemaFieldType;
 
     /**
      * Configures the graph objects that can be fetched primary key.
@@ -449,6 +453,7 @@ public GmsGraphQLEngine(final EntityClient entityClient, final GraphClient graph
         this.testType = new TestType(entityClient);
         this.dataHubPolicyType = new DataHubPolicyType(entityClient);
         this.dataHubRoleType = new DataHubRoleType(entityClient);
+        this.schemaFieldType = new SchemaFieldType();
         // Init Lists
         this.entityTypes = ImmutableList.of(
             datasetType,
@@ -476,7 +481,8 @@ public GmsGraphQLEngine(final EntityClient entityClient, final GraphClient graph
             accessTokenMetadataType,
             testType,
             dataHubPolicyType,
-            dataHubRoleType
+            dataHubRoleType,
+            schemaFieldType
         );
         this.loadableTypes = new ArrayList<>(entityTypes);
         this.ownerTypes = ImmutableList.of(corpUserType, corpGroupType);
@@ -535,6 +541,8 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
         configureAccessAccessTokenMetadataResolvers(builder);
         configureTestResultResolvers(builder);
         configureRoleResolvers(builder);
+        configureSchemaFieldResolvers(builder);
+        configureEntityPathResolvers(builder);
     }
 
     public GraphQLEngine.Builder builder() {
@@ -1008,6 +1016,20 @@ private void configureGlossaryNodeResolvers(final RuntimeWiring.Builder builder)
         );
     }
 
+    private void configureSchemaFieldResolvers(final RuntimeWiring.Builder builder) {
+        builder.type("SchemaFieldEntity", typeWiring -> typeWiring
+            .dataFetcher("parent", new EntityTypeResolver(entityTypes,
+                (env) -> ((SchemaFieldEntity) env.getSource()).getParent()))
+        );
+    }
+
+    private void configureEntityPathResolvers(final RuntimeWiring.Builder builder) {
+        builder.type("EntityPath", typeWiring -> typeWiring
+            .dataFetcher("path",  new BatchGetEntitiesResolver(entityTypes,
+                (env) -> ((EntityPath) env.getSource()).getPath()))
+        );
+    }
+
     /**
      * Configures resolvers responsible for resolving the {@link com.linkedin.datahub.graphql.generated.CorpUser} type.
      */

diff --git a/...re/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java b/...re/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java
@@ -25,6 +25,7 @@
 import com.linkedin.datahub.graphql.generated.MLModelGroup;
 import com.linkedin.datahub.graphql.generated.MLPrimaryKey;
 import com.linkedin.datahub.graphql.generated.Notebook;
+import com.linkedin.datahub.graphql.generated.SchemaFieldEntity;
 import com.linkedin.datahub.graphql.generated.Tag;
 import com.linkedin.datahub.graphql.generated.Test;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
@@ -163,6 +164,11 @@ public Entity apply(Urn input) {
       ((DataHubPolicy) partialEntity).setUrn(input.toString());
       ((DataHubPolicy) partialEntity).setType(EntityType.DATAHUB_POLICY);
     }
+    if (input.getEntityType().equals(SCHEMA_FIELD_ENTITY_NAME)) {
+      partialEntity = new SchemaFieldEntity();
+      ((SchemaFieldEntity) partialEntity).setUrn(input.toString());
+      ((SchemaFieldEntity) partialEntity).setType(EntityType.SCHEMA_FIELD);
+    }
     return partialEntity;
   }
 }
diff --git a/...n/java/com/linkedin/datahub/graphql/types/dataset/mappers/ForeignKeyConstraintMapper.java b/...n/java/com/linkedin/datahub/graphql/types/dataset/mappers/ForeignKeyConstraintMapper.java
@@ -5,9 +5,12 @@
 import com.linkedin.datahub.graphql.generated.ForeignKeyConstraint;
 import com.linkedin.datahub.graphql.generated.SchemaFieldEntity;
 import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
+import lombok.extern.slf4j.Slf4j;
+
 import java.util.stream.Collectors;
 
 
+@Slf4j
 public class ForeignKeyConstraintMapper {
   private ForeignKeyConstraintMapper() { }
 
@@ -34,7 +37,12 @@ public static ForeignKeyConstraint map(com.linkedin.schema.ForeignKeyConstraint
 
   private static SchemaFieldEntity mapSchemaFieldEntity(Urn schemaFieldUrn) {
     SchemaFieldEntity result = new SchemaFieldEntity();
-    result.setParent(schemaFieldUrn.getEntityKey().get(0));
+    try {
+      Urn resourceUrn = Urn.createFromString(schemaFieldUrn.getEntityKey().get(0));
+      result.setParent(UrnToEntityMapper.map(resourceUrn));
+    } catch (Exception e) {
+      throw new RuntimeException("Error converting schemaField parent urn string to Urn", e);
+    }
     result.setFieldPath(schemaFieldUrn.getEntityKey().get(1));
     return result;
   }

diff --git a/.../java/com/linkedin/datahub/graphql/types/mappers/UrnSearchAcrossLineageResultsMapper.java b/.../java/com/linkedin/datahub/graphql/types/mappers/UrnSearchAcrossLineageResultsMapper.java
@@ -1,9 +1,11 @@
 package com.linkedin.datahub.graphql.types.mappers;
 
+import com.linkedin.common.UrnArray;
 import com.linkedin.data.template.DoubleMap;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.datahub.graphql.generated.AggregationMetadata;
 import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.EntityPath;
 import com.linkedin.datahub.graphql.generated.FacetMetadata;
 import com.linkedin.datahub.graphql.generated.MatchedField;
 import com.linkedin.datahub.graphql.generated.SearchAcrossLineageResult;
@@ -46,11 +48,17 @@ private SearchAcrossLineageResult mapResult(LineageSearchEntity searchEntity) {
         .setEntity(UrnToEntityMapper.map(searchEntity.getEntity()))
         .setInsights(getInsightsFromFeatures(searchEntity.getFeatures()))
         .setMatchedFields(getMatchedFieldEntry(searchEntity.getMatchedFields()))
-        .setPath(searchEntity.getPath().stream().map(UrnToEntityMapper::map).collect(Collectors.toList()))
+        .setPaths(searchEntity.getPaths().stream().map(this::mapPath).collect(Collectors.toList()))
         .setDegree(searchEntity.getDegree())
         .build();
   }
 
+  private EntityPath mapPath(UrnArray path) {
+    EntityPath entityPath = new EntityPath();
+    entityPath.setPath(path.stream().map(UrnToEntityMapper::map).collect(Collectors.toList()));
+    return entityPath;
+  }
+
   private FacetMetadata mapFacet(com.linkedin.metadata.search.AggregationMetadata aggregationMetadata) {
     final FacetMetadata facetMetadata = new FacetMetadata();
     boolean isEntityTypeFilter = aggregationMetadata.getName().equals("entity");

diff --git a/...ql-core/src/main/java/com/linkedin/datahub/graphql/types/schemafield/SchemaFieldType.java b/...ql-core/src/main/java/com/linkedin/datahub/graphql/types/schemafield/SchemaFieldType.java
@@ -0,0 +1,71 @@
+package com.linkedin.datahub.graphql.types.schemafield;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.SchemaFieldEntity;
+import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
+import graphql.execution.DataFetcherResult;
+
+import javax.annotation.Nonnull;
+import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public class SchemaFieldType implements com.linkedin.datahub.graphql.types.EntityType<SchemaFieldEntity, String> {
+
+  public SchemaFieldType() { }
+
+  @Override
+  public EntityType type() {
+    return EntityType.SCHEMA_FIELD;
+  }
+
+  @Override
+  public Function<Entity, String> getKeyProvider() {
+    return Entity::getUrn;
+  }
+
+  @Override
+  public Class<SchemaFieldEntity> objectClass() {
+    return SchemaFieldEntity.class;
+  }
+
+  @Override
+  public List<DataFetcherResult<SchemaFieldEntity>> batchLoad(@Nonnull List<String> urns, @Nonnull QueryContext context) throws Exception {
+    final List<Urn> schemaFieldUrns = urns.stream()
+        .map(UrnUtils::getUrn)
+        .collect(Collectors.toList());
+
+    try {
+      return schemaFieldUrns.stream()
+          .map(this::mapSchemaFieldUrn)
+          .map(schemaFieldEntity -> DataFetcherResult.<SchemaFieldEntity>newResult()
+              .data(schemaFieldEntity)
+              .build()
+          )
+          .collect(Collectors.toList());
+
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to load schemaField entity", e);
+    }
+  }
+
+  private SchemaFieldEntity mapSchemaFieldUrn(Urn urn) {
+    try {
+      SchemaFieldEntity result = new SchemaFieldEntity();
+      result.setUrn(urn.toString());
+      result.setType(EntityType.SCHEMA_FIELD);
+      result.setFieldPath(urn.getEntityKey().get(1));
+      Urn parentUrn = Urn.createFromString(urn.getEntityKey().get(0));
+      result.setParent(UrnToEntityMapper.map(parentUrn));
+      return result;
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to load schemaField entity", e);
+    }
+  }
+
+}
+
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -2343,21 +2343,31 @@ type KeyValueSchema {
 Standalone schema field entity. Differs from the SchemaField struct because it is not directly nested inside a
 schema field
 """
-type SchemaFieldEntity {
+type SchemaFieldEntity implements Entity {
     """
     Primary key of the schema field
     """
     urn: String!
 
+    """
+    A standard Entity Type
+    """
+    type: EntityType!
+
     """
     Field path identifying the field in its dataset
     """
     fieldPath: String!
 
     """
-    The primary key of the field's parent.
+    The field's parent.
     """
-    parent: String!
+    parent: Entity!
+
+    """
+    Granular API for querying edges extending from this entity
+    """
+    relationships(input: RelationshipsInput!): EntityRelationshipsResult
 }
 
 """

diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql
@@ -318,14 +318,25 @@ type SearchAcrossLineageResult {
   """
   Optional list of entities between the source and destination node
   """
-  path: [Entity!]
+  paths: [EntityPath]
 
   """
   Degree of relationship (number of hops to get to entity)
   """
   degree: Int!
 }
 
+"""
+An overview of the field that was matched in the entity search document
+"""
+type EntityPath {
+    """
+    Path of entities between source and destination nodes
+    """
+    path: [Entity]
+}
+
+
 """
 An overview of the field that was matched in the entity search document
 """

diff --git a/datahub-web-react/src/app/entity/dataset/profile/stories/sampleSchema.ts b/datahub-web-react/src/app/entity/dataset/profile/stories/sampleSchema.ts
@@ -204,14 +204,16 @@ export const sampleSchemaWithPkFk: SchemaMetadata = {
             sourceFields: [
                 {
                     urn: 'datasetUrn',
-                    parent: 'dataset',
+                    type: EntityType.Dataset,
+                    parent: { urn: 'test', type: EntityType.Dataset },
                     fieldPath: 'shipping_address',
                 },
             ],
             foreignFields: [
                 {
                     urn: dataset3.urn,
-                    parent: dataset3.name,
+                    type: EntityType.Dataset,
+                    parent: { urn: dataset3.name, type: EntityType.Dataset },
                     fieldPath: 'address',
                 },
             ],

diff --git a/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx b/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx
@@ -31,7 +31,7 @@ export enum EntityMenuItems {
     MOVE,
 }
 
-const MenuIcon = styled(MoreOutlined)<{ fontSize?: number }>`
+export const MenuIcon = styled(MoreOutlined)<{ fontSize?: number }>`
     display: flex;
     justify-content: center;
     align-items: center;

diff --git a/...ub-web-react/src/app/entity/shared/components/styled/search/EmbeddedListSearchResults.tsx b/...ub-web-react/src/app/entity/shared/components/styled/search/EmbeddedListSearchResults.tsx
@@ -11,7 +11,7 @@ import { SearchFiltersSection } from '../../../../../search/SearchFiltersSection
 
 const SearchBody = styled.div`
     height: 100%;
-    overflow-y: scroll;
+    overflow-y: auto;
     display: flex;
 `;
 
@@ -127,6 +127,8 @@ export const EmbeddedListSearchResults = ({
                                     // when we add impact analysis, we will want to pipe the path to each element to the result this
                                     // eslint-disable-next-line @typescript-eslint/dot-notation
                                     degree: searchResult['degree'],
+                                    // eslint-disable-next-line @typescript-eslint/dot-notation
+                                    paths: searchResult['paths'],
                                 })) || []
                             }
                             isSelectMode={isSelectMode}

diff --git a/...ub-web-react/src/app/entity/shared/components/styled/search/EmbeddedListSearchSection.tsx b/...ub-web-react/src/app/entity/shared/components/styled/search/EmbeddedListSearchSection.tsx
@@ -40,9 +40,10 @@ export const EmbeddedListSearchSection = ({
 }: Props) => {
     const history = useHistory();
     const location = useLocation();
-    const baseParams = useEntityQueryParams();
+    const entityQueryParams = useEntityQueryParams();
 
     const params = QueryString.parse(location.search, { arrayFormat: 'comma' });
+    const baseParams = { ...params, ...entityQueryParams };
     const query: string = params?.query as string;
     const page: number = params.page && Number(params.page as string) > 0 ? Number(params.page as string) : 1;
     const unionType: UnionType = Number(params.unionType as any as UnionType) || UnionType.AND;

diff --git a/...hub-web-react/src/app/entity/shared/components/styled/search/navigateToEntitySearchUrl.ts b/...hub-web-react/src/app/entity/shared/components/styled/search/navigateToEntitySearchUrl.ts
@@ -30,11 +30,11 @@ export const navigateToEntitySearchUrl = ({
 
     const search = QueryString.stringify(
         {
+            ...baseParams,
             ...filtersToQueryStringParams(constructedFilters),
             query: newQuery,
             page: newPage,
             unionType,
-            ...baseParams,
         },
         { arrayFormat: 'comma' },
     );

diff --git a/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/SchemaTable.tsx b/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/SchemaTable.tsx
@@ -21,6 +21,7 @@ import { SchemaRow } from './components/SchemaRow';
 import { FkContext } from './utils/selectedFkContext';
 import useSchemaBlameRenderer from './utils/useSchemaBlameRenderer';
 import { ANTD_GRAY } from '../../../constants';
+import MenuColumn from './components/MenuColumn';
 
 const TableContainer = styled.div`
     &&& .ant-table-tbody > tr > .ant-table-cell-with-append {
@@ -163,6 +164,14 @@ export default function SchemaTable({
         render: usageStatsRenderer,
     };
 
+    const menuColumn = {
+        width: '5%',
+        title: '',
+        dataIndex: '',
+        key: 'menu',
+        render: (field: SchemaField) => <MenuColumn field={field} />,
+    };
+
     let allColumns: ColumnsType<ExtendedSchemaFields> = [fieldColumn, descriptionColumn, tagColumn, termColumn];
 
     if (hasUsageStats) {
@@ -173,6 +182,8 @@ export default function SchemaTable({
         allColumns = [...allColumns, blameColumn];
     }
 
+    allColumns = [...allColumns, menuColumn];
+
     const [expandedRows, setExpandedRows] = useState<Set<string>>(new Set());
 
     useEffect(() => {