mini-db-mcp/src/mini_db_mcp/db.py

import asyncio
import base64
import json
from time import monotonic
from typing import Any, Iterable, Sequence
from datetime import datetime, date, time
from decimal import Decimal
from uuid import UUID
from sqlalchemy import URL, inspect, Table, MetaData, select
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.exc import SQLAlchemyError

from mini_db_mcp.config import Config
from mini_db_mcp.utils.pagination_util import paginate_list


class QueryExecutionError(RuntimeError):
    def __init__(
        self,
        code: str,
        message: str | None = None,
        *,
        retryable: bool = False,
    ) -> None:
        self.code = code
        self.message = message or code
        self.retryable = retryable
        super().__init__(self.message)

class Database:
    def __init__(self, settings: Config):
        self.settings = settings
        self.database = settings.database
        url = URL.create(
            "mysql+aiomysql",
            username=self.database.user,
            password=self.database.password,
            host=self.database.host,
            port=self.database.port,
            database=self.database.schema or None,
        )
        self.engine: AsyncEngine = create_async_engine(
            url,
            pool_size=self.database.pool_size,
            max_overflow=self.database.max_overflow,
            pool_timeout=self.database.pool_timeout,
            pool_recycle=self.database.pool_recycle,
            pool_pre_ping=self.database.pool_pre_ping,
            connect_args={"connect_timeout": self.database.connect_timeout},
        )

    async def close(self) -> None:
        await self.engine.dispose()


    async def list_tables(
        self,
        *,
        keyword: str | None,
        limit: int,
        cursor: str | None,
    ) -> tuple[list[dict[str, Any]], str | None]:
        """Return one name-sorted page of tables and views.

        ``keyword`` filters object names case-insensitively. ``cursor`` must be
        the last object name returned by the previous page. The returned cursor
        is ``None`` when no additional objects remain.

        Raises:
            ValueError: If the cursor is not present in the filtered result.
        """
        schema = self.settings.database.schema or None

        def inspect_tables(sync_connection):
            inspector = inspect(sync_connection)
            table_names = set(inspector.get_table_names(schema=schema))
            try:
                view_names = set(inspector.get_view_names(schema=schema))
            except NotImplementedError:
                view_names = set()
            return table_names, view_names

        async with self.engine.connect() as connection:
            table_names, view_names = await connection.run_sync(inspect_tables)

        keyword_cf = keyword.casefold() if keyword else None
        sorted_table_names = sorted(
            name
            for name in table_names | view_names
            if keyword_cf is None or keyword_cf in name.casefold()
        )

        paged_table_names, next_cursor = paginate_list(
            sorted_table_names,
            cursor=cursor,
            page_size=limit,
            identifier=lambda name: name,
        )

        objects: list[dict[str, Any]] = []
        for name in paged_table_names:
            object_type = "view" if name in view_names else "table"
            objects.append(
                {
                    "schema_name": schema,
                    "object_name": name,
                    "object_type": object_type,
                    "description": None,
                }
            )
        return objects, next_cursor


    async def describe_table(
        self,
        *,
        table_name: str,
        schema: str | None = None,
        limit: int,
        cursor: str | None,
    ) -> dict[str, Any]:
        """Return one cursor-paginated page of a table's columns."""
        effective_schema = schema or self.settings.database.schema or None

        def inspect_table(sync_connection):
            inspector = inspect(sync_connection)
            columns = inspector.get_columns(table_name, schema=effective_schema)
            primary_keys = set(
                (
                    inspector.get_pk_constraint(
                        table_name,
                        schema=effective_schema,
                    )
                    or {}
                ).get("constrained_columns") or []
            )
            try:
                table_description = (
                    inspector.get_table_comment(
                        table_name,
                        schema=effective_schema,
                    )
                    or {}
                ).get("text")
            except NotImplementedError:
                table_description = None
            return columns, primary_keys, table_description

        try:
            async with self.engine.connect() as connection:
                columns, primary_keys, table_description = await connection.run_sync(
                    inspect_table
                )
        except SQLAlchemyError as exc:
            raise QueryExecutionError(f"Failed to read table metadata: {exc}") from exc

        paged_columns, next_cursor = paginate_list(
            columns,
            cursor=cursor,
            page_size=limit,
            identifier=lambda column: column["name"],
        )
        visible_columns: list[dict[str, Any]] = []
        for column in paged_columns:
            visible_columns.append(
                {
                    "name": column["name"],
                    "data_type": str(column.get("type")),
                    "nullable": column.get("nullable"),
                    "primary_key": column["name"] in primary_keys,
                    "description": column.get("comment"),
                }
            )

        return {
            "schema_name": effective_schema,
            "table_name": table_name,
            "table_description": table_description,
            "columns": visible_columns,
            "next_cursor": next_cursor,
            "truncated": next_cursor is not None,
        }


    async def sample_values(
        self,         
        *,
        table_name: str,
        schema: str | None = None,
        column: str,
        limit: int
    ) -> tuple[list[Any], bool]:
        effective_schema = schema or self.settings.database.schema or None
        metadata = MetaData()
        try:
            async with self.engine.connect() as connection:
                table_obj = await connection.run_sync(
                    lambda sync_connection: Table(
                        table_name,
                        metadata,
                        schema=effective_schema,
                        autoload_with=sync_connection,
                    )
                )
        except SQLAlchemyError as exc:
            raise QueryExecutionError("metadata_error", f"反射表结构失败：{exc}") from exc
        try:
            column_obj = table_obj.c[column]
        except KeyError as exc:
            raise QueryExecutionError(
                "metadata_error", f"列不存在：{table_name}.{column}"
            ) from exc
        statement = (
            select(column_obj)
            .where(column_obj.is_not(None))
            .distinct()
            .limit(limit + 1)
        )
        try:
            async with self.engine.connect() as connection:
                result = await connection.execute(statement)
                values = [row[0] for row in result]
        except SQLAlchemyError as exc:
            raise QueryExecutionError("query_failed", f"采样查询失败：{exc}") from exc
        truncated = len(values) > limit
        return [
            _serialize_scalar(
                value,
                max_chars=self.settings.query.max_cell_chars,
            )
            for value in values[:limit]
        ], truncated

    async def execute_guarded(self, sql: str) -> dict[str, Any]:
        start = monotonic()
        try:
            async with asyncio.timeout(self.settings.query.timeout_seconds):
                async with self.engine.connect() as connection:
                    result = await connection.exec_driver_sql(sql)
                    columns = list(result.keys())
                    raw_rows: Iterable[Any] = result
                    rows, truncated = bound_rows(
                        columns,
                        raw_rows,
                        max_rows=self.settings.query.max_rows,
                        max_bytes=self.settings.query.max_bytes,
                        max_cell_chars=self.settings.query.max_cell_chars,
                    )
        except SQLAlchemyError as exc:
            raise QueryExecutionError(
                "query_failed",
                "查询执行失败，请检查 SQL、权限及数据库状态。",
            ) from exc
        except TimeoutError as exc:
            raise QueryExecutionError(
                "query_timeout",
                "查询超过执行时间上限，已中断。",
                retryable=True,
            ) from exc

        return {
            "columns": columns,
            "rows": rows,
            "row_count": len(rows),
            "truncated": truncated,
            "elapsed_ms": int((monotonic() - start) * 1000),
        }


