nexus/tools/mysql_tool.py

"""
MySQL Tool - Quản lý MySQL qua pymysql.
Author: Hieu Louis (2026)
Operations: query, list_tables, describe_table, create_index, optimize, explain.
"""
from __future__ import annotations

import json
from datetime import date, datetime, decimal
from typing import Any, Dict, List, Optional

from .base import Tool, ToolResult, ToolContext, ToolCategory, ToolSafety


DDL_OPS = {"create_index"}
MAINTENANCE_OPS = {"optimize"}


def _json_default(o: Any) -> Any:
    if isinstance(o, (datetime, date)):
        return o.isoformat()
    if isinstance(o, decimal.Decimal):
        return float(o)
    if isinstance(o, bytes):
        try:
            return o.decode("utf-8")
        except Exception:
            return o.hex()
    return str(o)


def _parse_mysql_url(url: str) -> Dict[str, Any]:
    """Parse mysql://user:pw@host:3306/db // parse MySQL URL."""
    from urllib.parse import urlparse
    p = urlparse(url)
    return {
        "host": p.hostname or "localhost",
        "port": p.port or 3306,
        "user": p.username or "root",
        "password": p.password or "",
        "database": (p.path or "/").lstrip("/"),
    }


class MySQLTool(Tool):
    """Quản lý MySQL: query, list_tables, describe_table, create_index, optimize, explain."""

    category = ToolCategory.DATABASE
    safety = ToolSafety.DANGEROUS
    requires_confirmation = True

    @property
    def name(self) -> str:
        return "mysql"

    @property
    def description(self) -> str:
        return (
            "Quản lý MySQL qua pymysql: query (parameterized), list_tables, "
            "describe_table, create_index, optimize (ANALYZE TABLE), explain."
        )

    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "connection_string": {"type": "string", "description": "MySQL URL (mysql://user:pw@host:3306/db)"},
                "operation": {
                    "type": "string",
                    "enum": ["query", "list_tables", "describe_table", "create_index", "optimize", "explain"],
                    "description": "MySQL operation",
                },
                "query": {"type": "string", "description": "SQL query (operation=query/explain)"},
                "params": {
                    "oneOf": [{"type": "array", "items": {}}, {"type": "object"}],
                    "description": "Tham số bind (%s hoặc %(name)s)",
                },
                "table": {"type": "string", "description": "Table name (describe_table/create_index)"},
                "index_name": {"type": "string", "description": "Tên index (create_index)"},
                "columns": {"type": "array", "items": {"type": "string"}, "description": "Cột cho index"},
                "limit": {"type": "integer", "description": "Giới hạn rows (default 1000)"},
            },
            "required": ["connection_string", "operation"],
        }

    def validate_args(self, args: Dict[str, Any]) -> Optional[str]:
        if not args.get("connection_string"):
            return "Missing required arg: connection_string"
        op = args.get("operation")
        if not op:
            return "Missing required arg: operation"
        if op == "query" and not args.get("query"):
            return "Operation 'query' requires 'query' arg"
        if op in {"describe_table", "create_index"} and not args.get("table"):
            return f"Operation '{op}' requires 'table' arg"
        if op == "create_index" and not args.get("columns"):
            return "Operation 'create_index' requires 'columns' arg"
        return None

    def execute(self, args: Dict[str, Any], context: ToolContext) -> ToolResult:
        cs: str = args["connection_string"]
        op: str = args["operation"]

        # Lazy import pymysql (fallback mysql.connector) // lazy import
        try:
            import pymysql  # type: ignore
            import pymysql.cursors  # type: ignore
            driver = "pymysql"
        except ImportError:
            try:
                import mysql.connector  # type: ignore
                import mysql.connector.cursor  # type: ignore
                driver = "mysql.connector"
            except ImportError as e:
                return ToolResult(
                    success=False,
                    error=f"Neither pymysql nor mysql-connector installed: {e}",
                    return_code=127,
                )

        # Dry-run cho DDL // dry-run
        if context.dry_run and op in DDL_OPS:
            return ToolResult(
                success=True,
                output=f"[dry-run] Would run {op} on {args.get('table')}",
                metadata={"dry_run": True, "operation": op},
            )

        cfg = _parse_mysql_url(cs)
        try:
            if driver == "pymysql":
                conn = pymysql.connect(
                    host=cfg["host"], port=cfg["port"], user=cfg["user"],
                    password=cfg["password"], database=cfg["database"] or None,
                    cursorclass=pymysql.cursors.DictCursor, charset="utf8mb4",
                    autocommit=op in DDL_OPS,
                )
            else:
                conn = mysql.connector.connect(
                    host=cfg["host"], port=cfg["port"], user=cfg["user"],
                    password=cfg["password"], database=cfg["database"] or None,
                    charset="utf8mb4", autocommit=op in DDL_OPS,
                )
            cur = conn.cursor(dictionary=True) if driver == "mysql.connector" else conn.cursor()

            if op == "list_tables":
                cur.execute("SHOW TABLES")
                rows = cur.fetchall()
                payload = json.dumps(rows, default=_json_default, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"count": len(rows)})

            if op == "describe_table":
                cur.execute(f"DESCRIBE `{args['table']}`")
                rows = cur.fetchall()
                payload = json.dumps(rows, default=_json_default, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"table": args["table"], "columns": len(rows)})

            if op == "query":
                params = args.get("params") or ()
                cur.execute(args["query"], params)
                if cur.description:
                    rows = cur.fetchmany(int(args.get("limit") or 1000))
                    payload = json.dumps(rows, default=_json_default, ensure_ascii=False, indent=2)
                    return ToolResult(success=True, output=payload, metadata={"rowcount": len(rows)})
                conn.commit()
                return ToolResult(success=True, output=json.dumps({"rowcount": cur.rowcount}), metadata={"rowcount": cur.rowcount})

            if op == "create_index":
                idx = args["index_name"] or f"idx_{args['table']}_" + "_".join(args["columns"])
                cols = ", ".join(f"`{c}`" for c in args["columns"])
                cur.execute(f"CREATE INDEX `{idx}` ON `{args['table']}` ({cols})")
                return ToolResult(success=True, output=f"Index created: {idx}", metadata={"index": idx})

            if op == "optimize":
                cur.execute(f"ANALYZE TABLE `{args['table']}`" if args.get("table") else "ANALYZE TABLE")
                rows = cur.fetchall()
                payload = json.dumps(rows, default=_json_default, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"operation": "optimize"})

            if op == "explain":
                cur.execute("EXPLAIN " + args["query"])
                rows = cur.fetchall()
                payload = json.dumps(rows, default=_json_default, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"operation": "explain"})

            return ToolResult(success=False, error=f"Unknown operation: {op}", return_code=1)
        except Exception as e:
            return ToolResult(success=False, error=str(e), return_code=1)
        finally:
            try:
                cur.close(); conn.close()
            except Exception:
                pass