PyMTA-server/email_server/server_web_ui/database.py

"""
Database utilities for the SMTP server.
"""

import sys
import subprocess
from urllib.parse import urlparse
from email_server.tool_box import get_logger

logger = get_logger()

# Database driver mappings
DB_DRIVERS = {
    'mysql': {
        'package': 'pymysql',
        'import_name': 'pymysql',
        'friendly_name': 'MySQL'
    },
    'postgresql': {
        'package': 'psycopg2-binary',
        'import_name': 'psycopg2',
        'friendly_name': 'PostgreSQL'
    },
    'mssql': {
        'package': 'pyodbc',
        'import_name': 'pyodbc',
        'friendly_name': 'MSSQL'
    }
}

def install_package(package_name: str) -> bool:
    """
    Install a Python package using pip.
    
    Args:
        package_name: Name of the package to install
        
    Returns:
        bool: True if installation was successful, False otherwise
    """
    try:
        logger.info(f"Installing {package_name}...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to install {package_name}: {e}")
        return False

def import_or_install_driver(db_type: str) -> bool:
    """
    Import database driver, installing it if necessary.
    
    Args:
        db_type: Type of database ('mysql', 'postgresql', 'mssql')
        
    Returns:
        bool: True if driver is available (installed or already present), False otherwise
    """
    if db_type not in DB_DRIVERS:
        return True  # SQLite or unsupported type
        
    driver_info = DB_DRIVERS[db_type]
    try:
        __import__(driver_info['import_name'])
        return True
    except ImportError:
        logger.warning(f"{driver_info['friendly_name']} driver not found. Attempting to install...")
        if install_package(driver_info['package']):
            try:
                __import__(driver_info['import_name'])
                logger.info(f"Successfully installed {driver_info['friendly_name']} driver")
                return True
            except ImportError:
                logger.error(f"Failed to import {driver_info['friendly_name']} driver after installation")
                return False
        return False

def test_database_connection(url: str) -> tuple[bool, str]:
    """
    Test if a database connection can be established.
    
    Args:
        url: Database connection URL
        
    Returns:
        tuple: (success: bool, message: str)
        
    Supported URL formats:
        - sqlite:///path/to/file.db
        - mysql://user:password@host:port/dbname
        - postgresql://user:password@host:port/dbname
        - mssql+pyodbc://user:password@host:port/dbname?driver=ODBC+Driver+17+for+SQL+Server
    """
    try:
        # Parse the database URL
        parsed = urlparse(url)
        scheme = parsed.scheme.lower()
        
        # SQLite connection test
        if scheme == 'sqlite':
            import sqlite3
            conn = sqlite3.connect(url.replace('sqlite:///', ''))
            conn.close()
            return True, "Successfully connected to SQLite database"
            
        # Other database types
        db_type = scheme.split('+')[0]  # Handle mssql+pyodbc
        if db_type not in DB_DRIVERS and db_type != 'sqlite':
            return False, f"Unsupported database type: {scheme}"
            
        # Try to import/install the required driver
        if not import_or_install_driver(db_type):
            return False, f"Failed to install required driver for {DB_DRIVERS[db_type]['friendly_name']}"
            
        # MySQL connection test
        if db_type == 'mysql':
            import pymysql
            params = {
                'host': parsed.hostname,
                'port': parsed.port or 3306,
                'user': parsed.username,
                'password': parsed.password,
                'db': parsed.path.lstrip('/')
            }
            conn = pymysql.connect(**params)
            conn.close()
            return True, "Successfully connected to MySQL database"
            
        # PostgreSQL connection test
        elif db_type == 'postgresql':
            import psycopg2
            conn = psycopg2.connect(url)
            conn.close()
            return True, "Successfully connected to PostgreSQL database"
            
        # MSSQL connection test
        elif db_type == 'mssql':
            import pyodbc
            conn = pyodbc.connect(url.replace('mssql+pyodbc://', ''))
            conn.close()
            return True, "Successfully connected to MSSQL database"
            
    except Exception as e:
        error_msg = str(e)
        logger.error(f"Database connection error: {error_msg}")
        return False, f"Connection error: {error_msg}"
        
    return False, "Unknown error occurred"