postgres-mcp-jacob 0.3.0__py3-none-any.whl

postgres_mcp/__init__.py

@@ -0,0 +1,24 @@
+import asyncio
+import sys
+
+from . import server
+from . import top_queries
+
+
+def main():
+    """Main entry point for the package."""
+    # As of version 3.3.0 Psycopg on Windows is not compatible with the default
+    # ProactorEventLoop.
+    # See: https://www.psycopg.org/psycopg3/docs/advanced/async.html#async
+    if sys.platform == "win32":
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+    asyncio.run(server.main())
+
+
+# Optionally expose other important items at package level
+__all__ = [
+    "main",
+    "server",
+    "top_queries",
+]

postgres_mcp/artifacts.py

@@ -0,0 +1,323 @@
+"""Artifacts for the Database Tuning Advisor."""
+
+import difflib
+import json
+from typing import Any
+
+from attrs import define
+from attrs import field
+
+# If the recommendation cost is 0.0, we can't calculate the improvement multiple.
+# Return 1000000.0 to indicate infinite improvement.
+INFINITE_IMPROVEMENT_MULTIPLIER = 1000000.0
+
+
+class ErrorResult:
+    """Simple error result class."""
+
+    def to_text(self) -> str:
+        return self.value
+
+    def __init__(self, message: str):
+        self.value = message
+
+
+def calculate_improvement_multiple(base_cost: float, rec_cost: float) -> float:
+    """Calculate the improvement multiple from this recommendation."""
+    if base_cost <= 0.0:
+        # base_cost or rec_cost might be zero, but as they are floats, the might be
+        # represented as -0.0. That's why we compare to <= 0.0.
+        return 1.0
+    if rec_cost <= 0.0:
+        # If the recommendation cost is 0.0, we can't calculate the improvement multiple.
+        # Return INFINITE_IMPROVEMENT_MULTIPLIER to indicate infinite improvement.
+        return INFINITE_IMPROVEMENT_MULTIPLIER
+    return base_cost / rec_cost
+
+
+@define
+class PlanNode:
+    node_type: str
+    total_cost: float
+    startup_cost: float
+    plan_rows: int
+    plan_width: int
+
+    # Actual metrics from ANALYZE
+    actual_total_time: float | None = field(default=None)
+    actual_startup_time: float | None = field(default=None)
+    actual_rows: int | None = field(default=None)
+    actual_loops: int | None = field(default=None)
+
+    # Buffer info
+    shared_hit_blocks: int | None = field(default=None)
+    shared_read_blocks: int | None = field(default=None)
+    shared_written_blocks: int | None = field(default=None)
+
+    # Other common fields
+    relation_name: str | None = field(default=None)
+    filter: str | None = field(default=None)
+    children: list["PlanNode"] = field(factory=list)
+
+    @classmethod
+    def from_json_data(cls, json_node: dict[str, Any]) -> "PlanNode":
+        # Extract basic fields
+        node = cls(
+            node_type=json_node["Node Type"],
+            total_cost=json_node["Total Cost"],
+            startup_cost=json_node["Startup Cost"],
+            plan_rows=json_node["Plan Rows"],
+            plan_width=json_node["Plan Width"],
+        )
+
+        # Optional ANALYZE fields
+        if "Actual Total Time" in json_node:
+            node.actual_total_time = json_node["Actual Total Time"]
+            node.actual_startup_time = json_node["Actual Startup Time"]
+            node.actual_rows = json_node["Actual Rows"]
+            node.actual_loops = json_node["Actual Loops"]
+
+        # Optional BUFFERS fields
+        if "Shared Hit Blocks" in json_node:
+            node.shared_hit_blocks = json_node["Shared Hit Blocks"]
+            node.shared_read_blocks = json_node["Shared Read Blocks"]
+            node.shared_written_blocks = json_node["Shared Written Blocks"]
+
+        # Common optional fields
+        if "Relation Name" in json_node:
+            node.relation_name = json_node["Relation Name"]
+        if "Filter" in json_node:
+            node.filter = json_node["Filter"]
+
+        # Recursively process child plans
+        if "Plans" in json_node:
+            node.children = [cls.from_json_data(child) for child in json_node["Plans"]]
+
+        return node
+
+
+@define
+class ExplainPlanArtifact:
+    value: str
+    plan_tree: PlanNode
+    planning_time: float | None = field(default=None)
+    execution_time: float | None = field(default=None)
+
+    def __init__(
+        self,
+        value: str,
+        plan_tree: PlanNode,
+        planning_time: float | None = None,
+        execution_time: float | None = None,
+    ):
+        self.value = value
+        self.plan_tree = plan_tree
+        self.planning_time = planning_time
+        self.execution_time = execution_time
+
+    def to_text(self) -> str:
+        """Convert the explain plan to a text representation.
+
+        Returns:
+            str: A string representation of the execution plan with timing information.
+        """
+        result = []
+
+        # Add timing information if available
+        if self.planning_time is not None:
+            result.append(f"Planning Time: {self.planning_time:.3f} ms")
+        if self.execution_time is not None:
+            result.append(f"Execution Time: {self.execution_time:.3f} ms")
+
+        # Add plan tree representation
+        result.append(self._format_plan_node(self.plan_tree))
+
+        return "\n".join(result)
+
+    @staticmethod
+    def _format_plan_node(node: PlanNode, level: int = 0) -> str:
+        """Recursively format a plan node and its children.
+
+        Args:
+            node: The plan node to format
+            level: The current indentation level
+
+        Returns:
+            str: A formatted string representation of the node and its children
+        """
+        indent = "  " * level
+        output = f"{indent}→ {node.node_type} (Cost: {node.startup_cost:.2f}..{node.total_cost:.2f})"
+
+        # Add table name if present
+        if node.relation_name:
+            output += f" on {node.relation_name}"
+
+        # Add rows information
+        output += f" [Rows: {node.plan_rows}]"
+
+        # Add actual metrics if available in a compact form
+        if node.actual_total_time is not None:
+            output += (
+                f" [Actual: {node.actual_startup_time:.2f}..{node.actual_total_time:.2f} ms, Rows: {node.actual_rows}, Loops: {node.actual_loops}]"
+            )
+
+        # Add filter if present
+        if node.filter:
+            filter_text = node.filter
+            # Truncate long filters for readability
+            if len(filter_text) > 100:
+                filter_text = filter_text[:97] + "..."
+            output += f"\n{indent}  Filter: {filter_text}"
+
+        # Add buffer information if available in a compact form
+        if node.shared_hit_blocks is not None:
+            output += f"\n{indent}  Buffers - hit: {node.shared_hit_blocks}, read: {node.shared_read_blocks}, written: {node.shared_written_blocks}"
+
+        # Recursively format children
+        if node.children:
+            for child in node.children:
+                output += "\n" + ExplainPlanArtifact._format_plan_node(child, level + 1)
+
+        return output
+
+    @classmethod
+    def from_json_data(cls, plan_data: dict[str, Any]) -> "ExplainPlanArtifact":
+        if "Plan" not in plan_data:
+            raise ValueError("Missing 'Plan' field in explain plan data")
+
+        # Create plan tree from the "Plan" field
+        plan_tree = PlanNode.from_json_data(plan_data["Plan"])
+
+        # Extract optional timing information
+        planning_time = plan_data.get("Planning Time")
+        execution_time = plan_data.get("Execution Time")
+
+        return cls(
+            value=json.dumps(plan_data, indent=2),
+            plan_tree=plan_tree,
+            planning_time=planning_time,
+            execution_time=execution_time,
+        )
+
+    @staticmethod
+    def format_plan_summary(plan_data):
+        """Extract and format key information from a raw plan data."""
+        if not plan_data:
+            return "No plan data available"
+
+        try:
+            # Create a PlanNode from the raw JSON data
+            if "Plan" in plan_data:
+                plan_node = PlanNode.from_json_data(plan_data["Plan"])
+
+                # Use _format_plan_node to format the output
+                plan_tree = ExplainPlanArtifact._format_plan_node(plan_node, 0)
+
+                return f"{plan_tree}"
+            else:
+                return "Invalid plan data (missing Plan field)"
+
+        except Exception as e:
+            return f"Error summarizing plan: {e}"
+
+    @staticmethod
+    def create_plan_diff(before_plan: dict[str, Any], after_plan: dict[str, Any]) -> str:
+        """Generate a textual diff between two explain plans.
+
+        Args:
+            before_plan: The explain plan before changes
+            after_plan: The explain plan after changes
+
+        Returns:
+            A string containing a readable diff between the two plans
+        """
+        if not before_plan or not after_plan:
+            return "Cannot generate diff: Missing plan data"
+
+        try:
+            # Create PlanNode objects from the plans
+            before_tree = PlanNode.from_json_data(before_plan["Plan"]) if "Plan" in before_plan else None
+            after_tree = PlanNode.from_json_data(after_plan["Plan"]) if "Plan" in after_plan else None
+
+            if not before_tree or not after_tree:
+                return "Cannot generate diff: Invalid plan structure"
+
+            # Format the plans as text
+            before_lines = ExplainPlanArtifact._format_plan_node(before_tree).split("\n")
+            after_lines = ExplainPlanArtifact._format_plan_node(after_tree).split("\n")
+
+            # Generate a readable diff with context
+            diff_lines = []
+            diff_lines.append("PLAN CHANGES:")
+            diff_lines.append("------------")
+
+            # Extract cost information for a summary
+            before_cost = before_tree.total_cost
+            after_cost = after_tree.total_cost
+            improvement = calculate_improvement_multiple(before_cost, after_cost)
+
+            diff_lines.append(f"Cost: {before_cost:.2f} → {after_cost:.2f} ({improvement:.1f}x improvement)")
+            diff_lines.append("")
+
+            # Node type changes - a simplified structural diff
+            diff_lines.append("Operation Changes:")
+
+            # Helper function to extract node types with indentation
+            def extract_node_types(node, level=0, result=None):
+                if result is None:
+                    result = []
+                indent = "  " * level
+                node_info = f"{indent}→ {node.node_type}"
+                if node.relation_name:
+                    node_info += f" on {node.relation_name}"
+                result.append(node_info)
+                for child in node.children:
+                    extract_node_types(child, level + 1, result)
+                return result
+
+            before_structure = extract_node_types(before_tree)
+            after_structure = extract_node_types(after_tree)
+
+            # Generate the structural diff
+            structure_diff = list(
+                difflib.unified_diff(
+                    before_structure,
+                    after_structure,
+                    n=1,  # Context lines
+                    lineterm="",
+                )
+            )
+
+            # Add structural diff to output
+            if structure_diff:
+                diff_lines.extend(structure_diff)
+            else:
+                diff_lines.append("No structural changes detected")
+
+            # Add more specific details about key changes
+            diff_lines.append("")
+            diff_lines.append("Major Changes:")
+
+            # Look for significant changes like seq scan to index scan, changed filters, etc.
+            # This requires traversing both trees and comparing nodes
+
+            # For simplicity, we'll just list key changes in cost and rows
+            if before_tree.node_type != after_tree.node_type:
+                diff_lines.append(f"- Root operation changed: {before_tree.node_type} → {after_tree.node_type}")
+
+            # Compare scan methods used
+            before_scans = [line for line in before_lines if "Seq Scan" in line]
+            after_scans = [line for line in after_lines if "Seq Scan" in line]
+            if len(before_scans) > len(after_scans):
+                diff_lines.append(f"- {len(before_scans) - len(after_scans)} sequential scans replaced with more efficient access methods")
+
+            # Look for new index scans
+            before_idx_scans = [line for line in before_lines if "Index Scan" in line]
+            after_idx_scans = [line for line in after_lines if "Index Scan" in line]
+            if len(after_idx_scans) > len(before_idx_scans):
+                diff_lines.append(f"- {len(after_idx_scans) - len(before_idx_scans)} new index scans used")
+
+            return "\n".join(diff_lines)
+
+        except Exception as e:
+            return f"Error generating plan diff: {e}"

