supporting connection to mysql and pgsql

Author: Chenglong-MS

.env.template

@@ -3,4 +3,15 @@
 #   python -m data_formulator -p 5000 --exec-python-in-subprocess true --disable-display-keys true
 
 DISABLE_DISPLAY_KEYS=false # if true, the display keys will not be shown in the frontend
-EXEC_PYTHON_IN_SUBPROCESS=false # if true, the python code will be executed in a subprocess to avoid crashing the main app, but it will increase the time of response
+EXEC_PYTHON_IN_SUBPROCESS=false # if true, the python code will be executed in a subprocess to avoid crashing the main app, but it will increase the time of response
+
+# External atabase connection settings
+# check https://duckdb.org/docs/stable/extensions/mysql.html and https://duckdb.org/docs/stable/extensions/postgres.html
+USE_EXTERNAL_DB=false # if true, the app will use an external database instead of the one in the app
+DB_NAME=mysql_db # the name to refer to this database connection
+DB_TYPE=mysql # mysql or postgresql
+DB_HOST=localhost 
+DB_PORT=0 
+DB_DATABASE=mysql 
+DB_USER=root 
+DB_PASSWORD=

py-src/data_formulator/app.py

@@ -15,18 +15,20 @@
 
 import webbrowser
 import threading
+import numpy as np
+import datetime
+import time
 
 import logging
 
 import json
-import time
 from pathlib import Path
 
 from vega_datasets import data as vega_data
 
 from dotenv import load_dotenv
 import secrets
-
+import base64
 APP_ROOT = Path(os.path.join(Path(__file__).parent)).absolute()
 
 import os
@@ -38,6 +40,16 @@
 app = Flask(__name__, static_url_path='', static_folder=os.path.join(APP_ROOT, "dist"))
 app.secret_key = secrets.token_hex(16)  # Generate a random secret key for sessions
 
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.int64):
+            return int(obj)
+        if isinstance(obj, (bytes, bytearray)):
+            return base64.b64encode(obj).decode('ascii')
+        return super().default(obj)
+
+app.json_encoder = CustomJSONEncoder
+
 # Load env files early
 load_dotenv(os.path.join(APP_ROOT, "..", "..", 'api-keys.env'))
 load_dotenv(os.path.join(APP_ROOT, 'api-keys.env'))

py-src/data_formulator/db_manager.py

@@ -1,16 +1,22 @@
 import duckdb
 import pandas as pd
-from typing import Optional, Dict, List, ContextManager
+from typing import Optional, Dict, List, ContextManager, Any, Tuple
 import time
 from flask import session
 import tempfile
 import os
 from contextlib import contextmanager
+from dotenv import load_dotenv
 
 class DuckDBManager:
-    def __init__(self):
+    def __init__(self, external_db_connections: Dict[str, Dict[str, Any]]):
         # Store session db file paths
         self._db_files: Dict[str, str] = {}
+        # Track which extensions have been installed for which db files
+        self._installed_extensions: Dict[str, List[str]] = {}
+
+        # external db connections
+        self._external_db_connections: Dict[str, Dict[str, Any]] = external_db_connections
 
     @contextmanager
     def connection(self, session_id: str) -> ContextManager[duckdb.DuckDBPyConnection]:
@@ -31,13 +37,49 @@ def get_connection(self, session_id: str) -> duckdb.DuckDBPyConnection:
             db_file = os.path.join(tempfile.gettempdir(), f"df_{session_id}.db")
             print(f"=== Creating new db file: {db_file}")
             self._db_files[session_id] = db_file
+            # Initialize extension tracking for this file
+            self._installed_extensions[db_file] = []
         else:
             print(f"=== Using existing db file: {self._db_files[session_id]}")
             db_file = self._db_files[session_id]
 
         # Create a fresh connection to the database file
         conn = duckdb.connect(database=db_file)
+
+        if self._external_db_connections and self._external_db_connections['db_type'] in ['mysql', 'postgresql']:
+            db_name = self._external_db_connections['db_name']
+            db_type = self._external_db_connections['db_type']
+            
+            print(f"=== connecting to {db_type} extension")
+            # Only install if not already installed for this db file
+            if db_type not in self._installed_extensions.get(db_file, []):
+                conn.execute(f"INSTALL {db_type};")
+                self._installed_extensions[db_file].append(db_type)
+                
+            conn.execute(f"LOAD {db_type};")
+            conn.execute(f"""CREATE SECRET (
+                TYPE {db_type}, 
+                HOST '{self._external_db_connections['host']}', 
+                PORT '{self._external_db_connections['port']}', 
+                DATABASE '{self._external_db_connections['database']}', 
+                USER '{self._external_db_connections['user']}', 
+                PASSWORD '{self._external_db_connections['password']}');
+            """)
+            conn.execute(f"ATTACH '' AS {db_name} (TYPE {db_type});")
+            # result = conn.execute(f"SELECT * FROM {db_name}.information_schema.tables WHERE table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys');").fetch_df()
+            # print(f"=== result: {result}")
+                
         return conn
 
