From 922bd3a74156c7d458c2a1aa49e0fb9c62d7e915 Mon Sep 17 00:00:00 2001
From: "Han-Ru Chen (Future-Outlier)" <eric901201@gmail.com>
Date: Thu, 30 Jan 2025 11:39:19 -0800
Subject: [PATCH] feat: Streaming Decks (#890)

* "Use Yi-Cheng Lu's code"

Signed-off-by: Future-Outlier <eric901201@gmail.com>
Co-authored-by: Yicheng-Lu-llll <luyc58576@gmail.com>

* Add Refresh Bottom

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* refresh fix

Signed-off-by: Lyon Lu <lyon@union.ai>

* move refresh button out

Signed-off-by: Lyon Lu <lyon@union.ai>

* update my code

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* handle 404 vs 500 error rendering

Signed-off-by: Carina Ursu <carina@union.ai>

---------

Signed-off-by: Future-Outlier <eric901201@gmail.com>
Signed-off-by: Lyon Lu <lyon@union.ai>
Signed-off-by: Carina Ursu <carina@union.ai>
Co-authored-by: Yicheng-Lu-llll <luyc58576@gmail.com>
Co-authored-by: Lyon Lu <lyon@union.ai>
Co-authored-by: Carina Ursu <carina@union.ai>
---
 .../FlyteDeckButton.tsx                       | 22 ++++++++++------
 .../ExecutionDetails/ExecutionNodeDeck.tsx    | 26 +++++++++++++++++--
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionDetailsActions/FlyteDeckButton.tsx b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionDetailsActions/FlyteDeckButton.tsx
index dbff7814a..110699a94 100644
--- a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionDetailsActions/FlyteDeckButton.tsx
+++ b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionDetailsActions/FlyteDeckButton.tsx
@@ -8,6 +8,8 @@ import Grid from '@mui/material/Grid';
 import IconButton from '@mui/material/IconButton';
 import Typography from '@mui/material/Typography';
 import DialogContent from '@mui/material/DialogContent';
+import { useDownloadLink } from '@clients/oss-console/components/hooks/useDataProxy';
+import RefreshIcon from '@mui/icons-material/Refresh';
 import t from '../strings';
 import { WorkflowNodeExecution } from '../../contexts';
 import { NodeExecutionPhase } from '../../../../models/Execution/enums';
@@ -52,14 +54,11 @@ export const FlyteDeckButton: FC<FlyteDeckButtonProps> = ({
     setSetFullScreen(!fullScreen);
   };
 
+  const downloadLink = useDownloadLink(nodeExecution?.id);
+
   return nodeExecution?.closure?.deckUri ? (
     <>
-      <Button
-        variant="outlined"
-        color="primary"
-        onClick={() => setShowDeck(true)}
-        disabled={phase !== NodeExecutionPhase.SUCCEEDED}
-      >
+      <Button variant="outlined" color="primary" onClick={() => setShowDeck(true)}>
         {flyteDeckText || t('flyteDeck')}
       </Button>
       <Dialog
@@ -91,10 +90,17 @@ export const FlyteDeckButton: FC<FlyteDeckButtonProps> = ({
                 fontSize: '24px',
                 lineHeight: '32px',
                 marginBlock: 0,
+                display: 'flex',
+                justifyContent: 'center',
+                alignItems: 'center',
+                gap: 1,
               }}
               py={2}
             >
-              {t('flyteDeck')}
+              {flyteDeckText || t('flyteDeck')}
+              <IconButton onClick={() => downloadLink.fetch()}>
+                <RefreshIcon />
+              </IconButton>
             </Typography>
           </Grid>
           <Grid item>
@@ -109,7 +115,7 @@ export const FlyteDeckButton: FC<FlyteDeckButtonProps> = ({
             overflow: 'hidden',
           }}
         >
-          <ExecutionNodeDeck nodeExecutionId={nodeExecution.id} />
+          <ExecutionNodeDeck nodeExecutionId={nodeExecution.id} downloadLink={downloadLink} />
         </DialogContent>
       </Dialog>
     </>
diff --git a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeDeck.tsx b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeDeck.tsx
index 105c05f1a..dbafc5ca6 100644
--- a/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeDeck.tsx
+++ b/packages/oss-console/src/components/Executions/ExecutionDetails/ExecutionNodeDeck.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import Core from '@clients/common/flyteidl/core';
+import NotFoundError from '@clients/common/Errors/NotFoundError';
 import { LoadingSpinner } from '@clients/primitives/LoadingSpinner';
 import { useDownloadLink } from '../../hooks/useDataProxy';
 import { WaitForData } from '../../common/WaitForData';
@@ -8,8 +9,8 @@ import { WaitForData } from '../../common/WaitForData';
 export const ExecutionNodeDeck: React.FC<{
   nodeExecutionId: Core.NodeExecutionIdentifier;
   className?: string;
-}> = ({ nodeExecutionId, className = '' }) => {
-  const downloadLink = useDownloadLink(nodeExecutionId);
+  downloadLink: ReturnType<typeof useDownloadLink>;
+}> = ({ className = '', downloadLink }) => {
   const iFrameSrc = downloadLink?.value?.signedUrl?.[0];
 
   // Taken from https://developer.mozilla.org/en-US/docs/Web/HTML/Element/iframe#sandbox
@@ -27,6 +28,27 @@ export const ExecutionNodeDeck: React.FC<{
     'allow-downloads',
   ].join(' ');
 
+  /**
+   * if the download link query has an error other than 404, we show a 'pretty' message to the user.
+   * else: we show the built in DataError handler passed to the WaitForQuery component.
+   */
+  if (downloadLink?.lastError && !(downloadLink?.lastError instanceof NotFoundError)) {
+    return (
+      <div style={{ textAlign: 'center' }}>
+        <h1>The deck will be ready soon. Please try again later.</h1>
+        <p>
+          If you're using the real-time deck, it's because the 'publish' function has not been
+          invoked yet.
+        </p>
+        <p>
+          If you're not using the real-time deck, it's because the corresponding task is still in
+          progress.
+        </p>
+      </div>
+    );
+  }
+
+  // 404 or no error case
   return (
     <WaitForData {...downloadLink} loadingComponent={LoadingSpinner}>
       <iframe