postgres_mcp/database_health/__init__.py

@@ -0,0 +1,4 @@
+from .database_health import DatabaseHealthTool
+from .database_health import HealthType
+
+__all__ = ["DatabaseHealthTool", "HealthType"]

postgres_mcp/database_health/buffer_health_calc.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from ..sql import SqlDriver
+
+
+class BufferHealthCalc:
+    _cached_indexes: list[dict[str, Any]] | None = None
+
+    def __init__(self, sql_driver: SqlDriver):
+        self.sql_driver = sql_driver
+
+    async def index_hit_rate(self, threshold: float = 0.95) -> str:
+        """Calculate the index cache hit rate.
+
+        Returns:
+            String describing the index cache hit rate as a percentage and comparison to threshold
+        """
+        result = await self.sql_driver.execute_query("""
+            SELECT
+                (sum(idx_blks_hit)) / nullif(sum(idx_blks_hit + idx_blks_read), 0) AS rate
+            FROM
+                pg_statio_user_indexes
+        """)
+
+        result_list = [dict(x.cells) for x in result] if result else []
+
+        if not result_list or result_list[0]["rate"] is None:
+            return "No index cache statistics available."
+
+        hit_rate = float(result_list[0]["rate"]) * 100
+        threshold_pct = threshold * 100
+
+        if hit_rate >= threshold_pct:
+            return f"Index cache hit rate: {hit_rate:.1f}% (above {threshold_pct:.1f}% threshold)"
+        else:
+            return f"Index cache hit rate: {hit_rate:.1f}% (below {threshold_pct:.1f}% threshold)"
+
+    async def table_hit_rate(self, threshold: float = 0.95) -> str:
+        """Calculate the table cache hit rate.
+
+        Returns:
+            String describing the table cache hit rate as a percentage and comparison to threshold
+        """
+        result = await self.sql_driver.execute_query("""
+            SELECT
+                sum(heap_blks_hit) / nullif(sum(heap_blks_hit + heap_blks_read), 0) AS rate
+            FROM
+                pg_statio_user_tables
+        """)
+
+        result_list = [dict(x.cells) for x in result] if result else []
+
+        if not result_list or result_list[0]["rate"] is None:
+            return "No table cache statistics available."
+
+        hit_rate = float(result_list[0]["rate"]) * 100
+        threshold_pct = threshold * 100
+
+        if hit_rate >= threshold_pct:
+            return f"Table cache hit rate: {hit_rate:.1f}% (above {threshold_pct:.1f}% threshold)"
+        else:
+            return f"Table cache hit rate: {hit_rate:.1f}% (below {threshold_pct:.1f}% threshold)"

