Refactor SparkAppConfigModel, Notebook.js, and spark_app.py to update…

… Spark app configuration handling
xuwenyihust · Aug 22, 2024 · c2bebc7 · c2bebc7
1 parent 2c2164f
commit c2bebc7
Show file tree

Hide file tree

Showing 5 changed files with 206 additions and 105 deletions.
diff --git a/server/app/routes/spark_app.py b/server/app/routes/spark_app.py
@@ -21,4 +21,4 @@ def get_spark_app_config(notbook_path):
 def update_spark_app_config(notbook_path):
     logging.info(f"Updating spark app config for notebook path: {notbook_path}")
     data = request.get_json()
-    return SparkApp.update_spark_app_config_by_notebook_id(notbook_path, data)
+    return SparkApp.update_spark_app_config_by_notebook_path(notbook_path, data)
diff --git a/server/tests/routes/test_spark_app_route.py b/server/tests/routes/test_spark_app_route.py
@@ -0,0 +1,69 @@
+import unittest
+import json
+from flask_cors import CORS
+from flask import g
+from database import db
+from run import create_app
+from app.routes.notebook import notebook_blueprint
+from app.routes.login import login_blueprint
+from app.services.directory import Directory
+from app.models.user import UserModel
+from app.services.user import User
+from app.models.spark_app import SparkAppModel
+from app.models.notebook import NotebookModel
+
+class SparkAppRouteTestCase(unittest.TestCase):
+
+  def setUp(self):
+    self.app = create_app()
+    self.app.register_blueprint(notebook_blueprint)
+    self.app.register_blueprint(login_blueprint)
+    self.client = self.app.test_client()
+    with self.app.app_context():
+      db.create_all()
+      user = UserModel(name='test_user', email='test_email')
+      user.set_password('test_password')
+      db.session.add(user)
+      db.session.commit()
+
+  def tearDown(self):
+    with self.app.app_context():
+      db.session.remove()
+      db.drop_all()
+
+  def login_and_get_token(self):
+    with self.app.app_context():
+      response = self.client.post('/login', auth=('test_user', 'test_password'))
+      return json.loads(response.data)['access_token']
+
+  def test_create_spark_app(self):
+    with self.app.app_context():
+      # Create Notebook
+      notebook = NotebookModel(name='Test Notebook', path='/path/to/notebook', user_id=1)
+      db.session.add(notebook)
+      db.session.commit()
+
+      # Create Spark App
+      spark_app_id = '1234'
+      path = f'/spark-app/${spark_app_id}'
+
+      data = {
+        'notebookPath': notebook.path
+      }
+
+      token = self.login_and_get_token()
+      headers = {
+        'Authorization': f'Bearer {token}',
+      }
+
+      response = self.client.get(
+        path,
+        headers=headers,
+        data=json.dumps(data),
+      )
+
+      self.assertEqual(response.status_code, 200)
+      self.assertEqual(json.loads(response.data)['spark_app_id'], spark_app_id)
+      self.assertEqual(json.loads(response.data)['notebook_id'], notebook.id)
+      self.assertEqual(json.loads(response.data)['user_id'], notebook.user_id)
+
diff --git a/webapp/src/components/notebook/Notebook.js b/webapp/src/components/notebook/Notebook.js
@@ -324,7 +324,7 @@ function Notebook({
                             deleteNotebook={handleDeleteNotebook}
                             /> : contentType === ContentType.Config ?
                             <Config 
-                                notebook={notebook}
+                                notebookPath={notebook.path}
                                 /> :
                             <Runs 
                                 notebook={notebook}

diff --git a/webapp/src/components/notebook/content/Config.js b/webapp/src/components/notebook/content/Config.js
@@ -1,22 +1,32 @@
 
 import React, { useEffect, useState } from 'react';
 import { Box, Card, CardHeader, CardContent, CardActions, Button, Typography, List, ListItem, ListItemText, TextField, Select, MenuItem } from '@mui/material';
+import SparkAppConfigModel from '../../../models/SparkAppConfigModel';
 
-function Config({ notebook }) {
+function Config({ notebookPath }) {
 
-  const [executorCores, setExecutorCores] = useState('1');
-  const [executorInstances, setExecutorInstances] = useState('1');
+  const [loading, setLoading] = useState(true);
 
-  const [executorMemory, setExecutorMemory] = useState('512');
+  const [executorCores, setExecutorCores] = useState(null);
+  const [executorInstances, setExecutorInstances] = useState(null);
+
+  const [executorMemory, setExecutorMemory] = useState(null);
   const [executorMemoryUnit, setExecutorMemoryUnit] = useState('m');
 
   useEffect(() => {
-    // const fetchSparkConfig = async () => {
+    const fetchSparkConfig = async () => {
+      const config = await SparkAppConfigModel.getSparkAppConfigByNotebookPath(notebookPath);
+      console.log('config: ', config);
+
+      setExecutorCores(config.executor_cores);
+      setExecutorInstances(config.executor_instances);
+      setExecutorMemory(config.executor_memory);
+      setExecutorMemoryUnit(config.executor_memory_unit);
 
-    // };
-    // fetchSparkConfig();
-    console.log('notebook:', notebook);
-  }, [notebook]);
+      setLoading(false); 
+    };
+    fetchSparkConfig();
+  }, [notebookPath]);
 
   const handleExecutorMemoryUnitChange = (event) => {
     setExecutorMemoryUnit(event.target.value);
@@ -28,85 +38,89 @@ function Config({ notebook }) {
     console.log('executorMemory:', executorMemory + executorMemoryUnit);
   }
 
-  return (
-    <Box sx={{
-      marginTop: 5,
-      marginRight: 5,
-      marginLeft: 2,
-    }}>
-      <Card 
-        sx={{ 
-          display: 'flex',
-          flexDirection: 'column',
-        }}>
-        <CardHeader 
-          title={
-            <Typography 
-              variant="body1"
-              style={{ marginLeft: 10 }}
-              color="textSecondary">
-              Spark Configuration
-            </Typography>
-          }
-          sx={{
-            backgroundColor: '#f5f5f5',
-            borderBottom: 1,
-            borderBottomColor: '#f5f5f5',
-          }}/>
-        <CardContent>
-          <List>
-            <ListItem>
-              <ListItemText primary="executor.memory" />
-              <TextField 
-                defaultValue={executorMemory}
-                variant="outlined"
-                size="small" 
-                onInput={(e) => e.target.value = e.target.value.replace(/[^0-9]/g, '')}
-                onChange={(e) => {
-                  setExecutorMemory(e.target.value)
-                  }}/>
-              <Box m={1} />
-              <Select value={executorMemoryUnit}
-                size="small"
-                onChange={handleExecutorMemoryUnitChange}>
-                <MenuItem value={'m'}>MB</MenuItem>
-                <MenuItem value={'g'}>GB</MenuItem>
-              </Select>
-            </ListItem>
+  if (loading) {
+    return <div>Loading...</div>;
+  } else {
+    return (
+      <Box sx={{
+        marginTop: 5,
+        marginRight: 5,
+        marginLeft: 2,
+      }}>
+        <Card 
+          sx={{ 
+            display: 'flex',
+            flexDirection: 'column',
+          }}>
+          <CardHeader 
+            title={
+              <Typography 
+                variant="body1"
+                style={{ marginLeft: 10 }}
+                color="textSecondary">
+                Spark Configuration
+              </Typography>
+            }
+            sx={{
+              backgroundColor: '#f5f5f5',
+              borderBottom: 1,
+              borderBottomColor: '#f5f5f5',
+            }}/>
+          <CardContent>
+            <List>
+              <ListItem>
+                <ListItemText primary="spark.executor.memory" />
+                <TextField 
+                  defaultValue={executorMemory}
+                  variant="outlined"
+                  size="small" 
+                  onInput={(e) => e.target.value = e.target.value.replace(/[^0-9]/g, '')}
+                  onChange={(e) => {
+                    setExecutorMemory(e.target.value)
+                    }}/>
+                <Box m={1} />
+                <Select value={executorMemoryUnit}
+                  size="small"
+                  onChange={handleExecutorMemoryUnitChange}>
+                  <MenuItem value={'m'}>MB</MenuItem>
+                  <MenuItem value={'g'}>GB</MenuItem>
+                </Select>
+              </ListItem>
 
-            <ListItem>
-              <ListItemText primary="executor.cores" />
-              <TextField 
-                defaultValue={executorCores}
-                variant="outlined"
-                size="small"
-                onChange={(e) => setExecutorCores(e.target.value)} />
-            </ListItem>
+              <ListItem>
+                <ListItemText primary="spark.executor.cores" />
+                <TextField 
+                  defaultValue={executorCores}
+                  variant="outlined"
+                  size="small"
+                  onChange={(e) => setExecutorCores(e.target.value)} />
+              </ListItem>
 
-            <ListItem>
-              <ListItemText primary="spark.executor.instances" />
-              <TextField
-                defaultValue={executorInstances}
-                variant="outlined"
-                size="small" 
-                onChange={(e) => setExecutorInstances(e.target.value)}/>
-            </ListItem>
-          </List>
-        </CardContent>
-        <CardActions>
-          <Button 
-            variant="outlined" 
-            style={{ 
-              marginLeft: '20px',
-              borderColor: 'lightgrey', 
-              color: 'grey' }}
-            onClick={handleSave}>
-            Save
-          </Button>
-        </CardActions>
-      </Card>
-    </Box>
-  );
+              <ListItem>
+                <ListItemText primary="spark.executor.instances" />
+                <TextField
+                  defaultValue={executorInstances}
+                  variant="outlined"
+                  size="small" 
+                  onChange={(e) => setExecutorInstances(e.target.value)}/>
+              </ListItem>
+            </List>
+          </CardContent>
+          <CardActions>
+            <Button 
+              variant="outlined" 
+              style={{ 
+                marginLeft: '20px',
+                borderColor: 'lightgrey', 
+                color: 'grey' }}
+              onClick={handleSave}>
+              Save
+            </Button>
+          </CardActions>
+        </Card>
+      </Box>
+    );
+  }
 }
 
 export default Config;
diff --git a/webapp/src/models/SparkAppConfigModel.js b/webapp/src/models/SparkAppConfigModel.js
@@ -4,27 +4,45 @@ class SparkAppConfigModel {
   constructor() {
   }
 
-  static async getSparkAppConfigByNotebookId() {
-    const response = await fetch(`${config.serverBaseUrl}/spark_app/${}/config`);
+  static async getSparkAppConfigByNotebookPath(notebookPath) {
+    const response = await fetch(`${config.serverBaseUrl}/spark_app/${notebookPath}/config`);
     if (!response.ok) {
       throw new Error('Failed to fetch Spark app config');
     } else {
       const data = await response.json();
-      return data;
+      console.log('Spark app config:', data);
+
+      const data_transformed = {};
+
+      const executorMemory = data['spark.executor.memory'];
+      const memoryUnit = executorMemory.slice(-1);
+      const memoryValue = executorMemory.slice(0, -1);
+      data_transformed.executor_memory = memoryValue;
+      data_transformed.executor_memory_unit = memoryUnit;
+
+      const executorCores = data['spark.executor.cores'];
+      data_transformed.executor_cores = executorCores;
+
+      const executorInstances = data['spark.executor.instances'];
+      data_transformed.executor_instances = executorInstances;
+
+      return data_transformed;
+    }
+  }
+
+  static async updateSparkAppConfig(notebookPath, sparkAppConfig) {
+    const response = await fetch(`${config.serverBaseUrl}/spark_app/${notebookPath}/config`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify(sparkAppConfig),
+    });
+
+    if (!response.ok) {
+      throw new Error('Failed to update Spark app config');
     }
   }
+}
 
-  // static async updateSparkAppConfig(sparkAppConfig) {
-  //   const response = await fetch(`${config.serverBaseUrl}/spark_app_config`, {
-  //     method: 'PATCH',
-  //     headers: {
-  //       'Content-Type': 'application/json',
-  //     },
-  //     body: JSON.stringify(sparkAppConfig),
-  //   });
-
-  //   if (!response.ok) {
-  //     throw new Error('Failed to update Spark app config');
-  //   }
-  // }
-}
+export default SparkAppConfigModel;