fix various db manager issues, need to add a function to connect external DBs

Chenglong-MS · Chenglong-MS · commit 56c3ce7a462a · 2025-04-18T17:08:05.000-07:00
diff --git a/py-src/data_formulator/agents/agent_sql_data_transform.py b/py-src/data_formulator/agents/agent_sql_data_transform.py
@@ -55,65 +55,6 @@
 '''
 
 EXAMPLE='''
-
-For example:
-
-[CONTEXT]
-
-Here are our datasets, here are their field summaries and samples:
-
-table_0 (us_covid_cases) fields:
-	Date -- type: object, values: 1/1/2021, 1/1/2022, 1/1/2023, ..., 9/8/2022, 9/9/2020, 9/9/2021, 9/9/2022
-	Cases -- type: int64, values: -23999, -14195, -6940, ..., 1018935, 1032159, 1178403, 1433977
-
-table_0 (us_covid_cases) sample:
-```
-|Date|Cases
-0|1/21/2020|1
-1|1/22/2020|0
-2|1/23/2020|0
-3|1/24/2020|1
-4|1/25/2020|1
-......
-```
-
-[GOAL]
-
-{
-    "instruction": "calculate 7-day moving average",
-    "visualization_fields": ["Date", "7-day average cases"]
-}
-
-[OUTPUT]
-
-{  
-    "detailed_instruction": "Calculate the 7-day moving average of COVID-19 cases over time.",  
-    "output_fields": ["Date", "Cases", "7-day average cases"],  
-    "visualization_fields": ["Date", "7-day average cases"],  
-    "reason": "To calculate the 7-day moving average, the 'Cases' field is required, but it is not needed for visualization. The provided fields are sufficient to achieve the goal."  
-}  
-
-```python
-import pandas as pd  
-import collections  
-import numpy as np  
-  
-def transform_data(df):  
-    # Convert Date column to datetime  
-    df['Date'] = pd.to_datetime(df['Date'])  
-      
-    # Sort the dataframe by Date  
-    df = df.sort_values('Date')  
-      
-    # Calculate the 7-day moving average of cases  
-    df['7-day average cases'] = df['Cases'].rolling(window=7).mean()  
-      
-    # Select the output fields  
-    transformed_df = df[['Date', 'Cases', '7-day average cases']]  
-      
-    return transformed_df  
-```
-
 [CONTEXT]
 
 Here are our datasets, here are their field summaries and samples:
diff --git a/py-src/data_formulator/tables_routes.py b/py-src/data_formulator/tables_routes.py
@@ -287,6 +287,17 @@ def create_table():
             return jsonify({"status": "error", "message": "Invalid table name"}), 400
             
         with db_manager.connection(session['session_id']) as db:
+            # Check if table exists and generate unique name if needed
+            base_name = sanitized_table_name
+            counter = 1
+            while True:
+                # Check if table exists
+                exists = db.execute(f"SELECT COUNT(*) FROM duckdb_tables() WHERE table_name = '{sanitized_table_name}'").fetchone()[0] > 0
+                if not exists:
+                    break
+                # If exists, append counter to base name
+                sanitized_table_name = f"{base_name}_{counter}"
+                counter += 1
         
             # Read file based on extension
             if file.filename.endswith('.csv'):
@@ -297,7 +308,7 @@ def create_table():
                 df = pd.read_json(file)
             else:
                 return jsonify({"status": "error", "message": "Unsupported file format"}), 400
-            
+
             # Create table
             db.register('df_temp', df)
             db.execute(f"CREATE TABLE {sanitized_table_name} AS SELECT * FROM df_temp")
@@ -307,7 +318,9 @@ def create_table():
                 "status": "success",
                 "table_name": sanitized_table_name,
                 "row_count": len(df),
-                "columns": list(df.columns)
+                "columns": list(df.columns),
+                "original_name": base_name,  # Include the original name in response
+                "is_renamed": base_name != sanitized_table_name  # Flag indicating if name was changed
             })
     
     except Exception as e:
@@ -329,10 +342,21 @@ def drop_table():
             return jsonify({"status": "error", "message": "No table name provided"}), 400
             
         with db_manager.connection(session['session_id']) as db:
