MCP PostgreSQL Operations

import asyncpg
import logging
import os
from typing import Any, Dict, List, Optional, Union
import json
from datetime import datetime

# Logger configuration
logger = logging.getLogger(__name__)

# PostgreSQL connection configuration
POSTGRES_CONFIG = {
    "host": os.getenv("POSTGRES_HOST", "localhost"),
    "port": int(os.getenv("POSTGRES_PORT", "5432")),
    "user": os.getenv("POSTGRES_USER", "postgres"),
    "password": os.getenv("POSTGRES_PASSWORD", ""),
    "database": os.getenv("POSTGRES_DB", "postgres"),
}


async def get_current_database_name(database: str = None) -> str:
    """Get the name of the currently connected database.
    
    Args:
        database: Database name to connect to. If None, uses default from config.
        
    Returns:
        Current database name as string
    """
    try:
        query = "SELECT current_database() as database_name"
        result = await execute_query(query, [], database=database)
        return result[0]['database_name'] if result else "unknown"
    except Exception as e:
        logger.error(f"Failed to get current database name: {e}")
        return "unknown"


async def get_db_connection(database: str = None) -> asyncpg.Connection:
    """Create PostgreSQL database connection.
    
    Args:
        database: Database name to connect to. If None, uses default from config.
    """
    try:
        config = POSTGRES_CONFIG.copy()
        if database:
            config["database"] = database
            
        conn = await asyncpg.connect(**config)
        logger.debug(f"Connected to PostgreSQL at {config['host']}:{config['port']}/{config['database']}")
        return conn
    except Exception as e:
        logger.error(f"Failed to connect to PostgreSQL: {e}")
        raise


async def execute_query(query: str, params: Optional[List] = None, database: str = None) -> List[Dict[str, Any]]:
    """Execute query and return results.
    
    Args:
        query: SQL query to execute
        params: Query parameters
        database: Database name to connect to. If None, uses default from config.
    """
    conn = None
    try:
        conn = await get_db_connection(database)
        if params:
            rows = await conn.fetch(query, *params)
        else:
            rows = await conn.fetch(query)
        
        # Convert Record to Dict
        result = []
        for row in rows:
            result.append(dict(row))
        
        logger.debug(f"Query executed successfully, returned {len(result)} rows")
        return result
        
    except Exception as e:
        logger.error(f"Query execution failed: {e}")
        logger.debug(f"Failed query: {query}")
        raise
    finally:
        if conn:
            await conn.close()


async def execute_single_query(query: str, params: Optional[List] = None, database: str = None) -> Optional[Dict[str, Any]]:
    """Execute query that returns a single result.
    
    Args:
        query: SQL query to execute
        params: Query parameters  
        database: Database name to connect to. If None, uses default from config.
    """
    results = await execute_query(query, params, database)
    return results[0] if results else None


def format_bytes(bytes_value: Union[int, float, None]) -> str:
    """Format byte values into human-readable format."""
    if bytes_value is None:
        return "N/A"
    
    bytes_value = float(bytes_value)
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if bytes_value < 1024.0:
            return f"{bytes_value:.2f} {unit}"
        bytes_value /= 1024.0
    return f"{bytes_value:.2f} PB"


def format_duration(seconds: Union[int, float, None]) -> str:
    """Format seconds into human-readable format."""
    if seconds is None:
        return "N/A"
    
    seconds = float(seconds)
    if seconds < 60:
        return f"{seconds:.2f}s"
    elif seconds < 3600:
        return f"{seconds/60:.2f}m"
    else:
        return f"{seconds/3600:.2f}h"


