fix(ui) Properly display column-level lineage with v2 field paths (da…

…tahub-project#6217)
CybercentreCanada · Nov 18, 2022 · 23857f0 · 23857f0
1 parent 90637da
commit 23857f0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts b/datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts
@@ -1,5 +1,6 @@
 import { ColumnEdge, FetchedEntity, NodeData } from '../types';
 import { SchemaField } from '../../../types.generated';
+import { downgradeV2FieldPath } from '../../entity/dataset/profile/schema/utils/utils';
 
 export function getHighlightedColumnsForNode(highlightedEdges: ColumnEdge[], fields: SchemaField[], nodeUrn: string) {
     return highlightedEdges
@@ -35,6 +36,13 @@ export function sortRelatedLineageColumns(
     };
 }
 
+export function convertFieldsToV1FieldPath(fields: SchemaField[]) {
+    return fields.map((field) => ({
+        ...field,
+        fieldPath: downgradeV2FieldPath(field.fieldPath) || '',
+    }));
+}
+
 export function sortColumnsByDefault(
     columnsByUrn: Record<string, SchemaField[]>,
     fields: SchemaField[],
@@ -59,7 +67,10 @@ export function populateColumnsByUrn(
     let populatedColumnsByUrn = { ...columnsByUrn };
     Object.entries(fetchedEntities).forEach(([urn, fetchedEntity]) => {
         if (fetchedEntity.schemaMetadata && !columnsByUrn[urn]) {
-            populatedColumnsByUrn = { ...populatedColumnsByUrn, [urn]: fetchedEntity.schemaMetadata.fields };
+            populatedColumnsByUrn = {
+                ...populatedColumnsByUrn,
+                [urn]: convertFieldsToV1FieldPath(fetchedEntity.schemaMetadata.fields),
+            };
         }
     });
     setColumnsByUrn(populatedColumnsByUrn);
@@ -87,6 +98,9 @@ export function filterColumns(
 ) {
     const filteredFields = node.data.schemaMetadata?.fields.filter((field) => field.fieldPath.includes(filterText));
     if (filteredFields) {
-        setColumnsByUrn((colsByUrn) => ({ ...colsByUrn, [node.data.urn || 'noop']: filteredFields }));
+        setColumnsByUrn((colsByUrn) => ({
+            ...colsByUrn,
+            [node.data.urn || 'noop']: convertFieldsToV1FieldPath(filteredFields),
+        }));
     }
 }
diff --git a/datahub-web-react/src/app/lineage/utils/useSortColumnsBySelectedField.ts b/datahub-web-react/src/app/lineage/utils/useSortColumnsBySelectedField.ts
@@ -2,7 +2,12 @@ import { useContext, useEffect } from 'react';
 import usePrevious from '../../shared/usePrevious';
 import { NUM_COLUMNS_PER_PAGE } from '../constants';
 import { FetchedEntity } from '../types';
-import { getHighlightedColumnsForNode, sortColumnsByDefault, sortRelatedLineageColumns } from './columnLineageUtils';
+import {
+    convertFieldsToV1FieldPath,
+    getHighlightedColumnsForNode,
+    sortColumnsByDefault,
+    sortRelatedLineageColumns,
+} from './columnLineageUtils';
 import { LineageExplorerContext } from './LineageExplorerContext';
 
 export default function useSortColumnsBySelectedField(fetchedEntities: { [x: string]: FetchedEntity }) {
@@ -35,7 +40,7 @@ export default function useSortColumnsBySelectedField(fetchedEntities: { [x: str
                     updatedColumnsByUrn = sortColumnsByDefault(
                         updatedColumnsByUrn,
                         columns,
-                        fetchedEntity.schemaMetadata.fields,
+                        convertFieldsToV1FieldPath(fetchedEntity.schemaMetadata.fields),
                         urn,
                     );
                 }