postgres_mcp/database_health/connection_health_calc.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+
+from ..sql import SqlDriver
+
+
+@dataclass
+class ConnectionHealthMetrics:
+    total_connections: int
+    idle_connections: int
+    max_total_connections: int
+    max_idle_connections: int
+    is_total_connections_healthy: bool
+    is_idle_connections_healthy: bool
+
+    @property
+    def is_healthy(self) -> bool:
+        return self.is_total_connections_healthy and self.is_idle_connections_healthy
+
+
+class ConnectionHealthCalc:
+    def __init__(
+        self,
+        sql_driver: SqlDriver,
+        max_total_connections: int = 500,
+        max_idle_connections: int = 100,
+    ):
+        self.sql_driver = sql_driver
+        self.max_total_connections = max_total_connections
+        self.max_idle_connections = max_idle_connections
+
+    async def total_connections_check(self) -> str:
+        """Check if total number of connections is within healthy limits."""
+        total = await self._get_total_connections()
+
+        if total <= self.max_total_connections:
+            return f"Total connections healthy: {total}"
+        return f"High number of connections: {total} (max: {self.max_total_connections})"
+
+    async def idle_connections_check(self) -> str:
+        """Check if number of idle connections is within healthy limits."""
+        idle = await self._get_idle_connections()
+
+        if idle <= self.max_idle_connections:
+            return f"Idle connections healthy: {idle}"
+        return f"High number of idle connections: {idle} (max: {self.max_idle_connections})"
+
+    async def connection_health_check(self) -> str:
+        """Run all connection health checks and return combined results."""
+        total = await self._get_total_connections()
+        idle = await self._get_idle_connections()
+
+        if total > self.max_total_connections:
+            return f"High number of connections: {total}"
+        elif idle > self.max_idle_connections:
+            return f"High number of connections idle in transaction: {idle}"
+        else:
+            return f"Connections healthy: {total} total, {idle} idle"
+
+    async def _get_total_connections(self) -> int:
+        """Get the total number of database connections."""
+        result = await self.sql_driver.execute_query("""
+            SELECT COUNT(*) as count
+            FROM pg_stat_activity
+        """)
+        result_list = [dict(x.cells) for x in result] if result else []
+        return result_list[0]["count"] if result_list else 0
+
+    async def _get_idle_connections(self) -> int:
+        """Get the number of connections that are idle in transaction."""
+        result = await self.sql_driver.execute_query("""
+            SELECT COUNT(*) as count
+            FROM pg_stat_activity
+            WHERE state = 'idle in transaction'
+        """)
+        result_list = [dict(x.cells) for x in result] if result else []
+        return result_list[0]["count"] if result_list else 0

