336 fix spark badge display (#337)

xuwenyihust · web-flow · commit 3d346d0497ba · 2024-12-10T20:30:20.000+08:00
* Enhance notebook loading process with Spark integration

- Added SparkModel import to App.js for improved Spark session handling.
- Refactored handleExistingNotebookClick to use async/await for fetching notebook data and Spark application details concurrently.
- Implemented error handling for fetching notebook and Spark app, ensuring better user feedback and session management.
- Introduced a new method in SparkModel.js to retrieve Spark app details by notebook path, enhancing integration with the backend API.

* Enhance App.js with notebook reference and improved Spark app ID handling

- Added useRef for notebook reference to facilitate direct interaction with the Notebook component.
- Updated Spark app ID handling to ensure it is set correctly when an active Spark application is detected.
- Minor refactoring for improved clarity and maintainability of the code.

* Enhance App.js with additional logging for notebook and Spark app integration

- Added console log to display the associated Spark application after fetching notebook data.
- Improved visibility into the notebook and Spark app interaction for better debugging and user feedback.

* Refactor Notebook component to support forward refs and expose Spark app ID setter

- Updated Notebook.js to use forwardRef, allowing parent components to directly interact with the Notebook's internal state.
- Implemented useImperativeHandle to expose a method for setting the Spark app ID, enhancing integration with parent components.
- Improved code clarity and maintainability by restructuring the component's function signature.

* Enhance logging in App.js and Notebook.js for improved debugging

- Added useEffect in App.js to log the current notebook reference, aiding in tracking its state.
- Implemented useEffect in Notebook.js to log changes to sparkAppId, enhancing visibility into Spark application interactions.
- Updated setSparkAppId method in Notebook.js to include logging when setting the Spark app ID, improving traceability during state changes.

* Refactor Notebook.js to enhance Spark app ID handling and session management

- Updated the Notebook component to reset sparkAppId when switching notebooks and immediately fetch the associated Spark app.
- Improved error handling during the Spark app fetch process, logging any failures for better debugging.
- Streamlined the useEffect hook to ensure proper state management and clarity in the notebook's functionality.

* Update demo notebook to modify execution state and add Spark session termination cell

- Changed the execution state of an existing code cell from true to false, indicating it was not executed.
- Added a new code cell to the demo notebook for stopping the Spark session, including execution metadata and success status.
- Enhanced the notebook structure for better clarity and usability in Spark session management.
diff --git a/examples/user_0@gmail.com/demo.ipynb b/examples/user_0@gmail.com/demo.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "isExecuted": true,
+   "isExecuted": false,
    "lastExecutionResult": "success",
    "lastExecutionTime": "2024-12-10 10:26:03",
    "metadata": {},
@@ -43,6 +43,18 @@
     "spark = create_spark(\"work/user_0@gmail.com/demo.ipynb\")\n",
     "spark"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "isExecuted": true,
+   "lastExecutionResult": "success",
+   "lastExecutionTime": "2024-12-10 12:27:14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "spark.stop()"
+   ]
   }
  ],
  "metadata": {
diff --git a/webapp/src/App.js b/webapp/src/App.js
@@ -1,4 +1,4 @@
-import React, { useState, useEffect } from 'react';
+import React, { useState, useEffect, useRef } from 'react';
 import LoginForm from './components/auth/LoginForm';
 import Sidebar from './components/sidebar/Sidebar';
 import Notebook from './components/notebook/Notebook';
@@ -8,6 +8,7 @@ import { createTheme, ThemeProvider } from '@mui/material/styles';
 import config from './config';
 import NotebookModel from './models/NotebookModel';
 import DirectoryModel from './models/DirectoryModel';
+import SparkModel from './models/SparkModel';
 
 const theme = createTheme({
   components: {
@@ -59,6 +60,12 @@ const App = () => {
   const [workspaceFiles, setWorkspaceFiles] = useState([]);
   const [refreshKey, setRefreshKey] = useState(0);
 
+  const notebookRef = useRef(null);
+
+  useEffect(() => {
+    console.log('notebookRef current:', notebookRef.current);
+  }, [notebookRef.current]);
+
   // Auth
   useEffect(() => {
     const storedUsername = localStorage.getItem('username');
@@ -128,22 +135,36 @@ const App = () => {
     }
   };  
 
-  const handleExistingNotebookClick = (path) => {
+  const handleExistingNotebookClick = async (path) => {
     if (handleUnsavedChanges()) {
-      NotebookModel.fetchNotebook(`${path}`).then((data) => {
-        if (data.message == 'Token has expired') {
-          console.error('Token has expired, please log in again');
-          logout();
-        } else {
-          console.log('Fetched notebook:', data);
-          setNotebook(data);
-          setShowHistoryServer(false);
-          setShowScheduler(false);
-          setShowNotebook(true);
+        try {
+            const [notebookData, sparkApp] = await Promise.all([
+                NotebookModel.fetchNotebook(`${path}`),
+                SparkModel.getSparkAppByNotebookPath(path)
+            ]);
+
+            if (notebookData.message === 'Token has expired') {
+                console.error('Token has expired, please log in again');
+                logout();
+                return;
+            }
+
+            console.log('Fetched notebook:', notebookData);
+            setNotebook(notebookData);
+
+            console.log('Associated Spark app:', sparkApp);
+            
+            // Update Spark badge if there's an active Spark app
+            if (sparkApp && sparkApp.spark_app_id) {
+                notebookRef.current?.setSparkAppId(sparkApp.spark_app_id);
+            }
+
+            setShowHistoryServer(false);
+            setShowScheduler(false);
+            setShowNotebook(true);
+        } catch (error) {
+            console.error('Failed to fetch notebook or Spark app:', error);
         }
-      }).catch((error) => {
-        console.error('Failed to fetch notebook:', error);
-      });
     }
   }
 
@@ -211,6 +232,7 @@ const App = () => {
             username={username}
             useremail={useremail}/>
           <Notebook 
+            ref={notebookRef}
             showNotebook={showNotebook}
             notebook={notebook}
             notebookState={notebookState}
diff --git a/webapp/src/components/notebook/Notebook.js b/webapp/src/components/notebook/Notebook.js
@@ -1,4 +1,4 @@
-import React, { useState, useEffect } from 'react';
+import React, { useState, useEffect, forwardRef, useImperativeHandle } from 'react';
 import NotebookHeader from './header/NotebookHeader';
 import Code from './content/Code';
 import Config from './content/Config';
@@ -13,15 +13,15 @@ import config from '../../config';
 import { Box } from '@mui/material';
 
 
-function Notebook({ 
+const Notebook = forwardRef(({ 
     showNotebook, 
     notebook,
     notebookState,
     setNotebookState,
     isNotebookModified,
     setIsNotebookModified,
-    handleDeleteNotebook }) {
-
+    handleDeleteNotebook 
+}, ref) => {
     const jupyterBaseUrl= `${config.jupyterBaseUrl}`
     const baseUrl = `${jupyterBaseUrl}/api/contents/`
 
@@ -61,7 +61,7 @@ function Notebook({
                 isExecuted: cell.cell_type === 'code' ? false : cell.cell_type === 'markdown' ? true : cell.isExecuted,
                 lastExecutionResult: cell.lastExecutionResult === null? null : cell.lastExecutionResult, 
                 lastExecutionTime: cell.lastExecutionTime === null? null : cell.lastExecutionTime
-              }));
+            }));
             setNotebookState({
                 ...notebook,
                 content: {
@@ -70,12 +70,24 @@ function Notebook({
                 }
             });
             setCurrentName(notebook.name);
+
+            // Reset sparkAppId when switching notebooks, but immediately fetch the associated Spark app
+            setSparkAppId(null);
+            SparkModel.getSparkAppByNotebookPath(notebook.path)
+                .then(sparkApp => {
+                    if (sparkApp && sparkApp.spark_app_id) {
+                        setSparkAppId(sparkApp.spark_app_id);
+                    }
+                })
+                .catch(error => {
+                    console.error('Failed to fetch Spark app:', error);
+                });
         }
+        
         SessionModel.getSession(notebook.path)
             .then((kernelId) => {
                 setKernelId(kernelId);
             });
-        setSparkAppId(null);
     }, [notebook]);
 
     const clearOutputs = () => {
@@ -318,6 +330,19 @@ function Notebook({
         }
     };
 
+    // Add useEffect to log sparkAppId changes
+    useEffect(() => {
+        console.log('sparkAppId changed:', sparkAppId);
+    }, [sparkAppId]);
+
+    // Expose setSparkAppId to parent through ref
+    useImperativeHandle(ref, () => ({
+        setSparkAppId: (id) => {
+            console.log('setSparkAppId called with:', id);
+            setSparkAppId(id);
+        }
+    }));
+
     return (
         <div>
             {showNotebook && (
@@ -382,6 +407,6 @@ function Notebook({
             )}
         </div>
     );
-}
+});
 
 export default Notebook;
diff --git a/webapp/src/models/SparkModel.js b/webapp/src/models/SparkModel.js
@@ -94,6 +94,28 @@ spark`;
         throw error;
     }
   }
+
+  static async getSparkAppByNotebookPath(notebookPath) {
+    const token = sessionStorage.getItem('token');
+    try {
+        const response = await fetch(`${config.serverBaseUrl}/notebook/spark_app/${notebookPath}`, {
+            headers: {
+                'Authorization': `Bearer ${token}`,
+                'Content-Type': 'application/json'
+            }
+        });
+        
+        if (!response.ok) {
+            return null;
+        }
+        
+        const sparkApps = await response.json();
+        return sparkApps.length > 0 ? sparkApps[0] : null;
+    } catch (error) {
+        console.error('Failed to fetch Spark app:', error);
+        return null;
+    }
+  }
 }
 
 export default SparkModel;
