| #!/usr/bin/env python |
| # -*- coding: utf-8 -*- |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """ |
| Apache Cloudberry MCP Server Implementation |
| |
| A Model Communication Protocol server for Apache Cloudberry database interaction |
| providing resources, tools, and prompts for database management. |
| """ |
| |
| from typing import Annotated, Any, Dict, List, Optional |
| import logging |
| from fastmcp import FastMCP |
| from pydantic import Field |
| |
| from .config import DatabaseConfig, ServerConfig |
| from .database import DatabaseManager |
| from .prompt import ( |
| ANALYZE_QUERY_PERFORMANCE_PROMPT, |
| SUGGEST_INDEXES_PROMPT, |
| DATABASE_HEALTH_CHECK_PROMPT |
| ) |
| |
| logger = logging.getLogger(__name__) |
| |
| class CloudberryMCPServer: |
| """Apache Cloudberry MCP Server implementation""" |
| |
| def __init__(self, server_config: ServerConfig, db_config: DatabaseConfig): |
| self.server_config = server_config |
| self.db_config = db_config |
| self.mcp = FastMCP("Apache Cloudberry MCP Server") |
| self.db_manager = DatabaseManager(db_config) |
| |
| self._setup_resources() |
| self._setup_tools() |
| self._setup_prompts() |
| |
| |
| def _setup_resources(self): |
| """Setup MCP resources for database metadata""" |
| |
| @self.mcp.resource("postgres://schemas", mime_type="application/json") |
| async def list_schemas() -> List[str]: |
| """List all database schemas""" |
| logger.info("Listing schemas") |
| return await self.db_manager.list_schemas() |
| |
| @self.mcp.resource("postgres://database/info", mime_type="application/json") |
| async def database_info() -> Dict[str, str]: |
| """Get general database information""" |
| logger.info("Getting database info") |
| return await self.db_manager.get_database_info() |
| |
| @self.mcp.resource("postgres://database/summary", mime_type="application/json") |
| async def database_summary() -> Dict[str, dict]: |
| """Get comprehensive database summary""" |
| logger.info("Getting database summary") |
| return await self.db_manager.get_database_summary() |
| |
| |
| def _setup_tools(self): |
| """Setup MCP tools for database operations""" |
| |
| @self.mcp.tool() |
| async def list_tables( |
| schema: Annotated[str, Field(description="The schema name to list tables from")] |
| ) -> List[str]: |
| """List tables in a specific schema""" |
| logger.info(f"Listing tables in schema: {schema}") |
| try: |
| return await self.db_manager.list_tables(schema) |
| except Exception as e: |
| return f"Error listing tables: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_views( |
| schema: Annotated[str, Field(description="The schema name to list views from")] |
| ) -> List[str]: |
| """List views in a specific schema""" |
| logger.info(f"Listing views in schema: {schema}") |
| try: |
| return await self.db_manager.list_views(schema) |
| except Exception as e: |
| return f"Error listing views: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_indexes( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name to list indexes for")] |
| ) -> List[str]: |
| """List indexes for a specific table""" |
| logger.info(f"Listing indexes for table: {schema}.{table}") |
| try: |
| indexes = await self.db_manager.list_indexes(schema, table) |
| return [f"{idx['indexname']}: {idx['indexdef']}" for idx in indexes] |
| except Exception as e: |
| return f"Error listing indexes: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_columns( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name to list columns for")] |
| ) -> List[Dict[str, Any]]: |
| """List columns for a specific table""" |
| logger.info(f"Listing columns for table: {schema}.{table}") |
| try: |
| return await self.db_manager.list_columns(schema, table) |
| except Exception as e: |
| return f"Error listing columns: {str(e)}" |
| |
| @self.mcp.tool() |
| async def execute_query( |
| query: Annotated[str, Field(description="The SQL query to execute")], |
| params: Annotated[Optional[Dict[str, Any]], Field(description="The parameters for the query")] = None, |
| readonly: Annotated[bool, Field(description="Whether the query is read-only")] = True |
| ) -> Dict[str, Any]: |
| """ |
| Execute a safe SQL query with parameters |
| """ |
| logger.info(f"Executing query: {query}") |
| try: |
| return await self.db_manager.execute_query(query, params, readonly) |
| except Exception as e: |
| return {"error": f"Error executing query: {str(e)}"} |
| |
| @self.mcp.tool() |
| async def explain_query( |
| query: Annotated[str, Field(description="The SQL query to explain")], |
| params: Annotated[Optional[Dict[str, Any]], Field(description="The parameters for the query")] = None |
| ) -> str: |
| """ |
| Get the execution plan for a query |
| """ |
| logger.info(f"Explaining query: {query}") |
| try: |
| return await self.db_manager.explain_query(query, params) |
| except Exception as e: |
| return f"Error explaining query: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_table_stats( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")], |
| ) -> Dict[str, Any]: |
| """ |
| Get statistics for a table |
| """ |
| logger.info(f"Getting table stats for: {schema}.{table}") |
| try: |
| result = await self.db_manager.get_table_stats(schema, table) |
| if "error" in result: |
| return result["error"] |
| return result |
| except Exception as e: |
| return f"Error getting table stats: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_large_tables(limit: Annotated[int, Field(description="Number of tables to return")] = 10) -> List[Dict[str, Any]]: |
| """ |
| List the largest tables in the database |
| """ |
| logger.info(f"Listing large tables, limit: {limit}") |
| try: |
| return await self.db_manager.list_large_tables(limit) |
| except Exception as e: |
| return f"Error listing large tables: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_database_schemas() -> List[str]: |
| """Get database schemas""" |
| logger.info("Getting database schemas") |
| try: |
| return await self.db_manager.list_schemas() |
| except Exception as e: |
| return f"Error getting schemas: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_database_information() -> Dict[str, str]: |
| """Get general database information""" |
| logger.info("Getting database information") |
| try: |
| return await self.db_manager.get_database_info() |
| except Exception as e: |
| return f"Error getting database info: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_database_summary() -> Dict[str, dict]: |
| """Get detailed database summary""" |
| logger.info("Getting database summary") |
| try: |
| return await self.db_manager.get_database_summary() |
| except Exception as e: |
| return f"Error getting database summary: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_users() -> List[str]: |
| """List all database users""" |
| logger.info("Listing database users") |
| try: |
| return await self.db_manager.list_users() |
| except Exception as e: |
| return f"Error listing users: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_user_permissions( |
| username: Annotated[str, Field(description="The username to check permissions for")] |
| ) -> List[Dict[str, Any]]: |
| """List permissions for a specific user""" |
| logger.info(f"Listing permissions for user: {username}") |
| try: |
| return await self.db_manager.list_user_permissions(username) |
| except Exception as e: |
| return f"Error listing user permissions: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_table_privileges( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> List[Dict[str, Any]]: |
| """List privileges for a specific table""" |
| logger.info(f"Listing table privileges for: {schema}.{table}") |
| try: |
| return await self.db_manager.list_table_privileges(schema, table) |
| except Exception as e: |
| return f"Error listing table privileges: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_constraints( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> List[Dict[str, Any]]: |
| """List constraints for a specific table""" |
| logger.info(f"Listing constraints for table: {schema}.{table}") |
| try: |
| return await self.db_manager.list_constraints(schema, table) |
| except Exception as e: |
| return f"Error listing constraints: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_foreign_keys( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> List[Dict[str, Any]]: |
| """List foreign keys for a specific table""" |
| logger.info(f"Listing foreign keys for table: {schema}.{table}") |
| try: |
| return await self.db_manager.list_foreign_keys(schema, table) |
| except Exception as e: |
| return f"Error listing foreign keys: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_referenced_tables( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> List[Dict[str, Any]]: |
| """List tables that reference this table""" |
| logger.info(f"Listing referenced tables for: {schema}.{table}") |
| try: |
| return await self.db_manager.list_referenced_tables(schema, table) |
| except Exception as e: |
| return f"Error listing referenced tables: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_slow_queries( |
| limit: Annotated[int, Field(description="Number of slow queries to return")] = 10 |
| ) -> List[Dict[str, Any]]: |
| """Get slow queries from database statistics""" |
| logger.info(f"Getting slow queries, limit: {limit}") |
| try: |
| return await self.db_manager.get_slow_queries(limit) |
| except Exception as e: |
| return f"Error getting slow queries: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_index_usage() -> List[Dict[str, Any]]: |
| """Get index usage statistics""" |
| logger.info("Getting index usage statistics") |
| try: |
| return await self.db_manager.get_index_usage() |
| except Exception as e: |
| return f"Error getting index usage: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_table_bloat_info() -> List[Dict[str, Any]]: |
| """Get table bloat information""" |
| logger.info("Getting table bloat information") |
| try: |
| return await self.db_manager.get_table_bloat_info() |
| except Exception as e: |
| return f"Error getting table bloat info: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_database_activity() -> List[Dict[str, Any]]: |
| """Get current database activity""" |
| logger.info("Getting database activity") |
| try: |
| return await self.db_manager.get_database_activity() |
| except Exception as e: |
| return f"Error getting database activity: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_functions( |
| schema: Annotated[str, Field(description="The schema name")] |
| ) -> List[Dict[str, Any]]: |
| """List functions in a specific schema""" |
| logger.info(f"Listing functions in schema: {schema}") |
| try: |
| return await self.db_manager.list_functions(schema) |
| except Exception as e: |
| return f"Error listing functions: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_function_definition( |
| schema: Annotated[str, Field(description="The schema name")], |
| function_name: Annotated[str, Field(description="The function name")] |
| ) -> str: |
| """Get function definition""" |
| logger.info(f"Getting function definition: {schema}.{function_name}") |
| try: |
| return await self.db_manager.get_function_definition(schema, function_name) |
| except Exception as e: |
| return f"Error getting function definition: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_triggers( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> List[Dict[str, Any]]: |
| """List triggers for a specific table""" |
| logger.info(f"Listing triggers for table: {schema}.{table}") |
| try: |
| return await self.db_manager.list_triggers(schema, table) |
| except Exception as e: |
| return f"Error listing triggers: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_table_ddl( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> str: |
| """Get DDL statement for a table""" |
| logger.info(f"Getting table DDL: {schema}.{table}") |
| try: |
| return await self.db_manager.get_table_ddl(schema, table) |
| except Exception as e: |
| return f"Error getting table DDL: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_materialized_views( |
| schema: Annotated[str, Field(description="The schema name")] |
| ) -> List[str]: |
| """List materialized views in a specific schema""" |
| logger.info(f"Listing materialized views in schema: {schema}") |
| try: |
| return await self.db_manager.list_materialized_views(schema) |
| except Exception as e: |
| return f"Error listing materialized views: {str(e)}" |
| |
| @self.mcp.tool() |
| async def get_vacuum_info( |
| schema: Annotated[str, Field(description="The schema name")], |
| table: Annotated[str, Field(description="The table name")] |
| ) -> Dict[str, Any]: |
| """Get vacuum information for a table""" |
| logger.info(f"Getting vacuum info for table: {schema}.{table}") |
| try: |
| return await self.db_manager.get_vacuum_info(schema, table) |
| except Exception as e: |
| return f"Error getting vacuum info: {str(e)}" |
| |
| @self.mcp.tool() |
| async def list_active_connections() -> List[Dict[str, Any]]: |
| """List active database connections""" |
| logger.info("Listing active connections") |
| try: |
| return await self.db_manager.list_active_connections() |
| except Exception as e: |
| return f"Error listing active connections: {str(e)}" |
| |
| def _setup_prompts(self): |
| """Setup MCP prompts for common database tasks""" |
| |
| @self.mcp.prompt() |
| def analyze_query_performance( |
| sql: Annotated[str, Field(description="The SQL query to analyze")], |
| explain: Annotated[str, Field(description="The EXPLAIN ANALYZE output")], |
| table_info: Annotated[str, Field(description="The table schema information")], |
| ) -> str: |
| """Prompt for analyzing query performance""" |
| logger.info(f"Analyzing query performance for: {sql}") |
| return ANALYZE_QUERY_PERFORMANCE_PROMPT.format( |
| sql=sql, |
| explain=explain, |
| table_info=table_info |
| ) |
| @self.mcp.prompt() |
| def suggest_indexes( |
| query: Annotated[str, Field(description="The common query pattern")], |
| table_info: Annotated[str, Field(description="The table schema information")], |
| table_stats: Annotated[str, Field(description="The table statistics")], |
| ) -> str: |
| """Prompt for suggesting indexes""" |
| logger.info(f"Suggesting indexes for query: {query}") |
| return SUGGEST_INDEXES_PROMPT.format( |
| query=query, |
| table_info=table_info, |
| table_stats=table_stats |
| ) |
| |
| |
| @self.mcp.prompt() |
| def database_health_check() -> str: |
| """Prompt for database health check""" |
| logger.info(f"Checking database health") |
| return DATABASE_HEALTH_CHECK_PROMPT |
| |
| def run(self, mode: str="http"): |
| """Run the MCP server""" |
| if mode == "stdio": |
| return self.mcp.run( |
| transport="stdio", |
| ) |
| elif mode == "http": |
| return self.mcp.run( |
| transport="streamable-http", |
| host=self.server_config.host, |
| port=self.server_config.port, |
| path=self.server_config.path, |
| stateless_http=True |
| ) |
| |
| async def close(self): |
| """Close the server and cleanup resources""" |
| await self.db_manager.close() |
| |
| |
| def main(): |
| """Main entry point""" |
| import argparse |
| |
| parser = argparse.ArgumentParser( |
| description="Cloudberry MCP Server - Cloudberry database management tools", |
| formatter_class=argparse.RawDescriptionHelpFormatter, |
| epilog=""" |
| Examples: |
| %(prog)s --mode stdio |
| %(prog)s --mode http --host 0.0.0.0 --port 8080 |
| %(prog)s --mode http --log-level INFO |
| %(prog)s --help |
| """ |
| ) |
| |
| parser.add_argument( |
| "--mode", |
| choices=["stdio", "http"], |
| default="http", |
| help="Server mode: stdio for stdin/stdout communication, http for HTTP server (default: http)" |
| ) |
| |
| parser.add_argument( |
| "--host", |
| default=None, |
| help="HTTP server host (default: from CLOUDBERRY_MCP_HOST env var or 127.0.0.1)" |
| ) |
| |
| parser.add_argument( |
| "--port", |
| type=int, |
| default=None, |
| help="HTTP server port (default: from CLOUDBERRY_MCP_PORT env var or 8080)" |
| ) |
| |
| parser.add_argument( |
| "--path", |
| default=None, |
| help="HTTP server path (default: from CLOUDBERRY_MCP_PATH env var or /mcp)" |
| ) |
| |
| parser.add_argument( |
| "--log-level", |
| choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], |
| default="WARNING", |
| help="Logging level (default: WARNING)" |
| ) |
| |
| parser.add_argument( |
| "--version", |
| action="version", |
| version="Cloudberry MCP Server 1.0.0" |
| ) |
| |
| args = parser.parse_args() |
| |
| # Configure logging |
| log_level = getattr(logging, args.log_level.upper()) |
| logging.basicConfig( |
| level=log_level, |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', |
| handlers=[ |
| logging.StreamHandler() |
| ] |
| ) |
| |
| # Create configurations |
| server_config = ServerConfig.from_env() |
| db_config = DatabaseConfig.from_env() |
| |
| # Override with command line arguments |
| if args.host: |
| server_config.host = args.host |
| if args.port: |
| server_config.port = args.port |
| if args.path: |
| server_config.path = args.path |
| |
| server = CloudberryMCPServer(server_config, db_config) |
| |
| try: |
| logger.info(f"Starting server in {args.mode} mode...") |
| server.run(args.mode) |
| except KeyboardInterrupt: |
| logger.error("Server stopped by user") |
| except Exception as e: |
| logger.error(f"Server error: {e}") |
| sys.exit(1) |
| finally: |
| import asyncio |
| asyncio.run(server.close()) |
| |
| |
| if __name__ == "__main__": |
| import sys |
| |
| main() |
