feat(edit lineage): add edit lineage functionality to datahub (#12976)

datahub-web-react/src/app/entity/shared/utils.ts

@@ -140,3 +140,84 @@ export function getStructuredPropertyValue(value: PropertyValue) {
     }
     return null;
 }
+
+// Utility for formatting any casing of type to the expected casing for the API
+export function formatEntityType(type: string): string {
+    if (!type) return '';
+
+    switch (type.toLowerCase()) {
+        case 'dataset':
+            return EntityType.Dataset;
+        case 'role':
+            return EntityType.Role;
+        case 'corpuser':
+            return EntityType.CorpUser;
+        case 'corpgroup':
+            return EntityType.CorpGroup;
+        case 'dataplatform':
+            return EntityType.DataPlatform;
+        case 'dashboard':
+            return EntityType.Dashboard;
+        case 'chart':
+            return EntityType.Chart;
+        case 'tag':
+            return EntityType.Tag;
+        case 'dataflow':
+            return EntityType.DataFlow;
+        case 'datajob':
+            return EntityType.DataJob;
+        case 'glossaryterm':
+            return EntityType.GlossaryTerm;
+        case 'glossarynode':
+            return EntityType.GlossaryNode;
+        case 'mlmodel':
+            return EntityType.Mlmodel;
+        case 'mlmodelgroup':
+            return EntityType.MlmodelGroup;
+        case 'mlfeaturetable':
+            return EntityType.MlfeatureTable;
+        case 'mlfeature':
+            return EntityType.Mlfeature;
+        case 'mlprimarykey':
+            return EntityType.MlprimaryKey;
+        case 'container':
+            return EntityType.Container;
+        case 'domain':
+            return EntityType.Domain;
+        case 'notebook':
+            return EntityType.Notebook;
+        case 'dataplatforminstance':
+            return EntityType.DataPlatformInstance;
+        case 'test':
+            return EntityType.Test;
+        case 'schemafield':
+            return EntityType.SchemaField;
+
+        // these are const in the java app
+        case 'dataprocessinstance': // Constants.DATA_PROCESS_INSTANCE_ENTITY_NAME
+            return EntityType.DataProcessInstance;
+        case 'datahubview': // Constants.DATAHUB_VIEW_ENTITY_NAME
+            return EntityType.DatahubView;
+        case 'dataproduct': // Constants.DATA_PRODUCT_ENTITY_NAME
+            return EntityType.DataProduct;
+        case 'datahubconnection': // Constants.DATAHUB_CONNECTION_ENTITY_NAME
+            return EntityType.DatahubConnection;
+        case 'structuredproperty': // Constants.STRUCTURED_PROPERTY_ENTITY_NAME
+            return EntityType.StructuredProperty;
+        case 'assertion': // Constants.ASSERTION_ENTITY_NAME
+            return EntityType.Assertion;
+
+        default:
+            return '';
+    }
+}
+
+// Utility for getting entity type from urn if it's in the 3rd position
+export function extractTypeFromUrn(urn: string): EntityType {
+    const regex = /[^:]+:[^:]+:([^:]+):/;
+    const match = urn.match(regex);
+
+    if (match && match[1]) return formatEntityType(match[1]) as EntityType;
+
+    return '' as EntityType;
+}

datahub-web-react/src/app/entityV2/shared/components/styled/search/EntitySearchResults.tsx

@@ -1,8 +1,7 @@
 import { Checkbox, Empty, List } from 'antd';
 import React from 'react';
 import styled from 'styled-components';
-import { EntityPath, EntityType, SearchResult } from '../../../../../../types.generated';
-import { EntityAndType } from '../../../../../entity/shared/types';
+import { Entity, EntityPath, EntityType, SearchResult } from '../../../../../../types.generated';
 import { useSearchContext } from '../../../../../search/context/SearchContext';
 import { MATCHES_CONTAINER_HEIGHT } from '../../../../../searchV2/SearchResultList';
 import { MatchContextContainer } from '../../../../../searchV2/matches/MatchContextContainer';
