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

sql_query_mcp/introspection.py

@@ -0,0 +1,225 @@
+"""Metadata query 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 require_engine, resolve_namespace
+from .registry import ConnectionRegistry
+
+
+class MetadataService:
+    """Read-only metadata queries across PostgreSQL and MySQL."""
+
+    def __init__(
+        self,
+        registry: ConnectionRegistry,
+        settings: ServerSettings,
+        audit_logger: AuditLogger,
+    ):
+        self._registry = registry
+        self._settings = settings
+        self._audit = audit_logger
+
+    def list_schemas(self, connection_id: str) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            require_engine(config, "postgres", "list_schemas")
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                schemas = adapter.list_schemas(conn)
+            duration_ms = _elapsed_ms(started)
+            self._audit.log(
+                tool="list_schemas",
+                connection_id=connection_id,
+                success=True,
+                duration_ms=duration_ms,
+                row_count=len(schemas),
+                extra={"engine": config.engine},
+            )
+            return {
+                "connection_id": connection_id,
+                "engine": "postgres",
+                "schemas": schemas,
+            }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="list_schemas",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                error=sanitized,
+                extra=_build_audit_extra(config),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+    def list_databases(self, connection_id: str) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            require_engine(config, "mysql", "list_databases")
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                databases = adapter.list_databases(conn)
+            duration_ms = _elapsed_ms(started)
+            self._audit.log(
+                tool="list_databases",
+                connection_id=connection_id,
+                success=True,
+                duration_ms=duration_ms,
+                row_count=len(databases),
+                extra={"engine": config.engine},
+            )
+            return {
+                "connection_id": connection_id,
+                "engine": "mysql",
+                "databases": databases,
+            }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="list_databases",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                error=sanitized,
+                extra=_build_audit_extra(config),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+    def list_tables(
+        self,
+        connection_id: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+    ) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            namespace = resolve_namespace(config, schema=schema, database=database)
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                tables = adapter.list_tables(conn, namespace.value)
+            duration_ms = _elapsed_ms(started)
+            self._audit.log(
+                tool="list_tables",
+                connection_id=connection_id,
+                success=True,
+                duration_ms=duration_ms,
+                row_count=len(tables),
+                extra={"engine": config.engine, namespace.field_name: namespace.value},
+            )
+            return {
+                "connection_id": connection_id,
+                "engine": config.engine,
+                namespace.field_name: namespace.value,
+                "tables": tables,
+            }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="list_tables",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                error=sanitized,
+                extra=_build_audit_extra(config, schema=schema, database=database),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+    def describe_table(
+        self,
+        connection_id: str,
+        table_name: str,
+        schema: Optional[str] = None,
+        database: Optional[str] = None,
+    ) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        try:
+            config = self._registry.get_connection_config(connection_id)
+            namespace = resolve_namespace(config, schema=schema, database=database)
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(
+                    adapter, conn, self._settings.statement_timeout_ms
+                )
+                description = adapter.describe_table(conn, namespace.value, table_name)
+                if not description:
+                    raise QueryExecutionError(
+                        f"未找到表 {namespace.value}.{table_name}，或当前用户没有访问权限"
+                    )
+
+            duration_ms = _elapsed_ms(started)
+            self._audit.log(
+                tool="describe_table",
+                connection_id=connection_id,
+                success=True,
+                duration_ms=duration_ms,
+                row_count=len(description["columns"]),
+                extra={
+                    "engine": config.engine,
+                    namespace.field_name: namespace.value,
+                    "table_name": table_name,
+                },
+            )
+            return {
+                "connection_id": connection_id,
+                "engine": config.engine,
+                namespace.field_name: namespace.value,
+                "table_name": table_name,
+                "columns": description["columns"],
+                "indexes": description["indexes"],
+            }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="describe_table",
+                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,
+                ),
+            )
+            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

sql_query_mcp/namespace.py

@@ -0,0 +1,48 @@
+"""Namespace resolution helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+from .config import ConnectionConfig
+from .errors import SecurityError
+
+
+@dataclass(frozen=True)
+class NamespaceSelection:
+    field_name: str
+    value: str
+
+
+def resolve_namespace(
+    config: ConnectionConfig,
+    *,
+    schema: Optional[str] = None,
+    database: Optional[str] = None,
+) -> NamespaceSelection:
+    if schema and database:
+        raise SecurityError("schema 和 database 不能同时传入。")
+
+    if config.engine == "postgres":
+        if database:
+            raise SecurityError("PostgreSQL 连接不接受 database 参数。")
+        resolved = schema or config.default_schema
+        if not resolved:
+            raise SecurityError("PostgreSQL 连接必须显式传 schema，或在配置中设置 default_schema。")
+        return NamespaceSelection(field_name="schema", value=resolved)
+
+    if config.engine == "mysql":
+        if schema:
+            raise SecurityError("MySQL 连接不接受 schema 参数。")
+        resolved = database or config.default_database
+        if not resolved:
+            raise SecurityError("MySQL 连接必须显式传 database，或在配置中设置 default_database。")
+        return NamespaceSelection(field_name="database", value=resolved)
+
+    raise SecurityError(f"未知 engine: {config.engine}")
+
+
+def require_engine(config: ConnectionConfig, engine: str, tool_name: str) -> None:
+    if config.engine != engine:
+        raise SecurityError(f"{tool_name} 仅适用于 {engine} 连接，当前连接 engine={config.engine}")

