add query domain handler

src/mini_db_mcp/db.py

@@ -1,8 +1,11 @@
-from typing import Any
+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
-import base64
 from sqlalchemy import URL, inspect, Table, MetaData, select
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from sqlalchemy.exc import SQLAlchemyError
@@ -12,7 +15,17 @@ from mini_db_mcp.utils.pagination_util import paginate_list
 
 
 class QueryExecutionError(RuntimeError):
-    pass
+    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):
@@ -218,6 +231,76 @@ class Database:
             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)):
@@ -232,4 +315,4 @@ def _serialize_scalar(value: Any, *, max_chars: int) -> Any:
         text = str(value)
     if len(text) > max_chars:
         return text[:max_chars] + "…[truncated]"
-    return text
+    return text