feat(ui/homepage): Add drag-and-drop and other updates to Asset Collection Modal (#14168)

Co-authored-by: Chris Collins <chriscollins3456@gmail.com>

datahub-web-react/package.json

@@ -10,6 +10,7 @@
         "@ant-design/icons": "^4.3.0",
         "@apollo/client": "^3.3.19",
         "@dnd-kit/core": "^6.3.1",
+        "@dnd-kit/modifiers": "^9.0.0",
         "@dnd-kit/sortable": "^10.0.0",
         "@dnd-kit/utilities": "^3.2.2",
         "@fontsource/mulish": "^5.0.16",

datahub-web-react/src/app/homeV3/module/components/EntityItem.tsx

@@ -13,8 +13,9 @@ const StyledLink = styled(Link)`
 
 interface Props {
     entity: Entity;
-    customDetailsRenderer?: (entity: Entity) => void;
+    customDetailsRenderer?: (entity: Entity) => React.ReactNode;
     navigateOnlyOnNameClick?: boolean;
+    dragIconRenderer?: () => React.ReactNode;
     hideSubtitle?: boolean;
     hideMatches?: boolean;
     padding?: string;
@@ -24,6 +25,7 @@ export default function EntityItem({
     entity,
     customDetailsRenderer,
     navigateOnlyOnNameClick = false,
+    dragIconRenderer,
     hideSubtitle,
     hideMatches,
     padding,
@@ -41,6 +43,7 @@ export default function EntityItem({
                     hideMatches={hideMatches}
                     padding={padding}
                     navigateOnlyOnNameClick
+                    dragIconRenderer={dragIconRenderer}
                 />
             ) : (
                 <StyledLink to={entityRegistry.getEntityUrl(entity.type, entity.urn)}>
@@ -51,6 +54,7 @@ export default function EntityItem({
                         hideMatches={hideMatches}
                         padding={padding}
                         customDetailsRenderer={customDetailsRenderer}
+                        dragIconRenderer={dragIconRenderer}
                     />
                 </StyledLink>
             )}

datahub-web-react/src/app/homeV3/module/components/SmallModule.tsx

@@ -12,7 +12,7 @@ const DragIcon = styled(Icon)<{ isDragging: boolean }>`
     cursor: ${(props) => (props.isDragging ? 'grabbing' : 'grab')};
     display: none;
     position: absolute;
-    left: 4px;
+    left: 2px;
     top: 50%;
     transform: translateY(-50%);
 `;

datahub-web-react/src/app/homeV3/modules/assetCollection/AssetsSection.tsx

@@ -9,7 +9,7 @@ import SelectedAssetsSection from '@app/homeV3/modules/assetCollection/SelectedA
 const Container = styled.div`
     display: flex;
     width: 100%;
-    gap: 16px;
+    gap: 8px;
 `;
 
 const LeftSection = styled.div`

datahub-web-react/src/app/homeV3/modules/assetCollection/SelectAssetsSection.tsx

@@ -24,6 +24,10 @@ const ItemDetailsContainer = styled.div`
     align-items: center;
 `;
 
+const ResultsContainer = styled.div`
+    margin: 0 -16px 0 -8px;
+`;
+
 type Props = {
     selectedAssetUrns: string[];
     setSelectedAssetUrns: React.Dispatch<React.SetStateAction<string[]>>;
@@ -92,7 +96,7 @@ const SelectAssetsSection = ({ selectedAssetUrns, setSelectedAssetUrns }: Props)
                 appliedFilters={appliedFilters}
                 updateFieldFilters={updateFieldFilters}
             />
-            {content}
+            <ResultsContainer>{content}</ResultsContainer>
         </AssetsSection>
     );
 };

datahub-web-react/src/app/homeV3/modules/assetCollection/SelectedAssetsSection.tsx

@@ -1,8 +1,10 @@
 import { Text } from '@components';
-import React from 'react';
+import { isEqual } from 'lodash';
+import React, { useEffect, useMemo, useState } from 'react';
 import styled from 'styled-components';
 
