fix(ui): Handle encoded schemaField urns on the frontend (datahub-pro…

…ject#6321)
CybercentreCanada · Nov 18, 2022 · 92467c4 · 92467c4
1 parent 8fb19f5
commit 92467c4
Show file tree

Hide file tree

Showing 9 changed files with 70 additions and 11 deletions.
diff --git a/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/components/SchemaRow.tsx b/datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/components/SchemaRow.tsx
@@ -4,6 +4,7 @@ import { Link } from 'react-router-dom';
 import styled from 'styled-components';
 import { GetDatasetQuery } from '../../../../../../../graphql/dataset.generated';
 import { EntityType } from '../../../../../../../types.generated';
+import { decodeSchemaField } from '../../../../../../lineage/utils/columnLineageUtils';
 import CompactContext from '../../../../../../shared/CompactContext';
 import { useEntityRegistry } from '../../../../../../useEntityRegistry';
 import { ANTD_GRAY } from '../../../../constants';
@@ -130,7 +131,7 @@ export const SchemaRow = ({
                                     <TableTitle>{baseEntity.dataset?.name}</TableTitle>
                                     {selectedFk?.constraint?.sourceFields?.map((field) => (
                                         <div key={field?.fieldPath}>
-                                            <FieldBadge count={field?.fieldPath} />
+                                            <FieldBadge count={decodeSchemaField(field?.fieldPath || '')} />
                                         </div>
                                     ))}
                                 </div>
@@ -139,7 +140,7 @@ export const SchemaRow = ({
                                     <TableTitle>{selectedFk?.constraint?.foreignDataset?.name}</TableTitle>
                                     {selectedFk?.constraint?.foreignFields?.map((field) => (
                                         <div key={field?.fieldPath}>
-                                            <FieldBadge count={field?.fieldPath} />
+                                            <FieldBadge count={decodeSchemaField(field?.fieldPath || '')} />
                                         </div>
                                     ))}
                                 </div>

diff --git a/...-web-react/src/app/entity/shared/tabs/Dataset/Validations/DatasetAssertionDescription.tsx b/...-web-react/src/app/entity/shared/tabs/Dataset/Validations/DatasetAssertionDescription.tsx
@@ -9,6 +9,7 @@ import {
     DatasetAssertionScope,
     SchemaFieldRef,
 } from '../../../../../../types.generated';
+import { decodeSchemaField } from '../../../../../lineage/utils/columnLineageUtils';
 import { getFormattedParameterValue } from './assertionUtils';
 import { DatasetAssertionLogicModal } from './DatasetAssertionLogicModal';
 
@@ -37,7 +38,7 @@ const getSchemaAggregationText = (
         case AssertionStdAggregation.Columns:
             return <Typography.Text>Dataset columns are</Typography.Text>;
         case AssertionStdAggregation.Native: {
-            const fieldNames = fields?.map((field) => field.path) || [];
+            const fieldNames = fields?.map((field) => decodeSchemaField(field.path)) || [];
             return (
                 <Typography.Text>
                     Dataset columns <Typography.Text strong>{JSON.stringify(fieldNames)}</Typography.Text> are
@@ -76,7 +77,7 @@ const getColumnAggregationText = (
     aggregation: AssertionStdAggregation | undefined | null,
     field: SchemaFieldRef | undefined,
 ) => {
-    let columnText = field?.path;
+    let columnText = decodeSchemaField(field?.path || '');
     if (field === undefined) {
         columnText = 'undefined';
         console.error(`Invalid field provided for Dataset Assertion with scope Column ${JSON.stringify(field)}`);

diff --git a/datahub-web-react/src/app/entity/shared/tabs/Lineage/utils.ts b/datahub-web-react/src/app/entity/shared/tabs/Lineage/utils.ts
@@ -1,4 +1,6 @@
+import { encodeSchemaField } from '../../../../lineage/utils/columnLineageUtils';
+
 export function generateSchemaFieldUrn(fieldPath: string | undefined, resourceUrn: string) {
     if (!fieldPath) return null;
-    return `urn:li:schemaField:(${resourceUrn},${fieldPath})`;
+    return `urn:li:schemaField:(${resourceUrn},${encodeSchemaField(fieldPath)})`;
 }
diff --git a/datahub-web-react/src/app/lineage/utils/__tests__/columnLineageUtils.test.tsx b/datahub-web-react/src/app/lineage/utils/__tests__/columnLineageUtils.test.tsx
@@ -1,4 +1,9 @@
-import { getFieldPathFromSchemaFieldUrn, getSourceUrnFromSchemaFieldUrn } from '../columnLineageUtils';
+import {
+    decodeSchemaField,
+    encodeSchemaField,
+    getFieldPathFromSchemaFieldUrn,
+    getSourceUrnFromSchemaFieldUrn,
+} from '../columnLineageUtils';
 
 describe('getSourceUrnFromSchemaFieldUrn', () => {
     it('should get the source urn for a chart schemaField', () => {
@@ -43,3 +48,37 @@ describe('getFieldPathFromSchemaFieldUrn', () => {
         expect(sourceUrn).toBe('user.name.test');
     });
 });
+
+describe('decodeSchemaField', () => {
+    it('should decode a schemaField path when it has encoded reserved characters', () => {
+        const decodedSchemaField = decodeSchemaField('Test%2C test%2C test %28and test%29');
+        expect(decodedSchemaField).toBe('Test, test, test (and test)');
+    });
+
+    it('should return a regular schemaField path when it was not encoded', () => {
+        const schemaField = 'user.name';
+        const decodedSchemaField = decodeSchemaField(schemaField);
+        expect(decodedSchemaField).toBe(schemaField);
+    });
+});
+
+describe('encodeSchemaField', () => {
+    it('should encode a schemaField path when it has encoded reserved characters', () => {
+        const encodedSchemaField = encodeSchemaField('Test, test, test (and test)');
+        expect(encodedSchemaField).toBe('Test%2C test%2C test %28and test%29');
+    });
+
+    it('should return a regular schemaField path when it does not have reserved characters', () => {
+        const schemaField = 'user.name';
+        const encodedSchemaField = encodeSchemaField(schemaField);
+        expect(encodedSchemaField).toBe(schemaField);
+    });
+
+    it('should encode a decoded schemaField that we generate', () => {
+        const schemaField = 'Adults, men (% of pop)';
+        const encodedSchemaField = encodeSchemaField(schemaField);
+        const decodedSchemaField = decodeSchemaField(encodedSchemaField);
+        expect(encodedSchemaField).toBe('Adults%2C men %28% of pop%29');
+        expect(decodedSchemaField).toBe(schemaField);
+    });
+});
diff --git a/datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts b/datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts
@@ -116,9 +116,17 @@ export function filterColumns(
     }
 }
 
+export function decodeSchemaField(fieldPath: string) {
+    return fieldPath.replaceAll('%28', '(').replaceAll('%29', ')').replaceAll('%2C', ',');
+}
+
+export function encodeSchemaField(fieldPath: string) {
+    return fieldPath.replaceAll('(', '%28').replaceAll(')', '%29').replaceAll(',', '%2C');
+}
+
 export function getSourceUrnFromSchemaFieldUrn(schemaFieldUrn: string) {
     return schemaFieldUrn.replace('urn:li:schemaField:(', '').split(')')[0].concat(')');
 }
 export function getFieldPathFromSchemaFieldUrn(schemaFieldUrn: string) {
-    return schemaFieldUrn.replace('urn:li:schemaField:(', '').split(')')[1].replace(',', '');
+    return decodeSchemaField(schemaFieldUrn.replace('urn:li:schemaField:(', '').split(')')[1].replace(',', ''));
 }
diff --git a/datahub-web-react/src/app/lineage/utils/extendAsyncEntities.ts b/datahub-web-react/src/app/lineage/utils/extendAsyncEntities.ts
@@ -1,11 +1,15 @@
 import { SchemaFieldRef } from '../../../types.generated';
 import EntityRegistry from '../../entity/EntityRegistry';
 import { EntityAndType, FetchedEntities, FetchedEntity } from '../types';
-import { getFieldPathFromSchemaFieldUrn, getSourceUrnFromSchemaFieldUrn } from './columnLineageUtils';
+import {
+    decodeSchemaField,
+    getFieldPathFromSchemaFieldUrn,
+    getSourceUrnFromSchemaFieldUrn,
+} from './columnLineageUtils';
 
 const breakFieldUrn = (ref: SchemaFieldRef) => {
     const before = ref.urn;
-    const after = ref.path;
+    const after = decodeSchemaField(ref.path);
 
     return [before, after];
 };

diff --git a/datahub-web-react/src/app/preview/EntityPaths/ColumnsRelationshipText.tsx b/datahub-web-react/src/app/preview/EntityPaths/ColumnsRelationshipText.tsx
@@ -4,6 +4,7 @@ import styled from 'styled-components/macro';
 import { Entity, LineageDirection } from '../../../types.generated';
 import { downgradeV2FieldPath } from '../../entity/dataset/profile/schema/utils/utils';
 import { LineageTabContext } from '../../entity/shared/tabs/Lineage/LineageTabContext';
+import { decodeSchemaField } from '../../lineage/utils/columnLineageUtils';
 import DisplayedColumns from './DisplayedColumns';
 
 const ColumnNameWrapper = styled.span<{ isBlack?: boolean }>`
@@ -19,7 +20,7 @@ interface Props {
 export default function ColumnsRelationshipText({ displayedColumns }: Props) {
     const { selectedColumn, lineageDirection } = useContext(LineageTabContext);
 
-    const displayedFieldPath = downgradeV2FieldPath(selectedColumn);
+    const displayedFieldPath = decodeSchemaField(downgradeV2FieldPath(selectedColumn) || '');
 
     return (
         <>

diff --git a/datahub-web-react/src/app/preview/EntityPaths/DisplayedColumns.tsx b/datahub-web-react/src/app/preview/EntityPaths/DisplayedColumns.tsx
@@ -3,6 +3,7 @@ import React from 'react';
 import styled from 'styled-components/macro';
 import { Entity, EntityType, SchemaFieldEntity } from '../../../types.generated';
 import { downgradeV2FieldPath } from '../../entity/dataset/profile/schema/utils/utils';
+import { decodeSchemaField } from '../../lineage/utils/columnLineageUtils';
 import { useEntityRegistry } from '../../useEntityRegistry';
 
 const ColumnNameWrapper = styled.span<{ isBlack?: boolean }>`
@@ -25,7 +26,7 @@ export default function DisplayedColumns({ displayedColumns }: Props) {
                     return (
                         <ColumnNameWrapper>
                             {entity.type === EntityType.SchemaField
-                                ? downgradeV2FieldPath((entity as SchemaFieldEntity).fieldPath)
+                                ? decodeSchemaField(downgradeV2FieldPath((entity as SchemaFieldEntity).fieldPath) || '')
                                 : entityRegistry.getDisplayName(entity.type, entity)}
                             {index !== displayedColumns.length - 1 && ', '}
                         </ColumnNameWrapper>

diff --git a/metadata-ingestion/src/datahub/utilities/urn_encoder.py b/metadata-ingestion/src/datahub/utilities/urn_encoder.py
@@ -1,6 +1,8 @@
 import urllib.parse
 from typing import List
 
+# NOTE: Frontend relies on encoding these three characters. Specifically, we decode and encode schema fields for column level lineage.
+# If this changes, make appropriate changes to datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts
 RESERVED_CHARS = [",", "(", ")"]