postgres_mcp/database_health/constraint_health_calc.py

@@ -0,0 +1,101 @@
+from dataclasses import dataclass
+
+from ..sql import SqlDriver
+
+
+@dataclass
+class ConstraintMetrics:
+    schema: str
+    table: str
+    name: str
+    referenced_schema: str | None
+    referenced_table: str | None
+
+
+class ConstraintHealthCalc:
+    def __init__(self, sql_driver: SqlDriver):
+        self.sql_driver = sql_driver
+
+    async def invalid_constraints_check(self) -> str:
+        """Check for any invalid constraints in the database.
+
+        Returns:
+            String describing any invalid constraints found
+        """
+        metrics = await self._get_invalid_constraints()
+
+        if not metrics:
+            return "No invalid constraints found."
+
+        result = ["Invalid constraints found:"]
+        for metric in metrics:
+            if metric.referenced_table:
+                result.append(
+                    f"Constraint '{metric.name}' on table '{metric.schema}.{metric.table}' "
+                    f"referencing '{metric.referenced_schema}.{metric.referenced_table}' is invalid"
+                )
+            else:
+                result.append(f"Constraint '{metric.name}' on table '{metric.schema}.{metric.table}' is invalid")
+        return "\n".join(result)
+
+    async def _get_invalid_constraints(self) -> list[ConstraintMetrics]:
+        """Get all invalid constraints in the database."""
+        results = await self.sql_driver.execute_query("""
+            SELECT
+                nsp.nspname AS schema,
+                rel.relname AS table,
+                con.conname AS name,
+                fnsp.nspname AS referenced_schema,
+                frel.relname AS referenced_table
+            FROM
+                pg_catalog.pg_constraint con
+            INNER JOIN
+                pg_catalog.pg_class rel ON rel.oid = con.conrelid
+            LEFT JOIN
+                pg_catalog.pg_class frel ON frel.oid = con.confrelid
+            LEFT JOIN
+                pg_catalog.pg_namespace nsp ON nsp.oid = con.connamespace
+            LEFT JOIN
+                pg_catalog.pg_namespace fnsp ON fnsp.oid = frel.relnamespace
+            WHERE
+                con.convalidated = 'f'
+        """)
+
+        if not results:
+            return []
+
+        result_list = [dict(x.cells) for x in results]
+
+        return [
+            ConstraintMetrics(
+                schema=row["schema"],
+                table=row["table"],
+                name=row["name"],
+                referenced_schema=row["referenced_schema"],
+                referenced_table=row["referenced_table"],
+            )
+            for row in result_list
+        ]
+
+    async def _get_total_constraints(self) -> int:
+        """Get the total number of constraints."""
+        result = await self.sql_driver.execute_query("""
+            SELECT COUNT(*) as count
+            FROM information_schema.table_constraints
+        """)
+        if not result:
+            return 0
+        result_list = [dict(x.cells) for x in result]
+        return result_list[0]["count"] if result_list else 0
+
+    async def _get_active_constraints(self) -> int:
+        """Get the number of active constraints."""
+        result = await self.sql_driver.execute_query("""
+            SELECT COUNT(*) as count
+            FROM information_schema.table_constraints
+            WHERE is_deferrable = 'NO'
+        """)
+        if not result:
+            return 0
+        result_list = [dict(x.cells) for x in result]
+        return result_list[0]["count"] if result_list else 0