-            # First try to drop it as a view
-            db.execute(f"DROP VIEW IF EXISTS {table_name}")
-            # Then try to drop it as a table
-            db.execute(f"DROP TABLE IF EXISTS {table_name}")
+            # First check if it exists as a view
+            view_exists = db.execute(f"SELECT view_name FROM duckdb_views() WHERE view_name = '{table_name}'").fetchone() is not None
+            if view_exists:
+                db.execute(f"DROP VIEW IF EXISTS {table_name}")
+            
+            # Then check if it exists as a table
+            table_exists = db.execute(f"SELECT table_name FROM duckdb_tables() WHERE table_name = '{table_name}'").fetchone() is not None
+            if table_exists:
+                db.execute(f"DROP TABLE IF EXISTS {table_name}")
+
+            if not view_exists and not table_exists:
+                return jsonify({
+                    "status": "error",
+                    "message": f"Table/view '{table_name}' does not exist"
+                }), 404
         
             return jsonify({
                 "status": "success",
diff --git a/src/app/App.tsx b/src/app/App.tsx
@@ -618,7 +618,7 @@ export const AppFC: FC<AppFCProps> = function AppFC(appProps) {
                     <Divider orientation="vertical" variant="middle" flexItem />
                     <DBTableSelectionDialog buttonElement={
                         <Typography sx={{ display: 'flex', fontSize: 14, alignItems: 'center', gap: 1, textTransform: 'none' }}>
-                            <CloudQueueIcon fontSize="small" /> View Database
+                            <CloudQueueIcon fontSize="small" /> Database
                         </Typography>
                     } />
                     <Divider orientation="vertical" variant="middle" flexItem />
diff --git a/src/views/DBTableManager.tsx b/src/views/DBTableManager.tsx
@@ -226,10 +226,14 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
     const [dbTables, setDbTables] = useState<DBTable[]>([]);
     const [selectedTabIndex, setSelectedTabIndex] = useState(0);
     
-    const [errorMessage, setErrorMessage] = useState<string | null>(null);
+    const [errorMessage, setErrorMessage] = useState<{content: string, severity: "error" | "warning" | "info" | "success"} | null>(null);
     const [showError, setShowError] = useState(false);
     const [isUploading, setIsUploading] = useState<boolean>(false);
 
+    useEffect(() => {
+        fetchTables();
+    }, []);
+
     // Fetch list of tables
     const fetchTables = async () => {
         try {
@@ -267,7 +271,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
             }
         } catch (error) {
             console.error('Failed to upload table:', error);
-            setErrorMessage('Failed to upload table. The server may need to be restarted.');
+            setErrorMessage({content: 'Failed to upload table. The server may need to be restarted.', severity: "error"});
             setShowError(true);
         } finally {
             setIsUploading(false);
@@ -304,7 +308,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
             URL.revokeObjectURL(url);
         } catch (error) {
             console.error('Failed to download database:', error);
-            setErrorMessage('Failed to download database file');
+            setErrorMessage({content: 'Failed to download database file', severity: "error"});
             setShowError(true);
         }
     };
@@ -325,18 +329,25 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
             });
             const data = await response.json();
             if (data.status === 'success') {
+                if (data.is_renamed) {
+                    setErrorMessage({content: `Table ${data.original_name} already exists. Renamed to ${data.table_name}`, severity: "warning"});
+                    setShowError(true);
+                } 
                 fetchTables();  // Refresh table list
             } else {
-                // Handle error from server
-                setErrorMessage(data.error || 'Failed to upload table');
+                setErrorMessage({content: data.error || 'Failed to upload table', severity: "error"});
                 setShowError(true);
             }
         } catch (error) {
             console.error('Failed to upload table:', error);
-            setErrorMessage('Failed to upload table. The server may need to be restarted.');
+            setErrorMessage({content: 'Failed to upload table. The server may need to be restarted.', severity: "error"});
             setShowError(true);
         } finally {
             setIsUploading(false);
+            // Clear the file input value to allow uploading the same file again
+            if (event.target) {
+                event.target.value = '';
+            }
         }
     };
 
@@ -354,7 +365,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
             }
         } catch (error) {
             console.error('Failed to reset database:', error);
-            setErrorMessage('Failed to reset database');
+            setErrorMessage({content: 'Failed to reset database', severity: "error"});
             setShowError(true);
         }
     }
@@ -409,7 +420,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
             }
         } catch (error) {
             console.error('Failed to analyze table data:', error);
-            setErrorMessage('Failed to analyze table data');
+            setErrorMessage({content: 'Failed to analyze table data', severity: "error"});
             setShowError(true);
         }
     };
@@ -476,6 +487,44 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
         }
     }, [tableDialogOpen]);
 
+    let importButton = (buttonElement: React.ReactNode) => {
+        return <Tooltip title="import a duckdb .db file to the local database">
+            <Button variant="text" sx={{fontSize: "inherit", minWidth: "auto"}} component="label" disabled={isUploading}>
+                {buttonElement}
+                <input type="file" hidden onChange={handleDBUpload} accept=".db" disabled={isUploading} />
+            </Button>
+        </Tooltip>
+    }
+
+    let exportButton = 
+        <Tooltip title="save the local database to a duckdb .db file">
+            <Button variant="text" size="small" onClick={handleDBDownload} disabled={isUploading || dbTables.length === 0}>
+                export
+            </Button>
+        </Tooltip>
+
+    function uploadFileButton(element: React.ReactNode, buttonSx?: SxProps) {
+        return (
+            <Tooltip title="upload a csv/tsv file to the local database">
+                <Button
+                    variant="text"
+                    component="label"
+                    sx={{ fontSize: "inherit", ...buttonSx}}                    
+                    disabled={isUploading}
+                >
+                    {element}
+                    <input
+                        type="file"
+                        hidden
+                        onChange={handleFileUpload}
+                        accept=".csv,.xlsx,.json"
+                        disabled={isUploading}
+                    />
+                </Button>
+            </Tooltip>
+        );
+    }
+
     let mainContent = dbTables.length > 0 ? 
         <Box sx={{ flexGrow: 1, bgcolor: 'background.paper', display: 'flex', flexDirection: 'row', height: '100%' }}>
             <Box sx={{display: "flex", flexDirection: "column", width: "120px", borderRight: 1, borderColor: 'divider'}}>