+env = load_dotenv()
+
 # Initialize the DB manager
-db_manager = DuckDBManager()
+db_manager = DuckDBManager({
+    "db_name": os.getenv('DB_NAME'),
+    "db_type": os.getenv('DB_TYPE'),
+    "host": os.getenv('DB_HOST'),
+    "port": os.getenv('DB_PORT'),
+    "database": os.getenv('DB_DATABASE'),
+    "user": os.getenv('DB_USER'),
+    "password": os.getenv('DB_PASSWORD')
+} if os.getenv('USE_EXTERNAL_DB') == 'true' else None)

py-src/data_formulator/tables_routes.py

@@ -7,10 +7,12 @@
 import mimetypes
 mimetypes.add_type('application/javascript', '.js')
 mimetypes.add_type('application/javascript', '.mjs')
+import json
 
 from flask import request, send_from_directory, session, jsonify, Blueprint
 import pandas as pd
-
+import random
+import string
 from pathlib import Path
 
 from data_formulator.db_manager import db_manager
@@ -39,15 +41,18 @@ def list_tables():
     try:
         result = []
         with db_manager.connection(session['session_id']) as db:
-            tables = db.execute("SHOW TABLES").fetchall()
+            table_metadata_list = db.execute("SELECT database_name, schema_name, table_name, schema_name==current_schema() as is_current_schema FROM duckdb_tables() WHERE internal=False").fetchall()
 
-            for table in tables:
-                table_name = table[0]
+            print(f"table_metadata_list: {table_metadata_list}")
+            for table_metadata in table_metadata_list:
+                [database_name, schema_name, table_name, is_current_schema] = table_metadata
+
+                table_name = table_name if is_current_schema else '.'.join([database_name, schema_name, table_name])
                 # Get column information
                 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").fetchall()
+                sample_rows = db.execute(f"SELECT * FROM {table_name} LIMIT 1000").fetchdf()
 
                 # Check if this is a view or a table
                 try:
@@ -62,7 +67,7 @@ def list_tables():
                     "name": table_name,
                     "columns": [{"name": col[0], "type": col[1]} for col in columns],
                     "row_count": row_count,
-                    "sample_rows": [dict(zip([col[0] for col in columns], row)) for row in sample_rows],
+                    "sample_rows": json.loads(sample_rows.to_json(orient='records')),
                     "view_source": view_source
                 })
 
@@ -186,11 +191,11 @@ def sample_table():
                 else:
                     query += f" ORDER BY ROWID DESC LIMIT {sample_size}"
 
-            result = db.execute(query).fetchall()
+            result = db.execute(query).fetchdf()
 
         return jsonify({
             "status": "success",
-            "rows": [dict(zip(output_column_names, row)) for row in result],
+            "rows": json.loads(result.to_json(orient='records')),
             "total_row_count": total_row_count
         })
     except Exception as e:
@@ -437,6 +442,108 @@ def upload_db_file():
             "message": safe_msg
         }), status_code
 
