metadata tools

src/mini_db_mcp/db.py

@@ -1,11 +1,19 @@
 from typing import Any
-
-from sqlalchemy import URL, inspect
+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
@@ -63,24 +71,21 @@ class Database:
             table_names, view_names = await connection.run_sync(inspect_tables)
 
         keyword_cf = keyword.casefold() if keyword else None
-        names = sorted(
+        sorted_table_names = sorted(
             name
             for name in table_names | view_names
             if keyword_cf is None or keyword_cf in name.casefold()
         )
 
-        start = 0
-        if cursor is not None:
-            try:
-                start = names.index(cursor) + 1
-            except ValueError as exc:
-                raise ValueError("invalid pagination cursor") from exc
-
-        page_names = names[start : start + limit]
-        next_cursor = page_names[-1] if start + len(page_names) < len(names) else None
+        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 page_names:
+        for name in paged_table_names:
             object_type = "view" if name in view_names else "table"
             objects.append(
                 {
@@ -91,3 +96,140 @@ class Database:
                 }
             )
         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