Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
143 changes: 115 additions & 28 deletions src/postgres_mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -114,15 +114,15 @@ 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,
),
)
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:
Expand All @@ -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 []
)
Expand Down Expand Up @@ -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,
Expand All @@ -197,31 +241,61 @@ 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 = (
[
{
"column": r.cells["column_name"],
"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
]
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions src/postgres_mcp/sql/safe_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down