Repository: crystaldba/postgres-mcp
Pull request #5: DBA-693 Added pre-checks for DTA.
Author: login: mdashti
Created: 2025-03-29T00:01:54Z | Merged: 2025-03-29T07:31:53Z



Diff:
diff --git a/src/postgres_mcp/dta/dta_calc.py b/src/postgres_mcp/dta/dta_calc.py
index 4b6c90c3..d1c522d2 100644
--- a/src/postgres_mcp/dta/dta_calc.py
+++ b/src/postgres_mcp/dta/dta_calc.py
@@ -200,7 +200,6 @@ def __init__(
         self.max_index_width = max_index_width
         self.min_column_usage = min_column_usage
         self.seed_columns_count = seed_columns_count
-        self._setup_hypopg()
         self._analysis_start_time = 0.0
         self.pareto_alpha = pareto_alpha
         self.min_time_improvement = min_time_improvement
@@ -216,6 +215,70 @@ def __init__(
         # Add trace accumulator
         self._dta_traces: list[str] = []
 
+    def _run_prechecks(self, session: DTASession) -> DTASession | None:
+        """
+        Run pre-checks before analysis and return a session with error if any check fails.
+
+        Args:
+            session: The current DTASession object
+
+        Returns:
+            The DTASession with error information if any check fails, None if all checks pass
+        """
+        # Pre-check 1: Check HypoPG with more granular feedback
+        # First check if HypoPG is installed
+        installed_result = self.sql_driver.execute_query(
+            "SELECT 1 FROM pg_extension WHERE extname = 'hypopg'"
+        )
+        if installed_result:
+            # Case 1: HypoPG is installed and ready to use
+            pass
+        else:
+            # Check if HypoPG is available but not installed
+            available_result = self.sql_driver.execute_query(
+                "SELECT 1 FROM pg_available_extensions WHERE name = 'hypopg'"
+            )
+            if available_result:
+                # Case 2: Available but not installed
+                error_message = (
+                    "The HypoPG extension is available but not installed. "
+                    "Please connect to the database and run this statement: 'CREATE EXTENSION hypopg;'"
+                )
+                session.error = error_message
+                logger.error(error_message)
+                return session
+            else:
+                # Case 3: Not available at all
+                error_message = (
+                    "The HypoPG extension is not available on this PostgreSQL server. "
+                    "To install HypoPG:\n"
+                    "1. For Debian/Ubuntu: sudo apt-get install postgresql-hypopg\n"
+                    "2. For RHEL/CentOS: sudo yum install postgresql-hypopg\n"
+                    "3. For MacOS with Homebrew: brew install hypopg\n"
+                    "4. For other systems, build from source: git clone https://github.com/HypoPG/hypopg\n"
+                    "After installing the extension packages, connect to your database and run: CREATE EXTENSION hypopg;"
+                )
+                session.error = error_message
+                logger.error(error_message)
+                return session
+
+        # Pre-check 2: Check if ANALYZE has been run at least once
+        result = self.sql_driver.execute_query(
+            "SELECT s.last_analyze FROM pg_stat_user_tables s LIMIT 1"
+        )
+        if not result or all(row.cells.get("last_analyze") is None for row in result):
+            error_message = (
+                "Statistics are not up-to-date. The database needs to be analyzed first. "
+                "Please run 'ANALYZE;' on your database before using the tuning advisor. "
+                "Without up-to-date statistics, the index recommendations may be inaccurate."
+            )
+            session.error = error_message
+            logger.error(error_message)
+            return session
+
+        # All checks passed
+        return None
+
     def analyze_workload(
         self,
         workload: list[dict[str, Any]] | None = None,
@@ -264,6 +327,11 @@ def analyze_workload(
         )
 
         try:
+            # Run pre-checks
+            precheck_result = self._run_prechecks(session)
+            if precheck_result:
+                return precheck_result
+
             # First try to use explicit workload if provided
             if workload:
                 logger.debug(f"Using explicit workload with {len(workload)} queries")
@@ -930,20 +998,6 @@ def _get_query_stats_direct(
         )
         return [dict(row.cells) for row in result] if result else []
 
-    def _setup_hypopg(self) -> None:
-        """Ensure HypoPG extension is installed."""
-        result = self.sql_driver.execute_query(
-            "SELECT 1 FROM pg_extension WHERE extname = 'hypopg'"
-        )
-        if not result:
-            try:
-                self.sql_driver.execute_query("CREATE EXTENSION IF NOT EXISTS hypopg")
-            except Exception as e:
-                logger.error(f"Failed to create HypoPG extension: {e}", exc_info=True)
-                raise RuntimeError(
-                    "HypoPG extension required but could not be installed"
-                ) from e
-
     def _check_time(self) -> bool:
         """Return True if we have exceeded max_runtime_seconds."""
         if self.max_runtime_seconds <= 0:
diff --git a/src/postgres_mcp/dta/dta_tools.py b/src/postgres_mcp/dta/dta_tools.py
index f542b0ad..4b766ef9 100644
--- a/src/postgres_mcp/dta/dta_tools.py
+++ b/src/postgres_mcp/dta/dta_tools.py
@@ -39,6 +39,11 @@ def _create_recommendations_response(self, session: DTASession) -> Dict[str, Any
         Returns:
             Dictionary with summary and recommendations
         """
+        if session.error:
+            return {
+                "error": session.error,
+                "_langfuse_trace": session.dta_traces,
+            }
         if not session.recommendations:
             return {
                 "recommendations": "No index recommendations found.",
diff --git a/src/postgres_mcp/dta/safe_sql.py b/src/postgres_mcp/dta/safe_sql.py
index 4e5705e7..72329951 100644
--- a/src/postgres_mcp/dta/safe_sql.py
+++ b/src/postgres_mcp/dta/safe_sql.py
@@ -708,10 +708,14 @@ def _validate(self, query: str) -> None:
         except pglast.parser.ParseError as e:
             raise ValueError("Failed to parse SQL statement") from e
 
-    def execute_query(self, query: str) -> Optional[list[SqlDriver.RowResult]]:  # noqa: UP007
+    def execute_query(
+        self, query: str, force_readonly: bool = True
+    ) -> Optional[list[SqlDriver.RowResult]]:  # noqa: UP007
         """Execute a query after validating it is safe"""
         self._validate(query)
-        return self.sql_driver.execute_query(f"/* crystaldba */ {query}")
+        return self.sql_driver.execute_query(
+            f"/* crystaldba */ {query}", force_readonly=force_readonly
+        )
 
     @staticmethod
     def sql_to_query(sql: Composable) -> str:
diff --git a/src/postgres_mcp/dta/sql_driver.py b/src/postgres_mcp/dta/sql_driver.py
index 2fd09797..701503ce 100644
--- a/src/postgres_mcp/dta/sql_driver.py
+++ b/src/postgres_mcp/dta/sql_driver.py
@@ -33,7 +33,9 @@ def __init__(self, conn: Any = None, engine_url: str | None = None):
         else:
             raise ValueError("Either conn or engine_url must be provided")
 
-    def execute_query(self, query: str) -> Optional[List[RowResult]]:
+    def execute_query(
+        self, query: str, force_readonly: bool = True
+    ) -> Optional[List[RowResult]]:
         """
         Execute a query and return results.
 
@@ -46,15 +48,19 @@ def execute_query(self, query: str) -> Optional[List[RowResult]]:
         try:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                 # Start read-only transaction
-                cursor.execute("BEGIN TRANSACTION READ ONLY")
+                if force_readonly:
+                    cursor.execute("BEGIN TRANSACTION READ ONLY")
                 try:
                     cursor.execute(query)
                     if cursor.description is None:  # No results (like DDL statements)
                         return None
                     rows = cursor.fetchall()
+                    if not force_readonly:
+                        cursor.execute("COMMIT")
                     return [SqlDriver.RowResult(cells=dict(row)) for row in rows]
                 finally:
-                    cursor.execute("ROLLBACK")
+                    if force_readonly:
+                        cursor.execute("ROLLBACK")
         except Exception as e:
             logger.error(f"Error executing query: {e}")
             self.conn.rollback()
diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py
index f89f7e8b..2f4a0a68 100644
--- a/src/postgres_mcp/server.py
+++ b/src/postgres_mcp/server.py
@@ -4,7 +4,6 @@
 from urllib.parse import urlparse
 
 import psycopg2
-from psycopg2.extras import RealDictCursor
 from mcp.server.models import InitializationOptions
 import mcp.types as types
 from mcp.server import NotificationOptions, Server
@@ -33,25 +32,26 @@ def get_connection():
 async def handle_list_resources() -> list[types.Resource]:
     """List available database tables as resources."""
     try:
-        with get_connection().cursor() as cursor:
-            cursor.execute(
-                "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"
+        sql_driver = SafeSqlDriver(sql_driver=SqlDriver(conn=get_connection()))
+        rows = sql_driver.execute_query(
+            "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"
+        )
+        tables = [row.cells["table_name"] for row in rows] if rows else []
+
+        base_url = urlparse(os.environ.get("DATABASE_URL", ""))
+        base_url = base_url._replace(scheme="postgres", password="")
+        base_url = base_url.geturl()
+
+        return [
+            types.Resource(
+                uri=AnyUrl(f"{base_url}/{table}/{SCHEMA_PATH}"),
+                name=f'"{table}" database schema',
+                description=f"Schema for table {table}",
+                mimeType="application/json",
             )
-            tables = cursor.fetchall()
-
-            base_url = urlparse(os.environ.get("DATABASE_URL", ""))
-            base_url = base_url._replace(scheme="postgres", password="")
-            base_url = base_url.geturl()
+            for table in tables
+        ]
 
-            return [
-                types.Resource(
-                    uri=AnyUrl(f"{base_url}/{table[0]}/{SCHEMA_PATH}"),
-                    name=f'"{table[0]}" database schema',
-                    description=f"Schema for table {table[0]}",
-                    mimeType="application/json",
-                )
-                for table in tables
-            ]
     except Exception as e:
         print(f"Error listing resources: {e}", file=sys.stderr)
         return []
@@ -179,7 +179,9 @@ async def handle_call_tool(
             rows = sql_driver.execute_query(sql)
             if rows is None:
                 return [types.TextContent(type="text", text="No results")]
-            return [types.TextContent(type="text", text=str(list([r.cells for r in rows])))]
+            return [
+                types.TextContent(type="text", text=str(list([r.cells for r in rows])))
+            ]
         except Exception as e:
             print(f"Error executing query: {e}", file=sys.stderr)
             raise


Reviews:
user: login: jssmith
body: - Are we going to support creating the hypopg extension via CREATE EXTENSION? How does that work with Safe SQL?
- Can we give immediate feedback on whether hypopg is available? There are 3 cases: 1/ available+installed, 2/ available but not installed, 3/ not available.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-29T00:09:45Z
user: login: jssmith
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-29T00:47:45Z
user: login: mdashti
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-29T01:20:54Z
user: login: mdashti
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-29T07:30:46Z

Review comments:
diff_hunk: @@ -239,14 +239,36 @@ def _run_prechecks(self, session: DTASession) -> DTASession | None:
                 "SELECT 1 FROM pg_available_extensions WHERE name = 'hypopg'"
             )
             if available_result:
-                # Case 2: Available but not installed
-                error_message = (
-                    "The HypoPG extension is available but not installed. "
-                    "Please connect to your database and run: CREATE EXTENSION hypopg;"
+                # Case 2: Available but not installed - try to install it
+                self.dta_trace(
+                    "HypoPG extension is available but not installed. Attempting to install it..."
                 )
-                session.error = error_message
-                logger.error(error_message)
-                return session
+                try:
+                    install_result = self.sql_driver.execute_query(
path: src/postgres_mcp/dta/dta_calc.py
user: login: jssmith
body: We should loop back through the calling LLM, giving it a chance to ask the user whether they want to create the extension. Let's separate the read-only path (checking for extensions) from the write path (creating extensions).
created_at: 2025-03-29T00:47:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -239,14 +239,36 @@ def _run_prechecks(self, session: DTASession) -> DTASession | None:
                 "SELECT 1 FROM pg_available_extensions WHERE name = 'hypopg'"
             )
             if available_result:
-                # Case 2: Available but not installed
-                error_message = (
-                    "The HypoPG extension is available but not installed. "
-                    "Please connect to your database and run: CREATE EXTENSION hypopg;"
+                # Case 2: Available but not installed - try to install it
+                self.dta_trace(
+                    "HypoPG extension is available but not installed. Attempting to install it..."
                 )
-                session.error = error_message
-                logger.error(error_message)
-                return session
+                try:
+                    install_result = self.sql_driver.execute_query(
path: src/postgres_mcp/dta/dta_calc.py
user: login: mdashti
body: How do you suggest to do it? IMHO, let's do it in a separate follow-up PR.
created_at: 2025-03-29T01:20:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -239,14 +239,36 @@ def _run_prechecks(self, session: DTASession) -> DTASession | None:
                 "SELECT 1 FROM pg_available_extensions WHERE name = 'hypopg'"
             )
             if available_result:
-                # Case 2: Available but not installed
-                error_message = (
-                    "The HypoPG extension is available but not installed. "
-                    "Please connect to your database and run: CREATE EXTENSION hypopg;"
+                # Case 2: Available but not installed - try to install it
+                self.dta_trace(
+                    "HypoPG extension is available but not installed. Attempting to install it..."
                 )
-                session.error = error_message
-                logger.error(error_message)
-                return session
+                try:
+                    install_result = self.sql_driver.execute_query(
path: src/postgres_mcp/dta/dta_calc.py
user: login: mdashti
body: OK, done as discussed on Slack.
created_at: 2025-03-29T07:30:46Z
author_association: CONTRIBUTOR

PR comments:
user: login: mdashti
created_at: 2025-03-29T00:19:46Z
body: @jssmith 
> Are we going to support creating the hypopg extension via CREATE EXTENSION? How does that work with Safe SQL?

SafeSql currently allows adding extensions.
However, I removed the "CREATE EXTENSION" one. It's also maybe best to remove it from SafeSql. We'd check for HypoPG, and if it's not there, we suggest the user to install it.

> Can we give immediate feedback on whether hypopg is available? There are 3 cases: 1/ available+installed, 2/ available but not installed, 3/ not available.

Good idea. Done in [074a384](https://github.com/crystaldba/postgres-mcp/pull/5/commits/074a384666f52af1cddabb5096e88399bd46fa69)

author_association: CONTRIBUTOR
user: login: jssmith
created_at: 2025-03-29T00:28:26Z
body: We should allow creating the extensions needed by our tools. We might broaden this eventually but let's keep it limited right now. I think we just whitelist those in Safe SQL.
author_association: CONTRIBUTOR
user: login: mdashti
created_at: 2025-03-29T00:41:51Z
body: @jssmith Done in [8c8118b](https://github.com/crystaldba/postgres-mcp/pull/5/commits/8c8118b2bc4df2e17c2e87d9e13a2b2367d4d88d)
author_association: CONTRIBUTOR