refactor(ui): Separate entity lineage counts query from rest of entity query (#7569)

datahub-web-react/src/Mocks.tsx

@@ -232,8 +232,6 @@ export const dataset1 = {
     ],
     domain: null,
     container: null,
-    upstream: null,
-    downstream: null,
     health: [],
     assertions: null,
     deprecation: null,
@@ -326,8 +324,6 @@ export const dataset2 = {
     ],
     domain: null,
     container: null,
-    upstream: null,
-    downstream: null,
     health: [],
     assertions: null,
     status: null,
@@ -446,8 +442,6 @@ export const dataset3 = {
     },
     incoming: null,
     outgoing: null,
-    upstream: null,
-    downstream: null,
     institutionalMemory: {
         elements: [
             {

datahub-web-react/src/app/entity/mlFeature/MLFeatureEntity.tsx

@@ -6,7 +6,7 @@ import { Entity, EntityCapabilityType, IconStyleType, PreviewType } from '../Ent
 import { getDataForEntityType } from '../shared/containers/profile/utils';
 import { EntityProfile } from '../shared/containers/profile/EntityProfile';
 import { GenericEntityProperties } from '../shared/types';
-import { GetMlFeatureQuery, useGetMlFeatureQuery } from '../../../graphql/mlFeature.generated';
+import { useGetMlFeatureQuery } from '../../../graphql/mlFeature.generated';
 import { SidebarAboutSection } from '../shared/containers/profile/sidebar/AboutSection/SidebarAboutSection';
 import { SidebarTagsSection } from '../shared/containers/profile/sidebar/SidebarTagsSection';
 import { SidebarOwnerSection } from '../shared/containers/profile/sidebar/Ownership/SidebarOwnerSection';
@@ -82,15 +82,6 @@ export class MLFeatureEntity implements Entity<MlFeature> {
                 {
                     name: 'Lineage',
                     component: LineageTab,
-                    display: {
-                        visible: (_, _1) => true,
-                        enabled: (_, result: GetMlFeatureQuery) => {
-                            return (
-                                (result?.mlFeature?.upstream?.total || 0) > 0 ||
-                                (result?.mlFeature?.downstream?.total || 0) > 0
-                            );
-                        },
-                    },
                 },
             ]}
             sidebarSections={[

datahub-web-react/src/app/entity/mlPrimaryKey/MLPrimaryKeyEntity.tsx

@@ -5,7 +5,7 @@ import { Preview } from './preview/Preview';
 import { Entity, EntityCapabilityType, IconStyleType, PreviewType } from '../Entity';
 import { getDataForEntityType } from '../shared/containers/profile/utils';
 import { GenericEntityProperties } from '../shared/types';
-import { GetMlPrimaryKeyQuery, useGetMlPrimaryKeyQuery } from '../../../graphql/mlPrimaryKey.generated';
+import { useGetMlPrimaryKeyQuery } from '../../../graphql/mlPrimaryKey.generated';
 import { EntityProfile } from '../shared/containers/profile/EntityProfile';
 import { FeatureTableTab } from '../shared/tabs/ML/MlPrimaryKeyFeatureTableTab';
 import { DocumentationTab } from '../shared/tabs/Documentation/DocumentationTab';
@@ -80,15 +80,6 @@ export class MLPrimaryKeyEntity implements Entity<MlPrimaryKey> {
                 {
                     name: 'Lineage',
                     component: LineageTab,
-                    display: {
-                        visible: (_, _1) => true,
-                        enabled: (_, result: GetMlPrimaryKeyQuery) => {
-                            return (
-                                (result?.mlPrimaryKey?.upstream?.total || 0) > 0 ||
-                                (result?.mlPrimaryKey?.downstream?.total || 0) > 0
-                            );
-                        },
-                    },
                 },
             ]}
             sidebarSections={[

datahub-web-react/src/app/entity/shared/containers/profile/nav/EntityProfileNavBar.tsx

@@ -1,10 +1,9 @@
-import { Affix } from 'antd';
 import React from 'react';
 import styled from 'styled-components/macro';
+import { Affix } from 'antd';
 import { useGetBrowsePathsQuery } from '../../../../../../graphql/browse.generated';
 import { EntityType } from '../../../../../../types.generated';
 import { useEntityRegistry } from '../../../../../useEntityRegistry';
-import { useLineageData } from '../../../EntityContext';
 import { ProfileNavBrowsePath } from './ProfileNavBrowsePath';
 
 type Props = {
@@ -20,18 +19,15 @@ export const EntityProfileNavBar = ({ urn, entityType }: Props) => {
         fetchPolicy: 'cache-first',
     });
     const entityRegistry = useEntityRegistry();
-
     const isBrowsable = entityRegistry.getBrowseEntityTypes().includes(entityType);
-    const lineage = useLineageData();
 
     return (
         <AffixWithHeight offsetTop={60}>
             <ProfileNavBrowsePath
-                breadcrumbLinksEnabled={isBrowsable}
+                urn={urn}
                 type={entityType}
+                breadcrumbLinksEnabled={isBrowsable}
                 path={browseData?.browsePaths?.[0]?.path || []}
-                upstreams={lineage?.numUpstreamChildren || 0}
-                downstreams={lineage?.numDownstreamChildren || 0}
             />
         </AffixWithHeight>
     );

datahub-web-react/src/app/entity/shared/containers/profile/nav/LineageSelector.tsx

@@ -0,0 +1,157 @@
+import React from 'react';
+import styled from 'styled-components/macro';
+import { useHistory, useLocation } from 'react-router-dom';
+import { Badge } from 'antd';
+import { InfoCircleOutlined, PartitionOutlined } from '@ant-design/icons';
+import { grey, blue } from '@ant-design/colors';
+import { EntityType } from '../../../../../../types.generated';
+import { navigateToLineageUrl } from '../../../../../lineage/utils/navigateToLineageUrl';
+import { ANTD_GRAY, ENTITY_TYPES_WITH_MANUAL_LINEAGE } from '../../../constants';
+import useIsLineageMode from '../../../../../lineage/utils/useIsLineageMode';
+import { useGetLineageTimeParams } from '../../../../../lineage/utils/useGetLineageTimeParams';
+import { useIsSeparateSiblingsMode } from '../../../siblingUtils';
+import { useGetLineageCountsQuery } from '../../../../../../graphql/lineage.generated';
+
+const LineageIconGroup = styled.div`
+    width: 180px;
+    display: flex;
+    justify-content: space-between;
+    margin-right: 8px;
+`;
+
+const LineageIcon = styled(PartitionOutlined)`
+    font-size: 20px;
+    vertical-align: middle;
+    padding-right: 6px;
+`;
+
+const DetailIcon = styled(InfoCircleOutlined)`
+    font-size: 20px;
+    vertical-align: middle;
+    padding-right: 6px;
+`;
+
+const IconGroup = styled.div<{ isSelected: boolean; disabled?: boolean }>`
+    font-size: 14px;
+    color: ${(props) => {
+        if (props.disabled) {
+            return grey[2];
+        }
+        return !props.isSelected ? 'black' : props.theme.styles['primary-color'] || blue[4];
+    }};
+    &:hover {
+        color: ${(props) => (props.disabled ? grey[2] : props.theme.styles['primary-color'] || blue[4])};
+        cursor: ${(props) => (props.disabled ? 'not-allowed' : 'pointer')};
+    }
+`;
+
+const LineageNavContainer = styled.div`
+    display: inline-flex;
+    line-height: 24px;
+    align-items: center;
+`;
+
+const LineageSummary = styled.div`
+    margin-left: 16px;
+`;
+
+const LineageBadge = styled(Badge)`
+    &&& .ant-badge-count {
+        background-color: ${ANTD_GRAY[1]};
+        color: ${ANTD_GRAY[9]};
+        border: 1px solid ${ANTD_GRAY[5]};
+        font-size: 12px;
+        font-weight: 600;
+        height: 22px;
+    }
+`;
+
+type Props = {
+    urn: string;
+    type: EntityType;
+};
+
+/**
+ * Responsible for rendering a clickable browse path view.
+ */
+export const LineageSelector = ({ urn, type }: Props): JSX.Element => {
+    const history = useHistory();
+    const location = useLocation();
+    const isLineageMode = useIsLineageMode();
+    const isHideSiblingsMode = useIsSeparateSiblingsMode();
+    const { startTimeMillis, endTimeMillis } = useGetLineageTimeParams();
+
+    // Fetch the lineage counts for the entity.
+    const { data: lineageData, loading: lineageLoading } = useGetLineageCountsQuery({
+        variables: {
+            urn,
+            separateSiblings: isHideSiblingsMode,
+            startTimeMillis,
+            endTimeMillis,
+        },
+    });
+
+    const upstreamTotal = (lineageData?.entity as any)?.upstream?.total || 0;
+    const upstreamFiltered = (lineageData?.entity as any)?.upstream?.filtered || 0;
+    const upstreamCount = upstreamTotal - upstreamFiltered;
+
+    const downstreamTotal = (lineageData?.entity as any)?.downstream?.total || 0;
+    const downstreamFiltered = (lineageData?.entity as any)?.downstream?.filtered || 0;
+    const downstreamCount = downstreamTotal - downstreamFiltered;
+
+    const hasLineage = upstreamCount > 0 || downstreamCount > 0;
+    const canNavigateToLineage = hasLineage || ENTITY_TYPES_WITH_MANUAL_LINEAGE.has(type);
+
+    const upstreamText = upstreamCount === 100 ? '100+' : upstreamCount;
+    const downstreamText = downstreamCount === 100 ? '100+' : downstreamCount;
+
+    return (
+        <LineageNavContainer>
+            <LineageIconGroup>
+                <IconGroup
+                    disabled={!canNavigateToLineage}
+                    isSelected={!isLineageMode}
+                    onClick={() => {
+                        if (canNavigateToLineage) {
+                            navigateToLineageUrl({
+                                location,
+                                history,
+                                isLineageMode: false,
+                                startTimeMillis,
+                                endTimeMillis,
+                            });
+                        }
+                    }}
+                >
+                    <DetailIcon />
+                    Details
+                </IconGroup>
+                <IconGroup
+                    disabled={!canNavigateToLineage}
+                    isSelected={isLineageMode}
+                    onClick={() => {
+                        if (canNavigateToLineage) {
+                            navigateToLineageUrl({
+                                location,
+                                history,
+                                isLineageMode: true,
+                                startTimeMillis,
+                                endTimeMillis,
+                            });
+                        }
+                    }}
+                >
+                    <LineageIcon />
+                    Lineage
+                </IconGroup>
+            </LineageIconGroup>
+            <LineageSummary>
+                <LineageBadge
+                    count={`${lineageLoading ? '-' : upstreamText} upstream, ${
+                        lineageLoading ? '-' : downstreamText
+                    } downstream`}
+                />
+            </LineageSummary>
+        </LineageNavContainer>
+    );
+};

datahub-web-react/src/app/entity/shared/containers/profile/nav/ProfileNavBrowsePath.tsx

@@ -1,57 +1,12 @@
 import React from 'react';
-import { Link, useHistory, useLocation } from 'react-router-dom';
+import { Link } from 'react-router-dom';
-import { Badge, Breadcrumb, Row } from 'antd';
 import styled from 'styled-components/macro';
-import { InfoCircleOutlined, PartitionOutlined } from '@ant-design/icons';
+import { Breadcrumb, Row } from 'antd';
-import { grey, blue } from '@ant-design/colors';
 import { EntityType } from '../../../../../../types.generated';
 import { useEntityRegistry } from '../../../../../useEntityRegistry';
 import { PageRoutes } from '../../../../../../conf/Global';
-import { navigateToLineageUrl } from '../../../../../lineage/utils/navigateToLineageUrl';
+import { ANTD_GRAY } from '../../../constants';
-import useIsLineageMode from '../../../../../lineage/utils/useIsLineageMode';
+import { LineageSelector } from './LineageSelector';
-import { ANTD_GRAY, ENTITY_TYPES_WITH_MANUAL_LINEAGE } from '../../../constants';
-import { useGetLineageTimeParams } from '../../../../../lineage/utils/useGetLineageTimeParams';
-
-type Props = {
-    type: EntityType;
-    path: Array<string>;
-    upstreams: number;
-    downstreams: number;
-    breadcrumbLinksEnabled: boolean;
-};
-
-const LineageIconGroup = styled.div`
-    width: 180px;
-    display: flex;
-    justify-content: space-between;
-    margin-right: 8px;
-`;
-
-const LineageIcon = styled(PartitionOutlined)`
-    font-size: 20px;
-    vertical-align: middle;
-    padding-right: 6px;
-`;
-
-const DetailIcon = styled(InfoCircleOutlined)`
-    font-size: 20px;
-    vertical-align: middle;
-    padding-right: 6px;
-`;
-
-const IconGroup = styled.div<{ isSelected: boolean; disabled?: boolean }>`
-    font-size: 14px;
-    color: ${(props) => {
-        if (props.disabled) {
-            return grey[2];
-        }
-        return !props.isSelected ? 'black' : props.theme.styles['primary-color'] || blue[4];
-    }};
-    &:hover {
-        color: ${(props) => (props.disabled ? grey[2] : props.theme.styles['primary-color'] || blue[4])};
-        cursor: ${(props) => (props.disabled ? 'not-allowed' : 'pointer')};
-    }
-`;
 
 const BrowseRow = styled(Row)`
     padding: 10px 20px;
@@ -61,49 +16,24 @@ const BrowseRow = styled(Row)`
     border-bottom: 1px solid ${ANTD_GRAY[4.5]};
 `;
 
-const LineageNavContainer = styled.div`
-    display: inline-flex;
-    line-height: 24px;
-    align-items: center;
-`;
-
-const LineageSummary = styled.div`
-    margin-left: 16px;
-`;
-
-const LineageBadge = styled(Badge)`
-    &&& .ant-badge-count {
-        background-color: ${ANTD_GRAY[1]};
-        color: ${ANTD_GRAY[9]};
-        border: 1px solid ${ANTD_GRAY[5]};
-        font-size: 12px;
-        font-weight: 600;
-        height: 22px;
-    }
-`;
-
 export const BreadcrumbItem = styled(Breadcrumb.Item)<{ disabled?: boolean }>`
     &&& :hover {
         color: ${(props) => (props.disabled ? ANTD_GRAY[7] : props.theme.styles['primary-color'])};
     }
 `;
 
+type Props = {
+    urn: string;
+    type: EntityType;
+    path: Array<string>;
+    breadcrumbLinksEnabled: boolean;
+};
+
 /**
  * Responsible for rendering a clickable browse path view.
  */
-// TODO(Gabe): use this everywhere
+export const ProfileNavBrowsePath = ({ urn, type, path, breadcrumbLinksEnabled }: Props): JSX.Element => {
-export const ProfileNavBrowsePath = ({
-    type,
-    path,
-    upstreams,
-    downstreams,
-    breadcrumbLinksEnabled,
-}: Props): JSX.Element => {
     const entityRegistry = useEntityRegistry();
-    const history = useHistory();
-    const location = useLocation();
-    const isLineageMode = useIsLineageMode();
-    const { startTimeMillis, endTimeMillis } = useGetLineageTimeParams();
 
     const createPartialPath = (parts: Array<string>) => {
         return parts.join('/');
@@ -129,12 +59,6 @@ export const ProfileNavBrowsePath = ({
         </BreadcrumbItem>
     ));
 
-    const hasLineage = upstreams > 0 || downstreams > 0;
-    const canNavigateToLineage = hasLineage || ENTITY_TYPES_WITH_MANUAL_LINEAGE.has(type);
-
-    const upstreamText = upstreams === 100 ? '100+' : upstreams;
-    const downstreamText = downstreams === 100 ? '100+' : downstreams;
-
     return (
         <BrowseRow>
             <Breadcrumb style={{ fontSize: '16px' }} separator=">">
@@ -149,49 +73,7 @@ export const ProfileNavBrowsePath = ({
                 </BreadcrumbItem>
                 {pathCrumbs}
             </Breadcrumb>
-            <LineageNavContainer>
+            <LineageSelector urn={urn} type={type} />
-                <LineageIconGroup>
-                    <IconGroup
-                        disabled={!canNavigateToLineage}
-                        isSelected={!isLineageMode}
-                        onClick={() => {
-                            if (canNavigateToLineage) {
-                                navigateToLineageUrl({
-                                    location,
-                                    history,
-                                    isLineageMode: false,
-                                    startTimeMillis,
-                                    endTimeMillis,
-                                });
-                            }
-                        }}
-                    >
-                        <DetailIcon />
-                        Details
-                    </IconGroup>
-                    <IconGroup
-                        disabled={!canNavigateToLineage}
-                        isSelected={isLineageMode}
-                        onClick={() => {
-                            if (canNavigateToLineage) {
-                                navigateToLineageUrl({
-                                    location,
-                                    history,
-                                    isLineageMode: true,
-                                    startTimeMillis,
-                                    endTimeMillis,
-                                });
-                            }
-                        }}
-                    >
-                        <LineageIcon />
-                        Lineage
-                    </IconGroup>
-                </LineageIconGroup>
-                <LineageSummary>
-                    <LineageBadge count={`${upstreamText} upstream, ${downstreamText} downstream`} />
-                </LineageSummary>
-            </LineageNavContainer>
         </BrowseRow>
     );
 };

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

@@ -64,12 +64,6 @@ query getChart($urn: String!) {
         parentContainers {
             ...parentContainersFields
         }
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
         status {
             removed
         }

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

@@ -7,12 +7,6 @@ query getDashboard($urn: String!) {
         datasets: relationships(input: { types: ["Consumes"], direction: OUTGOING, start: 0, count: 100 }) {
             ...fullRelationshipResults
         }
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
     }
 }
 

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

@@ -1,12 +1,6 @@
 query getDataJob($urn: String!) {
     dataJob(urn: $urn) {
         ...dataJobFields
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
         runs(start: 0, count: 20) {
             count
             start

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

@@ -115,12 +115,6 @@ fragment nonSiblingDatasetFields on Dataset {
         timestampMillis
         lastUpdatedTimestamp
     }
-    upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-        ...partialLineageResults
-    }
-    downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-        ...partialLineageResults
-    }
     ...viewProperties
     autoRenderAspects: aspects(input: { autoRenderOnly: true }) {
         aspectName

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

@@ -413,3 +413,45 @@ fragment canEditLineageFragment on EntityWithRelationships {
         }
     }
 }
+
+query getLineageCounts(
+    $urn: String!
+    $separateSiblings: Boolean
+    $startTimeMillis: Long
+    $endTimeMillis: Long
+    $excludeUpstream: Boolean = false
+    $excludeDownstream: Boolean = false
+) {
+    entity(urn: $urn) {
+        urn
+        type
+        ... on EntityWithRelationships {
+            upstream: lineage(
+                input: {
+                    direction: UPSTREAM
+                    start: 0
+                    count: 100
+                    separateSiblings: $separateSiblings
+                    startTimeMillis: $startTimeMillis
+                    endTimeMillis: $endTimeMillis
+                }
+            ) @skip(if: $excludeUpstream) {
+                filtered
+                total
+            }
+            downstream: lineage(
+                input: {
+                    direction: DOWNSTREAM
+                    start: 0
+                    count: 100
+                    separateSiblings: $separateSiblings
+                    startTimeMillis: $startTimeMillis
+                    endTimeMillis: $endTimeMillis
+                }
+            ) @skip(if: $excludeDownstream) {
+                filtered
+                total
+            }
+        }
+    }
+}

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

@@ -1,12 +1,6 @@
 query getMLFeature($urn: String!) {
     mlFeature(urn: $urn) {
         ...nonRecursiveMLFeature
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
         featureTables: relationships(input: { types: ["Contains"], direction: INCOMING, start: 0, count: 100 }) {
             ...fullRelationshipResults
         }

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

@@ -1,11 +1,5 @@
 query getMLFeatureTable($urn: String!) {
     mlFeatureTable(urn: $urn) {
         ...nonRecursiveMLFeatureTable
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
     }
 }

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

@@ -1,12 +1,6 @@
 query getMLModel($urn: String!) {
     mlModel(urn: $urn) {
         ...nonRecursiveMLModel
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
         features: relationships(input: { types: ["Consumes"], direction: OUTGOING, start: 0, count: 100 }) {
             start
             count

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

@@ -21,11 +21,5 @@ query getMLModelGroup($urn: String!) {
         ) {
             ...fullRelationshipResults
         }
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
     }
 }

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

@@ -1,12 +1,6 @@
 query getMLPrimaryKey($urn: String!) {
     mlPrimaryKey(urn: $urn) {
         ...nonRecursiveMLPrimaryKey
-        upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
-        downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...partialLineageResults
-        }
         featureTables: relationships(input: { types: ["KeyedBy"], direction: INCOMING, start: 0, count: 100 }) {
             ...fullRelationshipResults
         }