feat(structured properties): use wider search select modal to edit structured properties (#13076)

datahub-web-react/src/app/entity/shared/components/styled/StructuredProperty/StructuredPropertyInput.tsx

@@ -1,12 +1,13 @@
-import React from 'react';
 import { PropertyCardinality, StdDataType, StructuredPropertyEntity } from '@src/types.generated';
-import SingleSelectInput from './SingleSelectInput';
+import React from 'react';
-import MultiSelectInput from './MultiSelectInput';
+import StructuredPropertySearchSelectUrnInput from '../../../entityForm/prompts/StructuredPropertyPrompt/UrnInput/StructuredPropertySearchSelectUrnInput';
-import StringInput from './StringInput';
-import RichTextInput from './RichTextInput';
-import DateInput from './DateInput';
-import NumberInput from './NumberInput';
 import UrnInput from '../../../entityForm/prompts/StructuredPropertyPrompt/UrnInput/UrnInput';
+import DateInput from './DateInput';
+import MultiSelectInput from './MultiSelectInput';
+import NumberInput from './NumberInput';
+import RichTextInput from './RichTextInput';
+import SingleSelectInput from './SingleSelectInput';
+import StringInput from './StringInput';
 
 interface Props {
     structuredProperty: StructuredPropertyEntity;
@@ -14,6 +15,7 @@ interface Props {
     selectSingleValue: (value: string | number) => void;
     toggleSelectedValue: (value: string | number) => void;
     updateSelectedValues: (value: (string | number | null)[]) => void;
+    canUseSearchSelectUrnInput?: boolean;
 }
 
 export default function StructuredPropertyInput({
@@ -22,6 +24,7 @@ export default function StructuredPropertyInput({
     selectedValues,
     toggleSelectedValue,
     updateSelectedValues,
+    canUseSearchSelectUrnInput = false,
 }: Props) {
     const { allowedValues, cardinality, valueType } = structuredProperty.definition;
 
@@ -66,10 +69,17 @@ export default function StructuredPropertyInput({
                     updateSelectedValues={updateSelectedValues}
                 />
             )}
-            {!allowedValues && valueType.info.type === StdDataType.Urn && (
+            {!allowedValues && valueType.info.type === StdDataType.Urn && canUseSearchSelectUrnInput && (
+                <StructuredPropertySearchSelectUrnInput
+                    structuredProperty={structuredProperty}
+                    selectedValues={selectedValues as string[]}
+                    updateSelectedValues={updateSelectedValues}
+                />
+            )}
+            {!allowedValues && valueType.info.type === StdDataType.Urn && !canUseSearchSelectUrnInput && (
                 <UrnInput
                     structuredProperty={structuredProperty}
-                    selectedValues={selectedValues}
+                    selectedValues={selectedValues as string[]}
                     updateSelectedValues={updateSelectedValues}
                 />
             )}

datahub-web-react/src/app/entity/shared/entityForm/prompts/StructuredPropertyPrompt/UrnInput/StructuredPropertySearchSelectUrnInput.tsx

@@ -0,0 +1,36 @@
+import { SearchSelectUrnInput } from '@src/app/entityV2/shared/components/styled/search/SearchSelectUrnInput';
+import React, { useMemo } from 'react';
+import { EntityType, PropertyCardinality, StructuredPropertyEntity } from '../../../../../../../types.generated';
+
+interface StructuredPropertySearchSelectUrnInputProps {
+    structuredProperty: StructuredPropertyEntity;
+    selectedValues: string[];
+    updateSelectedValues: (values: string[] | number[]) => void;
+}
+
+// Wrapper component that extracts information from StructuredProperty
+export default function StructuredPropertySearchSelectUrnInput({
+    structuredProperty,
+    selectedValues,
+    updateSelectedValues,
+}: StructuredPropertySearchSelectUrnInputProps) {
+    // Get the allowed entity types from the structured property
+    const allowedEntityTypes = useMemo(() => {
+        return (
+            structuredProperty.definition.typeQualifier?.allowedTypes?.map(
+                (allowedType) => allowedType.info.type as EntityType,
+            ) || []
+        );
+    }, [structuredProperty]);
+
+    const isMultiple = structuredProperty.definition.cardinality === PropertyCardinality.Multiple;
+
+    return (
+        <SearchSelectUrnInput
+            allowedEntityTypes={allowedEntityTypes}
+            isMultiple={isMultiple}
+            selectedValues={selectedValues}
+            updateSelectedValues={updateSelectedValues}
+        />
+    );
+}

datahub-web-react/src/app/entity/shared/entityForm/prompts/StructuredPropertyPrompt/UrnInput/UrnInput.tsx

@@ -3,8 +3,8 @@ import { Select } from 'antd';
 import React from 'react';
 import styled from 'styled-components';
 import { StructuredPropertyEntity } from '../../../../../../../types.generated';
-import useUrnInput from './useUrnInput';
 import SelectedEntity from './SelectedEntity';
+import useUrnInput from './useUrnInput';
 
 const EntitySelect = styled(Select)`
     width: 75%;

datahub-web-react/src/app/entity/shared/tabs/Properties/Edit/EditStructuredPropertyModal.tsx

@@ -1,11 +1,16 @@
+import { Button } from '@src/alchemy-components';
 import analytics, { EventType } from '@src/app/analytics';
+import { ModalButtonContainer } from '@src/app/shared/button/styledComponents';
 import { Modal, message } from 'antd';
 import React, { useEffect, useMemo } from 'react';
 import styled from 'styled-components';
-import { Button } from '@src/alchemy-components';
-import { ModalButtonContainer } from '@src/app/shared/button/styledComponents';
 import { useUpsertStructuredPropertiesMutation } from '../../../../../../graphql/structuredProperties.generated';
-import { EntityType, PropertyValueInput, StructuredPropertyEntity } from '../../../../../../types.generated';
+import {
+    EntityType,
+    PropertyValueInput,
+    StdDataType,
+    StructuredPropertyEntity,
+} from '../../../../../../types.generated';
 import handleGraphQLError from '../../../../../shared/handleGraphQLError';
 import StructuredPropertyInput from '../../../components/styled/StructuredProperty/StructuredPropertyInput';
 import { useEditStructuredProperty } from '../../../components/styled/StructuredProperty/useEditStructuredProperty';
@@ -27,6 +32,9 @@ interface Props {
     isAddMode?: boolean;
 }
 
+const SEARCH_SELECT_MODAL_WIDTH = 1400;
+const DEFAULT_MODAL_WIDTH = 650;
+
 export default function EditStructuredPropertyModal({
     isOpen,
     structuredProperty,
@@ -44,6 +52,7 @@ export default function EditStructuredPropertyModal({
     const { selectedValues, selectSingleValue, toggleSelectedValue, updateSelectedValues, setSelectedValues } =
         useEditStructuredProperty(initialValues);
     const [upsertStructuredProperties] = useUpsertStructuredPropertiesMutation();
+    const { allowedValues } = structuredProperty.definition;
 
     useEffect(() => {
         setSelectedValues(initialValues);
@@ -99,12 +108,15 @@ export default function EditStructuredPropertyModal({
             });
     }
 
+    const isUrnInput = structuredProperty.definition.valueType.info.type === StdDataType.Urn && !allowedValues;
+
     return (
         <Modal
             title={`${isAddMode ? 'Add property' : 'Edit property'} ${structuredProperty?.definition?.displayName}`}
             onCancel={closeModal}
             open={isOpen}
-            width={650}
+            // Urn input is a special case that requires a wider modal since it employs a search select component
+            width={isUrnInput ? SEARCH_SELECT_MODAL_WIDTH : DEFAULT_MODAL_WIDTH}
             footer={
                 <ModalButtonContainer>
                     <Button variant="text" onClick={closeModal} color="gray">
@@ -125,6 +137,7 @@ export default function EditStructuredPropertyModal({
                 <Description>{structuredProperty.definition.description}</Description>
             )}
             <StructuredPropertyInput
+                canUseSearchSelectUrnInput
                 structuredProperty={structuredProperty}
                 selectedValues={selectedValues}
                 selectSingleValue={selectSingleValue}

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

@@ -0,0 +1,161 @@
+import { Icon } 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 { useHydratedEntityMap } from '@src/app/entityV2/shared/tabs/Properties/useHydratedEntityMap';
+import { EntityLink } from '@src/app/homeV2/reference/sections/EntityLink';
+import { EntityType } from '@src/types.generated';
+import { Skeleton } from 'antd';
+import React, { useEffect, useMemo, useState } from 'react';
+import styled from 'styled-components';
+
+const Container = styled.div`
+    display: flex;
+    width: 100%;
+    min-height: 500px;
+    height: 70vh;
+    border: 1px solid #f0f0f0;
+    border-radius: 4px;
+`;
+
+const SearchSection = styled.div`
+    flex: 2;
+    height: 100%;
+    overflow: hidden;
+    display: flex;
+    flex-direction: column;
+`;
+
+const SubSearchSection = styled.div`
+    flex: 1;
+    height: 100%;
+    overflow: hidden;
+    display: flex;
+    flex-direction: column;
+`;
+
+const CurrentSection = styled.div`
+    flex: 1;
+    width: 40%;
+    border-left: 1px solid #f0f0f0;
+    display: flex;
+    flex-direction: column;
+`;
+
+const SectionHeader = styled.div`
+    padding-left: 20px;
+    margin-top: 10px;
+    font-size: 16px;
+    font-weight: 500;
+    color: #666;
+`;
+
+const ScrollableContent = styled.div`
+    flex: 1;
+    overflow: auto;
+    padding: 20px;
+`;
+
+const SelectedItem = styled.div`
+    display: flex;
+    padding: 8px;
+    border-radius: 4px;
+    margin-bottom: 8px;
+    align-items: center;
+    border: 1px solid #f0f0f0;
+    justify-content: space-between;
+
+    &:hover {
+        background-color: #fafafa;
+    }
+`;
+
+const IconWrapper = styled.div`
+    cursor: pointer;
+`;
+
+const INITIAL_SELECTED_ENTITIES = [] as EntityAndType[];
+
+// New interface for the generic component
+interface SearchSelectUrnInputProps {
+    allowedEntityTypes: EntityType[];
+    isMultiple: boolean;
+    selectedValues: string[];
+    updateSelectedValues: (values: string[] | number[]) => void;
+}
+
+// Generic component that doesn't depend on StructuredProperty
+export function SearchSelectUrnInput({
+    allowedEntityTypes,
+    isMultiple,
+    selectedValues,
+    updateSelectedValues,
+}: SearchSelectUrnInputProps) {
+    const [tempSelectedEntities, setTempSelectedEntities] = useState<EntityAndType[]>(INITIAL_SELECTED_ENTITIES);
+
+    // Convert the selected values (urns) to EntityAndType format for SearchSelect
+    const selectedEntities = useMemo(
+        () => selectedValues.map((urn) => ({ urn, type: extractTypeFromUrn(urn) })),
+        [selectedValues],
+    );
+
+    // only run this once
+    useEffect(() => {
+        if (tempSelectedEntities === INITIAL_SELECTED_ENTITIES) {
+            setTempSelectedEntities(selectedEntities);
+        }
+    }, [selectedEntities, setTempSelectedEntities, tempSelectedEntities]);
+
+    // call updateSelectedValues when tempSelectedEntities changes
+    useEffect(() => {
+        if (tempSelectedEntities !== INITIAL_SELECTED_ENTITIES) {
+            updateSelectedValues(tempSelectedEntities.map((entity) => entity.urn));
+        }
+        // this is excluded from the dependency array because updateSelectedValues is reconstructed
+        // by the parent component on every render
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [tempSelectedEntities]);
+
+    const removeEntity = (entity: EntityAndType) => {
+        setTempSelectedEntities(tempSelectedEntities.filter((e) => e.urn !== entity.urn));
+    };
+
+    const hydratedEntityMap = useHydratedEntityMap(tempSelectedEntities.map((entity) => entity.urn));
+
+    return (
+        <Container>
+            <SearchSection>
+                <SectionHeader>Search Values</SectionHeader>
+                <SubSearchSection>
+                    <SearchSelect
+                        fixedEntityTypes={allowedEntityTypes}
+                        selectedEntities={tempSelectedEntities}
+                        setSelectedEntities={setTempSelectedEntities}
+                        limit={isMultiple ? undefined : 1}
+                    />
+                </SubSearchSection>
+            </SearchSection>
+            <CurrentSection>
+                <SectionHeader>Selected Values</SectionHeader>
+                <ScrollableContent>
+                    {tempSelectedEntities.length === 0 ? (
+                        <div>No values selected</div>
+                    ) : (
+                        tempSelectedEntities.map((entity) => (
+                            <SelectedItem key={entity.urn}>
+                                {hydratedEntityMap[entity.urn] ? (
+                                    <EntityLink entity={hydratedEntityMap[entity.urn]} />
+                                ) : (
+                                    <Skeleton.Input active />
+                                )}
+                                <IconWrapper>
+                                    <Icon icon="X" source="phosphor" onClick={() => removeEntity(entity)} />
+                                </IconWrapper>
+                            </SelectedItem>
+                        ))
+                    )}
+                </ScrollableContent>
+            </CurrentSection>
+        </Container>
+    );
+}

datahub-web-react/src/app/entityV2/shared/components/styled/search/tests/SearchSelectUrnInput.test.tsx

@@ -0,0 +1,194 @@
+import { MockedProvider } from '@apollo/client/testing';
+import { act, fireEvent, render, screen } from '@testing-library/react';
+import React from 'react';
+import { EntityType } from '../../../../../../../types.generated';
+import TestPageContainer from '../../../../../../../utils/test-utils/TestPageContainer';
+import { EntityAndType } from '../../../../../../entity/shared/types';
+import { SearchSelectUrnInput } from '../SearchSelectUrnInput';
+
+// Mock the useHydratedEntityMap hook
+vi.mock('@src/app/entityV2/shared/tabs/Properties/useHydratedEntityMap', () => ({
+    useHydratedEntityMap: (urns: string[]) => {
+        const map: { [key: string]: any } = {};
+        urns?.forEach((urn) => {
+            // Simple mock entity structure
+            map[urn] = {
+                urn,
+                type: EntityType.Dataset, // Assuming dataset for simplicity
+                properties: { name: `Entity ${urn.split(':').pop()}` },
+                platform: {
+                    urn: 'urn:li:dataPlatform:test',
+                    name: 'test-platform',
+                    properties: {
+                        displayName: 'Test Platform',
+                        logoUrl: '',
+                    },
+                },
+                // Add other necessary fields if EntityLink mock needs them
+            };
+        });
+        return map;
+    },
+}));
+
+// Mock SearchSelect component - capture the setSelectedEntities callback
+const mockSetSelectedEntities = vi.fn(); // Spy will now store the *callback itself* when called
+vi.mock('@src/app/entityV2/shared/components/styled/search/SearchSelect', () => ({
+    SearchSelect: (props: { setSelectedEntities: (entities: EntityAndType[]) => void }) => {
+        // Call the spy with the actual callback function when the mock renders
+        if (props.setSelectedEntities) {
+            // Check if the prop exists before calling
+            mockSetSelectedEntities(props.setSelectedEntities);
+        }
+        return <div data-testid="mock-search-select">Mock SearchSelect</div>;
+    },
+}));
+
+// Mock EntityLink component
+vi.mock('@src/app/homeV2/reference/sections/EntityLink', () => ({
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    EntityLink: ({ entity }: { entity: any }) => (
+        <div data-testid={`entity-link-${entity.urn}`}>{entity.properties?.name || entity.urn}</div>
+    ),
+}));
+
+// Mock Icon component used for removal
+vi.mock('@src/alchemy-components', async () => {
+    // Import the actual module
+    const actual = await vi.importActual<typeof import('@src/alchemy-components')>('@src/alchemy-components');
+
+    return {
+        // Spread the actual module exports
+        ...actual,
+        // Override specific components with mocks
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        Icon: (props: { icon: string; onClick?: () => void; [key: string]: any }) => (
+            <button type="button" data-testid={`mock-icon-${props.icon.toLowerCase()}`} onClick={props.onClick}>
+                {props.icon}
+            </button>
+        ),
+        Tooltip: ({ children }: { children: React.ReactNode }) => <>{children}</>,
+        Button: ({ children, onClick }: { children: React.ReactNode; onClick?: () => void }) => (
+            <button type="button" onClick={onClick}>
+                {children}
+            </button>
+        ),
+        // No need to mock colors and typography anymore, we use the actual ones
+    };
+});
+
+describe('SearchSelectUrnInput', () => {
+    const mockUpdateSelectedValues = vi.fn();
+
+    beforeEach(() => {
+        mockUpdateSelectedValues.mockClear();
+        mockSetSelectedEntities.mockClear();
+    });
+
+    it('renders correctly with no initial selection', () => {
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <SearchSelectUrnInput
+                        allowedEntityTypes={[EntityType.Dataset]}
+                        isMultiple
+                        selectedValues={[]}
+                        updateSelectedValues={mockUpdateSelectedValues}
+                    />
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        expect(screen.getByText('Search Values')).toBeInTheDocument();
+        expect(screen.getByTestId('mock-search-select')).toBeInTheDocument();
+        expect(screen.getByText('Selected Values')).toBeInTheDocument();
+        expect(screen.getByText('No values selected')).toBeInTheDocument();
+    });
+
+    it('renders with initial selected values', () => {
+        const initialUrns = ['urn:li:dataset:1', 'urn:li:dataset:2'];
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <SearchSelectUrnInput
+                        allowedEntityTypes={[EntityType.Dataset]}
+                        isMultiple
+                        selectedValues={initialUrns}
+                        updateSelectedValues={mockUpdateSelectedValues}
+                    />
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        expect(screen.getByText('Selected Values')).toBeInTheDocument();
+        expect(screen.getByTestId('entity-link-urn:li:dataset:1')).toHaveTextContent('Entity 1');
+        expect(screen.getByTestId('entity-link-urn:li:dataset:2')).toHaveTextContent('Entity 2');
+        expect(screen.queryByText('No values selected')).not.toBeInTheDocument();
+    });
+
+    it('calls updateSelectedValues when a value is removed', () => {
+        const initialUrns = ['urn:li:dataset:1', 'urn:li:dataset:2'];
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <SearchSelectUrnInput
+                        allowedEntityTypes={[EntityType.Dataset]}
+                        isMultiple
+                        selectedValues={initialUrns}
+                        updateSelectedValues={mockUpdateSelectedValues}
+                    />
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        // Find the remove button for the first entity
+        const removeButton = screen.getAllByTestId('mock-icon-x')[0]; // Assuming order matches initialUrns
+        fireEvent.click(removeButton);
+
+        // Check if updateSelectedValues was called with the remaining URN
+        // It gets called twice: once on initial load (useEffect) and once on removal
+        expect(mockUpdateSelectedValues).toHaveBeenCalledTimes(2);
+        expect(mockUpdateSelectedValues).toHaveBeenLastCalledWith(['urn:li:dataset:2']);
+
+        // Check if the removed entity is no longer rendered
+        expect(screen.queryByTestId('entity-link-urn:li:dataset:1')).not.toBeInTheDocument();
+        expect(screen.getByTestId('entity-link-urn:li:dataset:2')).toBeInTheDocument();
+    });
+
+    it('calls updateSelectedValues when SearchSelect updates selection', () => {
+        const initialUrns: string[] = [];
+        const newEntity: EntityAndType = { urn: 'urn:li:dataset:3', type: EntityType.Dataset };
+        render(
+            <MockedProvider>
+                <TestPageContainer>
+                    <SearchSelectUrnInput
+                        allowedEntityTypes={[EntityType.Dataset]}
+                        isMultiple
+                        selectedValues={initialUrns}
+                        updateSelectedValues={mockUpdateSelectedValues}
+                    />
+                </TestPageContainer>
+            </MockedProvider>,
+        );
+
+        // Ensure the SearchSelect mock rendered and called our spy with the callback
+        expect(mockSetSelectedEntities).toHaveBeenCalled();
+
+        // Get the actual callback function from the spy's arguments
+        const capturedSetSelectedEntities = mockSetSelectedEntities.mock.calls[0][0];
+        expect(capturedSetSelectedEntities).toBeInstanceOf(Function); // Verify it's a function
+
+        // Simulate SearchSelect calling its setSelectedEntities prop within act
+        act(() => {
+            capturedSetSelectedEntities([newEntity]);
+        });
+
+        // Check if updateSelectedValues was called with the new URN
+        // It gets called once on initial load (useEffect) and once when the selection changes
+        expect(mockUpdateSelectedValues).toHaveBeenCalledTimes(2);
+        expect(mockUpdateSelectedValues).toHaveBeenLastCalledWith([newEntity.urn]);
+
+        // Check if the new entity is rendered
+        expect(screen.getByTestId('entity-link-urn:li:dataset:3')).toHaveTextContent('Entity 3');
+    });
+});