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

from typing import Any
from datetime import datetime, date, time
from decimal import Decimal
from uuid import UUID
import base64
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):
    pass

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


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