@@ -76,8 +75,8 @@ type Props = {
     additionalPropertiesList?: Array<AdditionalProperties>;
     searchResults: Array<SearchResult>;
     isSelectMode?: boolean;
-    selectedEntities?: EntityAndType[];
-    setSelectedEntities?: (entities: EntityAndType[]) => any;
+    selectedEntities?: Entity[];
+    setSelectedEntities?: (entities: Entity[]) => any;
     bordered?: boolean;
     entityAction?: React.FC<EntityActionProps>;
     compactUserSearchCardStyle?: boolean;
@@ -120,7 +119,7 @@ export const EntitySearchResults = ({
     /**
      * Invoked when a new entity is selected. Simply updates the state of the list of selected entities.
      */
-    const onSelectEntity = (selectedEntity: EntityAndType, selected: boolean) => {
+    const onSelectEntity = (selectedEntity: Entity, selected: boolean) => {
         if (selected) {
             setSelectedEntities?.([...selectedEntities, selectedEntity]);
         } else {
@@ -177,9 +176,7 @@ export const EntitySearchResults = ({
                                             selectedEntities.length >= selectLimit &&
                                             !selectedEntityUrns.includes(entity.urn)
                                         }
-                                        onChange={(e) =>
-                                            onSelectEntity({ urn: entity.urn, type: entity.type }, e.target.checked)
-                                        }
+                                        onChange={(e) => onSelectEntity(entity, e.target.checked)}
                                     />
                                 )}
                                 {entityRegistry.renderSearchResult(entity.type, searchResult)}

datahub-web-react/src/app/entityV2/shared/components/styled/search/SearchSelect.tsx

@@ -1,14 +1,14 @@
 import { FilterOutlined } from '@ant-design/icons';
 import { Button, message, Typography } from 'antd';
 import React, { useState } from 'react';
-import styled from 'styled-components';
 import { useDebounce } from 'react-use';
+import styled from 'styled-components';
 
-import useSortInput from '@src/app/searchV2/sorting/useSortInput';
 import SearchSortSelect from '@src/app/searchV2/sorting/SearchSortSelect';
+import useSortInput from '@src/app/searchV2/sorting/useSortInput';
 import { SearchCfg } from '../../../../../../conf';
 import { useGetSearchResultsForMultipleQuery } from '../../../../../../graphql/search.generated';
-import { EntityType, FacetFilterInput, FilterOperator } from '../../../../../../types.generated';
+import { Entity, EntityType, FacetFilterInput, FilterOperator } from '../../../../../../types.generated';
 import { EntityAndType } from '../../../../../entity/shared/types';
 import { SearchBar } from '../../../../../search/SearchBar';
 import { ENTITY_FILTER_NAME, UnionType } from '../../../../../search/utils/constants';
@@ -47,7 +47,7 @@ type Props = {
     fixedEntityTypes?: Array<EntityType> | null;
     placeholderText?: string | null;
     selectedEntities: EntityAndType[];
-    setSelectedEntities: (Entities: EntityAndType[]) => void;
+    setSelectedEntities: (Entities: Entity[]) => void;
     limit?: number;
 };
 

datahub-web-react/src/app/entityV2/shared/tabs/Lineage/LineageColumnView.tsx

@@ -1,28 +1,29 @@
-import { useIsSeparateSiblingsMode } from '@app/entity/shared/siblingUtils';
-import React, { useMemo, useState } from 'react';
-import { useLocation } from 'react-router';
-import styled from 'styled-components/macro';
-import * as QueryString from 'query-string';
 import {
     ArrowDownOutlined,
     ArrowUpOutlined,
     CaretDownFilled,
     CaretDownOutlined,
+    LoadingOutlined,
     ReloadOutlined,
     SubnodeOutlined,
-    LoadingOutlined,
 } from '@ant-design/icons';
-import { Button, Select, Typography } from 'antd';
+import { useIsSeparateSiblingsMode } from '@app/entity/shared/siblingUtils';
 import { Tooltip } from '@components';
 import { GenericEntityProperties } from '@src/app/entity/shared/types';
+import ManageLineageMenuForImpactAnalysis from '@src/app/entityV2/shared/tabs/Lineage/ManageLineageMenuFromImpactAnalysis';
+import { Direction } from '@src/app/lineage/types';
+import { Button, Select, Typography } from 'antd';
+import * as QueryString from 'query-string';
+import React, { useMemo, useState } from 'react';
+import { useLocation } from 'react-router';
+import styled from 'styled-components/macro';
 import { EntityType, LineageDirection } from '../../../../../types.generated';
-import ManageLineageMenu from '../../../../lineage/manage/ManageLineageMenu';
+import { useEntityData } from '../../../../entity/shared/EntityContext';
 import { useGetLineageTimeParams } from '../../../../lineage/utils/useGetLineageTimeParams';
 import { useEntityRegistry } from '../../../../useEntityRegistry';
 import { downgradeV2FieldPath } from '../../../dataset/profile/schema/utils/utils';
 import TabToolbar from '../../components/styled/TabToolbar';
 import { ANTD_GRAY } from '../../constants';
-import { useEntityData } from '../../../../entity/shared/EntityContext';
 import ColumnsLineageSelect from './ColumnLineageSelect';
 import { ImpactAnalysis } from './ImpactAnalysis';
 import { LineageTabContext } from './LineageTabContext';
@@ -73,9 +74,10 @@ interface SchemaFieldEntityData extends GenericEntityProperties {
 
 interface Props {
     defaultDirection: LineageDirection;
+    setVisualizeViewInEditMode: (view: boolean, direction: Direction) => void;
 }
 
-export function LineageColumnView({ defaultDirection }: Props) {
+export function LineageColumnView({ defaultDirection, setVisualizeViewInEditMode }: Props) {
     const { urn, entityType, entityData } = useEntityData();
     const location = useLocation();
     const entityRegistry = useEntityRegistry();
@@ -153,10 +155,8 @@ export function LineageColumnView({ defaultDirection }: Props) {
                     />
                 </LeftButtonsWrapper>
                 <RightButtonsWrapper>
-                    <ManageLineageMenu
-                        entityUrn={urn}
-                        refetchEntity={() => setShouldRefetch(true)}
-                        setUpdatedLineages={() => {}}
+                    <ManageLineageMenuForImpactAnalysis
+                        setVisualizeViewInEditMode={setVisualizeViewInEditMode}
                         menuIcon={
                             <Button type="text">
                                 <ManageLineageIcon />
@@ -166,9 +166,7 @@ export function LineageColumnView({ defaultDirection }: Props) {
                                 <StyledCaretDown />
                             </Button>
                         }
-                        showLoading
                         entityType={entityType}
-                        entityPlatform={entityData?.platform?.name}
                         canEditLineage={canEditLineage}
                         disableDropdown={!canEditLineage}
                     />

datahub-web-react/src/app/entityV2/shared/tabs/Lineage/LineageTab.tsx

@@ -1,15 +1,15 @@
 import LineageGraph from '@app/lineageV2/LineageGraph';
 import React, { useContext } from 'react';
 import styled from 'styled-components';
-import { useLineageV2 } from '../../../../lineageV2/useLineageV2';
 import { LineageDirection } from '../../../../../types.generated';
 import { useEntityData } from '../../../../entity/shared/EntityContext';
+import LineageExplorer from '../../../../lineage/LineageExplorer';
+import { useLineageV2 } from '../../../../lineageV2/useLineageV2';
+import TabFullsizedContext from '../../../../shared/TabFullsizedContext';
 import { TabRenderType } from '../../types';
 import { CompactLineageTab } from './CompactLineageTab';
-import LineageExplorer from '../../../../lineage/LineageExplorer';
-import { LineageColumnView } from './LineageColumnView';
-import TabFullsizedContext from '../../../../shared/TabFullsizedContext';
 import { useLineageViewState } from './hooks';
+import { LineageColumnView } from './LineageColumnView';
 
 const LINEAGE_SWITCH_WIDTH = 90;
 
@@ -68,7 +68,7 @@ function WideLineageTab({ defaultDirection }: { defaultDirection: LineageDirecti
     const { isTabFullsize } = useContext(TabFullsizedContext);
     const { urn, entityType } = useEntityData();
     const isLineageV2 = useLineageV2();
-    const { isVisualizeView, setVisualizeView } = useLineageViewState();
+    const { isVisualizeView, setVisualizeView, setVisualizeViewInEditMode } = useLineageViewState();
 
     return (
         <LineageTabWrapper>
@@ -84,7 +84,12 @@ function WideLineageTab({ defaultDirection }: { defaultDirection: LineageDirecti
                     </LineageSwitchWrapper>
                 </LineageTabHeader>
             )}
-            {!isVisualizeView && <LineageColumnView defaultDirection={defaultDirection} />}
+            {!isVisualizeView && (
+                <LineageColumnView
+                    defaultDirection={defaultDirection}
+                    setVisualizeViewInEditMode={setVisualizeViewInEditMode}
+                />
+            )}
             {isVisualizeView && !isLineageV2 && <LineageExplorer urn={urn} type={entityType} />}
             {isVisualizeView && isLineageV2 && (
                 <VisualizationWrapper>

datahub-web-react/src/app/entityV2/shared/tabs/Lineage/ManageLineageMenuFromImpactAnalysis.tsx

@@ -0,0 +1,138 @@
+import { ArrowDownOutlined, ArrowUpOutlined, MoreOutlined } from '@ant-design/icons';
+import { Popover, Tooltip } from '@components';
+import { Direction } from '@src/app/lineage/types';
+import { Dropdown } from 'antd';
+import React from 'react';
+import styled from 'styled-components';
+import { EntityType } from '../../../../../types.generated';
+import { ENTITY_TYPES_WITH_MANUAL_LINEAGE } from '../../../../entity/shared/constants';
+import { MenuItemStyle } from '../../../../entity/view/menu/item/styledComponent';
+
+const DROPDOWN_Z_INDEX = 100;
+const POPOVER_Z_INDEX = 101;
+const UNAUTHORIZED_TEXT = "You aren't authorized to edit lineage for this entity.";
+
+const UnderlineWrapper = styled.span`
+    text-decoration: underline;
+    cursor: pointer;
+`;
+
+const MenuItemContent = styled.div``;
+
+function PopoverContent({ centerEntity, direction }: { centerEntity?: () => void; direction: string }) {
+    return (
+        <div>
+            <UnderlineWrapper onClick={centerEntity}>Focus</UnderlineWrapper> on this entity to make {direction} edits.
+        </div>
+    );
+}
+
+function getDownstreamDisabledPopoverContent(canEditLineage: boolean, isDashboard: boolean, centerEntity?: () => void) {
+    if (!canEditLineage) {
+        return UNAUTHORIZED_TEXT;
+    }
+    if (isDashboard) {
+        return 'Dashboard entities have no downstream lineage';
+    }
+    return <PopoverContent centerEntity={centerEntity} direction="downstream" />;
+}
+
+interface Props {
+    disableUpstream?: boolean;
+    disableDownstream?: boolean;
+    centerEntity?: () => void;
+    menuIcon?: React.ReactNode;
+    entityType?: EntityType;
+    canEditLineage?: boolean;
+    disableDropdown?: boolean;
+    setVisualizeViewInEditMode: (view: boolean, direction: Direction) => void;
+}
+
+export default function ManageLineageMenuForImpactAnalysis({
+    disableUpstream,
+    disableDownstream,
+    centerEntity,
+    menuIcon,
+    entityType,
+    canEditLineage,
+    disableDropdown,
+    setVisualizeViewInEditMode,
+}: Props) {
+    function manageLineage(direction: Direction) {
+        setVisualizeViewInEditMode(true, direction);
+    }
+
+    const isCenterNode = !disableUpstream && !disableDownstream;
+    const isDashboard = entityType === EntityType.Dashboard;
+    const isDownstreamDisabled = disableDownstream || isDashboard || !canEditLineage;
+    const isUpstreamDisabled = disableUpstream || !canEditLineage;
+    const isManualLineageSupported = entityType && ENTITY_TYPES_WITH_MANUAL_LINEAGE.has(entityType);
+
+    // if we don't show manual lineage options or the center node option, this menu has no options
+    if (!isManualLineageSupported && isCenterNode) return null;
+
+    const items = [
+        isManualLineageSupported
+            ? {
+                  key: 0,
+                  label: (
+                      <MenuItemStyle onClick={() => manageLineage(Direction.Upstream)} disabled={isUpstreamDisabled}>
+                          <Popover
+                              content={
+                                  !canEditLineage ? (
+                                      UNAUTHORIZED_TEXT
+                                  ) : (
+                                      <PopoverContent centerEntity={centerEntity} direction="upstream" />
+                                  )
+                              }
+                              overlayStyle={isUpstreamDisabled ? { zIndex: POPOVER_Z_INDEX } : { display: 'none' }}
+                          >
+                              <MenuItemContent data-testid="edit-upstream-lineage">
+                                  <ArrowUpOutlined />
+                                  &nbsp; Edit Upstream
+                              </MenuItemContent>
+                          </Popover>
+                      </MenuItemStyle>
+                  ),
+              }
+            : null,
+        isManualLineageSupported
+            ? {
+                  key: 1,
+                  label: (
+                      <MenuItemStyle
+                          onClick={() => manageLineage(Direction.Downstream)}
+                          disabled={isDownstreamDisabled}
+                      >
+                          <Popover
+                              content={getDownstreamDisabledPopoverContent(!!canEditLineage, isDashboard, centerEntity)}
+                              overlayStyle={isDownstreamDisabled ? { zIndex: POPOVER_Z_INDEX } : { display: 'none' }}
+                          >
+                              <MenuItemContent data-testid="edit-downstream-lineage">
+                                  <ArrowDownOutlined />
+                                  &nbsp; Edit Downstream
+                              </MenuItemContent>
+                          </Popover>
+                      </MenuItemStyle>
+                  ),
+              }
+            : null,
+    ];
+
+    return (
+        <>
+            <Tooltip title={disableDropdown ? UNAUTHORIZED_TEXT : ''}>
+                <div data-testid="lineage-edit-menu-button">
+                    <Dropdown
+                        overlayStyle={{ zIndex: DROPDOWN_Z_INDEX }}
+                        disabled={disableDropdown}
+                        menu={{ items }}
+                        trigger={['click']}
+                    >
+                        {menuIcon || <MoreOutlined style={{ fontSize: 18 }} />}
+                    </Dropdown>
+                </div>
+            </Tooltip>
+        </>
+    );
+}

datahub-web-react/src/app/entityV2/shared/tabs/Lineage/hooks.ts

@@ -1,3 +1,4 @@
+import { Direction } from '@src/app/lineage/types';
 import { useCallback, useState } from 'react';
 import { useHistory, useLocation } from 'react-router-dom';
 
@@ -13,30 +14,51 @@ export function useLineageViewState() {
         initialView === 'impact' ? false : IS_VISUALIZE_VIEW_DEFAULT,
     );
 
+    // Helper function to update URL parameters
+    const updateUrlParam = useCallback((paramName: string, paramValue: string, currentSearch: string) => {
+        if (currentSearch.includes(`${paramName}=`)) {
+            // Replace the existing parameter
+            return currentSearch.replace(new RegExp(`(${paramName}=)[^&]+`), `${paramName}=${paramValue}`);
+        }
+        // Add the new parameter
+        return `${currentSearch + (currentSearch ? '&' : '?')}${paramName}=${paramValue}`;
+    }, []);
+
     const setVisualizeView = useCallback(
         (view: boolean) => {
             setIsVisualizeView(view);
 
             // Update the URL with the new view state
-            const newParam = `lineageView=${view ? 'explorer' : 'impact'}`;
-
-            let newSearch = location.search;
-            if (newSearch.includes('lineageView=')) {
-                // Replace the existing parameter
-                newSearch = newSearch.replace(/(lineageView=)[^&]+/, newParam);
-            } else {
-                // Add the new parameter
-                newSearch += (newSearch ? '&' : '?') + newParam;
-            }
+            const viewValue = view ? 'explorer' : 'impact';
+            const newSearch = updateUrlParam('lineageView', viewValue, location.search);
 
             // Update the URL without reloading the page
             history.replace({ search: newSearch });
         },
-        [location.search, history],
+        [location.search, history, updateUrlParam],
+    );
+
+    const setVisualizeViewInEditMode = useCallback(
+        (view: boolean, direction: Direction) => {
+            // First set isVisualizeView state value
+            setIsVisualizeView(view);
+
+            // First update lineageView parameter
+            const viewValue = view ? 'explorer' : 'impact';
+            let newSearch = updateUrlParam('lineageView', viewValue, location.search);
+
+            // Then update lineageEditDirection parameter
+            newSearch = updateUrlParam('lineageEditDirection', direction, newSearch);
+
+            // Update URL with both parameters in a single replace
+            history.replace({ search: newSearch });
+        },
+        [location.search, history, updateUrlParam],
     );
 
     return {
         isVisualizeView,
         setVisualizeView,
+        setVisualizeViewInEditMode,
     };
 }

datahub-web-react/src/app/lineageV2/LineageEntityNode/ManageLineageMenu.tsx

@@ -1,13 +1,15 @@
 import { ArrowLeftOutlined, ArrowRightOutlined, MoreOutlined } from '@ant-design/icons';
+import { Popover } from '@components';
 import Colors from '@components/theme/foundations/colors';
 import { Button, Dropdown, Menu } from 'antd';
-import { Popover } from '@components';
+import * as QueryString from 'query-string';
+import React, { useCallback, useContext, useEffect, useState } from 'react';
+import { useHistory, useLocation } from 'react-router-dom';
 import styled from 'styled-components';
-import React, { useCallback, useContext, useState } from 'react';
 import { EntityType, LineageDirection } from '../../../types.generated';
-import ManageLineageModal from '../manualLineage/ManageLineageModal';
-import { LineageDisplayContext, LineageEntity, onClickPreventSelect } from '../common';
 import { ENTITY_TYPES_WITH_MANUAL_LINEAGE } from '../../entityV2/shared/constants';
+import { LineageDisplayContext, LineageEntity, onClickPreventSelect } from '../common';
+import ManageLineageModal from '../manualLineage/ManageLineageModal';
 
 const DROPDOWN_Z_INDEX = 100;
 const POPOVER_Z_INDEX = 101;
@@ -63,13 +65,45 @@ const PopoverContent = styled.span`
 interface Props {
     node: LineageEntity;
     refetch: Record<LineageDirection, () => void>;
+    isRootUrn: boolean;
 }
 
-export default function ManageLineageMenu({ node, refetch }: Props) {
+export default function ManageLineageMenu({ node, refetch, isRootUrn }: Props) {
     const { displayedMenuNode, setDisplayedMenuNode } = useContext(LineageDisplayContext);
     const isMenuVisible = displayedMenuNode === node.urn;
     const [isModalVisible, setIsModalVisible] = useState(false);
     const [lineageDirection, setLineageDirection] = useState<LineageDirection>(LineageDirection.Upstream);
+    const location = useLocation();
+    const history = useHistory();
+
+    // Check for lineageEditDirection URL parameter when component mounts
+    useEffect(() => {
+        if (isRootUrn) {
+            const params = QueryString.parse(location.search);
+            const editDirection = params.lineageEditDirection as string;
+
+            if (editDirection) {
+                // Convert string parameter to LineageDirection enum
+                const direction =
+                    editDirection.toLowerCase() === 'downstream'
+                        ? LineageDirection.Downstream
+                        : LineageDirection.Upstream;
+
+                // Clear the parameter from URL
+                const newParams = { ...params };
+                delete newParams.lineageEditDirection;
+                const newSearch = QueryString.stringify(newParams);
+                history.replace({
+                    pathname: location.pathname,
+                    search: newSearch,
+                });
+
+                // Open the modal with the specified direction
+                setLineageDirection(direction);
+                setIsModalVisible(true);
+            }
+        }
+    }, [isRootUrn, location, history]);
 
     function manageLineage(direction: LineageDirection) {
         setLineageDirection(direction);
@@ -99,7 +133,7 @@ export default function ManageLineageMenu({ node, refetch }: Props) {
 
     return (
         <Wrapper>
-            <StyledButton onClick={handleMenuClick} type="text">
+            <StyledButton onClick={handleMenuClick} type="text" data-testid={`manage-lineage-menu-${node.urn}`}>
                 <Dropdown
                     open={isMenuVisible}
                     overlayStyle={{ zIndex: DROPDOWN_Z_INDEX }}
@@ -119,7 +153,7 @@ export default function ManageLineageMenu({ node, refetch }: Props) {
                                                 isUpstreamDisabled ? { zIndex: POPOVER_Z_INDEX } : { display: 'none' }
                                             }
                                         >
-                                            <MenuItemContent>
+                                            <MenuItemContent data-testid="edit-upstream-lineage">
                                                 <ArrowLeftOutlined />
                                                 &nbsp; Edit Upstream
                                             </MenuItemContent>
@@ -135,7 +169,7 @@ export default function ManageLineageMenu({ node, refetch }: Props) {
                                             content={getDownstreamDisabledPopoverContent(!!canEditLineage, isDashboard)}
                                             overlayStyle={!isDownstreamDisabled ? { display: 'none' } : undefined}
                                         >
-                                            <MenuItemContent>
+                                            <MenuItemContent data-testid="edit-downstream-lineage">
                                                 <ArrowRightOutlined />
                                                 &nbsp; Edit Downstream
                                             </MenuItemContent>

datahub-web-react/src/app/lineageV2/LineageEntityNode/NodeContents.tsx

@@ -4,11 +4,11 @@ import ContainerPath from '@app/lineageV2/LineageEntityNode/ContainerPath';
 import GhostEntityMenu from '@app/lineageV2/LineageEntityNode/GhostEntityMenu';
 import SchemaFieldNodeContents from '@app/lineageV2/LineageEntityNode/SchemaFieldNodeContents';
 import MatchTextSizeWrapper from '@app/sharedV2/text/MatchTextSizeWrapper';
-import { DeprecationIcon } from '@src/app/entityV2/shared/components/styled/DeprecationIcon';
+import { Tooltip } from '@components';
 import { KeyboardArrowDown, KeyboardArrowUp } from '@mui/icons-material';
+import { DeprecationIcon } from '@src/app/entityV2/shared/components/styled/DeprecationIcon';
 import StructuredPropertyBadge from '@src/app/entityV2/shared/containers/profile/header/StructuredPropertyBadge';
 import { Skeleton, Spin } from 'antd';
-import { Tooltip } from '@components';
 import React, { Dispatch, SetStateAction, useCallback } from 'react';
 import { Handle, Position } from 'reactflow';
 import styled from 'styled-components';
@@ -482,7 +482,11 @@ function NodeContents(props: Props & LineageEntity & DisplayedColumns) {
                                 {!showColumns && <KeyboardArrowDown fontSize="inherit" style={{ marginLeft: 3 }} />}
                             </ExpandColumnsWrapper>
                         )}
-                        {isGhost ? <GhostEntityMenu urn={urn} /> : <ManageLineageMenu node={props} refetch={refetch} />}
+                        {isGhost ? (
+                            <GhostEntityMenu urn={urn} />
+                        ) : (
+                            <ManageLineageMenu node={props} refetch={refetch} isRootUrn={urn === rootUrn} />
+                        )}
                         {entity && (
                             <PropertyBadgeWrapper>
                                 <StructuredPropertyBadge structuredProperties={entity.structuredProperties} />

datahub-web-react/src/app/lineageV2/manualLineage/AddEntityEdge.tsx

@@ -1,157 +0,0 @@
-import React, { useState } from 'react';
-import { useDebounce } from 'react-use';
-import styled from 'styled-components/macro';
-import { AutoComplete, Empty } from 'antd';
-import { LoadingOutlined, SubnodeOutlined } from '@ant-design/icons';
-import { toTitleCase } from '../../../graphql-mock/helper';
-import { useEntityRegistry } from '../../useEntityRegistry';
-import { useGetAutoCompleteMultipleResultsLazyQuery } from '../../../graphql/search.generated';
-import { Entity, EntityType, LineageDirection } from '../../../types.generated';
-import LineageEntityView from './LineageEntityView';
-import EntityRegistry from '../../entity/EntityRegistry';
-import { ANTD_GRAY } from '../../entity/shared/constants';
-import { getValidEntityTypes } from './utils';
-
-const DEBOUNCE_WAIT_MS = 200;
-
-const AddEdgeWrapper = styled.div`
-    padding: 15px 20px;
-    display: flex;
-    align-items: center;
-`;
-
-const AddLabel = styled.span`
-    font-size: 12px;
-    font-weight: bold;
-    display: flex;
-    align-items: center;
-    white-space: nowrap;
-`;
-
-const AddIcon = styled(SubnodeOutlined)`
-    margin-right: 5px;
-    font-size: 16px;
-`;
-
-const StyledAutoComplete = styled(AutoComplete<string>)`
-    margin-left: 10px;
-    flex: 1;
-`;
-
-const LoadingWrapper = styled.div`
-    padding: 8px;
-    display: flex;
-    justify-content: center;
-
-    svg {
-        height: 15px;
-        width: 15px;
-        color: ${ANTD_GRAY[8]};
-    }
-`;
-
-interface Props {
-    direction: LineageDirection;
-    setEntitiesToAdd: React.Dispatch<React.SetStateAction<Entity[]>>;
-    entitiesToAdd: Entity[];
-    entityUrn: string;
-    entityType?: EntityType;
-}
-
-export default function AddEntityEdge({ direction, setEntitiesToAdd, entitiesToAdd, entityUrn, entityType }: Props) {
-    const entityRegistry = useEntityRegistry();
-    const [getAutoCompleteResults, { data: autoCompleteResults, loading }] =
-        useGetAutoCompleteMultipleResultsLazyQuery();
-    const [queryText, setQueryText] = useState<string>('');
-
-    const validEntityTypes = getValidEntityTypes(direction, entityType);
-
-    useDebounce(
-        () => {
-            if (queryText.trim()) {
-                getAutoCompleteResults({
-                    variables: {
-                        input: {
-                            types: validEntityTypes,
-                            query: queryText,
-                            limit: 15,
-                        },
-                    },
-                });
-            }
-        },
-        DEBOUNCE_WAIT_MS,
-        [queryText],
-    );
-
-    const selectEntity = (urn: string) => {
-        const resultEntities = autoCompleteResults?.autoCompleteForMultiple?.suggestions?.flatMap(
-            (suggestion) => suggestion.entities || [],
-        );
-        const selectedEntity = resultEntities?.find((entity) => entity.urn === urn);
-        if (selectedEntity) {
-            setEntitiesToAdd((existingEntities) => [...existingEntities, selectedEntity]);
-        }
-    };
-
-    const renderSearchResult = (entity: Entity) => (
-        <AutoComplete.Option value={entity.urn} key={entity.urn}>
-            <LineageEntityView entity={entity} displaySearchResult />
-        </AutoComplete.Option>
-    );
-
-    const searchResults = autoCompleteResults?.autoCompleteForMultiple?.suggestions
-        .flatMap((suggestion) => suggestion.entities || [])
-        .filter((entity) => entity && !existsInEntitiesToAdd(entity, entitiesToAdd) && entity.urn !== entityUrn)
-        .map((entity) => renderSearchResult(entity));
-
-    const placeholderText = getPlaceholderText(validEntityTypes, entityRegistry);
-
-    return (
-        <AddEdgeWrapper>
-            <AddLabel>
-                <AddIcon />
-                Add {toTitleCase(direction.toLocaleLowerCase())}
-            </AddLabel>
-            <StyledAutoComplete
-                autoFocus
-                showSearch
-                placeholder={placeholderText}
-                value={queryText}
-                onSearch={(value) => setQueryText(value)}
-                onSelect={(urn: string) => selectEntity(urn)}
-                filterOption={false}
-                notFoundContent={queryText.length > 3 && <Empty description="No Assets Found" />}
-            >
-                {loading && (
-                    <AutoComplete.Option value="loading">
-                        <LoadingWrapper>
-                            <LoadingOutlined />
-                        </LoadingWrapper>
-                    </AutoComplete.Option>
-                )}
-                {!!queryText && searchResults}
-            </StyledAutoComplete>
-        </AddEdgeWrapper>
-    );
-}
-
-function existsInEntitiesToAdd(result: Entity, entitiesAlreadyAdded: Entity[]) {
-    return !!entitiesAlreadyAdded.find((entity) => entity.urn === result.urn);
-}
-
-function getPlaceholderText(validEntityTypes: EntityType[], entityRegistry: EntityRegistry) {
-    let placeholderText = 'Search for ';
-    if (!validEntityTypes.length) {
-        placeholderText = `${placeholderText} entities to add...`;
-    } else if (validEntityTypes.length === 1) {
-        placeholderText = `${placeholderText} ${entityRegistry.getCollectionName(validEntityTypes[0])}...`;
-    } else {
-        validEntityTypes.forEach((type, index) => {
-            placeholderText = `${placeholderText} ${entityRegistry.getCollectionName(type)}${
-                index !== validEntityTypes.length - 1 ? ', ' : '...'
-            }`;
-        });
-    }
-    return placeholderText;
-}

datahub-web-react/src/app/lineageV2/manualLineage/EntityEdge.tsx

@@ -1,4 +1,4 @@
-import { CloseOutlined } from '@ant-design/icons';
+import { Icon } from '@src/alchemy-components';
 import Text from 'antd/lib/typography/Text';
 import React from 'react';
 import styled from 'styled-components/macro';
@@ -25,15 +25,18 @@ const NameAndLogoWrapper = styled.span`
     max-width: 85%;
 `;
 
-const StyledClose = styled(CloseOutlined)`
-    cursor: pointer;
-`;
-
 const EntityName = styled(Text)`
     font-size: 14px;
     font-weight: bold;
 `;
 
+const AvatarWrapper = styled.div`
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    cursor: pointer;
+`;
+
 interface Props {
     entity: Entity;
     removeEntity: (removedEntity: Entity) => void;
@@ -56,10 +59,14 @@ export default function EntityEdge({ entity, removeEntity, createdOn, createdAct
                     {entityRegistry.getDisplayName(entity.type, entity)}
                 </EntityName>
             </NameAndLogoWrapper>
-            <span>
-                {shouldDisplayAvatar && <UserAvatar createdActor={createdActor} createdOn={createdOn} />}
-                <StyledClose onClick={() => removeEntity(entity)} />
-            </span>
+            <AvatarWrapper>
+                {shouldDisplayAvatar && (
+                    <div style={{ marginRight: '10px' }}>
+                        <UserAvatar createdActor={createdActor} createdOn={createdOn} />
+                    </div>
+                )}
+                <Icon icon="X" source="phosphor" onClick={() => removeEntity(entity)} />
+            </AvatarWrapper>
         </EntityItem>
     );
 }

datahub-web-react/src/app/lineageV2/manualLineage/LineageEdges.tsx

@@ -7,17 +7,17 @@ import { getEdgeId, LineageNodesContext, setDifference } from '../common';
 import EntityEdge from './EntityEdge';
 
 const LineageEdgesWrapper = styled.div`
-    height: 225px;
-    overflow: auto;
     padding: 0 20px 10px 20px;
+    height: 100%;
 `;
 
 const EmptyWrapper = styled.div`
     display: flex;
     align-items: center;
     justify-content: center;
-    height: 100%;
+    height: 95%;
     background-color: ${ANTD_GRAY[3]};
+    margin-top: 10px;
 `;
 
 interface Props {
@@ -25,18 +25,10 @@ interface Props {
     direction: LineageDirection;
     entitiesToAdd: Entity[];
     entitiesToRemove: Entity[];
-    setEntitiesToAdd: React.Dispatch<React.SetStateAction<Entity[]>>;
-    setEntitiesToRemove: React.Dispatch<React.SetStateAction<Entity[]>>;
+    onRemoveEntity: (entity: Entity) => void;
 }
 
-export default function LineageEdges({
-    parentUrn,
-    direction,
-    entitiesToAdd,
-    entitiesToRemove,
-    setEntitiesToAdd,
-    setEntitiesToRemove,
-}: Props) {
+export default function LineageEdges({ parentUrn, direction, entitiesToAdd, entitiesToRemove, onRemoveEntity }: Props) {
     const { nodes, edges, adjacencyList } = useContext(LineageNodesContext);
 
     const children = adjacencyList[direction].get(parentUrn) || new Set();
@@ -46,16 +38,6 @@ export default function LineageEdges({
     );
     const filteredChildren = setDifference(children, urnsToRemove);
 
-    function removeEntity(removedEntity: Entity) {
-        if (children.has(removedEntity.urn)) {
-            setEntitiesToRemove((existingEntitiesToRemove) => [...existingEntitiesToRemove, removedEntity]);
-        } else {
-            setEntitiesToAdd((existingEntitiesToAdd) =>
-                existingEntitiesToAdd.filter((addedEntity) => addedEntity.urn !== removedEntity.urn),
-            );
-        }
-    }
-
     return (
         <LineageEdgesWrapper>
             {!filteredChildren?.length && !entitiesToAdd.length && (
@@ -72,14 +54,14 @@ export default function LineageEdges({
                     <EntityEdge
                         key={childUrn}
                         entity={childNode.rawEntity || backupEntity}
-                        removeEntity={removeEntity}
+                        removeEntity={onRemoveEntity}
                         createdOn={edge?.created?.timestamp}
                         createdActor={edge?.created?.actor as CorpUser | undefined}
                     />
                 );
             })}
             {entitiesToAdd.map((addedEntity) => (
-                <EntityEdge key={addedEntity.urn} entity={addedEntity} removeEntity={removeEntity} />
+                <EntityEdge key={addedEntity.urn} entity={addedEntity} removeEntity={onRemoveEntity} />
             ))}
         </LineageEdgesWrapper>
     );

datahub-web-react/src/app/lineageV2/manualLineage/LineageEntityView.tsx

@@ -1,67 +0,0 @@
-import { Divider } from 'antd';
-import React from 'react';
-import styled from 'styled-components/macro';
-import { Entity } from '../../../types.generated';
-import { ANTD_GRAY } from '../../entity/shared/constants';
-import { getPlatformName } from '../../entity/shared/utils';
-import { capitalizeFirstLetterOnly } from '../../shared/textUtil';
-import { useEntityRegistry } from '../../useEntityRegistry';
-
-const EntityWrapper = styled.div<{ shrinkPadding?: boolean }>`
-    border-bottom: 1px solid ${ANTD_GRAY[4]};
-    padding: ${(props) => (props.shrinkPadding ? '4px 6px' : '12px 20px')};
-`;
-
-const PlatformContent = styled.div<{ removeMargin?: boolean }>`
-    display: flex;
-    align-items: center;
-    font-size: 10px;
-    color: ${ANTD_GRAY[7]};
-    margin-bottom: ${(props) => (props.removeMargin ? '0' : '5px')};
-`;
-
-const StyledDivider = styled(Divider)`
-    margin: 0 5px;
-`;
-
-const PlatformLogo = styled.img`
-    height: 14px;
-    margin-right: 5px;
-`;
-
-export const EntityName = styled.span<{ shrinkSize?: boolean }>`
-    font-size: ${(props) => (props.shrinkSize ? '12px' : '14px')};
-    font-weight: bold;
-`;
-
-interface Props {
-    entity: Entity;
-    displaySearchResult?: boolean;
-}
-
-export default function LineageEntityView({ entity, displaySearchResult }: Props) {
-    const entityRegistry = useEntityRegistry();
-    const genericProps = entityRegistry.getGenericEntityProperties(entity.type, entity);
-
-    const platformLogoUrl = genericProps?.platform?.properties?.logoUrl;
-    const platformName = getPlatformName(genericProps);
-
-    return (
-        <EntityWrapper shrinkPadding={displaySearchResult}>
-            <PlatformContent removeMargin={displaySearchResult}>
-                {platformLogoUrl && (
-                    <PlatformLogo src={platformLogoUrl} alt="platform logo" data-testid="platform-logo" />
-                )}
-                <span>{platformName}</span>
-                {platformName && <StyledDivider type="vertical" data-testid="divider" />}
-                <span>
-                    {capitalizeFirstLetterOnly(genericProps?.subTypes?.typeNames?.[0]) ||
-                        entityRegistry.getEntityName(entity.type)}
-                </span>
-            </PlatformContent>
-            <EntityName shrinkSize={displaySearchResult}>
-                {entityRegistry.getDisplayName(entity.type, entity)}
-            </EntityName>
-        </EntityWrapper>
-    );
-}

datahub-web-react/src/app/lineageV2/manualLineage/ManageLineageModal.tsx

@@ -1,45 +1,65 @@
-import { LoadingOutlined } from '@ant-design/icons';
+import { colors, Modal } from '@src/alchemy-components';
+import { EntityAndType } from '@src/app/entity/shared/types';
+import { extractTypeFromUrn } from '@src/app/entity/shared/utils';
+import { SearchSelect } from '@src/app/entityV2/shared/components/styled/search/SearchSelect';
+import ClickOutside from '@src/app/shared/ClickOutside';
+import { Modal as AntModal, message } from 'antd';
 import React, { useContext, useEffect, useState } from 'react';
-import { message, Modal } from 'antd';
 import styled from 'styled-components/macro';
-import { Button } from '@src/alchemy-components';
 import { toTitleCase } from '../../../graphql-mock/helper';
-import { EventType } from '../../analytics';
-import analytics from '../../analytics/analytics';
-import { useUserContext } from '../../context/useUserContext';
-import EntityRegistry from '../../entity/EntityRegistry';
-import { Direction } from '../../lineage/types';
-import { FetchStatus, LineageEntity, LineageNodesContext } from '../common';
-import AddEntityEdge from './AddEntityEdge';
-import LineageEntityView from './LineageEntityView';
-import LineageEdges from './LineageEdges';
-import { Entity, EntityType, LineageDirection, LineageEdge } from '../../../types.generated';
 import { useUpdateLineageMutation } from '../../../graphql/mutations.generated';
-import { useEntityRegistry } from '../../useEntityRegistry';
-import updateNodeContext from './updateNodeContext';
+import { Entity, EntityType, LineageDirection } from '../../../types.generated';
+import { useUserContext } from '../../context/useUserContext';
+import { useEntityRegistryV2 as useEntityRegistry } from '../../useEntityRegistry';
 import { useOnClickExpandLineage } from '../LineageEntityNode/useOnClickExpandLineage';
+import { FetchStatus, LineageEntity, LineageNodesContext } from '../common';
+import LineageEdges from './LineageEdges';
+import { buildUpdateLineagePayload } from './buildUpdateLineagePayload';
+import { recordAnalyticsEvents } from './recordManualLineageAnalyticsEvent';
+import updateNodeContext from './updateNodeContext';
+import { getValidEntityTypes } from './utils';
 
-const ModalFooter = styled.div`
-    display: flex;
-    justify-content: space-between;
-`;
-
-const TitleText = styled.div`
-    font-weight: bold;
-`;
+const MODAL_WIDTH_PX = 1400;
 
 const StyledModal = styled(Modal)`
-    .ant-modal-body {
-        padding: 0;
-    }
+    top: 30px;
+    padding: 0;
 `;
 
-const LoadingWrapper = styled.div`
+const ModalContentContainer = styled.div`
+    height: 75vh;
+    margin: -24px -20px;
     display: flex;
-    align-items: center;
-    justify-content: center;
-    height: 225px;
-    font-size: 30px;
+    flex-direction: row;
+`;
+
+const SearchSection = styled.div`
+    flex: 2;
+    height: 100%;
+    overflow: hidden;
+    display: flex;
+    flex-direction: column;
+`;
+
+const CurrentSection = styled.div`
+    flex: 1;
+    width: 40%;
+    border-left: 1px solid ${colors.gray[100]};
+    display: flex;
+    flex-direction: column;
+`;
+
+const SectionHeader = styled.div`
+    padding-left: 20px;
+    margin-top: 10px;
+    font-size: 16px;
+    font-weight: 500;
+    color: ${colors.gray[600]};
+`;
+
+const ScrollableContent = styled.div`
+    flex: 1;
+    overflow: auto;
 `;
 
 interface Props {
@@ -54,11 +74,16 @@ export default function ManageLineageModal({ node, direction, closeModal, refetc
     const expandOneLevel = useOnClickExpandLineage(node.urn, node.type, direction, false);
     const { user } = useUserContext();
     const entityRegistry = useEntityRegistry();
-    const [entitiesToAdd, setEntitiesToAdd] = useState<Entity[]>([]);
-    const [entitiesToRemove, setEntitiesToRemove] = useState<Entity[]>([]);
     const [updateLineage] = useUpdateLineageMutation();
     const fetchStatus = node.fetchStatus[direction];
-    const loading = node.fetchStatus[direction] === FetchStatus.LOADING;
+    const { adjacencyList } = nodeContext;
+    const validEntityTypes = getValidEntityTypes(direction, node.entity?.type);
+    const initialSetOfRelationshipsUrns = adjacencyList[direction].get(node.urn) || new Set();
+    const [isSaving, setIsSaving] = useState(false);
+
+    const [selectedEntities, setSelectedEntities] = useState<EntityAndType[]>(
+        Array.from(initialSetOfRelationshipsUrns).map((urn) => ({ urn, type: extractTypeFromUrn(urn) })) || [],
+    );
 
     useEffect(() => {
         if (fetchStatus === FetchStatus.UNFETCHED) {
@@ -66,8 +91,16 @@ export default function ManageLineageModal({ node, direction, closeModal, refetc
         }
     }, [fetchStatus, expandOneLevel]);
 
+    const entitiesToAdd = selectedEntities.filter((entity) => !initialSetOfRelationshipsUrns.has(entity.urn));
+    const entitiesToRemove = Array.from(initialSetOfRelationshipsUrns)
+        .filter((urn) => !selectedEntities.map((entity) => entity.urn).includes(urn))
+        .map((urn) => ({ urn } as Entity));
+
+    // save lineage changes will disable the button while its processing
     function saveLineageChanges() {
+        setIsSaving(true);
         const payload = buildUpdateLineagePayload(direction, entitiesToAdd, entitiesToRemove, node.urn);
+
         updateLineage({ variables: { input: payload } })
             .then((res) => {
                 if (res.data?.updateLineage) {
@@ -84,121 +117,85 @@ export default function ManageLineageModal({ node, direction, closeModal, refetc
                         entityType: node.type,
                         entityPlatform: entityRegistry.getDisplayName(EntityType.DataPlatform, node.entity?.platform),
                     });
+                    setIsSaving(false);
                 }
             })
             .catch((error) => {
                 message.error(error.message || 'Error updating lineage');
+                setIsSaving(false);
             });
     }
 
-    const isSaveDisabled = !entitiesToAdd.length && !entitiesToRemove.length;
+    const directionTitle = toTitleCase(direction.toLocaleLowerCase());
+
+    const onCancelSelect = () => {
+        if (entitiesToAdd.length > 0 || entitiesToRemove.length > 0) {
+            AntModal.confirm({
+                title: `Exit Lineage Management`,
+                content: `Are you sure you want to exit? ${
+                    entitiesToAdd.length + entitiesToRemove.length
+                } change(s) will be cleared.`,
+                onOk() {
+                    closeModal();
+                },
+                onCancel() {},
+                okText: 'Yes',
+                maskClosable: true,
+                closable: true,
+            });
+        } else {
+            closeModal();
+        }
+    };
 
     return (
-        <StyledModal
-            title={<TitleText>Manage {toTitleCase(direction.toLocaleLowerCase())} Lineage</TitleText>}
-            onCancel={closeModal}
-            keyboard
-            open
-            footer={
-                <ModalFooter>
-                    <Button onClick={closeModal} variant="text" color="gray">
-                        Cancel
-                    </Button>
-                    <Button onClick={saveLineageChanges} disabled={isSaveDisabled}>
-                        Save
-                    </Button>
-                </ModalFooter>
-            }
-        >
-            {node.entity && <LineageEntityView entity={node.entity} />}
-            <AddEntityEdge
-                direction={direction}
-                setEntitiesToAdd={setEntitiesToAdd}
-                entitiesToAdd={entitiesToAdd}
-                entityUrn={node.urn}
-                entityType={node.type}
-            />
-            {!loading && (
-                <LineageEdges
-                    parentUrn={node.urn}
-                    direction={direction}
-                    entitiesToAdd={entitiesToAdd}
-                    entitiesToRemove={entitiesToRemove}
-                    setEntitiesToAdd={setEntitiesToAdd}
-                    setEntitiesToRemove={setEntitiesToRemove}
-                />
-            )}
-            {loading && (
-                <LoadingWrapper>
-                    <LoadingOutlined />
-                </LoadingWrapper>
-            )}
-        </StyledModal>
+        <ClickOutside onClickOutside={onCancelSelect} wrapperClassName="search-select-modal">
+            <StyledModal
+                title={`Select the ${directionTitle}s to add to ${node.entity?.name}`}
+                width={MODAL_WIDTH_PX}
+                open
+                onCancel={onCancelSelect}
+                style={{ padding: 0 }}
+                buttons={[
+                    {
+                        text: 'Cancel',
+                        variant: 'text',
+                        onClick: onCancelSelect,
+                    },
+                    {
+                        text: isSaving ? 'Saving...' : `Set ${directionTitle}s`,
+                        onClick: saveLineageChanges,
+                        disabled: (entitiesToAdd.length === 0 && entitiesToRemove.length === 0) || isSaving,
+                    },
+                ]}
+            >
+                <ModalContentContainer>
+                    <SearchSection>
+                        <SectionHeader>Search and Add</SectionHeader>
+                        <ScrollableContent>
+                            <SearchSelect
+                                fixedEntityTypes={Array.from(validEntityTypes)}
+                                selectedEntities={selectedEntities}
+                                setSelectedEntities={setSelectedEntities}
+                            />
+                        </ScrollableContent>
+                    </SearchSection>
+                    <CurrentSection>
+                        <SectionHeader>Current {directionTitle}s</SectionHeader>
+                        <ScrollableContent>
+                            <LineageEdges
+                                parentUrn={node.urn}
+                                direction={direction}
+                                entitiesToAdd={entitiesToAdd}
+                                entitiesToRemove={entitiesToRemove}
+                                onRemoveEntity={(entity) => {
+                                    setSelectedEntities(selectedEntities.filter((e) => e.urn !== entity.urn));
+                                }}
+                            />
+                        </ScrollableContent>
+                    </CurrentSection>
+                </ModalContentContainer>
+            </StyledModal>
+        </ClickOutside>
     );
 }
-
-interface AnalyticsEventsProps {
-    direction: LineageDirection;
-    entitiesToAdd: Entity[];
-    entitiesToRemove: Entity[];
-    entityRegistry: EntityRegistry;
-    entityType?: EntityType;
-    entityPlatform?: string;
-}
-
-function recordAnalyticsEvents({
-    direction,
-    entitiesToAdd,
-    entitiesToRemove,
-    entityRegistry,
-    entityType,
-    entityPlatform,
-}: AnalyticsEventsProps) {
-    entitiesToAdd.forEach((entityToAdd) => {
-        const genericProps = entityRegistry.getGenericEntityProperties(entityToAdd.type, entityToAdd);
-        analytics.event({
-            type: EventType.ManuallyCreateLineageEvent,
-            direction: directionFromLineageDirection(direction),
-            sourceEntityType: entityType,
-            sourceEntityPlatform: entityPlatform,
-            destinationEntityType: entityToAdd.type,
-            destinationEntityPlatform: genericProps?.platform?.name,
-        });
-    });
-    entitiesToRemove.forEach((entityToRemove) => {
-        const genericProps = entityRegistry.getGenericEntityProperties(entityToRemove.type, entityToRemove);
-        analytics.event({
-            type: EventType.ManuallyDeleteLineageEvent,
-            direction: directionFromLineageDirection(direction),
-            sourceEntityType: entityType,
-            sourceEntityPlatform: entityPlatform,
-            destinationEntityType: entityToRemove.type,
-            destinationEntityPlatform: genericProps?.platform?.name,
-        });
-    });
-}
-
-function buildUpdateLineagePayload(
-    lineageDirection: LineageDirection,
-    entitiesToAdd: Entity[],
-    entitiesToRemove: Entity[],
-    entityUrn: string,
-) {
-    let edgesToAdd: LineageEdge[] = [];
-    let edgesToRemove: LineageEdge[] = [];
-
-    if (lineageDirection === LineageDirection.Upstream) {
-        edgesToAdd = entitiesToAdd.map((entity) => ({ upstreamUrn: entity.urn, downstreamUrn: entityUrn }));
-        edgesToRemove = entitiesToRemove.map((entity) => ({ upstreamUrn: entity.urn, downstreamUrn: entityUrn }));
-    }
-    if (lineageDirection === LineageDirection.Downstream) {
-        edgesToAdd = entitiesToAdd.map((entity) => ({ upstreamUrn: entityUrn, downstreamUrn: entity.urn }));
-        edgesToRemove = entitiesToRemove.map((entity) => ({ upstreamUrn: entityUrn, downstreamUrn: entity.urn }));
-    }
-
-    return { edgesToAdd, edgesToRemove };
-}
-
-function directionFromLineageDirection(lineageDirection: LineageDirection): Direction {
-    return lineageDirection === LineageDirection.Upstream ? Direction.Upstream : Direction.Downstream;
-}

datahub-web-react/src/app/lineageV2/manualLineage/UserAvatar.tsx

@@ -1,6 +1,5 @@
 import { PartitionOutlined } from '@ant-design/icons';
-import { Avatar } from 'antd';
-import { Popover } from '@components';
+import { Avatar, Popover } from '@components';
 import React from 'react';
 import styled from 'styled-components/macro';
 import { CorpUser, EntityType } from '../../../types.generated';
@@ -54,9 +53,9 @@ export default function UserAvatar({ createdActor, createdOn }: Props) {
                 </PopoverWrapper>
             }
         >
-            <StyledAvatar src={avatarPhotoUrl} $backgroundColor={getAvatarColor(userName)}>
-                {userName.charAt(0).toUpperCase()}
-            </StyledAvatar>
+            <div>
+                <StyledAvatar imageUrl={avatarPhotoUrl} $backgroundColor={getAvatarColor(userName)} name={userName} />
+            </div>
         </Popover>
     );
 }

datahub-web-react/src/app/lineageV2/manualLineage/__tests__/ManageLineageModal.test.tsx

@@ -0,0 +1,242 @@
+import { MockedProvider } from '@apollo/client/testing';
+import { fireEvent, render, screen } from '@testing-library/react';
+import React from 'react';
+import { EntityType, LineageDirection } from '../../../../types.generated';
+import TestPageContainer from '../../../../utils/test-utils/TestPageContainer';
+import { FetchStatus, LineageNodesContext } from '../../common';
+import ManageLineageModal from '../ManageLineageModal';
+
+// Mock the SearchSelect component to avoid rendering it
+vi.mock('../../../entityV2/shared/components/styled/search/SearchSelect', () => ({
+    SearchSelect: () => <div data-testid="mocked-search-select">Mocked SearchSelect</div>,
+}));
+
+// Mock the LineageEdges component
+vi.mock('../LineageEdges', () => ({
+    default: () => <div data-testid="mocked-lineage-edges">Mocked LineageEdges</div>,
+}));
+
+// Mock the useEntityRegistry hook
+vi.mock('../../../useEntityRegistry', () => ({
+    useEntityRegistry: () => ({
+        getDisplayName: vi.fn().mockImplementation((type, entity) => {
+            if (type === EntityType.DataPlatform) return 'Hive';
+            return entity?.name || 'Test Entity';
+        }),
+        getEntityUrl: vi.fn().mockImplementation((urn) => `/entity/${urn}`),
+        getSearchEntityTypes: vi.fn().mockReturnValue([EntityType.Dataset, EntityType.Dashboard]),
+    }),
+    useEntityRegistryV2: () => ({
+        getDisplayName: vi.fn().mockImplementation((type, entity) => {
+            if (type === EntityType.DataPlatform) return 'Hive';
+            return entity?.name || 'Test Entity';
+        }),
+        getEntityUrl: vi.fn().mockImplementation((urn) => `/entity/${urn}`),
+        getSearchEntityTypes: vi.fn().mockReturnValue([EntityType.Dataset, EntityType.Dashboard]),
+    }),
+}));
+
+// Mock the useUserContext hook
+vi.mock('../../../context/useUserContext', () => ({
+    useUserContext: () => ({
+        user: {
+            urn: 'urn:li:corpuser:testUser',
+            username: 'testUser',
+            type: EntityType.CorpUser,
+        },
+    }),
+}));
+
+// Mock the useOnClickExpandLineage hook
+const mockExpandOneLevel = vi.fn();
+vi.mock('../../LineageEntityNode/useOnClickExpandLineage', () => ({
+    useOnClickExpandLineage: () => mockExpandOneLevel,
+}));
+
+// Mock the updateLineageMutation
+vi.mock('../../../../graphql/mutations.generated', () => ({
+    useUpdateLineageMutation: () => [vi.fn().mockResolvedValue({ data: { updateLineage: true } })],
+}));
+
+describe('ManageLineageModal', () => {
+    // Mock DataPlatform entity
+    const mockPlatform = {
+        urn: 'urn:li:dataPlatform:hive',
+        type: EntityType.DataPlatform,
+        name: 'hive',
+        properties: {
+            displayName: 'Hive',
+            type: 'RELATIONAL_DB',
+            datasetNameDelimiter: '.',
+        },
+    };
+
+    // Mock LineageEntity
+    const mockNode = {
+        urn: 'urn:li:dataset:test',
+        type: EntityType.Dataset,
+        entity: {
+            urn: 'urn:li:dataset:test',
+            type: EntityType.Dataset,
+            name: 'Test Dataset',
+            platform: mockPlatform,
+        },
+        fetchStatus: {
+            [LineageDirection.Upstream]: FetchStatus.COMPLETE,
+            [LineageDirection.Downstream]: FetchStatus.COMPLETE,
+        },
+        filters: {
+            [LineageDirection.Upstream]: {},
+            [LineageDirection.Downstream]: {},
+        },
+        id: 'test-id',
+        isExpanded: true,
+    };
+
+    const mockCloseModal = vi.fn();
+    const mockRefetch = vi.fn();
+
+    // Mock LineageNodesContext
+    const mockLineageNodesContext = {
+        rootUrn: 'urn:li:dataset:test',
+        rootType: EntityType.Dataset,
+        nodes: new Map([['urn:li:dataset:test', mockNode]]),
+        edges: new Map(),
+        adjacencyList: {
+            [LineageDirection.Upstream]: new Map([['urn:li:dataset:test', new Set(['urn:li:dataset:upstream1'])]]),
+            [LineageDirection.Downstream]: new Map([['urn:li:dataset:test', new Set(['urn:li:dataset:downstream1'])]]),
+        },
+        nodeVersion: 0,
+        setNodeVersion: vi.fn(),
+        dataVersion: 0,
+        setDataVersion: vi.fn(),
+        displayVersion: [0, []],
+        setDisplayVersion: vi.fn(),
+        columnEdgeVersion: 0,
+        setColumnEdgeVersion: vi.fn(),
+        hideTransformations: false,
+        setHideTransformations: vi.fn(),
+        showDataProcessInstances: false,
+        setShowDataProcessInstances: vi.fn(),
+        showGhostEntities: false,
+        setShowGhostEntities: vi.fn(),
+    };
+
+    beforeEach(() => {
+        vi.clearAllMocks();
+    });
+
+    it('renders the modal with the correct title for upstream direction', () => {
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <LineageNodesContext.Provider value={mockLineageNodesContext as any}>
+                        <ManageLineageModal
+                            node={mockNode as any}
+                            direction={LineageDirection.Upstream}
+                            closeModal={mockCloseModal}
+                            refetch={mockRefetch}
+                        />
+                    </LineageNodesContext.Provider>
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        expect(screen.getByText('Select the Upstreams to add to Test Dataset')).toBeInTheDocument();
+        expect(screen.getByText('Search and Add')).toBeInTheDocument();
+        expect(screen.getByText('Current Upstreams')).toBeInTheDocument();
+        expect(screen.getByTestId('mocked-search-select')).toBeInTheDocument();
+        expect(screen.getByTestId('mocked-lineage-edges')).toBeInTheDocument();
+    });
+
+    it('renders the modal with the correct title for downstream direction', () => {
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <LineageNodesContext.Provider value={mockLineageNodesContext as any}>
+                        <ManageLineageModal
+                            node={mockNode as any}
+                            direction={LineageDirection.Downstream}
+                            closeModal={mockCloseModal}
+                            refetch={mockRefetch}
+                        />
+                    </LineageNodesContext.Provider>
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        expect(screen.getByText('Select the Downstreams to add to Test Dataset')).toBeInTheDocument();
+        expect(screen.getByText('Search and Add')).toBeInTheDocument();
+        expect(screen.getByText('Current Downstreams')).toBeInTheDocument();
+    });
+
+    it('calls closeModal when cancel button is clicked', () => {
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <LineageNodesContext.Provider value={mockLineageNodesContext as any}>
+                        <ManageLineageModal
+                            node={mockNode as any}
+                            direction={LineageDirection.Upstream}
+                            closeModal={mockCloseModal}
+                            refetch={mockRefetch}
+                        />
+                    </LineageNodesContext.Provider>
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        const cancelButton = screen.getByText('Cancel');
+        fireEvent.click(cancelButton);
+
+        expect(mockCloseModal).toHaveBeenCalledTimes(1);
+    });
+
+    it('expands lineage if fetchStatus is UNFETCHED', () => {
+        const unfetchedNode = {
+            ...mockNode,
+            fetchStatus: {
+                [LineageDirection.Upstream]: FetchStatus.UNFETCHED,
+                [LineageDirection.Downstream]: FetchStatus.COMPLETE,
+            },
+        };
+
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <LineageNodesContext.Provider value={mockLineageNodesContext as any}>
+                        <ManageLineageModal
+                            node={unfetchedNode as any}
+                            direction={LineageDirection.Upstream}
+                            closeModal={mockCloseModal}
+                            refetch={mockRefetch}
+                        />
+                    </LineageNodesContext.Provider>
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        // The expandOneLevel function should be called for unfetched status
+        expect(mockExpandOneLevel).toHaveBeenCalled();
+    });
+
+    it('does not expand lineage if fetchStatus is already COMPLETE', () => {
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <LineageNodesContext.Provider value={mockLineageNodesContext as any}>
+                        <ManageLineageModal
+                            node={mockNode as any}
+                            direction={LineageDirection.Upstream}
+                            closeModal={mockCloseModal}
+                            refetch={mockRefetch}
+                        />
+                    </LineageNodesContext.Provider>
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        // The expandOneLevel function should not be called for complete status
+        expect(mockExpandOneLevel).not.toHaveBeenCalled();
+    });
+});

datahub-web-react/src/app/lineageV2/manualLineage/buildUpdateLineagePayload.ts

@@ -0,0 +1,22 @@
+import { Entity, LineageDirection, LineageEdge } from '../../../types.generated';
+
+export function buildUpdateLineagePayload(
+    lineageDirection: LineageDirection,
+    entitiesToAdd: Entity[],
+    entitiesToRemove: Entity[],
+    entityUrn: string,
+) {
+    let edgesToAdd: LineageEdge[] = [];
+    let edgesToRemove: LineageEdge[] = [];
+
+    if (lineageDirection === LineageDirection.Upstream) {
+        edgesToAdd = entitiesToAdd.map((entity) => ({ upstreamUrn: entity.urn, downstreamUrn: entityUrn }));
+        edgesToRemove = entitiesToRemove.map((entity) => ({ upstreamUrn: entity.urn, downstreamUrn: entityUrn }));
+    }
+    if (lineageDirection === LineageDirection.Downstream) {
+        edgesToAdd = entitiesToAdd.map((entity) => ({ upstreamUrn: entityUrn, downstreamUrn: entity.urn }));
+        edgesToRemove = entitiesToRemove.map((entity) => ({ upstreamUrn: entityUrn, downstreamUrn: entity.urn }));
+    }
+
+    return { edgesToAdd, edgesToRemove };
+}

datahub-web-react/src/app/lineageV2/manualLineage/recordManualLineageAnalyticsEvent.ts

@@ -0,0 +1,50 @@
+import { Entity, EntityType, LineageDirection } from '../../../types.generated';
+import { EventType } from '../../analytics';
+import analytics from '../../analytics/analytics';
+import EntityRegistry from '../../entity/EntityRegistry';
+import { Direction } from '../../lineage/types';
+
+interface AnalyticsEventsProps {
+    direction: LineageDirection;
+    entitiesToAdd: Entity[];
+    entitiesToRemove: Entity[];
+    entityRegistry: EntityRegistry;
+    entityType?: EntityType;
+    entityPlatform?: string;
+}
+
+export function recordAnalyticsEvents({
+    direction,
+    entitiesToAdd,
+    entitiesToRemove,
+    entityRegistry,
+    entityType,
+    entityPlatform,
+}: AnalyticsEventsProps) {
+    entitiesToAdd.forEach((entityToAdd) => {
+        const genericProps = entityRegistry.getGenericEntityProperties(entityToAdd.type, entityToAdd);
+        analytics.event({
+            type: EventType.ManuallyCreateLineageEvent,
+            direction: directionFromLineageDirection(direction),
+            sourceEntityType: entityType,
+            sourceEntityPlatform: entityPlatform,
+            destinationEntityType: entityToAdd.type,
+            destinationEntityPlatform: genericProps?.platform?.name,
+        });
+    });
+    entitiesToRemove.forEach((entityToRemove) => {
+        const genericProps = entityRegistry.getGenericEntityProperties(entityToRemove.type, entityToRemove);
+        analytics.event({
+            type: EventType.ManuallyDeleteLineageEvent,
+            direction: directionFromLineageDirection(direction),
+            sourceEntityType: entityType,
+            sourceEntityPlatform: entityPlatform,
+            destinationEntityType: entityToRemove.type,
+            destinationEntityPlatform: genericProps?.platform?.name,
+        });
+    });
+}
+
+function directionFromLineageDirection(lineageDirection: LineageDirection): Direction {
+    return lineageDirection === LineageDirection.Upstream ? Direction.Upstream : Direction.Downstream;
+}

smoke-test/tests/cypress/cypress/e2e/lineageV2/v2_impact_analysis.js

@@ -158,4 +158,46 @@ describe("impact analysis", () => {
     );
     cy.contains("temperature_etl_2");
   });
+
+  it("editing upstream lineage will redirect to visual view with edit modal open", () => {
+    cy.login();
+    cy.visit(
+      `/dataset/${DATASET_URN}/Lineage?is_lineage_mode=false&lineageView=impact`,
+    );
+
+    // Wait for the page to load
+    cy.contains("SampleCypressKafkaDataset").should("be.visible");
+
+    // Click on a node that has the lineage edit menu
+    cy.get('[data-testid="lineage-edit-menu-button"]').first().click();
+
+    // Click on Edit Upstream option
+    cy.get('[data-testid="edit-upstream-lineage"]').click();
+
+    // Verify the edit modal is open
+    cy.contains(
+      "Select the Upstreams to add to SampleCypressKafkaDataset",
+    ).should("be.visible");
+  });
+
+  it("editing downstream lineage will redirect to visual view with edit modal open", () => {
+    cy.login();
+    cy.visit(
+      `/dataset/${DATASET_URN}/Lineage?is_lineage_mode=false&lineageView=impact`,
+    );
+
+    // Wait for the page to load
+    cy.contains("SampleCypressKafkaDataset").should("be.visible");
+
+    // Click on a node that has the lineage edit menu
+    cy.get('[data-testid="lineage-edit-menu-button"]').first().click();
+
+    // Click on Edit Downstream option
+    cy.get('[data-testid="edit-downstream-lineage"]').click();
+
+    // Verify the edit modal is open
+    cy.contains(
+      "Select the Downstreams to add to SampleCypressKafkaDataset",
+    ).should("be.visible");
+  });
 });

smoke-test/tests/cypress/cypress/e2e/lineageV2/v2_lineage_graph.js

@@ -126,4 +126,42 @@ describe("lineage_graph", () => {
     cy.contains("gdp");
     cy.contains("factor_income").should("not.exist");
   });
+
+  it("can edit upstream lineage", () => {
+    // note: this test does not make any mutations. This is to prevent introducing a flaky test.
+    // Ideally, we should refactor this test to verify that the mutations are correct as well. However, it needs to be done carefully to avoid introducing flakiness.
+
+    cy.login();
+    cy.goToEntityLineageGraphV2(DATASET_ENTITY_TYPE, DATASET_URN);
+    cy.get(
+      '[data-testid="manage-lineage-menu-urn:li:dataset:(urn:li:dataPlatform:kafka,SampleCypressKafkaDataset,PROD)"]',
+    ).click();
+    cy.get('[data-testid="edit-upstream-lineage"]').click();
+    cy.contains(
+      "Select the Upstreams to add to SampleCypressKafkaDataset",
+    ).should("be.visible");
+
+    // find the button that says "Set Upstreams" - not via test id
+    // verify that is is not disabled
+    cy.contains("Set Upstreams").should("be.disabled");
+
+    // there are multiple search inputs, we want to find the one for the search bar in the modal
+    cy.get('[role="dialog"] [data-testid="search-input"]').type(
+      "cypress_health_test",
+    );
+
+    // find the checkbox for cypress_health_test (data-testid="preview-urn:li:dataset:(urn:li:dataPlatform:hive,cypress_health_test,PROD)") is the
+    // test id for a sibling of the checkbox, we want to find that id, find its parent, and click the checkbox
+
+    cy.get(
+      '[data-testid="preview-urn:li:dataset:(urn:li:dataPlatform:hive,cypress_health_test,PROD)"]',
+    )
+      .parent()
+      .find("input")
+      .click();
+
+    // find the button that says "Set Upstreams" - not via test id
+    // verify that is is not disabled
+    cy.contains("Set Upstreams").should("not.be.disabled");
+  });
 });