srikanthbaskaran/autonomous_software_solutions

1from fastmcp import FastMCP2from database import DatabaseExplorer3import logging4import atexit5import asyncio6import pymysql.cursors  # Import pymysql.cursors for explicit DictCursor usage7 8# Configure logging for this module9logging.basicConfig(10    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"11)12logger = logging.getLogger(__name__)13 14# Initialize database connection globally15db = None16 17 18def initialize_database():19    """Initialize database connection"""20    global db21    logger.info("Attempting to initialize database connection...")22    try:23        db = DatabaseExplorer()24        logger.info("Database connection established successfully.")25    except Exception as e:26        logger.error(27            f"Database connection failed during initialization: {e}", exc_info=True28        )29        raise  # Re-raise to prevent the server from starting without a DB connection30 31 32def cleanup_database():33    """Cleanup database connection"""34    global db35    logger.info("Attempting to clean up database connection...")36    if db:37        try:38            # Check if the connection object exists and has a close method39            if hasattr(db, "connection") and db.connection:40                if db.connection.is_connected():41                    db.connection.close()42                    logger.info("Database connection closed successfully.")43                else:44                    logger.info(45                        "Database connection was already closed or not connected."46                    )47            else:48                logger.info("Database connection object not found or not initialized.")49        except Exception as e:50            logger.warning(51                f"Error closing database connection during cleanup: {e}", exc_info=True52            )53    else:54        logger.info("No database connection to clean up (db object is None).")55 56 57# Register cleanup function to run on program exit58atexit.register(cleanup_database)59 60# Initialize database at startup61try:62    initialize_database()63except Exception:64    # If database initialization fails, log and exit65    logger.critical("Failed to initialize database. Exiting application.")66    exit(1)  # Exit with an error code67 68mcp = FastMCP("sql_report_server")69 70 71@mcp.tool()72async def get_database_schema() -> list:  # Changed return type to list of strings73    """Get a list of all table names in the database."""74    logger.info("Tool 'get_database_schema' called to get table names.")75    try:76        if not db:77            logger.error("Database object is not initialized in get_database_schema.")78            raise RuntimeError("Database not initialized")79 80        conn = db.get_connection()81        cursor = conn.cursor()82 83        try:84            cursor.execute("SHOW TABLES")85            raw_tables = cursor.fetchall()86 87            tables = []88            if raw_tables:89                # Dynamically get the key name for the table name column90                table_column_name = cursor.description[0][0]91                tables = [table[table_column_name] for table in raw_tables]92 93            logger.info(94                f"Successfully fetched {len(tables)} table names from database: {tables}"95            )96            return tables  # Return only table names97        finally:98            cursor.close()99            logger.debug("Cursor closed in get_database_schema.")100 101    except Exception as e:102        logger.error(f"Error in 'get_database_schema': {str(e)}", exc_info=True)103        raise  # Re-raise the exception so the agent knows the tool failed104 105 106@mcp.tool()107async def get_table_details(table_name: str) -> dict:108    """109    Get detailed schema (columns and foreign keys) for a specific table.110 111    Args:112        table_name: The name of the table to get details for.113    """114    logger.info(f"Tool 'get_table_details' called for table: '{table_name}'.")115    try:116        if not db:117            logger.error("Database object is not initialized in get_table_details.")118            raise RuntimeError("Database not initialized")119 120        conn = db.get_connection()121 122        # Use DictCursor for detailed schema queries123        dict_cursor = conn.cursor(pymysql.cursors.DictCursor)124        try:125            # Get columns126            dict_cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")127            columns = {col["Field"]: col["Type"] for col in dict_cursor.fetchall()}128            logger.debug(f"Columns for {table_name}: {columns}")129 130            # Get foreign keys131            dict_cursor.execute(132                f"""133                SELECT COLUMN_NAME, REFERENCED_TABLE_NAME 134                FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE 135                WHERE TABLE_SCHEMA = DATABASE()136                AND TABLE_NAME = '{table_name}' 137                AND REFERENCED_TABLE_NAME IS NOT NULL138                """139            )140            foreign_keys = {141                fk["COLUMN_NAME"]: fk["REFERENCED_TABLE_NAME"]142                for fk in dict_cursor.fetchall()143            }144            logger.debug(f"Foreign keys for {table_name}: {foreign_keys}")145 146            table_schema = {147                "table_name": table_name,148                "columns": columns,149                "foreign_keys": foreign_keys,150            }151            logger.info(f"Successfully retrieved details for table: {table_name}.")152            return table_schema153        except pymysql.Error as e:154            logger.error(155                f"MySQL error fetching details for table '{table_name}': {e}",156                exc_info=True,157            )158            raise ValueError(159                f"Could not retrieve details for table '{table_name}'. It might not exist or there's a permission issue."160            )161        except Exception as e:162            logger.error(163                f"Unexpected error fetching details for table '{table_name}': {e}",164                exc_info=True,165            )166            raise167        finally:168            dict_cursor.close()169            logger.debug(f"DictCursor closed for {table_name}.")170 171    except Exception as e:172        logger.error(f"Error in 'get_table_details': {str(e)}", exc_info=True)173        raise  # Re-raise the exception so the agent knows the tool failed174 175 176@mcp.tool()177async def generate_sql_report(query: str, rationale: str) -> dict:178    """Execute SQL query and return formatted report179 180    Args:181        query: The SQL query to execute182        rationale: Explanation why this query provides the needed report183    """184    logger.info(185        f"Tool 'generate_sql_report' called with query: '{query}' and rationale: '{rationale}'"186    )187    try:188        if not db:189            logger.error("Database object is not initialized in generate_sql_report.")190            raise RuntimeError("Database not initialized")191 192        report = db.execute_report_query(query, rationale)193        logger.info("SQL report generated successfully.")194        return report.model_dump()195    except Exception as e:196        logger.error(f"Error generating SQL report: {str(e)}", exc_info=True)197        raise  # Re-raise the exception198 199 200@mcp.tool()201async def generate_standard_report(report_type: str) -> dict:202    """Generate standard report based on common business needs203 204    Args:205        report_type: Type of report (sales, inventory, financial, etc.)206    """207    logger.info(208        f"Tool 'generate_standard_report' called for report type: '{report_type}')"209    )210    try:211        if not db:212            logger.error(213                "Database object is not initialized in generate_standard_report."214            )215            raise RuntimeError("Database not initialized")216 217        report = db.generate_sample_report(report_type)218        logger.info(f"Standard report '{report_type}' generated successfully.")219        return report.model_dump()220    except Exception as e:221        logger.error(222            f"Error generating standard report for type '{report_type}': {str(e)}",223            exc_info=True,224        )225        raise  # Re-raise the exception226 227 228@mcp.tool()229async def get_report_history() -> list:230    """Get history of all generated reports"""231    logger.info("Tool 'get_report_history' called.")232    try:233        if not db:234            logger.error("Database object is not initialized in get_report_history.")235            raise RuntimeError("Database not initialized")236 237        history = db.get_report_history()238        logger.info(f"Successfully retrieved {len(history)} historical reports.")239        return history240    except Exception as e:241        logger.error(f"Error getting report history: {str(e)}", exc_info=True)242        raise  # Re-raise the exception243 244 245@mcp.resource("sqlreports://schema")246async def read_schema_resource() -> list:  # Changed return type to list247    """Resource endpoint for database schema (returns only table names)"""248    logger.info("Resource 'sqlreports://schema' accessed.")249    try:250        if not db:251            logger.error("Database object is not initialized for schema resource.")252            raise RuntimeError("Database not initialized")253 254        # Call the tool to get just the table names255        table_names = await get_database_schema()256        logger.info(257            f"Successfully prepared table names for schema resource: {len(table_names)} tables."258        )259        return table_names  # Return only table names260    except Exception as e:261        logger.error(f"Error reading schema resource: {str(e)}", exc_info=True)262        raise  # Re-raise the exception263 264 265@mcp.resource("sqlreports://history")266async def read_history_resource() -> list:267    """Resource endpoint for report history"""268    logger.info("Resource 'sqlreports://history' accessed.")269    try:270        if not db:271            logger.error("Database object is not initialized for history resource.")272            raise RuntimeError("Database not initialized")273 274        history = db.get_report_history()275        logger.info(276            f"Successfully prepared history for resource: {len(history)} entries."277        )278        return history279    except Exception as e:280        logger.error(f"Error reading history resource: {str(e)}", exc_info=True)281        raise  # Re-raise the exception282 283 284if __name__ == "__main__":285    logger.info("Starting SQL Report MCP Server")286    try:287        # FastMCP's run method handles the asyncio event loop internally288        mcp.run(transport="stdio")289    except KeyboardInterrupt:290        logger.info("Server interrupted by user (KeyboardInterrupt).")291    except Exception as e:292        logger.critical(293            f"SQL Report MCP Server crashed unexpectedly: {str(e)}", exc_info=True294        )295        # Re-raise to ensure the error is propagated if this is a critical failure296        raise297