sql-query-mcp 0.1.1__py3-none-any.whl

sql_query_mcp/__init__.py

@@ -0,0 +1,5 @@
+"""sql-query-mcp package."""
+
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"

sql_query_mcp/__main__.py

@@ -0,0 +1,5 @@
+from .app import main
+
+
+if __name__ == "__main__":
+    main()

sql_query_mcp/adapters/__init__.py

@@ -0,0 +1,15 @@
+"""Engine adapters for sql-query-mcp."""
+
+__all__ = ["MySQLAdapter", "PostgresAdapter"]
+
+
+def __getattr__(name: str):
+    if name == "MySQLAdapter":
+        from .mysql import MySQLAdapter
+
+        return MySQLAdapter
+    if name == "PostgresAdapter":
+        from .postgres import PostgresAdapter
+
+        return PostgresAdapter
+    raise AttributeError(name)

sql_query_mcp/adapters/mysql.py

@@ -0,0 +1,171 @@
+"""MySQL adapter."""
+
+from __future__ import annotations
+
+import json
+from contextlib import contextmanager
+from typing import Iterator, List
+from urllib.parse import parse_qs, unquote, urlparse
+
+try:
+    import pymysql
+    from pymysql.cursors import DictCursor
+except ImportError:  # pragma: no cover - runtime dependency
+    pymysql = None
+    DictCursor = None
+
+from ..errors import ConfigurationError, SecurityError
+
+
+class MySQLAdapter:
+    engine = "mysql"
+
+    @contextmanager
+    def connection(self, connection_id: str, dsn: str) -> Iterator[object]:
+        if pymysql is None or DictCursor is None:
+            raise ConfigurationError("缺少 PyMySQL 依赖，请先安装项目依赖。")
+
+        conn = pymysql.connect(
+            autocommit=True,
+            cursorclass=DictCursor,
+            **self._parse_dsn(dsn),
+        )
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    def close(self) -> None:
+        return None
+
+    def set_statement_timeout(self, conn: object, timeout_ms: int) -> None:
+        with conn.cursor() as cur:
+            cur.execute("SET SESSION max_execution_time = %s", (int(timeout_ms),))
+
+    def list_databases(self, conn: object) -> List[str]:
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT schema_name AS database_name
+                FROM information_schema.schemata
+                WHERE schema_name NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
+                ORDER BY schema_name
+                """
+            )
+            return [row["database_name"] for row in cur.fetchall()]
+
+    def list_tables(self, conn: object, database: str):
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT table_schema AS database_name, table_name, table_type
+                FROM information_schema.tables
+                WHERE table_schema = %s
+                ORDER BY table_name
+                """,
+                (database,),
+            )
+            return cur.fetchall()
+
+    def describe_table(self, conn: object, database: str, table_name: str):
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT column_name, column_type, is_nullable, column_default, extra, column_key, ordinal_position
+                FROM information_schema.columns
+                WHERE table_schema = %s AND table_name = %s
+                ORDER BY ordinal_position
+                """,
+                (database, table_name),
+            )
+            columns = cur.fetchall()
+            cur.execute(
+                """
+                SELECT index_name, non_unique, seq_in_index, column_name
+                FROM information_schema.statistics
+                WHERE table_schema = %s AND table_name = %s
+                ORDER BY index_name, seq_in_index
+                """,
+                (database, table_name),
+            )
+            index_rows = cur.fetchall()
+
+        if not columns:
+            return None
+
+        return {
+            "columns": [
+                {
+                    "column_name": row["column_name"],
+                    "data_type": row["column_type"],
+                    "udt_name": None,
+                    "nullable": row["is_nullable"] == "YES",
+                    "default": row["column_default"],
+                    "primary_key": row["column_key"] == "PRI",
+                    "extra": row["extra"],
+                }
+                for row in columns
+            ],
+            "indexes": self._normalize_indexes(index_rows),
+        }
+
+    def build_sample_query(self, database: str, table_name: str, sentinel_limit: int) -> str:
+        return (
+            f"SELECT * FROM {self._quote_identifier(database)}."
+            f"{self._quote_identifier(table_name)} LIMIT {int(sentinel_limit)}"
+        )
+
+    def build_explain_query(self, sql_text: str, analyze: bool = False) -> str:
+        if analyze:
+            raise SecurityError("MySQL 首版不支持 analyze=True。")
+        return f"EXPLAIN FORMAT=JSON {sql_text}"
+
+    def extract_plan(self, rows):
+        if not rows:
+            return []
+        plan = rows[0].get("EXPLAIN", [])
+        if isinstance(plan, str):
+            try:
+                return json.loads(plan)
+            except json.JSONDecodeError:
+                return plan
+        return plan
+
+    def column_names(self, description) -> List[str]:
+        return [column[0] for column in (description or [])]
+
+    def _parse_dsn(self, dsn: str) -> dict:
+        parsed = urlparse(dsn)
+        if parsed.scheme not in {"mysql", "mysql+pymysql"}:
+            raise ConfigurationError(f"MySQL DSN 必须使用 mysql:// 或 mysql+pymysql://，当前为 {parsed.scheme}")
+
+        query_params = {key: values[-1] for key, values in parse_qs(parsed.query).items()}
+        connect_args = {
+            "host": parsed.hostname or "localhost",
+            "user": unquote(parsed.username) if parsed.username else None,
+            "password": unquote(parsed.password) if parsed.password else None,
+            "port": parsed.port or 3306,
+            "database": parsed.path.lstrip("/") or None,
+            "charset": query_params.get("charset", "utf8mb4"),
+        }
+        return {key: value for key, value in connect_args.items() if value is not None}
+
+    def _quote_identifier(self, value: str) -> str:
+        return "`" + value.replace("`", "``") + "`"
+
+    def _normalize_indexes(self, rows: List[dict]) -> List[dict]:
+        grouped = {}
+        for row in rows:
+            index_name = row["index_name"]
+            item = grouped.setdefault(
+                index_name,
+                {
+                    "index_name": index_name,
+                    "columns": [],
+                    "unique": row["non_unique"] == 0,
+                    "primary_key": index_name == "PRIMARY",
+                    "definition": None,
+                },
+            )
+            item["columns"].append(row["column_name"])
+        return [grouped[name] for name in sorted(grouped)]

sql_query_mcp/adapters/postgres.py

@@ -0,0 +1,180 @@
+"""PostgreSQL adapter."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator, List
+
+try:
+    from psycopg import sql
+    from psycopg.rows import dict_row
+    from psycopg_pool import ConnectionPool
+except ImportError:  # pragma: no cover - runtime dependency
+    sql = None
+    dict_row = None
+    ConnectionPool = None
+
+from ..errors import ConfigurationError
+
+
+class PostgresAdapter:
+    engine = "postgres"
+
+    def __init__(self) -> None:
+        self._pools = {}
+
+    @contextmanager
+    def connection(self, connection_id: str, dsn: str) -> Iterator[object]:
+        pool = self._get_pool(connection_id, dsn)
+        with pool.connection() as conn:
+            yield conn
+
+    def close(self) -> None:
+        for pool in self._pools.values():
+            pool.close()
+
+    def set_statement_timeout(self, conn: object, timeout_ms: int) -> None:
+        with conn.cursor() as cur:
+            cur.execute("SELECT set_config('statement_timeout', %s, false)", (str(timeout_ms),))
+
+    def list_schemas(self, conn: object) -> List[str]:
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT schema_name
+                FROM information_schema.schemata
+                WHERE schema_name NOT IN ('information_schema')
+                  AND schema_name NOT LIKE 'pg_%'
+                ORDER BY schema_name
+                """
+            )
+            return [row["schema_name"] for row in cur.fetchall()]
+
+    def list_tables(self, conn: object, schema: str):
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT table_schema AS schema, table_name, table_type
+                FROM information_schema.tables
+                WHERE table_schema = %s
+                ORDER BY table_name
+                """,
+                (schema,),
+            )
+            return cur.fetchall()
+
+    def describe_table(self, conn: object, schema: str, table_name: str):
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT column_name, data_type, udt_name, is_nullable, column_default, ordinal_position
+                FROM information_schema.columns
+                WHERE table_schema = %s AND table_name = %s
+                ORDER BY ordinal_position
+                """,
+                (schema, table_name),
+            )
+            columns = cur.fetchall()
+            cur.execute(
+                """
+                SELECT kcu.column_name
+                FROM information_schema.table_constraints tc
+                JOIN information_schema.key_column_usage kcu
+                  ON tc.constraint_name = kcu.constraint_name
+                 AND tc.table_schema = kcu.table_schema
+                WHERE tc.constraint_type = 'PRIMARY KEY'
+                  AND tc.table_schema = %s
+                  AND tc.table_name = %s
+                ORDER BY kcu.ordinal_position
+                """,
+                (schema, table_name),
+            )
+            primary_keys = {row["column_name"] for row in cur.fetchall()}
+            cur.execute(
+                """
+                SELECT
+                    idx.relname AS index_name,
+                    ix.indisunique AS is_unique,
+                    ix.indisprimary AS is_primary,
+                    pg_get_indexdef(ix.indexrelid) AS definition,
+                    COALESCE(
+                        array_agg(att.attname ORDER BY keys.ordinality)
+                        FILTER (WHERE att.attname IS NOT NULL),
+                        ARRAY[]::text[]
+                    ) AS columns
+                FROM pg_class tbl
+                JOIN pg_namespace ns ON ns.oid = tbl.relnamespace
+                JOIN pg_index ix ON ix.indrelid = tbl.oid
+                JOIN pg_class idx ON idx.oid = ix.indexrelid
+                LEFT JOIN LATERAL unnest(ix.indkey) WITH ORDINALITY AS keys(attnum, ordinality) ON TRUE
+                LEFT JOIN pg_attribute att
+                  ON att.attrelid = tbl.oid
+                 AND att.attnum = keys.attnum
+                WHERE ns.nspname = %s
+                  AND tbl.relname = %s
+                GROUP BY idx.relname, ix.indisunique, ix.indisprimary, ix.indexrelid
+                ORDER BY idx.relname
+                """,
+                (schema, table_name),
+            )
+            index_rows = cur.fetchall()
+
+        if not columns:
+            return None
+
+        return {
+            "columns": [
+                {
+                    "column_name": row["column_name"],
+                    "data_type": row["data_type"],
+                    "udt_name": row["udt_name"],
+                    "nullable": row["is_nullable"] == "YES",
+                    "default": row["column_default"],
+                    "primary_key": row["column_name"] in primary_keys,
+                }
+                for row in columns
+            ],
+            "indexes": [
+                {
+                    "index_name": row["index_name"],
+                    "columns": row["columns"],
+                    "unique": row["is_unique"],
+                    "primary_key": row["is_primary"],
+                    "definition": row["definition"],
+                }
+                for row in index_rows
+            ],
+        }
+
+    def build_sample_query(self, schema: str, table_name: str, sentinel_limit: int):
+        if sql is None:
+            raise ConfigurationError("缺少 psycopg 依赖，请先安装项目依赖。")
+        return sql.SQL("SELECT * FROM {}.{} LIMIT {}").format(
+            sql.Identifier(schema),
+            sql.Identifier(table_name),
+            sql.Literal(sentinel_limit),
+        )
+
+    def build_explain_query(self, sql_text: str, analyze: bool = False) -> str:
+        return f"EXPLAIN (FORMAT JSON, ANALYZE {'TRUE' if analyze else 'FALSE'}) {sql_text}"
+
+    def extract_plan(self, rows):
+        return rows[0].get("QUERY PLAN", []) if rows else []
+
+    def column_names(self, description) -> List[str]:
+        return [column.name for column in (description or [])]
+
+    def _get_pool(self, connection_id: str, dsn: str) -> ConnectionPool:
+        if ConnectionPool is None or dict_row is None:
+            raise ConfigurationError("缺少 psycopg / psycopg-pool 依赖，请先安装项目依赖。")
+        pool = self._pools.get(connection_id)
+        if pool is None:
+            pool = ConnectionPool(
+                conninfo=dsn,
+                min_size=0,
+                max_size=4,
+                open=True,
+                kwargs={"autocommit": True, "row_factory": dict_row},
+            )
+            self._pools[connection_id] = pool
+        return pool

sql_query_mcp/app.py

@@ -0,0 +1,105 @@
+"""FastMCP application for stateless SQL queries."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from mcp.server.fastmcp import FastMCP
+
+from .audit import AuditLogger
+from .config import load_config
+from .errors import SqlQueryMCPError
+from .executor import QueryExecutor
+from .introspection import MetadataService
+from .registry import ConnectionRegistry
+
+
+def create_app() -> FastMCP:
+    app_config = load_config()
+    registry = ConnectionRegistry(app_config)
+    audit_logger = AuditLogger(app_config.settings.audit_log_path)
+    metadata = MetadataService(registry, app_config.settings, audit_logger)
+    executor = QueryExecutor(registry, app_config.settings, audit_logger)
+
+    mcp = FastMCP("sql-query-mcp", json_response=True)
+
+    @mcp.tool()
+    def list_connections() -> dict:
+        """List configured SQL connections by connection_id."""
+
+        return {"connections": registry.list_connections()}
+
+    @mcp.tool()
+    def list_schemas(connection_id: str) -> dict:
+        """List visible schemas for a PostgreSQL connection."""
+
+        return _run_tool(lambda: metadata.list_schemas(connection_id))
+
+    @mcp.tool()
+    def list_databases(connection_id: str) -> dict:
+        """List visible databases for a MySQL connection."""
+
+        return _run_tool(lambda: metadata.list_databases(connection_id))
+
+    @mcp.tool()
+    def list_tables(
+        connection_id: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+    ) -> dict:
+        """List tables and views for a resolved PostgreSQL schema or MySQL database."""
+
+        return _run_tool(lambda: metadata.list_tables(connection_id, schema, database))
+
+    @mcp.tool()
+    def describe_table(
+        connection_id: str,
+        table_name: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+    ) -> dict:
+        """Describe columns, keys, and indexes for a table."""
+
+        return _run_tool(lambda: metadata.describe_table(connection_id, table_name, schema, database))
+
+    @mcp.tool()
+    def run_select(connection_id: str, sql: str, limit: Optional[int] = None) -> dict:
+        """Run a read-only SELECT or CTE query."""
+
+        return _run_tool(lambda: executor.run_select(connection_id, sql, limit))
+
+    @mcp.tool()
+    def explain_query(connection_id: str, sql: str, analyze: bool = False) -> dict:
+        """Run EXPLAIN on a read-only SELECT or CTE query."""
+
+        return _run_tool(lambda: executor.explain_query(connection_id, sql, analyze))
+
+    @mcp.tool()
+    def get_table_sample(
+        connection_id: str,
+        table_name: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> dict:
+        """Fetch a small sample from a table for schema discovery."""
+
+        return _run_tool(lambda: executor.get_table_sample(connection_id, table_name, schema, database, limit))
+
+    return mcp
+
+
+def _run_tool(func):
+    try:
+        return func()
+    except SqlQueryMCPError as exc:
+        raise ValueError(str(exc)) from exc
+
+
+def main() -> None:
+    app = create_app()
+    app.run()
+
+
+if __name__ == "__main__":
+    main()

sql_query_mcp/audit.py

@@ -0,0 +1,42 @@
+"""Audit logging utilities."""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+class AuditLogger:
+    """Write audit records as JSON lines."""
+
+    def __init__(self, log_path: Path):
+        self._log_path = Path(log_path)
+
+    def log(
+        self,
+        *,
+        tool: str,
+        connection_id: Optional[str],
+        success: bool,
+        duration_ms: int,
+        row_count: Optional[int] = None,
+        sql_summary: Optional[str] = None,
+        error: Optional[str] = None,
+        extra: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        record = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "tool": tool,
+            "connection_id": connection_id,
+            "success": success,
+            "duration_ms": duration_ms,
+            "row_count": row_count,
+            "sql_summary": sql_summary,
+            "error": error,
+            "extra": extra or {},
+        }
+        self._log_path.parent.mkdir(parents=True, exist_ok=True)
+        with self._log_path.open("a", encoding="utf-8") as handle:
+            handle.write(json.dumps(record, ensure_ascii=False) + "\n")