80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
import hashlib
|
|
from typing import Any
|
|
|
|
from django.core.cache import cache
|
|
from django.db.models import Q
|
|
|
|
from apps.bi.models import Dataset
|
|
from apps.bi.services import execute_rendered_dataset, render_dataset_sql
|
|
from mcp_server.context import authenticated_factory_user, require_permission
|
|
from mcp_server.tools.common import json_safe_result, validate_result_size
|
|
|
|
|
|
def search_datasets(query: str = "", limit: int = 20) -> dict[str, Any]:
|
|
"""搜索可执行的数据集目录，不返回 SQL 等敏感配置。"""
|
|
authenticated_factory_user()
|
|
safe_limit = max(1, min(limit, 100))
|
|
queryset = Dataset.objects.filter(enabled=True)
|
|
if query.strip():
|
|
queryset = queryset.filter(
|
|
Q(name__icontains=query.strip())
|
|
| Q(code__icontains=query.strip())
|
|
| Q(description__icontains=query.strip())
|
|
)
|
|
rows = queryset.order_by("name", "code", "id").values(
|
|
"code",
|
|
"name",
|
|
"description",
|
|
"default_param",
|
|
"test_param",
|
|
)[:safe_limit]
|
|
return {"items": list(rows), "limit": safe_limit}
|
|
|
|
|
|
def execute_dataset(
|
|
code: str,
|
|
parameters: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""以当前 Factory 用户身份执行启用的数据集。需要 dataset.exec 权限。"""
|
|
user = authenticated_factory_user()
|
|
require_permission(user, "dataset.exec")
|
|
try:
|
|
dataset = Dataset.objects.get(code=code, enabled=True)
|
|
except Dataset.DoesNotExist as exc:
|
|
raise ValueError(f"未找到已启用的数据集：{code}") from exc
|
|
except Dataset.MultipleObjectsReturned as exc:
|
|
raise RuntimeError(f"数据集 code 不唯一，无法执行：{code}") from exc
|
|
|
|
query = dict(parameters or {})
|
|
query["r_user"] = user.id
|
|
query["r_dept"] = user.belong_dept_id or ""
|
|
full_sql = render_dataset_sql(dataset, query)
|
|
|
|
cache_key = None
|
|
response_data = None
|
|
if full_sql and dataset.cache_seconds:
|
|
digest = hashlib.sha256(full_sql.strip(";").encode()).hexdigest()
|
|
cache_key = f"mcp:dataset:{digest}"
|
|
response_data = cache.get(cache_key)
|
|
|
|
if response_data is None:
|
|
response_data, can_cache = execute_rendered_dataset(dataset, full_sql)
|
|
if cache_key and can_cache and response_data["data"]:
|
|
cache.set(cache_key, response_data, dataset.cache_seconds)
|
|
|
|
result = json_safe_result(
|
|
{
|
|
"code": dataset.code,
|
|
"name": dataset.name,
|
|
"description": dataset.description,
|
|
**response_data,
|
|
}
|
|
)
|
|
validate_result_size(result)
|
|
return result
|
|
|
|
|
|
def register_dataset_tools(server) -> None:
|
|
server.tool()(search_datasets)
|
|
server.tool()(execute_dataset)
|