update performance and enhance data loader

Chenglong-MS · Chenglong-MS · commit f72b47a60764 · 2025-07-01T13:30:01.000-07:00
diff --git a/README.md b/README.md
@@ -26,9 +26,10 @@ Any questions? Ask on the Discord channel! [![Discord](https://img.shields.io/ba
 
 ## News 🔥🔥🔥
 
-- [05-13-2025] Data Formulator 0.2.3: External Data Loader 
+- [05-13-2025] Data Formulator 0.2.3 / 0.2.4: External Data Loader 
   - We introduced external data loader class to make import data easier. [Readme](https://github.com/microsoft/data-formulator/tree/main/py-src/data_formulator/data_loader) and [Demo](https://github.com/microsoft/data-formulator/pull/155)
     - Current data loaders: MySQL, Azure Data Explorer (Kusto), Azure Blob and Amazon S3 (json, parquet, csv).
+    - [07-01-2025] Updated with: Postgresql, mssql.
   - Call for action [link](https://github.com/microsoft/data-formulator/issues/156):
     - Users: let us know which data source you'd like to load data from.
     - Developers: let's build more data loaders.
diff --git a/py-src/data_formulator/sse_routes.py b/py-src/data_formulator/sse_routes.py
@@ -52,7 +52,7 @@ def event_stream() -> Generator[str, None, None]:
         
         # Send initial connection confirmation
         yield format_sse_message({
-            "type": "notification",
+            "type": "heartbeat",
             "text": "agent connection ready",
             "timestamp": time.time()
         })
diff --git a/py-src/data_formulator/tables_routes.py b/py-src/data_formulator/tables_routes.py
@@ -37,61 +37,6 @@
 
 tables_bp = Blueprint('tables', __name__, url_prefix='/api/tables')
 
-
-def list_tables_util(db_conn):
-    """
-    List all tables in the current session
-    """
-    results = []
-    
-    table_metadata_list = db_conn.execute("""
-        SELECT database_name, schema_name, table_name, schema_name==current_schema() as is_current_schema, 'table' as object_type 
-        FROM duckdb_tables() 
-        WHERE internal=False AND database_name == current_database()
-        UNION ALL 
-        SELECT database_name, schema_name, view_name as table_name, schema_name==current_schema() as is_current_schema, 'view' as object_type 
-        FROM duckdb_views()
-        WHERE view_name NOT LIKE 'duckdb_%' AND view_name NOT LIKE 'sqlite_%' AND view_name NOT LIKE 'pragma_%' AND database_name == current_database()
-    """).fetchall()
-
-    for table_metadata in table_metadata_list:
-        [database_name, schema_name, table_name, is_current_schema, object_type] = table_metadata
-        table_name = table_name if is_current_schema else '.'.join([database_name, schema_name, table_name])
-        if database_name in ['system', 'temp']:
-            continue
-        
-        print(f"table_metadata: {table_metadata}")
-
-        try:
-            # Get column information
-            columns = db_conn.execute(f"DESCRIBE {table_name}").fetchall()
-                
-            # Get row count
-            row_count = db_conn.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0]
-            sample_rows = db_conn.execute(f"SELECT * FROM {table_name} LIMIT 1000").fetchdf() if row_count > 0 else pd.DataFrame()
-            
-            # Check if this is a view or a table
-            try:
-                # Get both view existence and source in one query
-                view_info = db_conn.execute(f"SELECT view_name, sql FROM duckdb_views() WHERE view_name = '{table_name}'").fetchone()
-                view_source = view_info[1] if view_info else None
-            except Exception as e:
-                # If the query fails, assume it's a regular table
-                view_source = None
-
-            results.append({
-                "name": table_name,
-                "columns": [{"name": col[0], "type": col[1]} for col in columns],
-                "row_count": row_count,
-                "sample_rows": json.loads(sample_rows.to_json(orient='records')),
-                "view_source": view_source
-            })
-        except Exception as e:
-            logger.error(f"Error getting table metadata for {table_name}: {str(e)}")
-            continue
-
-    return results
-
 @tables_bp.route('/list-tables', methods=['GET'])
 def list_tables():
     """List all tables in the current session"""
@@ -123,7 +68,7 @@ def list_tables():
                     columns = db.execute(f"DESCRIBE {table_name}").fetchall()
                     # Get row count
                     row_count = db.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0]
-                    sample_rows = db.execute(f"SELECT * FROM {table_name} LIMIT 1000").fetchdf()
+                    sample_rows = db.execute(f"SELECT * FROM {table_name} LIMIT 1000").fetchdf() if row_count > 0 else pd.DataFrame()
                     
                     # Check if this is a view or a table
                     try:
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "data_formulator"
-version = "0.2.1.3"
+version = "0.2.1.4"
 
 requires-python = ">=3.9"
 authors = [
diff --git a/src/app/SSEActions.tsx b/src/app/SSEActions.tsx
@@ -0,0 +1,140 @@
+import { createTableFromFromObjectArray } from "../data/utils";
+import { DataFormulatorState, generateFreshChart, getDataFieldItems, SSEMessage } from "./dfSlice";
+
+import { Channel, Chart, ChartTemplate, EncodingItem, EncodingMap, FieldItem, Trigger } from '../components/ComponentType'
+
+
+let actionList = [
+    {
+        type: "create_chart",
+        requiredFields: ['chart_type', 'table_ref', 'encodings']
+    },
+    {
+        type: "load_table_from_object",
+        requiredFields: ['table_name', 'rows']
+    },
+    {
+        type: "derive_data_in_progress",
+        requiredFields: ['action_id', 'source_table_ids', 'instruction', 'fields']
+    },
+    {
+        type: "derive_data_completed",
+        requiredFields: ['action_id','derived_table']
+    },
+    {
+        type: "derive_data_failed",
+        requiredFields: ['action_id']
+    }
+]
+
+let checkActionRequiredFields = (message: SSEMessage, state: DataFormulatorState) => {
+    let action = actionList.find(a => a.type == message.data?.type);
+    if (!action) {
+        state.messages = [...state.messages, {
+            component: "server",
+            type: "error",
+            timestamp: message.timestamp,
+            value: `Unknown action type: ${message.data?.type}`
+        }];
+        return false;
+    }
+    let missingFields = action.requiredFields.filter(field => !message.data?.[field]);
+    if (missingFields.length > 0) {
+        state.messages = [...state.messages, {
+            component: "server",
+            type: "error",
+            timestamp: message.timestamp,
+            value: `[action] ${message.data?.type} - missing required fields: ${missingFields.join(', ')}`,
+            detail: JSON.stringify(message.data).slice(0, 1000)
+        }];
+        return false;
+    }
+    return true;
+}
+
+export const handleSSEMessage = (state: DataFormulatorState, message: SSEMessage) => {
+
+    if (message.type == "heartbeat") {
+        return;
+    }
+
+    if (message.type == "notification") {
+        state.messages = [...state.messages, {
+            component: "server",
+            type: "info",
+            timestamp: message.timestamp,
+            value: message.text || "Unknown message" + " (no data provided, no action taken)"
+        }];
+        return;
+    }
+
+    // otherwise, it's an action
+    // if it has no data, it's an error
+    if (!message.data) {
+        state.messages = [...state.messages, {
+            component: "server",
+            type: "warning",
+            timestamp: message.timestamp,
+            value: message.text || "Unknown message" + " (no data provided, no action taken)"
+        }];
+        return;
+    }
+
+    let action = message.data;
+    let actionStatus : 'ok' | 'error' | 'in_progress' = 'ok';
+
+    if (!checkActionRequiredFields(message, state)) {
+        return;
+    }
+
+    if (action.type == "create_chart") {
+        let chartType = action.chart_type;
+        let encodings = action.encodings;
+        let tableRef = action.table_ref;
+
+        let chart = generateFreshChart(tableRef, chartType);
+        for (let [channel, fieldName] of Object.entries(encodings)) {
+            let field = state.conceptShelfItems.find(f => f.name == fieldName);
+            if (field) {
+                chart.encodingMap[channel as Channel] = { fieldID: field.id };
+            } else {
+                let newField = { id: `custom--${fieldName}--${Date.now()}`, name: fieldName as string,
+                    type: "auto", source: "custom", domain: [], tableRef: 'custom' } as FieldItem;
+                state.conceptShelfItems = [newField, ...state.conceptShelfItems];
+                chart.encodingMap[channel as Channel] = { fieldID: newField.id };
+            }
+        }
+        state.charts = [...state.charts, chart];
+    } else if (action.type == "load_table_from_object") {
+        let rows = action.rows;
+        let tableName = action.table_name;
+        let table = createTableFromFromObjectArray(tableName, rows, false);
+        if (state.tables.find(t => t.id == table.id)) {
+            table.id = `${tableName}--${Date.now()}`;
+        }
+        state.tables = [...state.tables, table];
+        state.conceptShelfItems = [...state.conceptShelfItems, ...getDataFieldItems(table)];
+        state.focusedTableId = table.id;
+    } else if (action.type == "derive_data_in_progress") {
+        actionStatus = 'in_progress';
+        state.pendingSSEActions = [...state.pendingSSEActions, message];
+    } else if (action.type == "derive_data_completed") {
+        let actionId = action.action_id;
+        state.tables = [...state.tables, action.derived_table];
+        state.pendingSSEActions = state.pendingSSEActions.filter(m => m.data?.action_id != actionId);
+    } else if (action.type == "derive_data_failed") {
+        let actionId = action.action_id;
+        actionStatus = 'error';
+        state.pendingSSEActions = state.pendingSSEActions.filter(m => m.data?.action_id != actionId);
+    } else {
+        actionStatus = 'error';
+    }
+    state.messages = [...state.messages, {
+        component: "server",
+        type: actionStatus == 'ok' ? "success" : actionStatus == 'in_progress' ? "info" : "error",
+        timestamp: message.timestamp,
+        value: `[action] ${action.type} - ${message.text}`,
+        detail: actionStatus == 'error' ? JSON.stringify(action).slice(0, 1000) : undefined
+    }];
+    
+}
diff --git a/src/views/DBTableManager.tsx b/src/views/DBTableManager.tsx
@@ -32,7 +32,8 @@ import {
   Collapse,
   styled,
   ToggleButtonGroup,
-  ToggleButton
+  ToggleButton,
+  useTheme
 } from '@mui/material';
 
 import DeleteIcon from '@mui/icons-material/Delete';
@@ -43,6 +44,7 @@ import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
 import TableRowsIcon from '@mui/icons-material/TableRows';
 import RefreshIcon from '@mui/icons-material/Refresh';
 import ArrowForwardIcon from '@mui/icons-material/ArrowForward';
+import SearchIcon from '@mui/icons-material/Search';
 
 import { getUrls } from '../app/utils';
 import { CustomReactTable } from './ReactTable';
@@ -920,7 +922,7 @@ export const DataLoaderForm: React.FC<{
 }> = ({dataLoaderType, paramDefs, authInstructions, onImport, onFinish}) => {
 
     const dispatch = useDispatch();
-
+    const theme = useTheme();
     const params = useSelector((state: DataFormulatorState) => state.dataLoaderConnectParams[dataLoaderType] ?? {});
 
     const [tableMetadata, setTableMetadata] = useState<Record<string, any>>({});    let [displaySamples, setDisplaySamples] = useState<Record<string, boolean>>({});
@@ -1017,7 +1019,7 @@ export const DataLoaderForm: React.FC<{
                             }}>Import</Button>
                         </TableCell>
                     </TableRow>,
-                    <TableRow >
+                    <TableRow key={`${tableName}-sample`}>
                         <TableCell colSpan={4} sx={{ paddingBottom: 0, paddingTop: 0, px: 0, maxWidth: 800, overflowX: "auto", 
                                          borderBottom: displaySamples[tableName] ? '1px solid rgba(0, 0, 0, 0.1)' : 'none' }}>
                         <Collapse in={displaySamples[tableName]} timeout="auto" unmountOnExit>
@@ -1086,6 +1088,31 @@ export const DataLoaderForm: React.FC<{
                         />
                     </Box>
                 ))}
+                <TextField
+                    size="small"
+                    color="secondary"
+                    sx={{width: "270px", 
+                        '& .MuiInputLabel-root': {fontSize: 14, color: theme.palette.secondary.main},
+                        '& .MuiInputBase-root': {fontSize: 14},
+                        '& .MuiInputBase-input::placeholder': {fontSize: 12, fontStyle: "italic"},
+                        '&:hover': {
+                            backgroundColor: alpha(theme.palette.secondary.main, 0.03),
+                        },
+                    }}
+                    key="table-filter"
+                    autoComplete="off"
+                    variant="standard"
+                    label={<Box sx={{display: "flex", flexDirection: "row", alignItems: "center", gap: 0.5}}>
+                        <SearchIcon sx={{ fontSize: 16, color: theme.palette.secondary.main }} />
+                        table filter
+                    </Box>}
+                    placeholder="load only tables containing keywords"
+                    value={tableFilter}
+                    onChange={(event) => setTableFilter(event.target.value)}
+                    slotProps={{
+                        inputLabel: {shrink: true},
+                    }}
+                />
                 {paramDefs.length > 0 && <ButtonGroup sx={{height: 32, mt: 'auto'}} size="small" 
                  variant="contained" color="primary">
                     <Button 
@@ -1121,7 +1148,7 @@ export const DataLoaderForm: React.FC<{
                             setIsConnecting(false);
                         });
                     }}>
