feat(data platform) adding data platform indexing & select platform modal in frontend (#5988)

* adding data platform indexing & select platform modal in frontend

* comments

datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Platform/SelectPlatformModal.tsx

@@ -0,0 +1,186 @@
+import { Button, Form, Modal, Select, Tag, Tooltip } from 'antd';
+import React, { ReactNode, useRef, useState } from 'react';
+import styled from 'styled-components/macro';
+import { useGetSearchResultsLazyQuery } from '../../../../../../../graphql/search.generated';
+import { DataPlatform, Entity, EntityType } from '../../../../../../../types.generated';
+import { useEnterKeyListener } from '../../../../../../shared/useEnterKeyListener';
+
+type Props = {
+    onCloseModal: () => void;
+    defaultValues?: { urn: string; entity?: Entity | null }[];
+    onOkOverride?: (result: string[]) => void;
+    titleOverride?: string;
+};
+
+type SelectedPlatform = {
+    entity?: Entity | null;
+    urn: string;
+};
+
+const StyleTag = styled(Tag)`
+    padding: 0px 7px;
+    margin-right: 3px;
+    display: flex;
+    justify-content: start;
+    align-items: center;
+`;
+
+const PreviewImage = styled.img`
+    max-height: 18px;
+    width: auto;
+    object-fit: contain;
+    background-color: transparent;
+    margin-right: 4px;
+`;
+
+export const SelectPlatformModal = ({ onCloseModal, defaultValues, onOkOverride, titleOverride }: Props) => {
+    const [platformSearch, { data: platforSearchData }] = useGetSearchResultsLazyQuery();
+    const platformSearchResults =
+        platforSearchData?.search?.searchResults?.map((searchResult) => searchResult.entity) || [];
+
+    const [selectedPlatforms, setSelectedPlatforms] = useState<SelectedPlatform[] | undefined>(defaultValues);
+
+    const inputEl = useRef(null);
+
+    const onModalClose = () => {
+        onCloseModal();
+    };
+
+    const handleSearch = (text: string) => {
+        platformSearch({
+            variables: {
+                input: {
+                    type: EntityType.DataPlatform,
+                    query: text,
+                    start: 0,
+                    count: 5,
+                },
+            },
+        });
+    };
+
+    // Renders a search result in the select dropdown.
+    const renderSearchResult = (entity: DataPlatform) => {
+        const displayName = entity.properties?.displayName || entity.name;
+        const truncatedDisplayName = displayName.length > 25 ? `${displayName.slice(0, 25)}...` : displayName;
+        return (
+            <Tooltip title={displayName}>
+                {!!entity.properties?.logoUrl && <PreviewImage src={entity?.properties?.logoUrl} alt={entity?.name} />}
+                <span>{truncatedDisplayName}</span>
+            </Tooltip>
+        );
+    };
+
+    const platformSearchOptions = platformSearchResults?.map((result) => {
+        return (
+            <Select.Option value={result.urn} key={result.urn}>
+                {renderSearchResult(result as DataPlatform)}
+            </Select.Option>
+        );
+    });
+
+    const onSelectPlatform = (newValue: { value: string; label: ReactNode }) => {
+        const newUrn = newValue.value;
+
+        if (inputEl && inputEl.current) {
+            (inputEl.current as any).blur();
+        }
+
+        const filteredPlatforms =
+            platformSearchResults?.filter((entity) => entity.urn === newUrn).map((entity) => entity) || [];
+
+        if (filteredPlatforms.length) {
+            const platform = filteredPlatforms[0] as DataPlatform;
+            setSelectedPlatforms([
+                ...(selectedPlatforms || []),
+                {
+                    entity: platform,
+                    urn: newUrn,
+                },
+            ]);
+        }
+    };
+
+    const onDeselectPlatform = (val) => {
+        setSelectedPlatforms(selectedPlatforms?.filter((platform) => platform.urn !== val));
+    };
+
+    const onOk = async () => {
+        if (!selectedPlatforms) {
+            return;
+        }
+
+        if (onOkOverride) {
+            onOkOverride(selectedPlatforms?.map((platform) => platform.urn));
+        }
+    };
+
+    // Handle the Enter press
+    useEnterKeyListener({
+        querySelectorToExecuteClick: '#setPlatformButton',
+    });
+
+    const tagRender = (props) => {
+        // eslint-disable-next-line react/prop-types
+        const { label, closable, onClose } = props;
+        const onPreventMouseDown = (event) => {
+            event.preventDefault();
+            event.stopPropagation();
+        };
+        return (
+            <StyleTag onMouseDown={onPreventMouseDown} closable={closable} onClose={onClose}>
+                {label}
+            </StyleTag>
+        );
+    };
+
+    return (
+        <Modal
+            title={titleOverride || 'Select Platform'}
+            visible
+            onCancel={onModalClose}
+            footer={
+                <>
+                    <Button onClick={onModalClose} type="text">
+                        Cancel
+                    </Button>
+                    <Button id="setPlatformButton" disabled={selectedPlatforms?.length === 0} onClick={onOk}>
+                        Add
+                    </Button>
+                </>
+            }
+        >
+            <Form component={false}>
+                <Form.Item>
+                    <Select
+                        autoFocus
+                        filterOption={false}
+                        showSearch
+                        mode="multiple"
+                        defaultActiveFirstOption={false}
+                        placeholder="Search for Platforms..."
+                        onSelect={(platformUrn: any) => onSelectPlatform(platformUrn)}
+                        onDeselect={onDeselectPlatform}
+                        onSearch={(value: string) => {
+                            // eslint-disable-next-line react/prop-types
+                            handleSearch(value.trim());
+                        }}
+                        ref={inputEl}
+                        labelInValue
+                        value={selectedPlatforms?.map((platform) => ({
+                            value: platform.urn,
+                            label: platform.entity ? (
+                                renderSearchResult(platform.entity as DataPlatform)
+                            ) : (
+                                <span>{platform.urn}</span>
+                            ),
+                        }))}
+                        tagRender={tagRender}
+                    >
+                        {platformSearchOptions}
+                    </Select>
+                </Form.Item>
+            </Form>
+        </Modal>
+    );
+};

datahub-web-react/src/graphql/search.graphql

@@ -162,6 +162,9 @@ fragment autoCompleteFields on Entity {
             ...dataPlatformInstanceFields
         }
     }
