319 lines
11 KiB
Python
319 lines
11 KiB
Python
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
|