def bound_rows(
    columns: Sequence[str],
    source_rows: Iterable[Sequence[Any]],
    *,
    max_rows: int,
    max_bytes: int,
    max_cell_chars: int,
) -> tuple[list[list[Any]], bool]:
    """Bound model-facing output by rows and serialized byte size.

    This only limits what is returned to the model. It does not limit the amount of
    data the database may scan; query timeouts and database workload governance are
    separate controls.
    """

    rows: list[list[Any]] = []
    truncated = False
    base_size = len(json.dumps({"columns": list(columns), "rows": []}, ensure_ascii=False).encode("utf-8"))
    current_size = base_size

    for raw_row in source_rows:
        if len(rows) >= max_rows:
            truncated = True
            break
        row = [_serialize_scalar(value, max_chars=max_cell_chars) for value in raw_row]
        row_size = len(json.dumps(row, ensure_ascii=False, default=str).encode("utf-8"))
        if current_size + row_size > max_bytes:
            truncated = True
            break
        rows.append(row)
        current_size += row_size

    return rows, truncated


def _serialize_scalar(value: Any, *, max_chars: int) -> Any:
    if value is None or isinstance(value, (bool, int, float)):
        return value
    if isinstance(value, Decimal):
        return str(value)
    if isinstance(value, (datetime, date, time, UUID)):
        return value.isoformat() if hasattr(value, "isoformat") else str(value)
    if isinstance(value, (bytes, bytearray, memoryview)):
        text = base64.b64encode(bytes(value)).decode("ascii")
    else:
        text = str(value)
    if len(text) > max_chars:
        return text[:max_chars] + "…[truncated]"
    return text