-                        {Object.keys(tableMetadata).length > 0 ? "refresh" : "connect"}
+                        {Object.keys(tableMetadata).length > 0 ? "refresh" : "connect"} {tableFilter.trim() ? "with filter" : ""}
                     </Button>
                     <Button 
                         disabled={Object.keys(tableMetadata).length === 0}
@@ -1131,66 +1158,12 @@ export const DataLoaderForm: React.FC<{
                         }}>
                         disconnect
                     </Button>
-                </ButtonGroup>                }
-                
+                </ButtonGroup>}
+
             </Box>            
-            {Object.keys(tableMetadata).length === 0 && (
-                <Box sx={{display: "flex", flexDirection: "row", alignItems: "center", gap: 1, ml: 4, mt: 4}}>
-                    <TextField
-                        size="small"
-                        sx={{
-                            width: "200px",
-                            '& .MuiInputLabel-root': {fontSize: 12},
-                            '& .MuiInputBase-root': {fontSize: 12, height: 32},
-                        }}
-                        variant="outlined"
-                        label="Filter tables"
-                        placeholder="Type to filter tables..."
-                        value={tableFilter}
-                        onChange={(event) => setTableFilter(event.target.value)}
-                        InputProps={{
-                            startAdornment: (
-                                <SearchIcon sx={{ fontSize: 16, color: 'text.secondary', mr: 0.5 }} />
-                            ),
-                        }}
-                    />
-                    <Button 
-                        size="small"
-                        sx={{textTransform: "none", height: 32}}
-                        onClick={() => {
-                            setIsConnecting(true);
-                            fetch(getUrls().DATA_LOADER_LIST_TABLES, {
-                                method: 'POST',
-                                headers: {
-                                    'Content-Type': 'application/json',
-                                },
-                                body: JSON.stringify({
-                                    data_loader_type: dataLoaderType, 
-                                    data_loader_params: params,
-                                    table_filter: tableFilter.trim() || null
-                                })
-                            }).then(response => response.json())
-                            .then(data => {
-                                if (data.status === "success") {
-                                    console.log(data.tables);
-                                    setTableMetadata(Object.fromEntries(data.tables.map((table: any) => {
-                                        return [table.name, table.metadata];
-                                    })));
-                                } else {
-                                    console.error('Failed to fetch data loader tables: {}', data.message);
-                                    onFinish("error", `Failed to fetch data loader tables: ${data.message}`);
-                                }
-                                setIsConnecting(false);
-                            })
-                            .catch(error => {
-                                onFinish("error", `Failed to fetch data loader tables, please check the server is running`);
-                                setIsConnecting(false);
-                            });
-                        }}>
-                        apply filter
-                    </Button>
-                </Box>
-            )}
+            <Box  sx={{display: "flex", flexDirection: "row", alignItems: "center", gap: 1, ml: 4, mt: 4}}>
+                
+            </Box>
             <Button
                 variant="text" 
                 size="small" 
diff --git a/src/views/SSEClient.tsx b/src/views/SSEClient.tsx