def format_table_data(data: List[Dict[str, Any]], title: str = "") -> str:
    """Convert table data into formatted string."""
    if not data:
        return f"No data found{' for ' + title if title else ''}"
    
    result = []
    if title:
        result.append(f"=== {title} ===\n")
    
    # Format as table
    if len(data) == 1:
        # Display single record as key-value pairs
        for key, value in data[0].items():
            if isinstance(value, (int, float)) and key.endswith(('_bytes', '_size')):
                value = format_bytes(value)
            elif isinstance(value, (int, float)) and key.endswith(('_time', '_duration')):
                value = format_duration(value)
            result.append(f"{key}: {value}")
    else:
        # Display multiple records as table format
        headers = list(data[0].keys())
        
        # Add headers
        result.append(" | ".join(headers))
        result.append("-" * (sum(len(h) for h in headers) + len(headers) * 3 - 1))
        
        # Add data rows
        for row in data:
            formatted_row = []
            for key, value in row.items():
                if isinstance(value, (int, float)) and key.endswith(('_bytes', '_size')):
                    formatted_row.append(format_bytes(value))
                elif isinstance(value, (int, float)) and key.endswith(('_time', '_duration')):
                    formatted_row.append(format_duration(value))
                else:
                    formatted_row.append(str(value) if value is not None else "NULL")
            result.append(" | ".join(formatted_row))
    
    return "\n".join(result)


async def get_server_version() -> str:
    """Return PostgreSQL server version."""
    try:
        result = await execute_single_query("SELECT version()")
        return result["version"] if result else "Unknown"
    except Exception as e:
        logger.error(f"Failed to get server version: {e}")
        return f"Error: {e}"


async def check_extension_exists(extension_name: str) -> bool:
    """Check if extension is installed."""
    try:
        query = "SELECT 1 FROM pg_extension WHERE extname = $1"
        result = await execute_single_query(query, [extension_name])
        return result is not None
    except Exception:
        return False


# pg_stat_statements related functions
async def get_pg_stat_statements_data(limit: int = 20, database: str = None) -> List[Dict[str, Any]]:
    """Get query statistics from pg_stat_statements.
    
    Args:
        limit: Maximum number of results to return
        database: Database name to query (uses default if omitted)
    """
    from .version_compat import get_pg_stat_statements_query
    
    try:
        # Get version-compatible query
        base_query = await get_pg_stat_statements_query(database)
        query = f"{base_query} LIMIT $1"
        
        return await execute_query(query, [limit], database=database)
    except Exception as e:
        logger.error(f"Failed to fetch pg_stat_statements data: {e}")
        raise Exception(f"Failed to fetch pg_stat_statements data: {e}")


# pg_stat_monitor related functions
async def get_pg_stat_monitor_data(limit: int = 20, database: str = None) -> List[Dict[str, Any]]:
    """Get query statistics from pg_stat_monitor.
    
    Args:
        limit: Maximum number of results to return
        database: Database name to query (uses default if omitted)
    """
    from .version_compat import get_pg_stat_monitor_query
    
    try:
        # Get version-compatible query
        base_query = await get_pg_stat_monitor_query(database)
        query = f"{base_query} LIMIT $1"
        
        return await execute_query(query, [limit], database=database)
    except Exception as e:
        logger.error(f"Failed to fetch pg_stat_monitor data: {e}")
        raise Exception(f"Failed to fetch pg_stat_monitor data: {e}")


def sanitize_connection_info() -> Dict[str, Any]:
    """Remove sensitive information from connection info."""
    config = POSTGRES_CONFIG.copy()
    config["password"] = "***"
    return config


def read_prompt_template(path: str) -> str:
    """
    Reads the MCP prompt template file and returns its content as a string.
    """
    with open(path, "r", encoding="utf-8") as f:
        return f.read()


def parse_prompt_sections(template: str):
    """
    Parses the prompt template into section headings and sections.
    Returns (headings, sections).
    """
    lines = template.splitlines()
    sections = []
    current = []
    headings = []
    for line in lines:
        if line.startswith("## "):
            if current:
                sections.append("\n".join(current))
                current = []
            headings.append(line[3:].strip())
            current.append(line)
        else:
            current.append(line)
    if current:
        sections.append("\n".join(current))
    return headings, sections