-import EntityItem from '@app/homeV3/module/components/EntityItem';
+import DraggableEntityItem from '@app/homeV3/modules/assetCollection/dragAndDrop/DraggableEntityItem';
+import VerticalDragAndDrop from '@app/homeV3/modules/assetCollection/dragAndDrop/VerticalDragAndDrop';
 import { EmptyContainer, StyledIcon } from '@app/homeV3/styledComponents';
 import { useGetEntities } from '@app/sharedV2/useGetEntities';
 
@@ -15,13 +17,41 @@ const SelectedAssetsContainer = styled.div`
     height: 100%;
 `;
 
+const ResultsContainer = styled.div`
+    margin: 0 -12px 0 -8px;
+    height: 100%;
+`;
+
 type Props = {
     selectedAssetUrns: string[];
     setSelectedAssetUrns: React.Dispatch<React.SetStateAction<string[]>>;
 };
 
 const SelectedAssetsSection = ({ selectedAssetUrns, setSelectedAssetUrns }: Props) => {
-    const { entities, loading } = useGetEntities(selectedAssetUrns);
+    const [orderedUrns, setOrderedUrns] = useState(selectedAssetUrns);
+
+    useEffect(() => {
+        if (!isEqual(selectedAssetUrns, orderedUrns)) {
+            setOrderedUrns(selectedAssetUrns);
+        }
+    }, [orderedUrns, selectedAssetUrns]);
+
+    const onChangeOrder = (urns: string[]) => {
+        setOrderedUrns(urns);
+        setSelectedAssetUrns(urns);
+    };
+
+    // To prevent refetch on only order change
+    const stableUrns = useMemo(() => [...selectedAssetUrns].sort(), [selectedAssetUrns]);
+    const { entities } = useGetEntities(stableUrns);
+
+    const entitiesMap = useMemo(() => {
+        const map: Record<string, Entity> = {};
+        entities.forEach((entity) => {
+            map[entity.urn] = entity;
+        });
+        return map;
+    }, [entities]);
 
     const handleRemoveAsset = (entity: Entity) => {
         const newUrns = selectedAssetUrns.filter((urn) => !(entity.urn === urn));
@@ -44,16 +74,14 @@ const SelectedAssetsSection = ({ selectedAssetUrns, setSelectedAssetUrns }: Prop
     };
 
     let content;
-    if (entities && entities.length > 0) {
-        content = entities.map((entity) => (
-            <EntityItem
-                entity={entity}
-                key={entity.urn}
-                customDetailsRenderer={renderRemoveAsset}
-                navigateOnlyOnNameClick
-            />
-        ));
-    } else if (!loading && entities.length === 0) {
+    if (selectedAssetUrns.length > 0) {
+        content = selectedAssetUrns
+            .map((urn) => entitiesMap[urn])
+            .filter(Boolean)
+            .map((entity) => (
+                <DraggableEntityItem key={entity.urn} entity={entity} customDetailsRenderer={renderRemoveAsset} />
+            ));
+    } else {
         content = (
             <EmptyContainer>
                 <Text color="gray">No assets selected.</Text>
@@ -66,7 +94,9 @@ const SelectedAssetsSection = ({ selectedAssetUrns, setSelectedAssetUrns }: Prop
             <Text color="gray" weight="bold">
                 Selected Assets
             </Text>
-            {content}
+            <VerticalDragAndDrop items={orderedUrns} onChange={onChangeOrder}>
+                <ResultsContainer>{content}</ResultsContainer>
+            </VerticalDragAndDrop>
         </SelectedAssetsContainer>
     );
 };

datahub-web-react/src/app/homeV3/modules/assetCollection/__tests__/useGetAssetResults.test.ts

@@ -60,7 +60,7 @@ describe('useGetAssetResults', () => {
                 input: {
                     query: searchQuery,
                     start: 0,
-                    count: 10,
+                    count: 20,
                     orFilters: mockOrFilters,
                     searchFlags: { skipCache: true },
                 },

datahub-web-react/src/app/homeV3/modules/assetCollection/__tests__/useVerticalDragAndDrop.test.ts

@@ -0,0 +1,99 @@
+import { DragEndEvent } from '@dnd-kit/core';
+import { act, renderHook } from '@testing-library/react-hooks';
+import { describe, expect, it, vi } from 'vitest';
+
+import { useVerticalDragAndDrop } from '@app/homeV3/modules/assetCollection/dragAndDrop/useVerticalDragAndDrop';
+
+describe('useVerticalDragAndDrop', () => {
+    const initialItems = ['a', 'b', 'c', 'd'];
+
+    it('should return correct sensors, strategy, collisionDetection, and modifiers', () => {
+        const { result } = renderHook(() => useVerticalDragAndDrop({ items: initialItems, onChange: vi.fn() }));
+        const { sensors, strategy, collisionDetection, modifiers } = result.current;
+
+        expect(sensors).toBeDefined();
+        expect(strategy).toBeDefined();
+        expect(collisionDetection).toBeDefined();
+        expect(Array.isArray(modifiers)).toBe(true);
+        expect(modifiers.length).toBeGreaterThan(0);
+    });
+
+    describe('handleDragEnd callback behavior', () => {
+        it('should do nothing if over is null', () => {
+            const onChange = vi.fn();
+            const { result } = renderHook(() => useVerticalDragAndDrop({ items: initialItems, onChange }));
+
+            const event = { active: { id: 'a' }, over: null } as unknown as DragEndEvent;
+            act(() => {
+                result.current.handleDragEnd(event);
+            });
+
+            expect(onChange).not.toHaveBeenCalled();
+        });
+
+        it('should do nothing if active.id is equal to over.id', () => {
+            const onChange = vi.fn();
+            const { result } = renderHook(() => useVerticalDragAndDrop({ items: initialItems, onChange }));
+
+            const event = { active: { id: 'b' }, over: { id: 'b' } } as unknown as DragEndEvent;
+            act(() => {
+                result.current.handleDragEnd(event);
+            });
+
+            expect(onChange).not.toHaveBeenCalled();
+        });
+
+        it('should call onChange with new order if active and over ids are different and in items', () => {
+            const onChange = vi.fn();
+            // items in order: ['a','b','c','d']
+            const { result } = renderHook(() => useVerticalDragAndDrop({ items: initialItems, onChange }));
+
+            // simulate dragging 'a' over 'c'
+            const event = { active: { id: 'a' }, over: { id: 'c' } } as unknown as DragEndEvent;
+            act(() => {
+                result.current.handleDragEnd(event);
+            });
+
+            // Expected new order: ['b', 'c', 'a', 'd']
+            expect(onChange).toHaveBeenCalledTimes(1);
+            expect(onChange).toHaveBeenCalledWith(['b', 'c', 'a', 'd']);
+        });
+
+        it('should do nothing if active.id or over.id not in items', () => {
+            const onChange = vi.fn();
+            const { result } = renderHook(() => useVerticalDragAndDrop({ items: initialItems, onChange }));
+
+            // active.id not in items
+            let event = { active: { id: 'x' }, over: { id: 'a' } } as unknown as DragEndEvent;
+            act(() => {
+                result.current.handleDragEnd(event);
+            });
+            expect(onChange).not.toHaveBeenCalled();
+
+            // over.id not in items
+            event = { active: { id: 'a' }, over: { id: 'x' } } as unknown as DragEndEvent;
+            act(() => {
+                result.current.handleDragEnd(event);
+            });
+            expect(onChange).not.toHaveBeenCalled();
+        });
+    });
+
+    it('should update handleDragEnd callback when items or onChange changes', () => {
+        const onChange1 = vi.fn();
+        const onChange2 = vi.fn();
+        const { result, rerender } = renderHook(({ items, onChange }) => useVerticalDragAndDrop({ items, onChange }), {
+            initialProps: { items: initialItems, onChange: onChange1 },
+        });
+
+        const firstHandleDragEnd = result.current.handleDragEnd;
+
+        // Re-render with new items
+        rerender({ items: ['a', 'c', 'b', 'd'], onChange: onChange1 });
+        expect(result.current.handleDragEnd).not.toBe(firstHandleDragEnd);
+
+        // Re-render with new onChange callback
+        rerender({ items: initialItems, onChange: onChange2 });
+        expect(result.current.handleDragEnd).not.toBe(firstHandleDragEnd);
+    });
+});

datahub-web-react/src/app/homeV3/modules/assetCollection/dragAndDrop/DragHandle.tsx

@@ -0,0 +1,26 @@
+import { Icon } from '@components';
+import { SyntheticListenerMap } from '@dnd-kit/core/dist/hooks/utilities';
+import React from 'react';
+import styled from 'styled-components';
+
+const DragIcon = styled(Icon)<{ $isDragging?: boolean }>`
+    cursor: ${(props) => (props.$isDragging ? 'grabbing' : 'grab')};
+`;
+
+type Props = {
+    isDragging?: boolean;
+    listeners?: SyntheticListenerMap;
+};
+
+export default function DragHandle({ isDragging, listeners }: Props) {
+    return (
+        <DragIcon
+            {...listeners}
+            size="lg"
+            color="gray"
+            icon="DotsSixVertical"
+            source="phosphor"
+            $isDragging={isDragging}
+        />
+    );
+}

datahub-web-react/src/app/homeV3/modules/assetCollection/dragAndDrop/DraggableEntityItem.tsx

@@ -0,0 +1,51 @@
+import { colors } from '@components';
+import { useSortable } from '@dnd-kit/sortable';
+import { CSS } from '@dnd-kit/utilities';
+import React from 'react';
+import styled from 'styled-components';
+
+import EntityItem from '@app/homeV3/module/components/EntityItem';
+import DragHandle from '@app/homeV3/modules/assetCollection/dragAndDrop/DragHandle';
+
+import type { Entity } from '@types';
+
+const DraggableWrapper = styled.div<{ $isDragging: boolean; $transform?: string; $transition?: string }>`
+    background-color: ${colors.white};
+    box-shadow: ${(props) => (props.$isDragging ? '0px 4px 12px 0px rgba(9, 1, 61, 0.12)' : 'none')};
+    cursor: ${(props) => (props.$isDragging ? 'grabbing' : 'inherit')};
+    z-index: ${(props) => (props.$isDragging ? '999' : 'auto')};
+    transform: ${(props) => props.$transform};
+    transition: ${(props) => props.$transition};
+    position: ${({ $isDragging }) => ($isDragging ? 'relative' : 'static')};
+    border-radius: 8px;
+`;
+
+type Props = {
+    entity: Entity;
+    customDetailsRenderer?: (entity: Entity) => React.ReactNode;
+};
+
+export default function DraggableEntityItem({ entity, customDetailsRenderer }: Props) {
+    const { attributes, listeners, setNodeRef, transform, transition, isDragging } = useSortable({ id: entity.urn });
+
+    const dragHandle = () => {
+        return <DragHandle listeners={listeners} isDragging={isDragging} />;
+    };
+
+    return (
+        <DraggableWrapper
+            ref={setNodeRef}
+            {...attributes}
+            $isDragging={isDragging}
+            $transform={CSS.Transform.toString(transform)}
+            $transition={transition}
+        >
+            <EntityItem
+                entity={entity}
+                customDetailsRenderer={customDetailsRenderer}
+                navigateOnlyOnNameClick
+                dragIconRenderer={dragHandle}
+            />
+        </DraggableWrapper>
+    );
+}

datahub-web-react/src/app/homeV3/modules/assetCollection/dragAndDrop/VerticalDragAndDrop.tsx

@@ -0,0 +1,30 @@
+import { DndContext } from '@dnd-kit/core';
+import { SortableContext } from '@dnd-kit/sortable';
+import React from 'react';
+
+import { useVerticalDragAndDrop } from '@app/homeV3/modules/assetCollection/dragAndDrop/useVerticalDragAndDrop';
+
+type Props = {
+    items: string[];
+    onChange: (newItems: string[]) => void;
+    children: React.ReactNode;
+};
+
+export default function VerticalDragAndDrop({ items, onChange, children }: Props) {
+    const { sensors, handleDragEnd, strategy, collisionDetection, modifiers } = useVerticalDragAndDrop({
+        items,
+        onChange,
+    });
+    return (
+        <DndContext
+            sensors={sensors}
+            collisionDetection={collisionDetection}
+            onDragEnd={handleDragEnd}
+            modifiers={modifiers}
+        >
+            <SortableContext items={items} strategy={strategy}>
+                {children}
+            </SortableContext>
+        </DndContext>
+    );
+}

datahub-web-react/src/app/homeV3/modules/assetCollection/dragAndDrop/useVerticalDragAndDrop.ts

@@ -0,0 +1,36 @@
+import { DragEndEvent, PointerSensor, closestCenter, useSensor, useSensors } from '@dnd-kit/core';
+import { restrictToParentElement, restrictToVerticalAxis } from '@dnd-kit/modifiers';
+import { arrayMove, verticalListSortingStrategy } from '@dnd-kit/sortable';
+import { useCallback } from 'react';
+
+type Props = {
+    items: string[];
+    onChange: (newItems: string[]) => void;
+};
+
+export function useVerticalDragAndDrop({ items, onChange }: Props) {
+    const sensors = useSensors(useSensor(PointerSensor));
+
+    const handleDragEnd = useCallback(
+        (event: DragEndEvent) => {
+            const { active, over } = event;
+            if (active.id !== over?.id && over) {
+                const oldIndex = items.indexOf(String(active.id));
+                const newIndex = items.indexOf(String(over.id));
+                if (oldIndex !== -1 && newIndex !== -1) {
+                    const newItems = arrayMove(items, oldIndex, newIndex);
+                    onChange(newItems);
+                }
+            }
+        },
+        [items, onChange],
+    );
+
+    return {
+        sensors,
+        handleDragEnd,
+        strategy: verticalListSortingStrategy,
+        collisionDetection: closestCenter,
+        modifiers: [restrictToVerticalAxis, restrictToParentElement],
+    };
+}

datahub-web-react/src/app/homeV3/modules/assetCollection/useGetAssetResults.ts

@@ -21,7 +21,7 @@ export default function useGetAssetResults({ searchQuery, appliedFilters }: Prop
             input: {
                 query: searchQuery || '*',
                 start: 0,
-                count: 10,
+                count: 20,
                 orFilters,
                 searchFlags: {
                     skipCache: true,

datahub-web-react/src/app/homeV3/styledComponents.ts

@@ -79,13 +79,13 @@ export const StyledIcon = styled(Icon)`
 export const LoaderContainer = styled.div`
     display: flex;
     height: 100%;
-    min-height: 200px;
+    height: 200px;
 `;
 
 export const EmptyContainer = styled.div`
     display: flex;
-    height: 50%;
     width: 100%;
+    height: 200px;
     justify-content: center;
     align-items: center;
 `;

datahub-web-react/src/app/homeV3/template/components/addModuleMenu/components/MenuItem.tsx

@@ -1,5 +1,5 @@
-import { Icon, IconNames, Text } from '@components';
-import React from 'react';
+import { Icon, IconNames, Text, Tooltip } from '@components';
+import React, { useMemo } from 'react';
 import styled from 'styled-components';
 
 import spacing from '@components/theme/foundations/spacing';
@@ -32,10 +32,20 @@ interface Props {
     title: string;
     description?: string;
     hasChildren?: boolean;
+    isDisabled?: boolean;
+    isSmallModule?: boolean;
 }
 
-export default function MenuItem({ icon, title, description, hasChildren }: Props) {
-    return (
+export default function MenuItem({ icon, title, description, hasChildren, isDisabled, isSmallModule }: Props) {
+    const tooltipText = useMemo(() => {
+        if (!isDisabled) return undefined;
+        if (isSmallModule) {
+            return 'Cannot add small widget to large widget row';
+        }
+        return 'Cannot add large widget to small widget row';
+    }, [isDisabled, isSmallModule]);
+
+    const content = (
         <Wrapper>
             <IconWrapper>
                 <Icon icon={icon} source="phosphor" color="gray" size="2xl" />
@@ -44,7 +54,7 @@ export default function MenuItem({ icon, title, description, hasChildren }: Prop
             <Container>
                 <Text weight="semiBold">{title}</Text>
                 {description && (
-                    <Text color="gray" colorLevel={1700} size="sm">
+                    <Text color="gray" colorLevel={isDisabled ? 300 : 1700} size="sm">
                         {description}
                     </Text>
                 )}
@@ -55,4 +65,10 @@ export default function MenuItem({ icon, title, description, hasChildren }: Prop
             {hasChildren && <Icon icon="CaretRight" source="phosphor" color="gray" size="lg" />}
         </Wrapper>
     );
+
+    if (isDisabled && tooltipText) {
+        return <Tooltip title={tooltipText}>{content}</Tooltip>;
+    }
+
+    return content;
 }

datahub-web-react/src/app/homeV3/template/components/addModuleMenu/useAddModuleMenu.tsx

@@ -79,9 +79,18 @@ export default function useAddModuleMenu(position: ModulePositionInput, closeMen
         const items: MenuProps['items'] = [];
 
         const quickLink = {
-            title: 'Quick Link',
+            name: 'Quick Link',
             key: 'quick-link',
-            label: <MenuItem description="Choose links that are important" title="Quick Link" icon="LinkSimple" />,
+            label: (
+                <MenuItem
+                    description="Choose links that are important"
+                    title="Quick Link"
+                    icon="LinkSimple"
+                    isDisabled={isLargeModuleRow}
+                    isSmallModule
+                />
+            ),
+
             onClick: () => {
                 handleOpenCreateModuleModal(DataHubPageModuleType.Link);
             },
@@ -89,9 +98,18 @@ export default function useAddModuleMenu(position: ModulePositionInput, closeMen
         };
 
         const documentation = {
-            title: 'Documentation',
+            name: 'Documentation',
             key: 'documentation',
-            label: <MenuItem description="Pin docs for your DataHub users" title="Documentation" icon="TextT" />,
+            label: (
+                <MenuItem
+                    description="Pin docs for your DataHub users"
+                    title="Documentation"
+                    icon="TextT"
+                    isDisabled={isSmallModuleRow}
+                    isSmallModule={false}
+                />
+            ),
+
             onClick: () => {
                 handleOpenCreateModuleModal(DataHubPageModuleType.RichText);
             },
@@ -106,9 +124,18 @@ export default function useAddModuleMenu(position: ModulePositionInput, closeMen
         });
 
         const yourAssets = {
-            title: 'Your Assets',
+            name: 'Your Assets',
             key: 'your-assets',
-            label: <MenuItem description="Assets the current user owns" title="Your Assets" icon="Database" />,
+            label: (
+                <MenuItem
+                    description="Assets the current user owns"
+                    title="Your Assets"
+                    icon="Database"
+                    isDisabled={isSmallModuleRow}
+                    isSmallModule={false}
+                />
+            ),
+
             onClick: () => {
                 handleAddExistingModule(YOUR_ASSETS_MODULE);
             },
@@ -116,9 +143,18 @@ export default function useAddModuleMenu(position: ModulePositionInput, closeMen
         };
 
         const domains = {
-            title: 'Domains',
+            name: 'Domains',
             key: 'domains',
-            label: <MenuItem description="Most used domains in your organization" title="Domains" icon="Globe" />,
+            label: (
+                <MenuItem
+                    description="Most used domains in your organization"
+                    title="Domains"
+                    icon="Globe"
+                    isDisabled={isSmallModuleRow}
+                    isSmallModule={false}
+                />
+            ),
+
             onClick: () => {
                 handleAddExistingModule(DOMAINS_MODULE);
             },
@@ -126,13 +162,15 @@ export default function useAddModuleMenu(position: ModulePositionInput, closeMen
         };
 
         const assetCollection = {
-            title: 'Asset Collection',
+            name: 'Asset Collection',
             key: 'asset-collection',
             label: (
                 <MenuItem
                     description="A curated list of assets of your choosing"
                     title="Asset Collection"
                     icon="Stack"
+                    isDisabled={isSmallModuleRow}
+                    isSmallModule={false}
                 />
             ),
             onClick: () => {

datahub-web-react/src/app/searchV2/autoCompleteV2/AutoCompleteEntityItem.tsx

@@ -92,18 +92,23 @@ const TypeContainer = styled.div`
     align-items: center;
 `;
 
+const Icons = styled.div`
+    display: flex;
+    align-items: center;
+    gap: 8px;
+`;
+
 interface EntityAutocompleteItemProps {
     entity: Entity;
     query?: string;
     siblings?: Entity[];
     matchedFields?: MatchedField[];
     variant?: EntityItemVariant;
-    customDetailsRenderer?: (entity: Entity) => void;
+    customDetailsRenderer?: (entity: Entity) => React.ReactNode;
     navigateOnlyOnNameClick?: boolean;
-
+    dragIconRenderer?: () => React.ReactNode;
     hideSubtitle?: boolean;
     hideMatches?: boolean;
-
     padding?: string;
 }
 
@@ -115,6 +120,7 @@ export default function AutoCompleteEntityItem({
     variant,
     customDetailsRenderer,
     navigateOnlyOnNameClick,
+    dragIconRenderer,
     hideSubtitle,
     hideMatches,
     padding,
@@ -154,9 +160,18 @@ export default function AutoCompleteEntityItem({
     return (
         <Container $navigateOnlyOnNameClick={navigateOnlyOnNameClick} $padding={padding}>
             <ContentContainer>
-                <IconContainer $variant={variant}>
-                    <EntityIcon entity={entity} siblings={siblings} />
-                </IconContainer>
+                {dragIconRenderer ? (
+                    <Icons>
+                        {dragIconRenderer()}
+                        <IconContainer $variant={variant}>
+                            <EntityIcon entity={entity} siblings={siblings} />
+                        </IconContainer>
+                    </Icons>
+                ) : (
+                    <IconContainer $variant={variant}>
+                        <EntityIcon entity={entity} siblings={siblings} />
+                    </IconContainer>
+                )}
 
                 <DescriptionContainer>
                     <HoverEntityTooltip

datahub-web-react/src/app/sharedV2/useGetEntities.ts

@@ -1,3 +1,5 @@
+import { useEffect, useMemo, useState } from 'react';
+
 import { useGetEntitiesQuery } from '@graphql/entity.generated';
 import { Entity } from '@types';
 
@@ -5,13 +7,26 @@ export function useGetEntities(urns: string[]): {
     entities: Entity[];
     loading: boolean;
 } {
-    const verifiedUrns = urns.filter((urn) => typeof urn === 'string' && urn.startsWith('urn:li:'));
+    const verifiedUrns = useMemo(
+        () => urns.filter((urn) => typeof urn === 'string' && urn.startsWith('urn:li:')),
+        [urns],
+    );
 
     const { data, loading } = useGetEntitiesQuery({
         variables: { urns: verifiedUrns },
         skip: !verifiedUrns.length,
         fetchPolicy: 'cache-first',
     });
-    const entities: Entity[] = Array.isArray(data?.entities) ? (data?.entities.filter(Boolean) as Entity[]) : [];
+
+    const [entities, setEntities] = useState<Entity[]>([]);
+
+    useEffect(() => {
+        if (data?.entities && data.entities.length > 0) {
+            setEntities(data.entities as Entity[]);
+        } else if (!loading && (!data?.entities || data.entities.length === 0)) {
+            setEntities([]);
+        }
+    }, [data, loading]);
+
     return { entities, loading };
 }

datahub-web-react/yarn.lock

@@ -1101,6 +1101,14 @@
     "@dnd-kit/utilities" "^3.2.2"
     tslib "^2.0.0"
 
+"@dnd-kit/modifiers@^9.0.0":
+  version "9.0.0"
+  resolved "https://registry.yarnpkg.com/@dnd-kit/modifiers/-/modifiers-9.0.0.tgz#96a0280c77b10c716ef79d9792ce7ad04370771d"
+  integrity sha512-ybiLc66qRGuZoC20wdSSG6pDXFikui/dCNGthxv4Ndy8ylErY0N3KVxY2bgo7AWwIbxDmXDg3ylAFmnrjcbVvw==
+  dependencies:
+    "@dnd-kit/utilities" "^3.2.2"
+    tslib "^2.0.0"
+
 "@dnd-kit/sortable@^10.0.0":
   version "10.0.0"
   resolved "https://registry.npmjs.org/@dnd-kit/sortable/-/sortable-10.0.0.tgz"