+
+def validate_db_connection_params(db_type: str, db_host: str, db_port: int, 
+                                db_database: str, db_user: str, db_password: str):
+    """Validate database connection parameters"""
+    # Validate db_type
+    valid_db_types = ['postgresql', 'mysql']
+    if not db_type or db_type.lower() not in valid_db_types:
+        raise ValueError(f"Invalid database type. Must be one of: {', '.join(valid_db_types)}")
+    
+    # Validate host (basic DNS/IP format check)
+    if not db_host or not re.match(r'^[a-zA-Z0-9.-]+$', db_host):
+        raise ValueError("Invalid host format")
+    
+    # Validate port
+    try:
+        port = int(db_port)
+        if not (1 <= port <= 65535):
+            raise ValueError()
+    except (ValueError, TypeError):
+        raise ValueError("Port must be a number between 1 and 65535")
+    
+    # Validate database name (alphanumeric and underscores only)
+    if not db_database or not re.match(r'^[a-zA-Z0-9_]+$', db_database):
+        raise ValueError("Invalid database name format")
+    
+    # Validate username (alphanumeric and some special chars)
+    if not db_user or not re.match(r'^[a-zA-Z0-9@._-]+$', db_user):
+        raise ValueError("Invalid username format")
+    
+    # Validate password exists
+    if not db_password:
+        raise ValueError("Password cannot be empty")
+
+@tables_bp.route('/attach-external-db', methods=['POST'])
+def attach_external_db():
+    """Attach an external db to the session"""
+    try:
+        data = request.get_json()
+        db_type = data.get('db_type')
+        db_host = data.get('db_host')
+        db_port = data.get('db_port')
+        db_database = data.get('db_database')
+        db_user = data.get('db_user')
+        db_password = data.get('db_password')
+
+        # Generate a random suffix for the database name
+        suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=2))
+        db_name = f"{db_type.lower()}_{suffix}"
+
+        if 'session_id' not in session:
+            return jsonify({"status": "error", "message": "No session ID found"}), 400
+        
+        with db_manager.connection(session['session_id']) as conn:
+            # Create secret using parameterized query
+
+            # Install and load the extension
+            if db_type == 'mysql':
+                conn.install_extension("mysql")
+                conn.load_extension("mysql")
+            elif db_type == 'postgresql':
+                conn.install_extension("postgres")
+                conn.load_extension("postgres")
+                
+            connect_query = f"""CREATE SECRET (
+                TYPE {db_type},
+                HOST '{db_host}',
+                PORT '{db_port}',
+                DATABASE '{db_database}',
+                USER '{db_user}',
+                PASSWORD '{db_password}'
+            );"""
+            conn.execute(connect_query)
+            
+            # Attach the database
+            conn.execute(f"ATTACH '' AS {db_name} (TYPE {db_type});")
+            
+            result = conn.execute(f"SELECT * FROM {db_name}.information_schema.tables WHERE table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys');").fetch_df()
+
+            print(f"result: {result}")
+
+            result = conn.execute(f"SELECT * FROM {db_name}.sakila.actor LIMIT 10;").fetchdf()
+
+            print(f"result: {result}")
+   
+            # Log what we found for debugging
+            logger.info(f"Found {len(result)} tables: {result}")
+
+            return jsonify({
+                "status": "success",
+                "message": "External database attached successfully",
+                "result": result
+            })
+        
+    except Exception as e:
+        logger.error(f"Error attaching external database: {str(e)}")
+        safe_msg, status_code = sanitize_db_error_message(e)
+        return jsonify({
+            "status": "error", 
+            "message": safe_msg
+        }), status_code
+
+
 @tables_bp.route('/download-db-file', methods=['GET'])
 def download_db_file():
     """Download the db file for a session"""

src/app/utils.tsx

@@ -57,6 +57,7 @@ export function getUrls() {
         UPLOAD_DB_FILE: `/api/tables/upload-db-file`,
         DOWNLOAD_DB_FILE: `/api/tables/download-db-file`,
         RESET_DB_FILE: `/api/tables/reset-db-file`,
+        ATTACH_EXTERNAL_DB: `/api/tables/attach-external-db`,
 
         LIST_TABLES: `/api/tables/list-tables`,
         TABLE_DATA: `/api/tables/get-table`,

src/views/DBTableManager.tsx

@@ -29,7 +29,8 @@ import {
   TableRow,
   CircularProgress,
   ButtonGroup,
-  Tooltip
+  Tooltip,
+  MenuItem
 } from '@mui/material';
 import DeleteIcon from '@mui/icons-material/Delete';
 import UploadFileIcon from '@mui/icons-material/UploadFile';
@@ -548,7 +549,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
                             {...a11yProps(i)} 
                         />
                     ))}
-                </Tabs> 
+                </Tabs>
                 <Divider sx={{my: 1}} textAlign='left'> <TuneIcon sx={{fontSize: 12, color: "text.secondary"}} /></Divider>
                 {uploadFileButton(<Typography component="span" fontSize={12}>{isUploading ? 'uploading...' : 'upload file'}</Typography>)}
             </Box>
@@ -595,7 +596,11 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
                                 />
                             ) : (
                                 <CustomReactTable 
-                                    rows={currentTable.sample_rows.slice(0, 9)} 
+                                    rows={currentTable.sample_rows.map((row: any) => {
+                                        return Object.fromEntries(Object.entries(row).map(([key, value]: [string, any]) => {
+                                            return [key, String(value)];
+                                        }));
+                                    }).slice(0, 9)} 
                                     columnDefs={currentTable.columns.map(col => ({
                                         id: col.name,
                                         label: col.name,