fix and tested mongodb data loader

Author: Chenglong-MS

py-src/data_formulator/data_loader/mongodb_data_loader.py

@@ -24,7 +24,8 @@ def list_params() -> bool:
             {"name": "username", "type": "string", "required": False, "default": "", "description": ""},
             {"name": "password", "type": "string", "required": False, "default": "", "description": ""},
             {"name": "database", "type": "string", "required": True, "default": "", "description": ""},
-            {"name": "collection", "type": "string", "required": False, "default": "", "description": "If specified, only this collection will be accessed"}
+            {"name": "collection", "type": "string", "required": False, "default": "", "description": "If specified, only this collection will be accessed"},
+            {"name": "authSource", "type": "string", "required": False, "default": "", "description": "Authentication database (defaults to target database if empty)"}
         ]
         return params_list
 
@@ -67,9 +68,17 @@ def __init__(self, params: Dict[str, Any], duck_db_conn: duckdb.DuckDBPyConnecti
             password = self.params.get("password", "")
             database = self.params.get("database", "")
             collection = self.params.get("collection", "")
+            auth_source = self.params.get("authSource", "") or database  # Default to target database
 
             if username and password:
-                self.mongo_client = pymongo.MongoClient(host=host, port=port, username=username, password=password)
+                # Use authSource to specify which database contains user credentials
+                self.mongo_client = pymongo.MongoClient(
+                    host=host, 
+                    port=port, 
+                    username=username, 
+                    password=password,
+                    authSource=auth_source
+                )
             else:
                 self.mongo_client = pymongo.MongoClient(host=host, port=port)
 
@@ -81,6 +90,28 @@ def __init__(self, params: Dict[str, Any], duck_db_conn: duckdb.DuckDBPyConnecti
         except Exception as e:
             raise Exception(f"Failed to connect to MongoDB: {e}")
 
+    def close(self):
+        """Close the MongoDB connection"""
+        if hasattr(self, 'mongo_client') and self.mongo_client is not None:
+            try:
+                self.mongo_client.close()
+                self.mongo_client = None
+            except Exception as e:
+                print(f"Warning: Failed to close MongoDB connection: {e}")
+
+    def __enter__(self):
+        """Context manager entry"""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures connection is closed"""
+        self.close()
+        return False
+
+    def __del__(self):
+        """Destructor to ensure connection is closed"""
+        self.close()
+    
     @staticmethod
     def _flatten_document(doc: Dict[str, Any], parent_key: str = '', sep: str = '_') -> Dict[str, Any]:
         """
@@ -237,7 +268,7 @@ def ingest_data(self, table_name: str, name_as: Optional[str] = None, size: int
         return
 
 
-    def view_query_sample(self, query: str) -> str:
+    def view_query_sample(self, query: str) -> List[Dict[str, Any]]:
 
         self._existed_collections_in_duckdb()
         self._difference_collections()
@@ -271,17 +302,27 @@ def ingest_data_from_query(self, query: str, name_as: str) -> pd.DataFrame:
         self._difference_collections()
         self._preload_all_collections(self.collection.name if self.collection else "")
 
-        df = self.duck_db_conn.execute(query).df()
+        query_result_df = self.duck_db_conn.execute(query).df()
 
         self._drop_all_loaded_tables()
 
-        for collection_name, df in self.existed_collections.items():
-            self._load_dataframe_to_duckdb(df, collection_name)
+        for collection_name, existing_df in self.existed_collections.items():
+            self._load_dataframe_to_duckdb(existing_df, collection_name)
 
-        self._load_dataframe_to_duckdb(df, name_as)
+        self._load_dataframe_to_duckdb(query_result_df, name_as)
 
-        return df
+        return query_result_df
 
+    @staticmethod
+    def _quote_identifier(name: str) -> str:
+        """
+        Safely quote a SQL identifier to prevent SQL injection.
+        Double quotes are escaped by doubling them.
+        """
+        # Escape any double quotes in the identifier by doubling them
+        escaped = name.replace('"', '""')
+        return f'"{escaped}"'
+
     def _existed_collections_in_duckdb(self):
         """
         Return the names and contents of tables already loaded into DuckDB
@@ -290,7 +331,8 @@ def _existed_collections_in_duckdb(self):
         duckdb_tables = self.duck_db_conn.execute("SHOW TABLES").df()
         for _, row in duckdb_tables.iterrows():
             collection_name = row['name']
-            df = self.duck_db_conn.execute(f"SELECT * FROM {collection_name}").df()
+            quoted_name = self._quote_identifier(collection_name)
+            df = self.duck_db_conn.execute(f"SELECT * FROM {quoted_name}").df()
             self.existed_collections[collection_name] = df
 
 
@@ -311,7 +353,8 @@ def _drop_all_loaded_tables(self):
         """
         for table_name in self.loaded_tables.values():
             try:
-                self.duck_db_conn.execute(f"DROP TABLE IF EXISTS main.{table_name}")
+                quoted_name = self._quote_identifier(table_name)
+                self.duck_db_conn.execute(f"DROP TABLE IF EXISTS main.{quoted_name}")
                 print(f"Dropped loaded table: {table_name}")
             except Exception as e:
                 print(f"Warning: Failed to drop table '{table_name}': {e}")
@@ -366,5 +409,10 @@ def _load_dataframe_to_duckdb(self, df: pd.DataFrame, table_name: str, size: int
 
         self.duck_db_conn.register(temp_view_name, df)
         # Use CREATE OR REPLACE to directly replace existing table
-        self.duck_db_conn.execute(f"CREATE OR REPLACE TABLE main.{table_name} AS SELECT * FROM {temp_view_name} LIMIT {size}")
-        self.duck_db_conn.execute(f"DROP VIEW {temp_view_name}")
+        # Quote identifiers to prevent SQL injection
+        quoted_table_name = self._quote_identifier(table_name)
+        quoted_temp_view = self._quote_identifier(temp_view_name)
+        # Ensure size is an integer to prevent injection via size parameter
+        safe_size = int(size)
+        self.duck_db_conn.execute(f"CREATE OR REPLACE TABLE main.{quoted_table_name} AS SELECT * FROM {quoted_temp_view} LIMIT {safe_size}")
+        self.duck_db_conn.execute(f"DROP VIEW {quoted_temp_view}")

src/views/DBTableManager.tsx

@@ -34,7 +34,8 @@ import {
   ToggleButtonGroup,
   ToggleButton,
   useTheme,
-  Link
+  Link,
+  Checkbox
 } from '@mui/material';
 
 import DeleteIcon from '@mui/icons-material/Delete';
@@ -820,9 +821,14 @@ export const DBTableSelectionDialog: React.FC<{
                         onImport={() => {
                             setIsUploading(true);
                         }} 
-                        onFinish={(status, message) => {
+                        onFinish={(status, message, importedTables) => {
                             setIsUploading(false);
-                            fetchTables();
+                            fetchTables().then(() => {
+                                // Navigate to the first imported table after tables are fetched
+                                if (status === "success" && importedTables && importedTables.length > 0) {
+                                    setSelectedTabKey(importedTables[0]);
+                                }
+                            });
                             if (status === "error") {
                                 setSystemMessage(message, "error");
                             }
@@ -1030,7 +1036,7 @@ export const DataLoaderForm: React.FC<{
     paramDefs: {name: string, default: string, type: string, required: boolean, description: string}[],
     authInstructions: string,
     onImport: () => void,
-    onFinish: (status: "success" | "error", message: string) => void
+    onFinish: (status: "success" | "error", message: string, importedTables?: string[]) => void
 }> = ({dataLoaderType, paramDefs, authInstructions, onImport, onFinish}) => {
 
     const dispatch = useDispatch();
@@ -1039,6 +1045,7 @@ export const DataLoaderForm: React.FC<{
 
     const [tableMetadata, setTableMetadata] = useState<Record<string, any>>({});    let [displaySamples, setDisplaySamples] = useState<Record<string, boolean>>({});
     let [tableFilter, setTableFilter] = useState<string>("");
+    const [selectedTables, setSelectedTables] = useState<Set<string>>(new Set());
 
     const [displayAuthInstructions, setDisplayAuthInstructions] = useState(false);
 
@@ -1086,7 +1093,12 @@ export const DataLoaderForm: React.FC<{
                     return [
                     <TableRow
                         key={tableName}
-                        sx={{ '&:last-child td, &:last-child th': { border: 0 }, '& .MuiTableCell-root': { padding: 0.25, wordWrap: 'break-word', whiteSpace: 'normal' }}}
+                        sx={{ 
+                            '&:last-child td, &:last-child th': { border: 0 }, 
+                            '& .MuiTableCell-root': { padding: 0.25, wordWrap: 'break-word', whiteSpace: 'normal' },
+                            backgroundColor: selectedTables.has(tableName) ? 'action.selected' : 'inherit',
+                            '&:hover': { backgroundColor: selectedTables.has(tableName) ? 'action.selected' : 'action.hover' }
+                        }}
                     >
                         <TableCell sx={{borderBottom: displaySamples[tableName] ? 'none' : '1px solid rgba(0, 0, 0, 0.1)'}}>
                             <IconButton size="small" onClick={() => toggleDisplaySamples(tableName)}>
@@ -1103,33 +1115,20 @@ export const DataLoaderForm: React.FC<{
                                 <Chip key={column.name} label={column.name} sx={{fontSize: 11, margin: 0.25, height: 20}} size="small" />
                             ))}
                         </TableCell>
-                        <TableCell sx={{width: 60}}>
-                            <Button size="small" onClick={() => {
-                                onImport();
-                                fetch(getUrls().DATA_LOADER_INGEST_DATA, {
-                                    method: 'POST',
-                                    headers: {
-                                        'Content-Type': 'application/json',
-                                    },
-                                    body: JSON.stringify({
-                                        data_loader_type: dataLoaderType, 
-                                        data_loader_params: params, table_name: tableName
-                                    })
-                                })
-                                .then(response => response.json())
-                                .then(data => {
-                                    
-                                    if (data.status === "success") {
-                                        onFinish("success", "Data ingested successfully");
+                        <TableCell sx={{width: 40}} padding="checkbox">
+                            <Checkbox
+                                size="small"
+                                checked={selectedTables.has(tableName)}
+                                onChange={(e) => {
+                                    const newSelected = new Set(selectedTables);
+                                    if (e.target.checked) {
+                                        newSelected.add(tableName);
                                     } else {
-                                        onFinish("error", data.error);
+                                        newSelected.delete(tableName);
                                     }
-                                })
-                                .catch(error => {
-                                    console.error('Failed to ingest data:', error);
-                                    onFinish("error", `Failed to ingest data: ${error}`);
-                                });
-                            }}>Import</Button>
+                                    setSelectedTables(newSelected);
+                                }}
+                            />
                         </TableCell>
                     </TableRow>,
                     <TableRow key={`${tableName}-sample`}>
@@ -1155,6 +1154,49 @@ export const DataLoaderForm: React.FC<{
                 </TableBody>
                 </Table>
             </TableContainer>,
+        mode === "view tables" && Object.keys(tableMetadata).length > 0 && <Box sx={{ display: 'flex', justifyContent: 'flex-end', mt: 1 }}>
+            <Button 
+                variant="contained" 
+                size="small"
+                disabled={selectedTables.size === 0}
+                onClick={() => {
+                    const tablesToImport = Array.from(selectedTables);
+                    onImport();
+                    
+                    // Import all selected tables sequentially
+                    const importPromises = tablesToImport.map(tableName => 
+                        fetch(getUrls().DATA_LOADER_INGEST_DATA, {
+                            method: 'POST',
+                            headers: {
+                                'Content-Type': 'application/json',
+                            },
+                            body: JSON.stringify({
+                                data_loader_type: dataLoaderType, 
+                                data_loader_params: params, 
+                                table_name: tableName
+                            })
+                        }).then(response => response.json())
+                    );
+                    
+                    Promise.all(importPromises)
+                        .then(results => {
+                            const errors = results.filter(r => r.status !== "success");
+                            if (errors.length === 0) {
+                                setSelectedTables(new Set());
+                                onFinish("success", `Successfully imported ${tablesToImport.length} table(s)`, tablesToImport);
+                            } else {
+                                onFinish("error", `Failed to import some tables: ${errors.map(e => e.error).join(", ")}`);
+                            }
+                        })
+                        .catch(error => {
+                            console.error('Failed to ingest data:', error);
+                            onFinish("error", `Failed to ingest data: ${error}`);
+                        });
+                }}
+            >
+                Import Selected ({selectedTables.size})
+            </Button>
+        </Box>,
         mode === "query" && <DataQueryForm 
             dataLoaderType={dataLoaderType} 
             availableTables={Object.keys(tableMetadata).map(t => ({name: t, fields: tableMetadata[t].columns.map((c: any) => c.name)}))}