From 25b9576114a2cb83e362918c3465665e6220471c Mon Sep 17 00:00:00 2001
From: Karan Hotchandani <33024356+karanh37@users.noreply.github.com>
Date: Mon, 20 May 2024 12:22:49 +0530
Subject: [PATCH] create export data on every click (#16344)

---
 .../LineageProvider/LineageProvider.tsx       | 21 ++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
index 1eb60593ac9..339f051797d 100644
--- a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
@@ -175,7 +175,6 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   const [childMap, setChildMap] = useState<EntityReferenceChild>();
   const [paginationData, setPaginationData] = useState({});
   const { showModal } = useEntityExportModalProvider();
-  const [exportResult, setExportResult] = useState<string>('');
 
   const initLineageChildMaps = useCallback(
     (
@@ -228,11 +227,10 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
             entityType !== EntityType.PIPELINE &&
             entityType !== EntityType.STORED_PROCEDURE
           ) {
-            const { map: childMapObj, exportResult } = getChildMap(
+            const { map: childMapObj } = getChildMap(
               { ...res, nodes: allNodes },
               decodedFqn
             );
-            setExportResult(exportResult);
             setChildMap(childMapObj);
             const { nodes: newNodes, edges: newEdges } = getPaginatedChildMap(
               {
@@ -250,8 +248,6 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
               edges: [...(res.edges ?? []), ...newEdges],
             });
           } else {
-            const csv = getExportData(allNodes);
-            setExportResult(csv);
             setEntityLineage({
               ...res,
               nodes: allNodes,
@@ -281,9 +277,20 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
 
   const exportLineageData = useCallback(
     async (_: string) => {
+      if (
+        entityType === EntityType.PIPELINE ||
+        entityType === EntityType.STORED_PROCEDURE
+      ) {
+        // Since pipeline is an edge, we cannot create a tree, hence we take the nodes from the lineage response
+        // to get the export data.
+        return getExportData(entityLineage.nodes ?? []);
+      }
+
+      const { exportResult } = getChildMap(entityLineage, decodedFqn);
+
       return exportResult;
     },
-    [exportResult]
+    [entityType, decodedFqn, entityLineage]
   );
 
   const onExportClick = useCallback(() => {
@@ -293,7 +300,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
         onExport: exportLineageData,
       });
     }
-  }, [decodedFqn, exportResult]);
+  }, [entityType, decodedFqn, entityLineage]);
 
   const loadChildNodesHandler = useCallback(
     async (node: SourceType, direction: EdgeTypeEnum) => {