+    ... on DataPlatform {
+        ...nonConflictingPlatformFields
+    }
 }
 
 query getAutoCompleteResults($input: AutoCompleteInput!) {
@@ -672,6 +675,9 @@ fragment searchResultFields on Entity {
         }
         description
     }
+    ... on DataPlatform {
+        ...nonConflictingPlatformFields
+    }
 }
 
 fragment facetFields on FacetMetadata {

metadata-models/src/main/pegasus/com/linkedin/dataplatform/DataPlatformInfo.pdl

@@ -14,11 +14,21 @@ record DataPlatformInfo {
    * Name of the data platform
    */
   @validate.strlen.max = 15
+  @Searchable = {
+    "fieldType": "TEXT_PARTIAL",
+    "enableAutocomplete": false,
+    "boostScore": 10.0
+  }
   name: string
 
   /**
    * The name that will be used for displaying a platform type.
    */
+  @Searchable = {
+    "fieldType": "TEXT_PARTIAL",
+    "enableAutocomplete": true,
+    "boostScore": 10.0
+  }
   displayName: optional string
 
   /**
@@ -85,4 +95,4 @@ record DataPlatformInfo {
    * The URL for a logo associated with the platform
    */
   logoUrl: optional Url
-}
+}

metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/BootstrapManagerFactory.java

@@ -7,6 +7,7 @@ import com.linkedin.gms.factory.search.EntitySearchServiceFactory;
 import com.linkedin.gms.factory.search.SearchDocumentTransformerFactory;
 import com.linkedin.metadata.boot.BootstrapManager;
 import com.linkedin.metadata.boot.BootstrapStep;
+import com.linkedin.metadata.boot.steps.IndexDataPlatformsStep;
 import com.linkedin.metadata.boot.steps.IngestDataPlatformInstancesStep;
 import com.linkedin.metadata.boot.steps.IngestDataPlatformsStep;
 import com.linkedin.metadata.boot.steps.IngestPoliciesStep;