sql_query_mcp/registry.py

@@ -0,0 +1,67 @@
+"""Connection registry and adapter routing."""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from typing import Any, Iterator, Tuple
+
+from .adapters import MySQLAdapter, PostgresAdapter
+from .config import AppConfig, ConnectionConfig
+from .errors import ConfigurationError, ConnectionNotFoundError
+
+
+class ConnectionRegistry:
+    """Resolve connection config and route to the correct engine adapter."""
+
+    def __init__(self, app_config: AppConfig):
+        self._config = app_config
+        self._adapters = {
+            "postgres": PostgresAdapter(),
+            "mysql": MySQLAdapter(),
+        }
+
+    def list_connections(self):
+        return [item.summary for item in self._config.connections]
+
+    def get_connection_config(self, connection_id: str) -> ConnectionConfig:
+        try:
+            config = self._config.connection_map[connection_id]
+        except KeyError as exc:
+            raise ConnectionNotFoundError(
+                f"未知 connection_id: {connection_id}"
+            ) from exc
+        if not config.enabled:
+            raise ConnectionNotFoundError(f"connection_id 已被禁用: {connection_id}")
+        return config
+
+    @contextmanager
+    def connection(
+        self, connection_id: str
+    ) -> Iterator[Tuple[Any, ConnectionConfig, Any]]:
+        config = self.get_connection_config(connection_id)
+        with self.connection_from_config(config) as (conn, adapter):
+            yield conn, config, adapter
+
+    @contextmanager
+    def connection_from_config(
+        self, config: ConnectionConfig
+    ) -> Iterator[Tuple[Any, Any]]:
+        dsn = os.environ.get(config.dsn_env)
+        if not dsn:
+            raise ConfigurationError(
+                f"{config.connection_id} 缺少环境变量 {config.dsn_env}，无法建立数据库连接"
+            )
+        adapter = self.get_adapter(config)
+        with adapter.connection(config.connection_id, dsn) as conn:
+            yield conn, adapter
+
+    def close(self) -> None:
+        for adapter in self._adapters.values():
+            adapter.close()
+
+    def get_adapter(self, config: ConnectionConfig):
+        try:
+            return self._adapters[config.engine]
+        except KeyError as exc:
+            raise ConfigurationError(f"不支持的 engine: {config.engine}") from exc

sql_query_mcp/release_metadata.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+import re
+
+try:
+    import tomllib
+except ModuleNotFoundError:  # pragma: no cover - Python 3.10 fallback
+    import tomli as tomllib  # type: ignore[import-not-found]
+
+
+TAG_PATTERN = re.compile(r"^v\d+\.\d+\.\d+$")
+VERSION_PATTERN = re.compile(r"^\d+\.\d+\.\d+$")
+
+
+@dataclass(frozen=True)
+class ReleaseContext:
+    tag: str
+    version: str
+    release_branch: str
+
+
+def resolve_effective_tag(
+    event_name: str, github_ref_name: str, input_tag: str | None
+) -> str:
+    if event_name == "workflow_dispatch":
+        if not input_tag:
+            raise ValueError("workflow_dispatch requires an explicit tag")
+        return input_tag
+    return github_ref_name
+
+
+def parse_version_tag(tag: str) -> str:
+    if not TAG_PATTERN.match(tag):
+        raise ValueError("release tags must match vX.Y.Z")
+    return tag[1:]
+
+
+def read_project_version(pyproject_path: Path) -> str:
+    data = tomllib.loads(pyproject_path.read_text(encoding="utf-8"))
+    version = data.get("project", {}).get("version")
+    if not version:
+        raise ValueError("missing project.version")
+    if not VERSION_PATTERN.match(version):
+        raise ValueError("project.version must match X.Y.Z")
+    return version
+
+
+def build_release_context(tag: str, pyproject_path: Path) -> ReleaseContext:
+    version = parse_version_tag(tag)
+    project_version = read_project_version(pyproject_path)
+    if version != project_version:
+        raise ValueError("tag version does not match pyproject version")
+    return ReleaseContext(tag=tag, version=version, release_branch=f"release/{tag}")
+
+
+def should_skip_pypi_upload(
+    is_recovery_run: bool,
+    pypi_version_exists: bool,
+    recovery_confirmed: bool,
+) -> bool:
+    return is_recovery_run and pypi_version_exists and recovery_confirmed
+
+
+def decide_backmerge_action(target: str, has_open_pr: bool, has_diff: bool) -> str:
+    if has_open_pr:
+        return "reuse"
+    if target == "main":
+        return "create"
+    if has_diff:
+        return "create"
+    return "skip"
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Resolve release metadata.")
+    parser.add_argument("--tag", required=True)
+    parser.add_argument("--pyproject", required=True)
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    context = build_release_context(args.tag, Path(args.pyproject))
+    print(f"tag={context.tag}")
+    print(f"version={context.version}")
+    print(f"release_branch={context.release_branch}")
+
+
+if __name__ == "__main__":
+    main()