@@ -485,7 +534,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
                     value={selectedTabIndex}
                     onChange={handleTabChange}
                     aria-label="Database tables"
-                    sx={{ width: '120px' }}
+                    sx={{ width: '120px', maxHeight: '360px' }}
                 >
                     {dbTables.map((t, i) => (
                         <Tab 
@@ -501,25 +550,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
                     ))}
                 </Tabs> 
                 <Divider sx={{my: 1}} textAlign='left'> <TuneIcon sx={{fontSize: 12, color: "text.secondary"}} /></Divider>
-                <Tooltip title="upload a csv/tsv file to the local database">
-                    <Button
-                        variant="text"
-                        component="label"
-                        size="small"
-                        sx={{ width: '100%', textTransform: "none"}}                    
-                        disabled={isUploading}
-                        startIcon={<UploadIcon fontSize="small"/>}
-                    >
-                        {isUploading ? 'uploading...' : 'upload file'}
-                        <input
-                            type="file"
-                            hidden
-                            onChange={handleFileUpload}
-                            accept=".csv,.xlsx,.json"
-                            disabled={isUploading}
-                        />
-                    </Button>
-                </Tooltip>
+                {uploadFileButton(<Typography component="span" fontSize={12}>{isUploading ? 'uploading...' : 'upload file'}</Typography>)}
             </Box>
             {dbTables.map((t, i) => {
                 const currentTable = t;
@@ -582,8 +613,12 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
         </Box> : 
         <Box sx={{ p: 3, textAlign: 'center', height: '100%', display: 'flex', flexDirection: 'column', justifyContent: 'center', alignItems: 'center' }}>
             <StorageIcon sx={{ fontSize: 60, color: 'text.secondary', mb: 2 }} />
-            <Typography variant="h6" gutterBottom>
-                Please upload a file to create a table
+            <Typography variant="caption"> Database is currently empty. </Typography>
+            <Typography>
+                {uploadFileButton(<Typography component="span">Upload a csv dataset </Typography>)}
+                or
+                {importButton(<Typography component="span">Import a db file</Typography>)}
+                <Typography component="span"> to get started.</Typography>
             </Typography>
         </Box>
 
@@ -602,8 +637,8 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
                 onClose={handleCloseError}
                 anchorOrigin={{ vertical: 'top', horizontal: 'center' }}
             >
-                <Alert onClose={handleCloseError} severity="error" sx={{ width: '100%' }}>
-                    {errorMessage}
+                <Alert onClose={handleCloseError} severity={errorMessage?.severity} sx={{ width: '100%' }}>
+                    {errorMessage?.content}
                 </Alert>
             </Snackbar>
             
@@ -649,30 +684,9 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
                 </DialogContent>
                 <DialogActions>
                     <Typography variant="caption" sx={{ mr: 'auto', '& .MuiButton-root': { minWidth: 'auto',  textTransform: "none" } }}>
-                        <Button
-                            variant="text" 
-                            size="small"
-                            component="label"
-                            disabled={isUploading}
-                        >
-                            Reload
-                            <input
-                                type="file"
-                                hidden
-                                onChange={handleDBUpload}
-                                accept=".db"
-                                disabled={isUploading}
-                            />
-                        </Button>
+                        {importButton(<Typography component="span" fontSize="inherit">Import</Typography>)}
                         ,
-                        <Button
-                            variant="text" size="small"
-                           // endIcon={<DownloadIcon />}
-                            onClick={handleDBDownload}
-                            disabled={isUploading || dbTables.length === 0}
-                        >
-                            download
-                        </Button>
+                        {exportButton}
                         or
                         <Button
                             variant="text" size="small"
diff --git a/src/views/TableSelectionView.tsx b/src/views/TableSelectionView.tsx
@@ -284,6 +284,17 @@ export const TableUploadDialog: React.FC<TableUploadDialogProps> = ({ buttonElem
                     file.name.endsWith('.csv') || 
                     file.name.endsWith('.tsv') || 
                     file.name.endsWith('.json')) {
+
+                    // Check if file is larger than 5MB
+                    const MAX_FILE_SIZE = 5 * 1024 * 1024; // 5MB in bytes
+                    if (file.size > MAX_FILE_SIZE) {
+                        dispatch(dfActions.addMessages({
+                            "timestamp": Date.now(),
+                            "type": "error",
+                            "value": `File ${file.name} is too large (${(file.size / (1024 * 1024)).toFixed(2)}MB), upload it via DATABASE option instead.`
+                        }));
+                        continue; // Skip this file and process the next one
+                    }
                     
                     // Handle text files
                     file.text().then((text) => {