@@ -79,13 +80,15 @@ public class BootstrapManagerFactory {
         new IngestDataPlatformInstancesStep(_entityService, _migrationsDao);
     final RestoreGlossaryIndices restoreGlossaryIndicesStep =
         new RestoreGlossaryIndices(_entityService, _entitySearchService, _entityRegistry);
+    final IndexDataPlatformsStep indexDataPlatformsStep =
+        new IndexDataPlatformsStep(_entityService, _entitySearchService, _entityRegistry);
     final RestoreDbtSiblingsIndices restoreDbtSiblingsIndices =
         new RestoreDbtSiblingsIndices(_entityService, _entityRegistry);
     final RemoveClientIdAspectStep removeClientIdAspectStep = new RemoveClientIdAspectStep(_entityService);
 
     final List<BootstrapStep> finalSteps = new ArrayList<>(ImmutableList.of(ingestRootUserStep, ingestPoliciesStep, ingestRolesStep,
         ingestDataPlatformsStep, ingestDataPlatformInstancesStep, _ingestRetentionPoliciesStep, restoreGlossaryIndicesStep,
-        removeClientIdAspectStep, restoreDbtSiblingsIndices));
+        removeClientIdAspectStep, restoreDbtSiblingsIndices, indexDataPlatformsStep));
 
     if (_upgradeDefaultBrowsePathsEnabled) {
       finalSteps.add(new UpgradeDefaultBrowsePathsStep(_entityService));

metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IndexDataPlatformsStep.java

@@ -0,0 +1,113 @@
+package com.linkedin.metadata.boot.steps;
+
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.dataplatform.DataPlatformInfo;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspectMap;
+import com.linkedin.events.metadata.ChangeType;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.boot.UpgradeStep;
+import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.models.AspectSpec;
+import com.linkedin.metadata.models.registry.EntityRegistry;
+import com.linkedin.metadata.query.ListUrnsResult;
+import com.linkedin.metadata.search.EntitySearchService;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import lombok.extern.slf4j.Slf4j;
+
+
+@Slf4j
+public class IndexDataPlatformsStep extends UpgradeStep {
+  private static final String VERSION = "1";
+  private static final String UPGRADE_ID = "index-data-platforms";
+  private static final Integer BATCH_SIZE = 1000;
+
+  private final EntitySearchService _entitySearchService;
+  private final EntityRegistry _entityRegistry;
+
+  public IndexDataPlatformsStep(EntityService entityService, EntitySearchService entitySearchService,
+      EntityRegistry entityRegistry) {
+    super(entityService, VERSION, UPGRADE_ID);
+    _entitySearchService = entitySearchService;
+    _entityRegistry = entityRegistry;
+  }
+
+  @Override
+  public void upgrade() throws Exception {
+    final AspectSpec dataPlatformSpec = _entityRegistry.getEntitySpec(Constants.DATA_PLATFORM_ENTITY_NAME)
+        .getAspectSpec(Constants.DATA_PLATFORM_INFO_ASPECT_NAME);
+
+    final AuditStamp auditStamp =
+        new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis());
+
+    getAndReIndexDataPlatforms(auditStamp, dataPlatformSpec);
+
+    log.info("Successfully indexed data platform aspects");
+  }
+
+  @Nonnull
+  @Override
+  public ExecutionMode getExecutionMode() {
+    return ExecutionMode.ASYNC;
+  }
+
+  private int getAndReIndexDataPlatforms(AuditStamp auditStamp, AspectSpec dataPlatformInfoAspectSpec)
+      throws Exception {
+    ListUrnsResult listResult =
+        _entityService.listUrns(Constants.DATA_PLATFORM_ENTITY_NAME, 0, BATCH_SIZE);
+
+    List<Urn> dataPlatformUrns = listResult.getEntities();
+
+    if (dataPlatformUrns.size() == 0) {
+      return 0;
+    }
+
+    final Map<Urn, EntityResponse> dataPlatformInfoResponses =
+        _entityService.getEntitiesV2(Constants.DATA_PLATFORM_ENTITY_NAME, new HashSet<>(dataPlatformUrns),
+        Collections.singleton(Constants.DATA_PLATFORM_INFO_ASPECT_NAME)
+    );
+
+    //  Loop over Data platforms and produce changelog
+    for (Urn dpUrn : dataPlatformUrns) {
+      EntityResponse dataPlatformEntityResponse = dataPlatformInfoResponses.get(dpUrn);
+      if (dataPlatformEntityResponse == null) {
+        log.warn("Data Platform not in set of entity responses {}", dpUrn);
+        continue;
+      }
+
+      DataPlatformInfo dpInfo = mapDpInfo(dataPlatformEntityResponse);
+      if (dpInfo == null) {
+        log.warn("Received null dataPlatformInfo aspect for urn {}", dpUrn);
+        continue;
+      }
+
+      _entityService.produceMetadataChangeLog(
+          dpUrn,
+          Constants.DATA_PLATFORM_ENTITY_NAME,
+          Constants.DATA_PLATFORM_INFO_ASPECT_NAME,
+          dataPlatformInfoAspectSpec,
+          null,
+          dpInfo,
+          null,
+          null,
+          auditStamp,
+          ChangeType.RESTATE);
+    }
+
+    return listResult.getTotal();
+  }
+
+  private DataPlatformInfo mapDpInfo(EntityResponse entityResponse) {
+    EnvelopedAspectMap aspectMap = entityResponse.getAspects();
+    if (!aspectMap.containsKey(Constants.DATA_PLATFORM_INFO_ASPECT_NAME)) {
+      return null;
+    }
+
+    return new DataPlatformInfo(aspectMap.get(Constants.DATA_PLATFORM_INFO_ASPECT_NAME).getValue().data());
+  }
+}