db_handler.py

from typing import Any, Dict, List, Optional
import mysql.connector
from mysql.connector import Error
import psycopg2
import psycopg2.extras

# Global database connection storage
_connections: Dict[str, Any] = {}

def connect_database(connection_name: str, db_type: str, host: str, port: int,
                     user: str, password: str, database: str) -> Dict[str, Any]:
    """Connect to a database (MySQL or PostgreSQL)."""
    try:
        if connection_name in _connections:
            try:
                _connections[connection_name]['connection'].close()
            except:
                pass

        connection = None
        if db_type == "mysql":
            connection = mysql.connector.connect(
                host=host,
                port=port,
                user=user,
                password=password,
                database=database,
                autocommit=True
            )
        elif db_type == "postgresql":
            connection = psycopg2.connect(
                host=host,
                port=port,
                user=user,
                password=password,
                dbname=database
            )
            connection.autocommit = True
        else:
            return {"success": False, "error": "Unsupported database type"}

        if connection:
            _connections[connection_name] = {'connection': connection, 'db_type': db_type}
            return {
                "success": True,
                "message": f"Connected to {db_type} database '{database}' on {host}:{port}",
                "connection_name": connection_name
            }
        else:
            return {"success": False, "error": "Failed to establish connection"}

    except (Error, psycopg2.Error) as e:
        return {"success": False, "error": f"Database connection failed: {str(e)}"}
    except Exception as e:
        return {"success": False, "error": f"Unexpected error: {str(e)}"}

def execute_query(connection_name: str, query: str, parameters: Optional[List] = None) -> Dict[str, Any]:
    """Execute a READ-ONLY SQL query on the specified connection."""
    raw = query.strip()

    if ";" in raw.strip().rstrip(";"):
        return {"success": False, "error": "Multiple statements are not allowed."}

    first_token = (raw.split(None, 1)[0] if raw else "").upper()
    if first_token == "DESC":
        first_token = "DESCRIBE"

    allowed = {"SELECT", "SHOW", "DESCRIBE", "EXPLAIN"}
    if first_token not in allowed:
        return {"success": False, "error": f"'{first_token}' is not allowed. Only SELECT, SHOW, DESCRIBE, EXPLAIN are permitted."}

    if connection_name not in _connections:
        return {"success": False, "error": f"No connection found with name '{connection_name}'"}

    connection_info = _connections[connection_name]
    connection = connection_info['connection']
    db_type = connection_info['db_type']

    try:
        if db_type == "mysql":
            cursor = connection.cursor(dictionary=True)
        elif db_type == "postgresql":
            cursor = connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)

        cursor.execute(raw, parameters or [])

        result = cursor.fetchall()
        return {
            "success": True,
            "data": result,
            "row_count": len(result)
        }

    except (Error, psycopg2.Error) as e:
        return {"success": False, "error": f"Query execution failed: {str(e)}"}
    except Exception as e:
        return {"success": False, "error": f"Unexpected error: {str(e)}"}
    finally:
        if 'cursor' in locals():
            cursor.close()

def list_tables(connection_name: str) -> Dict[str, Any]:
    """List all tables in the database."""
    if connection_name not in _connections:
        return {"success": False, "error": f"No connection found with name '{connection_name}'"}

    db_type = _connections[connection_name]['db_type']
    query = ""
    if db_type == "mysql":
        query = "SHOW TABLES"
    elif db_type == "postgresql":
        query = "SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname != 'pg_catalog' AND schemaname != 'information_schema'"

    return execute_query(connection_name, query)

def describe_table(connection_name: str, table_name: str) -> Dict[str, Any]:
    """Get the structure of a specific table."""
    if connection_name not in _connections:
        return {"success": False, "error": f"No connection found with name '{connection_name}'"}

    db_type = _connections[connection_name]['db_type']

    if db_type == "mysql":
        query = f"DESCRIBE `{table_name}`" # MySQL's DESCRIBE doesn't support parameters for table names
        return execute_query(connection_name, query)
    elif db_type == "postgresql":
        query = "SELECT column_name, data_type, character_maximum_length FROM information_schema.columns WHERE table_name = %s"
        return execute_query(connection_name, query, [table_name])

    return {"success": False, "error": "Unsupported database for describe_table"}

def close_all_connections():
    """Close all active database connections."""
    for name, conn_info in _connections.items():
        try:
            conn_info['connection'].close()
        except:
            pass