CockroachDB MCP Server

Instructions

Args:
    sql: SQL query to explain.
    analyze: If True, actually execute to get runtime stats.

Returns:
    Query execution plan.

Implementation Reference

• src/cockroachdb_mcp/server.py:100-114 (registration)

  MCP tool registration for 'explain_query' using @mcp.tool() decorator. This wrapper function handles errors and delegates to the core implementation in query.py.

  @mcp.tool()
  async def explain_query(sql: str, analyze: bool = False) -> dict[str, Any]:
      """Get the execution plan for a query.
  
      Args:
          sql: SQL query to explain.
          analyze: If True, actually execute to get runtime stats.
  
      Returns:
          Query execution plan.
      """
      try:
          return await query.explain_query(sql, analyze)
      except Exception as e:
          return {"status": "error", "error": str(e)}

• src/cockroachdb_mcp/tools/query.py:130-175 (handler)

  Core handler function that implements the logic for explaining a SQL query: validates input, constructs EXPLAIN(ANALYZE) statement, executes it, and formats the execution plan output.

  async def explain_query(query: str, analyze: bool = False) -> dict[str, Any]:
      """Get the execution plan for a query.
  
      Args:
          query: SQL query to explain.
          analyze: If True, actually execute to get runtime stats.
  
      Returns:
          Execution plan.
      """
      # Validate the underlying query
      validation = await validate_query(query)
      if not validation["is_valid"]:
          return {
              "status": "error",
              "error": "Query validation failed",
              "issues": validation["issues"],
          }
  
      # Build EXPLAIN query
      if analyze:
          explain_query_str = f"EXPLAIN ANALYZE {query}"
      else:
          explain_query_str = f"EXPLAIN {query}"
  
      result = await connection_manager.execute_query(explain_query_str)
  
      if result.get("status") == "error":
          return result
  
      # Format the plan output
      plan_lines = []
      for row in result.get("rows", []):
          # CockroachDB returns plan in 'info' column
          if "info" in row:
              plan_lines.append(row["info"])
          else:
              # Fallback for different column names
              plan_lines.append(str(list(row.values())[0]) if row else "")
  
      return {
          "status": "success",
          "query": query,
          "analyzed": analyze,
          "plan": "\n".join(plan_lines),
      }

• src/cockroachdb_mcp/tools/query.py:52-102 (helper)

  Helper function validate_query used by explain_query to check for blocked commands, read-only mode compliance, and classify query type before execution.

  async def validate_query(query: str) -> dict[str, Any]:
      """Validate a SQL query without executing it.
  
      Args:
          query: SQL query to validate.
  
      Returns:
          Validation result with is_valid and any issues.
      """
      issues: list[str] = []
  
      # Check for empty query
      if not query or not query.strip():
          return {
              "is_valid": False,
              "issues": ["Query is empty"],
              "query_type": None,
          }
  
      # Check for blocked commands
      is_blocked, blocked_cmd = _is_blocked_command(query)
      if is_blocked:
          issues.append(f"Blocked command: {blocked_cmd}")
  
      # Check read-only mode
      if settings.read_only and not _is_read_only_query(query):
          issues.append("Server is in read-only mode; only SELECT/SHOW/EXPLAIN allowed")
  
      # Determine query type
      query_upper = query.strip().upper()
      if query_upper.startswith("SELECT") or query_upper.startswith("WITH"):
          query_type = "SELECT"
      elif query_upper.startswith("INSERT"):
          query_type = "INSERT"
      elif query_upper.startswith("UPDATE"):
          query_type = "UPDATE"
      elif query_upper.startswith("DELETE"):
          query_type = "DELETE"
      elif query_upper.startswith("SHOW"):
          query_type = "SHOW"
      elif query_upper.startswith("EXPLAIN"):
          query_type = "EXPLAIN"
      else:
          query_type = "OTHER"
  
      return {
          "is_valid": len(issues) == 0,
          "issues": issues,
          "query_type": query_type,
          "is_read_only": _is_read_only_query(query),
      }