diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index f3ba8f8b..5fdcffcb 100644 --- a/src/postgres_mcp/server.py +++ b/src/postgres_mcp/server.py @@ -82,7 +82,7 @@ def format_error_response(error: str) -> ResponseType: @mcp.tool( - description="List all schemas in the database", + description="List all schemas in the PostgreSQL database with their owners and types.", annotations=ToolAnnotations( title="List Schemas", readOnlyHint=True, @@ -114,7 +114,7 @@ async def list_schemas() -> ResponseType: @mcp.tool( - description="List objects in a schema", + description="List tables, views, materialized views, sequences, or extensions in a PostgreSQL schema. Returns names and comments/descriptions for each object.", annotations=ToolAnnotations( title="List Objects", readOnlyHint=True, @@ -122,7 +122,7 @@ async def list_schemas() -> ResponseType: ) async def list_objects( schema_name: str = Field(description="Schema name"), - object_type: str = Field(description="Object type: 'table', 'view', 'sequence', or 'extension'", default="table"), + object_type: str = Field(description="Object type: 'table', 'view', 'materialized_view', 'sequence', or 'extension'", default="table"), ) -> ResponseType: """List objects of a given type in a schema.""" try: @@ -133,15 +133,59 @@ async def list_objects( rows = await SafeSqlDriver.execute_param_query( sql_driver, """ - SELECT table_schema, table_name, table_type - FROM information_schema.tables - WHERE table_schema = {} AND table_type = {} - ORDER BY table_name + SELECT + t.table_schema, + t.table_name, + t.table_type, + obj_description( + (quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::regclass, + 'pg_class' + ) AS comment + FROM information_schema.tables t + WHERE t.table_schema = {} AND t.table_type = {} + ORDER BY t.table_name """, [schema_name, table_type], ) objects = ( - [{"schema": row.cells["table_schema"], "name": row.cells["table_name"], "type": row.cells["table_type"]} for row in rows] + [ + { + "schema": row.cells["table_schema"], + "name": row.cells["table_name"], + "type": row.cells["table_type"], + "comment": row.cells["comment"], + } + for row in rows + ] + if rows + else [] + ) + + elif object_type == "materialized_view": + rows = await SafeSqlDriver.execute_param_query( + sql_driver, + """ + SELECT + n.nspname AS schema_name, + c.relname AS view_name, + obj_description(c.oid, 'pg_class') AS comment + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relkind = 'm' AND n.nspname = {} + ORDER BY c.relname + """, + [schema_name], + ) + objects = ( + [ + { + "schema": row.cells["schema_name"], + "name": row.cells["view_name"], + "type": "MATERIALIZED VIEW", + "comment": row.cells["comment"], + } + for row in rows + ] if rows else [] ) @@ -188,7 +232,7 @@ async def list_objects( @mcp.tool( - description="Show detailed information about a database object", + description="Show columns (with types, nullability, defaults, and comments), constraints, indexes, and the table/view comment for a PostgreSQL object. Use this BEFORE writing queries to understand the schema.", annotations=ToolAnnotations( title="Get Object Details", readOnlyHint=True, @@ -197,24 +241,53 @@ async def list_objects( async def get_object_details( schema_name: str = Field(description="Schema name"), object_name: str = Field(description="Object name"), - object_type: str = Field(description="Object type: 'table', 'view', 'sequence', or 'extension'", default="table"), + object_type: str = Field(description="Object type: 'table', 'view', 'materialized_view', 'sequence', or 'extension'", default="table"), ) -> ResponseType: """Get detailed information about a database object.""" try: sql_driver = await get_sql_driver() - if object_type in ("table", "view"): - # Get columns - col_rows = await SafeSqlDriver.execute_param_query( - sql_driver, - """ - SELECT column_name, data_type, is_nullable, column_default - FROM information_schema.columns - WHERE table_schema = {} AND table_name = {} - ORDER BY ordinal_position - """, - [schema_name, object_name], - ) + if object_type in ("table", "view", "materialized_view"): + # Get columns with comments + if object_type == "materialized_view": + # Materialized views are not in information_schema.columns — use pg_catalog + col_rows = await SafeSqlDriver.execute_param_query( + sql_driver, + """ + SELECT + a.attname AS column_name, + pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type, + CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable, + NULL AS column_default, + col_description(c.oid, a.attnum) AS comment + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + JOIN pg_attribute a ON a.attrelid = c.oid + WHERE n.nspname = {} AND c.relname = {} + AND c.relkind = 'm' AND a.attnum > 0 AND NOT a.attisdropped + ORDER BY a.attnum + """, + [schema_name, object_name], + ) + else: + col_rows = await SafeSqlDriver.execute_param_query( + sql_driver, + """ + SELECT + c.column_name, + c.data_type, + c.is_nullable, + c.column_default, + col_description( + (quote_ident(c.table_schema) || '.' || quote_ident(c.table_name))::regclass, + c.ordinal_position + ) AS comment + FROM information_schema.columns c + WHERE c.table_schema = {} AND c.table_name = {} + ORDER BY c.ordinal_position + """, + [schema_name, object_name], + ) columns = ( [ { @@ -222,6 +295,7 @@ async def get_object_details( "data_type": r.cells["data_type"], "is_nullable": r.cells["is_nullable"], "default": r.cells["column_default"], + "comment": r.cells["comment"], } for r in col_rows ] @@ -270,8 +344,21 @@ async def get_object_details( indexes = [{"name": r.cells["indexname"], "definition": r.cells["indexdef"]} for r in idx_rows] if idx_rows else [] + # Get table/view/materialized view comment + comment_rows = await SafeSqlDriver.execute_param_query( + sql_driver, + """ + SELECT obj_description(c.oid, 'pg_class') AS comment + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = {} AND c.relname = {} + """, + [schema_name, object_name], + ) + object_comment = comment_rows[0].cells["comment"] if comment_rows else None + result = { - "basic": {"schema": schema_name, "name": object_name, "type": object_type}, + "basic": {"schema": schema_name, "name": object_name, "type": object_type, "comment": object_comment}, "columns": columns, "constraints": constraints_list, "indexes": indexes, @@ -327,7 +414,7 @@ async def get_object_details( @mcp.tool( - description="Explains the execution plan for a SQL query, showing how the database will execute it and provides detailed cost estimates.", + description="Explains the PostgreSQL execution plan for a SQL query, showing how the database will execute it and provides detailed cost estimates.", annotations=ToolAnnotations( title="Explain Query", readOnlyHint=True, @@ -428,7 +515,7 @@ async def execute_sql( @mcp.tool( - description="Analyze frequently executed queries in the database and recommend optimal indexes", + description="Analyze frequently executed PostgreSQL queries and recommend optimal indexes.", annotations=ToolAnnotations( title="Analyze Workload Indexes", readOnlyHint=True, @@ -455,7 +542,7 @@ async def analyze_workload_indexes( @mcp.tool( - description="Analyze a list of (up to 10) SQL queries and recommend optimal indexes", + description="Analyze a list of (up to 10) PostgreSQL queries and recommend optimal indexes.", annotations=ToolAnnotations( title="Analyze Query Indexes", readOnlyHint=True, @@ -488,7 +575,7 @@ async def analyze_query_indexes( @mcp.tool( - description="Analyzes database health. Here are the available health checks:\n" + description="Analyzes PostgreSQL database health. Here are the available health checks:\n" "- index - checks for invalid, duplicate, and bloated indexes\n" "- connection - checks the number of connection and their utilization\n" "- vacuum - checks vacuum health for transaction id wraparound\n" @@ -616,7 +703,7 @@ async def main(): else: mcp.add_tool( execute_sql, - description="Execute a read-only SQL query", + description="Execute a read-only SQL query against the PostgreSQL database and return the results.", annotations=ToolAnnotations( title="Execute SQL (Read-Only)", readOnlyHint=True, diff --git a/src/postgres_mcp/sql/safe_sql.py b/src/postgres_mcp/sql/safe_sql.py index 37382f0b..e2207751 100644 --- a/src/postgres_mcp/sql/safe_sql.py +++ b/src/postgres_mcp/sql/safe_sql.py @@ -417,6 +417,10 @@ class SafeSqlDriver(SqlDriver): # Validity Checking "pg_input_is_valid", "pg_input_error_info", + # Object Description/Comment Functions + "obj_description", + "col_description", + "shobj_description", # Object Definition/Information Functions "pg_get_serial_sequence", "pg_get_viewdef",