From ae274439302e8eb109f6017af345725be22cbf01 Mon Sep 17 00:00:00 2001
From: Curie <t-curiekim@microsoft.com>
Date: Wed, 6 Jul 2022 16:34:11 -0700
Subject: [PATCH 1/6] show feature lineage in feature details page

---
 ui/src/api/api.tsx                      |  4 +-
 ui/src/pages/feature/featureDetails.tsx | 90 ++++++++++++++++++++++++-
 2 files changed, 90 insertions(+), 4 deletions(-)

diff --git a/ui/src/api/api.tsx b/ui/src/api/api.tsx
index bcb25bd21..777947440 100644
--- a/ui/src/api/api.tsx
+++ b/ui/src/api/api.tsx
@@ -75,10 +75,10 @@ export const fetchProjectLineages = async (project: string) => {
     });
 };
 
-export const fetchFeatureLineages = async (project: string) => {
+export const fetchFeatureLineages = async (feature: string) => {
   const axios = await authAxios(msalInstance);
   return axios
-    .get<FeatureLineage>(`${getApiBaseUrl()}/features/lineage/${project}`, {})
+    .get<FeatureLineage>(`${getApiBaseUrl()}/features/${feature}/lineage`, {})
     .then((response) => {
       return response.data;
     });
diff --git a/ui/src/pages/feature/featureDetails.tsx b/ui/src/pages/feature/featureDetails.tsx
index 896a59c76..e98519200 100644
--- a/ui/src/pages/feature/featureDetails.tsx
+++ b/ui/src/pages/feature/featureDetails.tsx
@@ -1,11 +1,16 @@
-import React from 'react';
+import React, { useEffect, useState } from 'react';
 import { Alert, Button, Card, Col, Row, Space, Spin, Typography } from 'antd';
 import { LoadingOutlined } from '@ant-design/icons';
-import { useNavigate, useParams } from "react-router-dom";
+import { useNavigate, useParams, useSearchParams } from "react-router-dom";
 import { QueryStatus, useQuery } from "react-query";
 import { AxiosError } from 'axios';
 import { fetchFeature } from '../../api';
 import { Feature } from "../../models/model";
+import { FeatureLineage } from "../../models/model";
+import { fetchFeatureLineages } from "../../api";
+import { generateEdge, generateNode } from "../../components/graph/utils";
+import { Elements } from 'react-flow-renderer';
+import Graph from "../../components/graph/graph";
 
 const { Title } = Typography;
 
@@ -118,6 +123,86 @@ function InputDerivedFeatures(props: { project: string, feature: Feature }) {
   </>;
 }
 
+function FeatureLineageGraph() {
+  const [searchParams] = useSearchParams();
+  const { featureId } = useParams() as Params;
+  const nodeId = searchParams.get('nodeId') as string;
+
+  const [lineageData, setLineageData] = useState<FeatureLineage>({ guidEntityMap: null, relations: null });
+  const [elements, setElements] = useState<Elements>([]);
+  const [loading, setLoading] = useState<boolean>(false);
+
+  useEffect(() => {
+    const fetchLineageData = async () => {
+      setLoading(true);
+      const data = await fetchFeatureLineages(featureId);
+      setLineageData(data);
+      setLoading(false);
+    };
+    
+    fetchLineageData();
+  }, [featureId]);
+
+  // Generate graph data on client side, invoked after graphData or featureType is changed
+  useEffect(() => {
+    const generateGraphData = async () => {
+      if (lineageData.guidEntityMap === null && lineageData.relations === null) {
+        return;
+      }
+
+      const elements: Elements = [];
+      const elementObj: Record<string, string> = {};
+
+      for (let index = 0; index < Object.values(lineageData.guidEntityMap).length; index++) {
+        const currentNode: any = Object.values(lineageData.guidEntityMap)[index];
+
+        const nodeId = currentNode.guid;
+
+        const node = generateNode({
+          index,
+          nodeId,
+          currentNode
+        });
+
+        elementObj[nodeId] = index?.toString();
+        elements.push(node);
+      }
+
+      console.log(lineageData.relations.length);
+      for (let index = 0; index < lineageData.relations.length; index++) {
+        const { fromEntityId: to, toEntityId: from, relationshipType } = lineageData.relations[index];
+        const edge = generateEdge({ obj: elementObj, from, to });
+        if (edge?.source && edge?.target) {
+          if (relationshipType === "Consumes") {
+            elements.push(edge);
+          }
+        }
+      }
+
+      setElements(elements);
+    };
+
+    generateGraphData();
+  }, [lineageData])
+
+  return <>
+  {
+    loading
+    ? (
+      <Spin indicator={ <LoadingOutlined style={ { fontSize: 24 } } spin /> } />
+    )
+    : (
+      <Col span={ 24 }>
+        <Card className="card">
+          <title>Codestin Search App</title>
+          <Graph data={ elements } nodeId={ nodeId } />
+        </Card>
+      </Col>
+    )
+  }
+  </>;
+}
+
 type Params = {
   project: string;
   featureId: string;
@@ -193,6 +278,7 @@ const FeatureDetails: React.FC = () => {
                     <FeatureTransformation feature={ data } />
                     <FeatureKey feature={ data } />
                     <FeatureType feature={ data } />
+                    <FeatureLineageGraph />
                   </Row>
                 </div>
               </Card>

From 23e7087d5187a445203f982c7ca6b85d326679f8 Mon Sep 17 00:00:00 2001
From: Curie <t-curiekim@microsoft.com>
Date: Wed, 6 Jul 2022 17:16:29 -0700
Subject: [PATCH 2/6] show both produces and consumes relationship in lineage

---
 ui/src/pages/feature/featureDetails.tsx | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/ui/src/pages/feature/featureDetails.tsx b/ui/src/pages/feature/featureDetails.tsx
index e98519200..76c481140 100644
--- a/ui/src/pages/feature/featureDetails.tsx
+++ b/ui/src/pages/feature/featureDetails.tsx
@@ -170,10 +170,11 @@ function FeatureLineageGraph() {
 
       console.log(lineageData.relations.length);
       for (let index = 0; index < lineageData.relations.length; index++) {
-        const { fromEntityId: to, toEntityId: from, relationshipType } = lineageData.relations[index];
+        var { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
+        if (relationshipType == "Consumes") [from, to] = [to, from];
         const edge = generateEdge({ obj: elementObj, from, to });
         if (edge?.source && edge?.target) {
-          if (relationshipType === "Consumes") {
+          if (relationshipType === "Consumes" || relationshipType == "Produces") {
             elements.push(edge);
           }
         }

From 2437231b1edfe3af1e47982bd1c115e4864eb44b Mon Sep 17 00:00:00 2001
From: Curie <t-curiekim@microsoft.com>
Date: Fri, 8 Jul 2022 16:12:55 -0700
Subject: [PATCH 3/6] make the height of the graph responsive

---
 ui/src/components/graph/graph.tsx       | 15 ++++++++++++++-
 ui/src/pages/feature/featureDetails.tsx | 10 ++++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/ui/src/components/graph/graph.tsx b/ui/src/components/graph/graph.tsx
index ab625761c..a8031be87 100644
--- a/ui/src/components/graph/graph.tsx
+++ b/ui/src/components/graph/graph.tsx
@@ -64,6 +64,19 @@ const Graph: React.FC<Props> = ({ data, nodeId }) => {
     setElements(values);
   };
 
+  // calculate the height of the graph by finding the maximum y position of the nodes and adding some padding to take the size of the node into account
+  const calculateHeight = () => {
+    var padding = 200;
+    var max = 0;
+    for (let index = 0; index < elements.length; index++) {
+      const element = elements[index];
+      if (isNode(element) && element.position.y > max) {
+        max = element.position.y
+      }
+    }
+    return max + padding;
+  }
+
   // Highlight path of selected node, including all linked up and down stream nodes
   const highlightPath = (node: Node, check: boolean): void => {
     const checkElements = check ? layoutedElements : elements;
@@ -146,7 +159,7 @@ const Graph: React.FC<Props> = ({ data, nodeId }) => {
     <div className="lineage-graph">
         <ReactFlowProvider>
           <ReactFlow
-            style={ { height: "700px", width: "100%" } }
+            style={ { height: calculateHeight(), width: "100%" } }
             elements={ elements }
             snapToGrid
             snapGrid={ [15, 15] }
diff --git a/ui/src/pages/feature/featureDetails.tsx b/ui/src/pages/feature/featureDetails.tsx
index 76c481140..364f8647b 100644
--- a/ui/src/pages/feature/featureDetails.tsx
+++ b/ui/src/pages/feature/featureDetails.tsx
@@ -126,7 +126,6 @@ function InputDerivedFeatures(props: { project: string, feature: Feature }) {
 function FeatureLineageGraph() {
   const [searchParams] = useSearchParams();
   const { featureId } = useParams() as Params;
-  const nodeId = searchParams.get('nodeId') as string;
 
   const [lineageData, setLineageData] = useState<FeatureLineage>({ guidEntityMap: null, relations: null });
   const [elements, setElements] = useState<Elements>([]);
@@ -143,7 +142,6 @@ function FeatureLineageGraph() {
     fetchLineageData();
   }, [featureId]);
 
-  // Generate graph data on client side, invoked after graphData or featureType is changed
   useEffect(() => {
     const generateGraphData = async () => {
       if (lineageData.guidEntityMap === null && lineageData.relations === null) {
@@ -168,13 +166,13 @@ function FeatureLineageGraph() {
         elements.push(node);
       }
 
-      console.log(lineageData.relations.length);
+
       for (let index = 0; index < lineageData.relations.length; index++) {
         var { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
-        if (relationshipType == "Consumes") [from, to] = [to, from];
+        if (relationshipType === "Consumes") [from, to] = [to, from];
         const edge = generateEdge({ obj: elementObj, from, to });
         if (edge?.source && edge?.target) {
-          if (relationshipType === "Consumes" || relationshipType == "Produces") {
+          if (relationshipType === "Consumes" || relationshipType === "Produces") {
             elements.push(edge);
           }
         }
@@ -196,7 +194,7 @@ function FeatureLineageGraph() {
       <Col span={ 24 }>
         <Card className="card">
           <title>Codestin Search App</title>
-          <Graph data={ elements } nodeId={ nodeId } />
+          <Graph data={ elements } nodeId={ featureId }/>
         </Card>
       </Col>
     )

From 7bb020d3c2ee458f241642225aa68acf6ca0db41 Mon Sep 17 00:00:00 2001
From: Curie <t-curiekim@microsoft.com>
Date: Wed, 13 Jul 2022 14:15:07 -0700
Subject: [PATCH 4/6] refactor code

---
 ui/src/api/api.tsx                      |  4 +-
 ui/src/components/graph/graph.tsx       | 23 ++++++----
 ui/src/pages/feature/featureDetails.tsx | 52 +++------------------
 ui/src/pages/feature/lineageGraph.tsx   | 56 ++---------------------
 ui/src/pages/feature/useGraphHook.tsx   | 61 +++++++++++++++++++++++++
 5 files changed, 87 insertions(+), 109 deletions(-)
 create mode 100644 ui/src/pages/feature/useGraphHook.tsx

diff --git a/ui/src/api/api.tsx b/ui/src/api/api.tsx
index 777947440..67bc3456d 100644
--- a/ui/src/api/api.tsx
+++ b/ui/src/api/api.tsx
@@ -75,10 +75,10 @@ export const fetchProjectLineages = async (project: string) => {
     });
 };
 
-export const fetchFeatureLineages = async (feature: string) => {
+export const fetchFeatureLineages = async (featureId: string) => {
   const axios = await authAxios(msalInstance);
   return axios
-    .get<FeatureLineage>(`${getApiBaseUrl()}/features/${feature}/lineage`, {})
+    .get<FeatureLineage>(`${getApiBaseUrl()}/features/${featureId}/lineage`, {})
     .then((response) => {
       return response.data;
     });
diff --git a/ui/src/components/graph/graph.tsx b/ui/src/components/graph/graph.tsx
index a8031be87..98e61d403 100644
--- a/ui/src/components/graph/graph.tsx
+++ b/ui/src/components/graph/graph.tsx
@@ -22,8 +22,9 @@ const nodeTypes = {
 type Props = {
   data: Elements;
   nodeId: string;
+  isFeatureGraph: boolean;
 }
-const Graph: React.FC<Props> = ({ data, nodeId }) => {
+const Graph: React.FC<Props> = ({ data, nodeId, isFeatureGraph }) => {
   const [, setURLSearchParams] = useSearchParams();
 
   const { layoutedElements, elementMapping } = getLayoutedElements(data);
@@ -66,17 +67,21 @@ const Graph: React.FC<Props> = ({ data, nodeId }) => {
 
   // calculate the height of the graph by finding the maximum y position of the nodes and adding some padding to take the size of the node into account
   const calculateHeight = () => {
-    var padding = 200;
-    var max = 0;
-    for (let index = 0; index < elements.length; index++) {
-      const element = elements[index];
-      if (isNode(element) && element.position.y > max) {
-        max = element.position.y
+    if (isFeatureGraph) {
+      var padding = 200;
+      var max = 0;
+      for (let index = 0; index < elements.length; index++) {
+        const element = elements[index];
+        if (isNode(element) && element.position.y > max) {
+          max = element.position.y
+        }
       }
+      return max + padding;
+    } else {
+      return window.innerHeight - 250;
     }
-    return max + padding;
   }
-
+  
   // Highlight path of selected node, including all linked up and down stream nodes
   const highlightPath = (node: Node, check: boolean): void => {
     const checkElements = check ? layoutedElements : elements;
diff --git a/ui/src/pages/feature/featureDetails.tsx b/ui/src/pages/feature/featureDetails.tsx
index 364f8647b..87d320fd8 100644
--- a/ui/src/pages/feature/featureDetails.tsx
+++ b/ui/src/pages/feature/featureDetails.tsx
@@ -8,9 +8,9 @@ import { fetchFeature } from '../../api';
 import { Feature } from "../../models/model";
 import { FeatureLineage } from "../../models/model";
 import { fetchFeatureLineages } from "../../api";
-import { generateEdge, generateNode } from "../../components/graph/utils";
 import { Elements } from 'react-flow-renderer';
 import Graph from "../../components/graph/graph";
+import useGraphHook from "./useGraphHook"
 
 const { Title } = Typography;
 
@@ -123,12 +123,11 @@ function InputDerivedFeatures(props: { project: string, feature: Feature }) {
   </>;
 }
 
+
 function FeatureLineageGraph() {
-  const [searchParams] = useSearchParams();
   const { featureId } = useParams() as Params;
-
   const [lineageData, setLineageData] = useState<FeatureLineage>({ guidEntityMap: null, relations: null });
-  const [elements, setElements] = useState<Elements>([]);
+  const [elements, SetElements] = useState<Elements>([]);
   const [loading, setLoading] = useState<boolean>(false);
 
   useEffect(() => {
@@ -142,47 +141,8 @@ function FeatureLineageGraph() {
     fetchLineageData();
   }, [featureId]);
 
-  useEffect(() => {
-    const generateGraphData = async () => {
-      if (lineageData.guidEntityMap === null && lineageData.relations === null) {
-        return;
-      }
-
-      const elements: Elements = [];
-      const elementObj: Record<string, string> = {};
-
-      for (let index = 0; index < Object.values(lineageData.guidEntityMap).length; index++) {
-        const currentNode: any = Object.values(lineageData.guidEntityMap)[index];
-
-        const nodeId = currentNode.guid;
-
-        const node = generateNode({
-          index,
-          nodeId,
-          currentNode
-        });
-
-        elementObj[nodeId] = index?.toString();
-        elements.push(node);
-      }
-
-
-      for (let index = 0; index < lineageData.relations.length; index++) {
-        var { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
-        if (relationshipType === "Consumes") [from, to] = [to, from];
-        const edge = generateEdge({ obj: elementObj, from, to });
-        if (edge?.source && edge?.target) {
-          if (relationshipType === "Consumes" || relationshipType === "Produces") {
-            elements.push(edge);
-          }
-        }
-      }
-
-      setElements(elements);
-    };
-
-    generateGraphData();
-  }, [lineageData])
+  // Generate graph data on client side, invoked after graphData or featureType is changed
+  useGraphHook(lineageData, "all_nodes", SetElements);
 
   return <>
   {
@@ -194,7 +154,7 @@ function FeatureLineageGraph() {
       <Col span={ 24 }>
         <Card className="card">
           <title>Codestin Search App</title>
-          <Graph data={ elements } nodeId={ featureId }/>
+          <Graph data={ elements } nodeId={ featureId } isFeatureGraph={true}/>
         </Card>
       </Col>
     )
diff --git a/ui/src/pages/feature/lineageGraph.tsx b/ui/src/pages/feature/lineageGraph.tsx
index 0ba733550..54fe9a61c 100644
--- a/ui/src/pages/feature/lineageGraph.tsx
+++ b/ui/src/pages/feature/lineageGraph.tsx
@@ -3,11 +3,11 @@ import { Card, Col, Radio, Row, Spin, Tabs, Typography } from 'antd';
 import { useParams, useSearchParams } from "react-router-dom";
 import { Elements } from 'react-flow-renderer';
 import Graph from "../../components/graph/graph";
-import { generateEdge, generateNode } from "../../components/graph/utils";
 import { fetchProjectLineages } from "../../api";
 import { FeatureLineage } from "../../models/model";
 import { LoadingOutlined } from "@ant-design/icons";
 import GraphNodeDetails from "../../components/graph/graphNodeDetails";
+import useGraphHook from "./useGraphHook"
 
 const { Title } = Typography;
 const { TabPane } = Tabs;
@@ -37,57 +37,9 @@ const LineageGraph: React.FC = () => {
     fetchLineageData();
   }, [project]);
 
+  
   // Generate graph data on client side, invoked after graphData or featureType is changed
-  useEffect(() => {
-    const generateGraphData = async () => {
-      if (lineageData.guidEntityMap === null && lineageData.relations === null) {
-        return;
-      }
-
-      const elements: Elements = [];
-      const elementObj: Record<string, string> = {};
-
-      for (let index = 0; index < Object.values(lineageData.guidEntityMap).length; index++) {
-        const currentNode: any = Object.values(lineageData.guidEntityMap)[index];
-
-        if (currentNode.typeName === "feathr_workspace_v1") {
-          continue; // Open issue: should project node get displayed as well?
-        }
-
-        const nodeId = currentNode.guid;
-
-        // If toggled feature type exists, skip other types
-        if (featureType && featureType !== "all_nodes" && currentNode.typeName !== featureType) {
-          continue;
-        }
-
-        const node = generateNode({
-          index,
-          nodeId,
-          currentNode
-        });
-
-        elementObj[nodeId] = index?.toString();
-
-        elements.push(node);
-      }
-
-      for (let index = 0; index < lineageData.relations.length; index++) {
-        const { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
-        const edge = generateEdge({ obj: elementObj, from, to });
-        if (edge?.source && edge?.target) {
-          // Currently, API returns all relationships, filter out Contains, Consumes, etc
-          if (relationshipType === "Produces") {
-            elements.push(edge);
-          }
-        }
-      }
-
-      SetElements(elements);
-    };
-
-    generateGraphData();
-  }, [lineageData, featureType])
+  useGraphHook(lineageData, featureType, SetElements);
 
   const toggleFeatureType = (type: string) => {
     setFeatureType((prevType: string | null) => {
@@ -118,7 +70,7 @@ const LineageGraph: React.FC = () => {
               : (
                 <Row>
                   <Col flex="2">
-                    <Graph data={ elements } nodeId={ nodeId } />
+                    <Graph data={ elements } nodeId={ nodeId } isFeatureGraph={false} />
                   </Col>
                   <Col flex="1">
                     <Tabs defaultActiveKey="1">
diff --git a/ui/src/pages/feature/useGraphHook.tsx b/ui/src/pages/feature/useGraphHook.tsx
new file mode 100644
index 000000000..774f59a31
--- /dev/null
+++ b/ui/src/pages/feature/useGraphHook.tsx
@@ -0,0 +1,61 @@
+import { useEffect } from "react";
+import { Elements } from 'react-flow-renderer';
+import { generateEdge, generateNode } from "../../components/graph/utils";
+import { FeatureLineage } from "../../models/model";
+
+const useGraphHook = (lineageData: FeatureLineage, featureType: string|null, setElements: Function) => {
+  // Generate graph data on client side, invoked after graphData or featureType is changed
+  useEffect(() => {
+    const generateGraphData = async () => {
+      if (lineageData.guidEntityMap === null && lineageData.relations === null) {
+        return;
+      }
+
+      const elements: Elements = [];
+      const elementObj: Record<string, string> = {};
+
+      for (let index = 0; index < Object.values(lineageData.guidEntityMap).length; index++) {
+        const currentNode: any = Object.values(lineageData.guidEntityMap)[index];
+
+        if (currentNode.typeName === "feathr_workspace_v1") {
+          continue; // Open issue: should project node get displayed as well?
+        }
+
+        const nodeId = currentNode.guid;
+
+        // If toggled feature type exists, skip other types
+        if (featureType && featureType !== "all_nodes" && currentNode.typeName !== featureType) {
+          continue;
+        }
+
+        const node = generateNode({
+          index,
+          nodeId,
+          currentNode
+        });
+
+        elementObj[nodeId] = index?.toString();
+
+        elements.push(node);
+      }
+
+      for (let index = 0; index < lineageData.relations.length; index++) {
+        var { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
+        if (relationshipType === "Consumes") [from, to] = [to, from];
+        const edge = generateEdge({ obj: elementObj, from, to });
+        if (edge?.source && edge?.target) {
+          if (relationshipType === "Consumes" || relationshipType === "Produces") {
+            elements.push(edge);
+          }
+        }
+      }
+
+      setElements(elements);
+    };
+
+    generateGraphData();
+
+  }, [lineageData, featureType]);
+}
+
+export default useGraphHook;
\ No newline at end of file

From 70d2c1e71bc40a1f15faa04a8fabe61c6c7dbf1e Mon Sep 17 00:00:00 2001
From: Curie <t-curiekim@microsoft.com>
Date: Thu, 14 Jul 2022 08:06:44 -0700
Subject: [PATCH 5/6] refactor code

---
 ui/src/components/graph/graph.tsx       | 23 ++--------
 ui/src/components/graph/utils.ts        | 49 ++++++++++++++++++++
 ui/src/pages/feature/featureDetails.tsx | 29 ++++++++++--
 ui/src/pages/feature/lineageGraph.tsx   | 13 ++++--
 ui/src/pages/feature/useGraphHook.tsx   | 61 -------------------------
 5 files changed, 86 insertions(+), 89 deletions(-)
 delete mode 100644 ui/src/pages/feature/useGraphHook.tsx

diff --git a/ui/src/components/graph/graph.tsx b/ui/src/components/graph/graph.tsx
index 98e61d403..474118cba 100644
--- a/ui/src/components/graph/graph.tsx
+++ b/ui/src/components/graph/graph.tsx
@@ -22,9 +22,9 @@ const nodeTypes = {
 type Props = {
   data: Elements;
   nodeId: string;
-  isFeatureGraph: boolean;
+  height: number;
 }
-const Graph: React.FC<Props> = ({ data, nodeId, isFeatureGraph }) => {
+const Graph: React.FC<Props> = ({ data, nodeId, height }) => {
   const [, setURLSearchParams] = useSearchParams();
 
   const { layoutedElements, elementMapping } = getLayoutedElements(data);
@@ -65,23 +65,6 @@ const Graph: React.FC<Props> = ({ data, nodeId, isFeatureGraph }) => {
     setElements(values);
   };
 
-  // calculate the height of the graph by finding the maximum y position of the nodes and adding some padding to take the size of the node into account
-  const calculateHeight = () => {
-    if (isFeatureGraph) {
-      var padding = 200;
-      var max = 0;
-      for (let index = 0; index < elements.length; index++) {
-        const element = elements[index];
-        if (isNode(element) && element.position.y > max) {
-          max = element.position.y
-        }
-      }
-      return max + padding;
-    } else {
-      return window.innerHeight - 250;
-    }
-  }
-  
   // Highlight path of selected node, including all linked up and down stream nodes
   const highlightPath = (node: Node, check: boolean): void => {
     const checkElements = check ? layoutedElements : elements;
@@ -164,7 +147,7 @@ const Graph: React.FC<Props> = ({ data, nodeId, isFeatureGraph }) => {
     <div className="lineage-graph">
         <ReactFlowProvider>
           <ReactFlow
-            style={ { height: calculateHeight(), width: "100%" } }
+            style={ { height: height, width: "100%" } }
             elements={ elements }
             snapToGrid
             snapGrid={ [15, 15] }
diff --git a/ui/src/components/graph/utils.ts b/ui/src/components/graph/utils.ts
index 7d8a0510d..880efdbb0 100644
--- a/ui/src/components/graph/utils.ts
+++ b/ui/src/components/graph/utils.ts
@@ -1,5 +1,6 @@
 import dagre from 'dagre';
 import { ArrowHeadType, Edge, Elements, isNode, Node, Position, } from 'react-flow-renderer';
+import { FeatureLineage } from "../../models/model";
 
 const DEFAULT_WIDTH = 172;
 const DEFAULT_HEIGHT = 36;
@@ -9,6 +10,53 @@ type getLayoutElementsRet = {
   elementMapping: Record<string, number>,
 };
 
+const getElements = (lineageData: FeatureLineage, featureType: string|null) => {
+  if (lineageData.guidEntityMap === null && lineageData.relations === null) {
+    return;
+  }
+
+  const elements: Elements = [];
+  const elementObj: Record<string, string> = {};
+
+  for (let index = 0; index < Object.values(lineageData.guidEntityMap).length; index++) {
+    const currentNode: any = Object.values(lineageData.guidEntityMap)[index];
+
+    if (currentNode.typeName === "feathr_workspace_v1") {
+      continue; // Open issue: should project node get displayed as well?
+    }
+
+    const nodeId = currentNode.guid;
+
+    // If toggled feature type exists, skip other types
+    if (featureType && featureType !== "all_nodes" && currentNode.typeName !== featureType) {
+      continue;
+    }
+
+    const node = generateNode({
+      index,
+      nodeId,
+      currentNode
+    });
+
+    elementObj[nodeId] = index?.toString();
+
+    elements.push(node);
+  }
+
+  for (let index = 0; index < lineageData.relations.length; index++) {
+    var { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
+    if (relationshipType === "Consumes") [from, to] = [to, from];
+    const edge = generateEdge({ obj: elementObj, from, to });
+    if (edge?.source && edge?.target) {
+      if (relationshipType === "Consumes" || relationshipType === "Produces") {
+        elements.push(edge);
+      }
+    }
+  }
+
+  return elements;
+};
+
 const getLayoutedElements = (elements: Elements, direction = 'LR'): getLayoutElementsRet => {
   const dagreGraph = new dagre.graphlib.Graph();
   dagreGraph.setDefaultEdgeLabel(() => ({}));
@@ -112,6 +160,7 @@ export {
   generateEdge,
   generateNode,
   getLayoutedElements,
+  getElements
 };
 
 export const findNodeInElement = (nodeId: string | null, elements: Elements): Node | null => {
diff --git a/ui/src/pages/feature/featureDetails.tsx b/ui/src/pages/feature/featureDetails.tsx
index 87d320fd8..6f63a6ada 100644
--- a/ui/src/pages/feature/featureDetails.tsx
+++ b/ui/src/pages/feature/featureDetails.tsx
@@ -8,9 +8,10 @@ import { fetchFeature } from '../../api';
 import { Feature } from "../../models/model";
 import { FeatureLineage } from "../../models/model";
 import { fetchFeatureLineages } from "../../api";
-import { Elements } from 'react-flow-renderer';
+import { Elements, isNode } from 'react-flow-renderer';
 import Graph from "../../components/graph/graph";
-import useGraphHook from "./useGraphHook"
+import { getElements } from "../../components/graph/utils"
+import { getLayoutedElements } from "../../components/graph/utils"
 
 const { Title } = Typography;
 
@@ -142,7 +143,27 @@ function FeatureLineageGraph() {
   }, [featureId]);
 
   // Generate graph data on client side, invoked after graphData or featureType is changed
-  useGraphHook(lineageData, "all_nodes", SetElements);
+  useEffect(() => {
+    const generateGraphData = async () => {
+      SetElements(getElements(lineageData, "all_nodes")!);
+    };
+
+    generateGraphData();
+  }, [lineageData]);
+
+
+  const calculateHeight = () => {
+    const { layoutedElements } = getLayoutedElements(elements);
+    var padding = 200;
+    var max = 0;
+    for (let index = 0; index < layoutedElements.length; index++) {
+      const element = layoutedElements[index];
+      if (isNode(element) && element.position.y > max) {
+        max = element.position.y
+      }
+    }
+    return max + padding;
+  }
 
   return <>
   {
@@ -154,7 +175,7 @@ function FeatureLineageGraph() {
       <Col span={ 24 }>
         <Card className="card">
           <title>Codestin Search App</title>
-          <Graph data={ elements } nodeId={ featureId } isFeatureGraph={true}/>
+          <Graph data={ elements } nodeId={ featureId } height={calculateHeight()}/>
         </Card>
       </Col>
     )
diff --git a/ui/src/pages/feature/lineageGraph.tsx b/ui/src/pages/feature/lineageGraph.tsx
index 54fe9a61c..6906c0cee 100644
--- a/ui/src/pages/feature/lineageGraph.tsx
+++ b/ui/src/pages/feature/lineageGraph.tsx
@@ -7,7 +7,7 @@ import { fetchProjectLineages } from "../../api";
 import { FeatureLineage } from "../../models/model";
 import { LoadingOutlined } from "@ant-design/icons";
 import GraphNodeDetails from "../../components/graph/graphNodeDetails";
-import useGraphHook from "./useGraphHook"
+import { getElements } from "../../components/graph/utils"
 
 const { Title } = Typography;
 const { TabPane } = Tabs;
@@ -37,9 +37,14 @@ const LineageGraph: React.FC = () => {
     fetchLineageData();
   }, [project]);
 
-  
   // Generate graph data on client side, invoked after graphData or featureType is changed
-  useGraphHook(lineageData, featureType, SetElements);
+  useEffect(() => {
+    const generateGraphData = async () => {
+      SetElements(getElements(lineageData, featureType)!);
+    };
+
+    generateGraphData();
+  }, [lineageData, featureType])
 
   const toggleFeatureType = (type: string) => {
     setFeatureType((prevType: string | null) => {
@@ -70,7 +75,7 @@ const LineageGraph: React.FC = () => {
               : (
                 <Row>
                   <Col flex="2">
-                    <Graph data={ elements } nodeId={ nodeId } isFeatureGraph={false} />
+                    <Graph data={ elements } nodeId={ nodeId } height={window.innerHeight - 250}/>
                   </Col>
                   <Col flex="1">
                     <Tabs defaultActiveKey="1">
diff --git a/ui/src/pages/feature/useGraphHook.tsx b/ui/src/pages/feature/useGraphHook.tsx
deleted file mode 100644
index 774f59a31..000000000
--- a/ui/src/pages/feature/useGraphHook.tsx
+++ /dev/null
@@ -1,61 +0,0 @@
-import { useEffect } from "react";
-import { Elements } from 'react-flow-renderer';
-import { generateEdge, generateNode } from "../../components/graph/utils";
-import { FeatureLineage } from "../../models/model";
-
-const useGraphHook = (lineageData: FeatureLineage, featureType: string|null, setElements: Function) => {
-  // Generate graph data on client side, invoked after graphData or featureType is changed
-  useEffect(() => {
-    const generateGraphData = async () => {
-      if (lineageData.guidEntityMap === null && lineageData.relations === null) {
-        return;
-      }
-
-      const elements: Elements = [];
-      const elementObj: Record<string, string> = {};
-
-      for (let index = 0; index < Object.values(lineageData.guidEntityMap).length; index++) {
-        const currentNode: any = Object.values(lineageData.guidEntityMap)[index];
-
-        if (currentNode.typeName === "feathr_workspace_v1") {
-          continue; // Open issue: should project node get displayed as well?
-        }
-
-        const nodeId = currentNode.guid;
-
-        // If toggled feature type exists, skip other types
-        if (featureType && featureType !== "all_nodes" && currentNode.typeName !== featureType) {
-          continue;
-        }
-
-        const node = generateNode({
-          index,
-          nodeId,
-          currentNode
-        });
-
-        elementObj[nodeId] = index?.toString();
-
-        elements.push(node);
-      }
-
-      for (let index = 0; index < lineageData.relations.length; index++) {
-        var { fromEntityId: from, toEntityId: to, relationshipType } = lineageData.relations[index];
-        if (relationshipType === "Consumes") [from, to] = [to, from];
-        const edge = generateEdge({ obj: elementObj, from, to });
-        if (edge?.source && edge?.target) {
-          if (relationshipType === "Consumes" || relationshipType === "Produces") {
-            elements.push(edge);
-          }
-        }
-      }
-
-      setElements(elements);
-    };
-
-    generateGraphData();
-
-  }, [lineageData, featureType]);
-}
-
-export default useGraphHook;
\ No newline at end of file

From a346a8b6ba8fa656a2642ed130b4e7379f69cd18 Mon Sep 17 00:00:00 2001
From: Curie <t-curiekim@microsoft.com>
Date: Thu, 14 Jul 2022 10:44:10 -0700
Subject: [PATCH 6/6] Set height to be based on the window height

---
 ui/src/components/graph/graph.tsx       |  5 ++---
 ui/src/pages/feature/featureDetails.tsx | 16 +---------------
 ui/src/pages/feature/lineageGraph.tsx   |  2 +-
 3 files changed, 4 insertions(+), 19 deletions(-)

diff --git a/ui/src/components/graph/graph.tsx b/ui/src/components/graph/graph.tsx
index 474118cba..55e5f9746 100644
--- a/ui/src/components/graph/graph.tsx
+++ b/ui/src/components/graph/graph.tsx
@@ -22,9 +22,8 @@ const nodeTypes = {
 type Props = {
   data: Elements;
   nodeId: string;
-  height: number;
 }
-const Graph: React.FC<Props> = ({ data, nodeId, height }) => {
+const Graph: React.FC<Props> = ({ data, nodeId }) => {
   const [, setURLSearchParams] = useSearchParams();
 
   const { layoutedElements, elementMapping } = getLayoutedElements(data);
@@ -147,7 +146,7 @@ const Graph: React.FC<Props> = ({ data, nodeId, height }) => {
     <div className="lineage-graph">
         <ReactFlowProvider>
           <ReactFlow
-            style={ { height: height, width: "100%" } }
+            style={ { height: window.innerHeight - 250, width: "100%" } }
             elements={ elements }
             snapToGrid
             snapGrid={ [15, 15] }
diff --git a/ui/src/pages/feature/featureDetails.tsx b/ui/src/pages/feature/featureDetails.tsx
index 6f63a6ada..f3d8cf9ec 100644
--- a/ui/src/pages/feature/featureDetails.tsx
+++ b/ui/src/pages/feature/featureDetails.tsx
@@ -151,20 +151,6 @@ function FeatureLineageGraph() {
     generateGraphData();
   }, [lineageData]);
 
-
-  const calculateHeight = () => {
-    const { layoutedElements } = getLayoutedElements(elements);
-    var padding = 200;
-    var max = 0;
-    for (let index = 0; index < layoutedElements.length; index++) {
-      const element = layoutedElements[index];
-      if (isNode(element) && element.position.y > max) {
-        max = element.position.y
-      }
-    }
-    return max + padding;
-  }
-
   return <>
   {
     loading
@@ -175,7 +161,7 @@ function FeatureLineageGraph() {
       <Col span={ 24 }>
         <Card className="card">
           <title>Codestin Search App</title>
-          <Graph data={ elements } nodeId={ featureId } height={calculateHeight()}/>
+          <Graph data={ elements } nodeId={ featureId }/>
         </Card>
       </Col>
     )
diff --git a/ui/src/pages/feature/lineageGraph.tsx b/ui/src/pages/feature/lineageGraph.tsx
index 6906c0cee..051030d28 100644
--- a/ui/src/pages/feature/lineageGraph.tsx
+++ b/ui/src/pages/feature/lineageGraph.tsx
@@ -75,7 +75,7 @@ const LineageGraph: React.FC = () => {
               : (
                 <Row>
                   <Col flex="2">
-                    <Graph data={ elements } nodeId={ nodeId } height={window.innerHeight - 250}/>
+                    <Graph data={ elements } nodeId={ nodeId }/>
                   </Col>
                   <Col flex="1">
                     <Tabs defaultActiveKey="1">