postgres_mcp/database_health/database_health.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import logging
+from enum import Enum
+from typing import List
+
+import mcp.types as types
+
+from .buffer_health_calc import BufferHealthCalc
+from .connection_health_calc import ConnectionHealthCalc
+from .constraint_health_calc import ConstraintHealthCalc
+from .index_health_calc import IndexHealthCalc
+from .replication_calc import ReplicationCalc
+from .sequence_health_calc import SequenceHealthCalc
+from .vacuum_health_calc import VacuumHealthCalc
+
+ResponseType = List[types.TextContent | types.ImageContent | types.EmbeddedResource]
+
+logger = logging.getLogger(__name__)
+
+
+class HealthType(str, Enum):
+    INDEX = "index"
+    CONNECTION = "connection"
+    VACUUM = "vacuum"
+    SEQUENCE = "sequence"
+    REPLICATION = "replication"
+    BUFFER = "buffer"
+    CONSTRAINT = "constraint"
+    ALL = "all"
+
+
+class DatabaseHealthTool:
+    """Tool for analyzing database health metrics."""
+
+    def __init__(self, sql_driver):
+        self.sql_driver = sql_driver
+
+    async def health(self, health_type: str) -> str:
+        """Run database health checks for the specified components.
+
+        Args:
+            health_type: Comma-separated list of health check types to perform
+                         Valid values: index, connection, vacuum, sequence, replication, buffer, constraint, all
+
+        Returns:
+            A string with the health check results
+        """
+        try:
+            result = ""
+            try:
+                health_types = {HealthType(x.strip()) for x in health_type.split(",")}
+            except ValueError:
+                return (
+                    f"Invalid health types provided: '{health_type}'. "
+                    + f"Valid values are: {', '.join(sorted([t.value for t in HealthType]))}. "
+                    + "Please try again with a comma-separated list of valid health types."
+                )
+
+            if HealthType.ALL in health_types:
+                health_types = [t.value for t in HealthType if t != HealthType.ALL]
+
+            if HealthType.INDEX in health_types:
+                index_health = IndexHealthCalc(self.sql_driver)
+                result += "Invalid index check: " + await index_health.invalid_index_check() + "\n"
+                result += "Duplicate index check: " + await index_health.duplicate_index_check() + "\n"
+                result += "Index bloat: " + await index_health.index_bloat() + "\n"
+                result += "Unused index check: " + await index_health.unused_indexes() + "\n"
+
+            if HealthType.CONNECTION in health_types:
+                connection_health = ConnectionHealthCalc(self.sql_driver)
+                result += "Connection health: " + await connection_health.connection_health_check() + "\n"
+
+            if HealthType.VACUUM in health_types:
+                vacuum_health = VacuumHealthCalc(self.sql_driver)
+                result += "Vacuum health: " + await vacuum_health.transaction_id_danger_check() + "\n"
+
+            if HealthType.SEQUENCE in health_types:
+                sequence_health = SequenceHealthCalc(self.sql_driver)
+                result += "Sequence health: " + await sequence_health.sequence_danger_check() + "\n"
+
+            if HealthType.REPLICATION in health_types:
+                replication_health = ReplicationCalc(self.sql_driver)
+                result += "Replication health: " + await replication_health.replication_health_check() + "\n"
+
+            if HealthType.BUFFER in health_types:
+                buffer_health = BufferHealthCalc(self.sql_driver)
+                result += "Buffer health for indexes: " + await buffer_health.index_hit_rate() + "\n"
+                result += "Buffer health for tables: " + await buffer_health.table_hit_rate() + "\n"
+
+            if HealthType.CONSTRAINT in health_types:
+                constraint_health = ConstraintHealthCalc(self.sql_driver)
+                result += "Constraint health: " + await constraint_health.invalid_constraints_check() + "\n"
+
+            return result if result else "No health checks were performed."
+        except Exception as e:
+            logger.error(f"Error calculating database health: {e}", exc_info=True)
+            return f"Error calculating database health: {e}"