sql_query_mcp/validator.py

@@ -0,0 +1,128 @@
+"""SQL validation and normalization."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Optional, Tuple
+
+from sqlglot import exp, parse_one
+from sqlglot.errors import ParseError
+
+from .errors import SecurityError
+
+COMMENT_TOKENS = ("--", "/*", "*/")
+READ_ONLY_ROOT_TYPES = tuple(
+    expr_type
+    for expr_type in (
+        getattr(exp, name, None) for name in ("Select", "Union", "Except", "Intersect")
+    )
+    if isinstance(expr_type, type)
+)
+MUTATING_EXPRESSION_TYPES = tuple(
+    expr_type
+    for expr_type in (
+        getattr(exp, name, None)
+        for name in (
+            "Insert",
+            "Update",
+            "Delete",
+            "Merge",
+            "Create",
+            "Drop",
+            "Alter",
+            "Command",
+            "Copy",
+            "Call",
+            "Set",
+            "Pragma",
+            "Use",
+            "Grant",
+            "Revoke",
+            "Transaction",
+            "Commit",
+            "Rollback",
+            "TruncateTable",
+            "Vacuum",
+            "Into",
+        )
+    )
+    if isinstance(expr_type, type)
+)
+DIALECT_BY_ENGINE = {"postgres": "postgres", "mysql": "mysql"}
+
+
+def validate_select_sql(sql: str, engine: str) -> str:
+    cleaned = _clean_sql(sql)
+    lowered = cleaned.lstrip().lower()
+    if lowered.startswith("explain"):
+        raise SecurityError(
+            "explain_query 会自动包装 SQL，请直接传 SELECT 或 WITH 查询。"
+        )
+    if not (lowered.startswith("select") or lowered.startswith("with")):
+        raise SecurityError("仅允许 SELECT 或 WITH ... SELECT 语句。")
+    statement = _parse_statement(cleaned, engine)
+    _ensure_read_only_statement(statement)
+    return cleaned
+
+
+def clamp_limit(limit: Optional[int], default_limit: int, max_limit: int) -> int:
+    value = default_limit if limit is None else int(limit)
+    if value <= 0:
+        raise SecurityError("limit 必须大于 0。")
+    return min(value, max_limit)
+
+
+def build_limited_query(sql: str, row_limit: int) -> Tuple[str, int]:
+    sentinel_limit = row_limit + 1
+    wrapped = f"SELECT * FROM ({sql}) AS _pq_result LIMIT {sentinel_limit}"
+    return wrapped, sentinel_limit
+
+
+def summarize_sql(sql: str, max_chars: int = 160) -> str:
+    one_line = re.sub(r"\s+", " ", sql).strip()
+    if len(one_line) <= max_chars:
+        return one_line
+    return one_line[: max_chars - 3] + "..."
+
+
+def _clean_sql(sql: str) -> str:
+    if not sql or not sql.strip():
+        raise SecurityError("SQL 不能为空。")
+
+    cleaned = sql.strip()
+    for token in COMMENT_TOKENS:
+        if token in cleaned:
+            raise SecurityError("不允许使用 SQL 注释。")
+
+    semicolon_count = cleaned.count(";")
+    if semicolon_count > 1:
+        raise SecurityError("只允许单条 SQL 语句。")
+    if semicolon_count == 1 and not cleaned.endswith(";"):
+        raise SecurityError("只允许单条 SQL 语句。")
+    if cleaned.endswith(";"):
+        cleaned = cleaned[:-1].rstrip()
+
+    if not cleaned:
+        raise SecurityError("SQL 不能为空。")
+    return cleaned
+
+
+def _parse_statement(sql: str, engine: str) -> Any:
+    try:
+        dialect = DIALECT_BY_ENGINE[engine]
+    except KeyError as exc:
+        raise SecurityError(f"不支持的 SQL 方言: {engine}") from exc
+
+    try:
+        return parse_one(sql, dialect=dialect)
+    except ParseError as exc:
+        raise SecurityError(f"SQL 解析失败，已拒绝执行: {exc}") from exc
+
+
+def _ensure_read_only_statement(statement: Any) -> None:
+    if not isinstance(statement, READ_ONLY_ROOT_TYPES):
+        raise SecurityError("仅允许 SELECT 或 WITH ... SELECT 语句。")
+
+    for node in statement.walk():
+        if isinstance(node, MUTATING_EXPRESSION_TYPES):
+            raise SecurityError(f"仅允许只读查询，检测到写操作: {node.key.upper()}")