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

sql_query_mcp/config.py

@@ -0,0 +1,255 @@
+"""Configuration loading for sql-query-mcp."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+from .errors import ConfigurationError
+
+PACKAGE_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_CONFIG_ENV = "SQL_QUERY_MCP_CONFIG"
+DEFAULT_CONFIG_PATH = PACKAGE_ROOT / "config" / "connections.json"
+DEFAULT_AUDIT_LOG_PATH = PACKAGE_ROOT / "logs" / "audit.jsonl"
+CONNECTION_ID_RE = re.compile(r"^[a-z0-9]+(?:_[a-z0-9]+){3,}$")
+SUPPORTED_ENGINES = {"postgres", "mysql"}
+
+
+@dataclass(frozen=True)
+class ServerSettings:
+    default_limit: int = 200
+    max_limit: int = 1000
+    statement_timeout_ms: Optional[int] = None
+    audit_log_path: Path = DEFAULT_AUDIT_LOG_PATH
+
+
+@dataclass(frozen=True)
+class ConnectionConfig:
+    connection_id: str
+    engine: str
+    env: str
+    tenant: str
+    role: str
+    dsn_env: str
+    enabled: bool = True
+    label: Optional[str] = None
+    description: Optional[str] = None
+    default_schema: Optional[str] = None
+    default_database: Optional[str] = None
+
+    @property
+    def summary(self) -> Dict[str, object]:
+        return {
+            "connection_id": self.connection_id,
+            "engine": self.engine,
+            "label": self.label or self.connection_id,
+            "env": self.env,
+            "tenant": self.tenant,
+            "role": self.role,
+            "enabled": self.enabled,
+            "default_schema": self.default_schema,
+            "default_database": self.default_database,
+            "description": self.description,
+        }
+
+
+@dataclass(frozen=True)
+class AppConfig:
+    settings: ServerSettings
+    connections: List[ConnectionConfig]
+
+    @property
+    def connection_map(self) -> Dict[str, ConnectionConfig]:
+        return {item.connection_id: item for item in self.connections}
+
+    def enabled_connections(self) -> Iterable[ConnectionConfig]:
+        return (item for item in self.connections if item.enabled)
+
+
+def resolve_config_path(config_path: Optional[str] = None) -> Path:
+    raw_path = config_path or os.environ.get(DEFAULT_CONFIG_ENV)
+    return Path(raw_path).expanduser().resolve() if raw_path else DEFAULT_CONFIG_PATH
+
+
+def load_config(config_path: Optional[str] = None) -> AppConfig:
+    path = resolve_config_path(config_path)
+    if not path.exists():
+        return AppConfig(settings=ServerSettings(), connections=[])
+
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise ConfigurationError(f"配置文件不是有效 JSON: {path}") from exc
+
+    settings = _parse_settings(payload.get("settings", {}), path)
+    connections = _parse_connections(payload.get("connections", []))
+    return AppConfig(settings=settings, connections=connections)
+
+
+def _parse_settings(data: Dict[str, object], path: Path) -> ServerSettings:
+    default_limit = _required_int(data.get("default_limit", 200), "default_limit")
+    max_limit = _required_int(data.get("max_limit", 1000), "max_limit")
+    statement_timeout_ms = _optional_positive_int(
+        data.get("statement_timeout_ms"), "statement_timeout_ms"
+    )
+    audit_log_raw = str(
+        data.get("audit_log_path", DEFAULT_AUDIT_LOG_PATH.relative_to(PACKAGE_ROOT))
+    )
+    audit_log_path = Path(audit_log_raw)
+    if not audit_log_path.is_absolute():
+        audit_log_path = (path.parent / audit_log_path).resolve()
+
+    if default_limit <= 0:
+        raise ConfigurationError("default_limit 必须大于 0")
+    if max_limit < default_limit:
+        raise ConfigurationError("max_limit 不能小于 default_limit")
+    return ServerSettings(
+        default_limit=default_limit,
+        max_limit=max_limit,
+        statement_timeout_ms=statement_timeout_ms,
+        audit_log_path=audit_log_path,
+    )
+
+
+def _parse_connections(items: object) -> List[ConnectionConfig]:
+    if not isinstance(items, list):
+        raise ConfigurationError("connections 必须是数组")
+
+    result: List[ConnectionConfig] = []
+    seen = set()
+    for item in items:
+        if not isinstance(item, dict):
+            raise ConfigurationError("connections 数组中的每一项都必须是对象")
+
+        connection_id = str(item.get("connection_id", "")).strip()
+        engine = str(item.get("engine", "")).strip()
+        label = _required_string(item, "label", connection_id or "connection")
+        enabled = _required_bool(item, "enabled", connection_id or "connection")
+        if not CONNECTION_ID_RE.match(connection_id):
+            raise ConfigurationError(
+                "connection_id 必须符合 <system>_<env>_<tenant>_<role> 风格，且只包含小写字母、数字、下划线"
+            )
+        if engine not in SUPPORTED_ENGINES:
+            raise ConfigurationError(
+                f"{connection_id} 缺少合法 engine，必须是 postgres 或 mysql"
+            )
+        if connection_id in seen:
+            raise ConfigurationError(f"重复的 connection_id: {connection_id}")
+        seen.add(connection_id)
+
+        dsn_env = str(item.get("dsn_env", "")).strip()
+        env = str(item.get("env", "")).strip()
+        tenant = str(item.get("tenant", "")).strip()
+        role = str(item.get("role", "")).strip()
+        if not all((dsn_env, env, tenant, role)):
+            raise ConfigurationError(
+                f"{connection_id} 缺少必要字段，必须提供 env / tenant / role / dsn_env"
+            )
+
+        _reject_legacy_namespace_fields(item, connection_id)
+        default_schema = _optional_string(item.get("default_schema"))
+        default_database = _optional_string(item.get("default_database"))
+
+        if engine == "postgres" and "default_database" in item:
+            raise ConfigurationError(
+                f"{connection_id} 是 PostgreSQL 连接，不能配置 default_database"
+            )
+        if engine == "mysql" and "default_schema" in item:
+            raise ConfigurationError(
+                f"{connection_id} 是 MySQL 连接，不能配置 default_schema"
+            )
+
+        result.append(
+            ConnectionConfig(
+                connection_id=connection_id,
+                engine=engine,
+                label=label,
+                description=(
+                    str(item["description"]).strip()
+                    if item.get("description")
+                    else None
+                ),
+                env=env,
+                tenant=tenant,
+                role=role,
+                dsn_env=dsn_env,
+                enabled=enabled,
+                default_schema=default_schema,
+                default_database=default_database,
+            )
+        )
+
+    return result
+
+
+def _reject_legacy_namespace_fields(
+    item: Dict[str, object], connection_id: str
+) -> None:
+    if "default_namespace" in item:
+        raise ConfigurationError(
+            f"{connection_id} 仍在使用 default_namespace，请改为 default_schema 或 default_database"
+        )
+    if "default_schemas" in item:
+        raise ConfigurationError(
+            f"{connection_id} 仍在使用 default_schemas，请收敛为单值字段 default_schema"
+        )
+
+
+def _optional_string(value: object) -> Optional[str]:
+    if value is None:
+        return None
+    text = str(value).strip()
+    return text or None
+
+
+def _optional_positive_int(value: Any, field_name: str) -> Optional[int]:
+    if value is None:
+        return None
+    parsed = _required_int(value, field_name)
+    if parsed <= 0:
+        raise ConfigurationError(f"{field_name} 必须是大于 0 的整数")
+    return parsed
+
+
+def _required_int(value: Any, field_name: str) -> int:
+    if isinstance(value, bool):
+        raise ConfigurationError(f"{field_name} 必须是整数")
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        if value.is_integer():
+            return int(value)
+        raise ConfigurationError(f"{field_name} 必须是整数")
+    if isinstance(value, str):
+        try:
+            return int(value)
+        except ValueError as exc:
+            raise ConfigurationError(f"{field_name} 必须是整数") from exc
+    try:
+        return value.__int__()
+    except (TypeError, ValueError) as exc:
+        raise ConfigurationError(f"{field_name} 必须是整数") from exc
+
+
+def _required_string(
+    item: Dict[str, object], field_name: str, connection_id: str
+) -> str:
+    value = _optional_string(item.get(field_name))
+    if value is None:
+        raise ConfigurationError(f"{connection_id} 缺少必要字段 {field_name}")
+    return value
+
+
+def _required_bool(
+    item: Dict[str, object], field_name: str, connection_id: str
+) -> bool:
+    if field_name not in item:
+        raise ConfigurationError(f"{connection_id} 缺少必要字段 {field_name}")
+    value = item[field_name]
+    if not isinstance(value, bool):
+        raise ConfigurationError(f"{connection_id} 的 {field_name} 必须是布尔值")
+    return value

sql_query_mcp/errors.py

@@ -0,0 +1,34 @@
+"""Error types for sql-query-mcp."""
+
+from __future__ import annotations
+
+import re
+
+
+class SqlQueryMCPError(Exception):
+    """Base error for this package."""
+
+
+class ConfigurationError(SqlQueryMCPError):
+    """Raised when local configuration is invalid."""
+
+
+class ConnectionNotFoundError(SqlQueryMCPError):
+    """Raised when the requested connection_id does not exist."""
+
+
+class SecurityError(SqlQueryMCPError):
+    """Raised when SQL validation rejects a query."""
+
+
+class QueryExecutionError(SqlQueryMCPError):
+    """Raised when the database execution layer fails."""
+
+
+_DSN_CREDENTIALS_RE = re.compile(r"://([^:@/\s]+):([^@/\s]+)@")
+
+
+def sanitize_error_message(message: str) -> str:
+    """Mask DSN credentials before surfacing an error to the model."""
+
+    return _DSN_CREDENTIALS_RE.sub(r"://\1:***@", message)

sql_query_mcp/executor.py

@@ -0,0 +1,243 @@
+"""Query execution helpers."""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, Optional
+
+from .audit import AuditLogger
+from .config import ServerSettings
+from .errors import QueryExecutionError, sanitize_error_message
+from .namespace import resolve_namespace
+from .registry import ConnectionRegistry
+from .validator import (
+    build_limited_query,
+    clamp_limit,
+    summarize_sql,
+    validate_select_sql,
+)
+
+
+class QueryExecutor:
+    """Execute validated read-only SQL."""
+
+    def __init__(
+        self,
+        registry: ConnectionRegistry,
+        settings: ServerSettings,
+        audit_logger: AuditLogger,
+    ):
+        self._registry = registry
+        self._settings = settings
+        self._audit = audit_logger
+
+    def run_select(
+        self, connection_id: str, sql_text: str, limit: Optional[int] = None
+    ) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        row_limit = clamp_limit(
+            limit, self._settings.default_limit, self._settings.max_limit
+        )
+
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            cleaned_sql = validate_select_sql(sql_text, config.engine)
+            limited_sql, _ = build_limited_query(cleaned_sql, row_limit)
+            sql_summary = summarize_sql(cleaned_sql)
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                with conn.cursor() as cur:
+                    cur.execute(limited_sql)
+                    columns = adapter.column_names(cur.description)
+                    rows = cur.fetchall()
+
+                truncated = len(rows) > row_limit
+                trimmed_rows = rows[:row_limit]
+                duration_ms = _elapsed_ms(started)
+                self._audit.log(
+                    tool="run_select",
+                    connection_id=connection_id,
+                    success=True,
+                    duration_ms=duration_ms,
+                    row_count=len(trimmed_rows),
+                    sql_summary=sql_summary,
+                    extra={
+                        "engine": config.engine,
+                        "limit": row_limit,
+                        "truncated": truncated,
+                    },
+                )
+                return {
+                    "connection_id": connection_id,
+                    "engine": config.engine,
+                    "columns": columns,
+                    "rows": trimmed_rows,
+                    "row_count": len(trimmed_rows),
+                    "truncated": truncated,
+                    "duration_ms": duration_ms,
+                    "applied_limit": row_limit,
+                }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sql_summary = summarize_sql(sql_text)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="run_select",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                sql_summary=sql_summary,
+                error=sanitized,
+                extra=_build_audit_extra(config, limit=row_limit),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+    def explain_query(
+        self, connection_id: str, sql_text: str, analyze: bool = False
+    ) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            cleaned_sql = validate_select_sql(sql_text, config.engine)
+            sql_summary = summarize_sql(cleaned_sql)
+            adapter = self._registry.get_adapter(config)
+            explain_sql = adapter.build_explain_query(cleaned_sql, analyze=analyze)
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                with conn.cursor() as cur:
+                    cur.execute(explain_sql)
+                    rows = cur.fetchall()
+
+                duration_ms = _elapsed_ms(started)
+                plan = adapter.extract_plan(rows)
+                self._audit.log(
+                    tool="explain_query",
+                    connection_id=connection_id,
+                    success=True,
+                    duration_ms=duration_ms,
+                    row_count=1 if rows else 0,
+                    sql_summary=sql_summary,
+                    extra={"engine": config.engine, "analyze": analyze},
+                )
+                return {
+                    "connection_id": connection_id,
+                    "engine": config.engine,
+                    "plan": plan,
+                    "duration_ms": duration_ms,
+                    "analyze": analyze,
+                }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sql_summary = summarize_sql(sql_text)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="explain_query",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                sql_summary=sql_summary,
+                error=sanitized,
+                extra=_build_audit_extra(config, analyze=analyze),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+    def get_table_sample(
+        self,
+        connection_id: str,
+        table_name: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> Dict[str, object]:
+        row_limit = clamp_limit(
+            limit, self._settings.default_limit, self._settings.max_limit
+        )
+        started = time.perf_counter()
+        config = None
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            namespace = resolve_namespace(config, schema=schema, database=database)
+            adapter = self._registry.get_adapter(config)
+            query = adapter.build_sample_query(
+                namespace.value, table_name, row_limit + 1
+            )
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                with conn.cursor() as cur:
+                    cur.execute(query)
+                    columns = adapter.column_names(cur.description)
+                    rows = cur.fetchall()
+
+                truncated = len(rows) > row_limit
+                trimmed_rows = rows[:row_limit]
+                duration_ms = _elapsed_ms(started)
+                self._audit.log(
+                    tool="get_table_sample",
+                    connection_id=connection_id,
+                    success=True,
+                    duration_ms=duration_ms,
+                    row_count=len(trimmed_rows),
+                    sql_summary=f"sample {namespace.value}.{table_name}",
+                    extra={
+                        "engine": config.engine,
+                        namespace.field_name: namespace.value,
+                        "table_name": table_name,
+                        "limit": row_limit,
+                    },
+                )
+                return {
+                    "connection_id": connection_id,
+                    "engine": config.engine,
+                    namespace.field_name: namespace.value,
+                    "table_name": table_name,
+                    "columns": columns,
+                    "rows": trimmed_rows,
+                    "row_count": len(trimmed_rows),
+                    "truncated": truncated,
+                    "duration_ms": duration_ms,
+                    "applied_limit": row_limit,
+                }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="get_table_sample",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                error=sanitized,
+                extra=_build_audit_extra(
+                    config,
+                    schema=schema,
+                    database=database,
+                    table_name=table_name,
+                    limit=row_limit,
+                ),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+
+def _elapsed_ms(started: float) -> int:
+    return int((time.perf_counter() - started) * 1000)
+
+
+def _apply_statement_timeout(
+    adapter: Any, conn: Any, timeout_ms: Optional[int]
+) -> None:
+    if timeout_ms is not None:
+        getattr(adapter, "set_statement_timeout")(conn, timeout_ms)
+
+
+def _build_audit_extra(config, **kwargs: object) -> Dict[str, object]:
+    extra = dict(kwargs)
+    if config is not None:
+        extra["engine"] = config.engine
+    return extra