role_id 权限限制+查询

chenfeng559 · chenfeng559 · commit 67ffb0bbd955 · 2025-07-03T00:55:22.000+08:00
diff --git a/README.md b/README.md
@@ -26,6 +26,7 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
 - `api_key`：AskTable API 密钥（必需，环境变量）
 - `datasource_id`：数据源ID（必需，环境变量）
 - `base_url`：本地IP服务地址（可选，填写则走本地部署，不填则走SaaS）
+- `role_id` ：角色id（可选，填写则只能访问该角色被允许的数据，不填则即可查询所有数据）
 
 ---
 
@@ -36,7 +37,7 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
  - gen_conclusion ， 根据用户的问题，直接返回数据结果
    - 输入：用户问题，如："请给我出销售额前10的产品"
    - 输出：对应的数据结果
- - list_available_datasources ， 获取当前用户apikey下的可用的所有数据库（数据源）信息
+ - list_available_datasources ， 获取当前APIKEY下的用户（role_id）所有可用数据库（数据源）信息
    - 输入：我数据库中有哪些数据源？
    - 输出：对应的数据源信息，包括数据源id、数据库引擎、数据库描述
 
@@ -76,8 +77,9 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
       "command": "uvx",
       "args": ["asktable-mcp-server@latest", "--transport", "stdio"],
       "env": {
-        "api_key": "your_api_key",
-        "datasource_id": "your_datasource_id"
+        "api_key": "your_api_key",            // 必填
+        "datasource_id": "your_datasource_id", // 必填
+        // "role_id": "your_role_id"           // 可选：如需限定角色权限，请填写
       }
     }
   }
@@ -95,9 +97,10 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
       "command": "uvx",
       "args": ["asktable-mcp-server@latest", "--transport", "stdio"],
       "env": {
-        "api_key": "your_api_key",
-        "datasource_id": "your_datasource_id",
-        "base_url": "http://your_local_ip:port/api"
+        "api_key": "your_api_key",           // 必填
+        "datasource_id": "your_datasource_id",// 必填
+        "base_url": "http://your_local_ip:port/api", // 必填
+        // "role_id": "your_role_id"           // 可选：如需限定角色权限，请填写
       }
     }
   }
@@ -112,7 +115,8 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
 {
   "mcpServers": {    
     "asktable-mcp-server": {
-      "url": "http://localhost:8095/sse/?apikey=your_apikey&datasouce_id=your_datasouce_id",
+      // role_id 为可选参数，不指定则使用默认权限
+      "url": "http://localhost:8095/sse/?apikey=your_apikey&datasouce_id=your_datasouce_id&role_id=your_role_id",
       "headers": {},
       "timeout": 300,
       "sse_read_timeout": 300
diff --git a/src/asktable_mcp_server/server.py b/src/asktable_mcp_server/server.py
@@ -38,11 +38,13 @@ async def gen_sql(query: str) -> str:
         'question': query
     }
     
-    # 如果环境变量中有base_url，添加到参数中
+    # 如果环境变量中有base_url、role_id，添加到参数中
     base_url = os.getenv('base_url')
+    role_id = os.getenv('role_id')
     if base_url:
         params['base_url'] = base_url
-    
+    if role_id:
+        params['role_id'] = role_id
     message = await get_asktable_sql(**params)
     return message
 
@@ -74,11 +76,13 @@ async def gen_conclusion(query: str) -> str:
         'question': query
     }
     
-    # 如果环境变量中有base_url，添加到参数中
+    # 如果环境变量中有base_url、role_id，添加到参数中
     base_url = os.getenv('base_url')
+    role_id = os.getenv('role_id')
     if base_url:
         params['base_url'] = base_url
-
+    if role_id:
+        params['role_id'] = role_id
 
     message = await get_asktable_data(**params)
     return message
@@ -126,8 +130,9 @@ async def list_available_datasources() -> str:
     """
     api_key =  os.getenv('api_key')
     base_url = os.getenv('base_url') or None
+    role_id = os.getenv('role_id') or None
 
-    result = await get_datasources_info(api_key=api_key, base_url=base_url)
+    result = await get_datasources_info(api_key=api_key, base_url=base_url,role_id=role_id)
     logging.info(result['status'])
     return result['data']
 
diff --git a/src/asktable_mcp_server/sse_server.py b/src/asktable_mcp_server/sse_server.py
@@ -34,8 +34,7 @@ async def lifespan(fastmcp_instance):
     # 启动逻辑
     logger.info("服务器正在初始化...")
 
-    # 在这里可以进行一些初始化工作
-    await asyncio.sleep(3)  # 模拟初始化时间
+    await asyncio.sleep(3)
 
     server_ready = True
     logger.info("服务器初始化完成，准备接受请求")
@@ -65,14 +64,15 @@ async def health_check(request: Request):
 @mcp.custom_route("/sse/", methods=["GET"])
 async def sse_endpoint(request: Request):
     """自定义SSE端点，检查服务器是否准备就绪"""
-    global API_KEY, DATASOURCE_ID, server_ready
+    global API_KEY, DATASOURCE_ID, server_ready,ROLE_ID
 
     if not server_ready:
         return {"error": "Server is still initializing, please wait"}
 
     # 从URL参数获取配置
     API_KEY = request.query_params.get('apikey')
     DATASOURCE_ID = request.query_params.get('datasource_id')
+    ROLE_ID = request.query_params.get('role_id')
 
     if not API_KEY or not DATASOURCE_ID:
         logging.info("error: Missing required parameters: apikey and datasource_id")
@@ -102,7 +102,7 @@ async def gen_sql(query: str) -> str:
         - 需要将自然语言转化为SQL查询
         - 仅需要SQL文本而不需要执行结果
     """
-    global API_KEY, DATASOURCE_ID, server_ready
+    global API_KEY, DATASOURCE_ID, server_ready,ROLE_ID
 
     if not server_ready:
         return "Server is still initializing, please wait"
@@ -119,13 +119,25 @@ async def gen_sql(query: str) -> str:
         api_key = API_KEY
         datasource_id = DATASOURCE_ID
 
+    try:
+        request = get_http_request()
+        role_id = request.query_params.get('role_id', None)
+    except RuntimeError:
+        role_id = None
+
+    if not role_id:
+        role_id = ROLE_ID
+
     logging.info(f"api_key:{api_key}")
     logging.info(f"datasource_id:{datasource_id}")
+    logging.info(f"role_id:{role_id}")
 
     params = {
         'api_key': api_key,
         'datasource_id': datasource_id,
-        'question': query
+        'question': query,
+        'role_id':role_id
+
     }
     if args.base_url:
         params['base_url'] = args.base_url
@@ -159,7 +171,6 @@ async def gen_conclusion(query: str) -> str:
     if not server_ready:
         return "Server is still initializing, please wait"
 
-    # 其余代码保持不变...
     if not API_KEY or not DATASOURCE_ID:
         try:
             request = get_http_request()
@@ -172,16 +183,26 @@ async def gen_conclusion(query: str) -> str:
         api_key = API_KEY
         datasource_id = DATASOURCE_ID
 
+    try:
+        request = get_http_request()
+        role_id = request.query_params.get('role_id', None)
+    except RuntimeError:
+        role_id = None
+
+    if not role_id:
+        role_id = ROLE_ID
+
     logging.info(f"api_key:{api_key}")
     logging.info(f"datasource_id:{datasource_id}")
+    logging.info(f"role_id:{role_id}")
 
     params = {
-        'api_key': str(api_key),
-        'datasource_id': str(datasource_id),
-        'question': query
+        'api_key': api_key,
+        'datasource_id': datasource_id,
+        'question': query,
+        'role_id':role_id
+
     }
-    if args.base_url:
-        params['base_url'] = args.base_url
 
     message = await get_asktable_data(**params)
     return message
diff --git a/src/asktable_mcp_server/tools.py b/src/asktable_mcp_server/tools.py
@@ -1,69 +1,63 @@
 from asktable import Asktable
 import json
 import asyncio
+from utils import AskTableHelper
+
+async def get_asktable_data(api_key, datasource_id, question, base_url=None,role_id=None):
+    """
+    获取asktable数据
+    :param api_key:
+    :param datasource_id:
+    :param question:
+    :param base_url:
+    :param role_id:
+    :return:
+    """
 
-async def get_asktable_data(api_key, datasource_id, question, base_url=None):
     # 如果没有传入base_url或传入None，使用默认值
     if base_url is None:
         base_url = "https://api.asktable.com"
     
     asktable_client = Asktable(base_url=base_url, api_key=api_key)
-    answer_response = asktable_client.answers.create(datasource_id=datasource_id, question=question)
+    answer_response = asktable_client.answers.create(datasource_id=datasource_id, question=question,role_id=role_id)
     if answer_response.answer is None:
         return "没有查询到相关信息"
     return answer_response.answer.text
 
-async def get_asktable_sql(api_key, datasource_id, question, base_url=None):
+async def get_asktable_sql(api_key, datasource_id, question, base_url=None,role_id=None):
     # 如果没有传入base_url或传入None，使用默认值
     if base_url is None:
         base_url = "https://api.asktable.com"
     
     asktable_client = Asktable(base_url=base_url, api_key=api_key)
-    query_response = asktable_client.sqls.create(datasource_id=datasource_id, question=question)
+    query_response = asktable_client.sqls.create(datasource_id=datasource_id, question=question,role_id=role_id)
     if query_response.query.sql is None: 
         return "没有查询到相关信息"
     return query_response.query.sql
 
 
-async def get_datasources_info(api_key, base_url=None):
+async def get_datasources_info(api_key, base_url=None,role_id=None):
     """"
-    返回用户数据库meta data
+    返回用户数据库meta_data;
+    若输入roleid，则返回该角色可访问的数据库meta_data
+    若不输入roleid，则返回所有数据库meta_data
     args:
         api_key: str
         base_url: str
+        role_id:str
     return:
         {
             "status": "success" or "failure",
             "data": json.dumps(result, ensure_ascii=False, indent=2)
         }
     """
-    if base_url is None:
-        base_url = "https://api.asktable.com"
-
-    asktable_client = Asktable(base_url=base_url, api_key=api_key)
-    meta_data_list = asktable_client.datasources.list()
-    # 最前面判断
-    if not meta_data_list.items:
-        return {
-            "status": "failure",
-            "data": "该用户还没有创建任何数据库"
-        }
-
-    # 提取指定字段
-    result = [
-        {
-            "datasource_id": ds.id,
-            "数据库引擎": ds.engine,
-            "数据库描述": ds.desc,
-        }
-        for ds in meta_data_list.items
-    ]
+    helper = AskTableHelper(api_key=api_key,base_url=base_url)
+    if role_id is None:
+        return helper.get_datasources_info()
+    else:
+        return helper.get_datasources_info_by_role(role_id=role_id)
 
-    return {
-        "status": "success",
-        "data": json.dumps(result, ensure_ascii=False, indent=2)
-    }
 
 
 if __name__ == "__main__":
-    pass
+    pass
diff --git a/src/asktable_mcp_server/utils.py b/src/asktable_mcp_server/utils.py
@@ -0,0 +1,92 @@
+import json
+from asktable import Asktable
+
+class AskTableHelper:
+    def __init__(self, api_key, base_url=None):
+        if base_url is None:
+            base_url = "https://api.asktable.com"
+        self.asktable_client = Asktable(base_url=base_url, api_key=api_key)
+
+    def get_datasource_ids_by_role(self, role_id):
+        """
+        根据 role_id 获取该角色下所有 policy 涉及的 datasource_id 列表
+        :param role_id: 角色ID
+        :return: list, 包含所有 datasource_id
+        """
+        role_source = self.asktable_client.roles.get_polices(role_id=role_id)
+        datasource_ids = set()
+        for policy in role_source:
+            if hasattr(policy, 'dataset_config') and hasattr(policy.dataset_config, 'datasource_ids'):
+                datasource_ids.update(policy.dataset_config.datasource_ids)
+        return list(datasource_ids)
+
+    def get_role_name_by_role(self, role_id):
+        """
+        根据 role_id 获取角色名称
+        :param role_id: 角色ID
+        :return: 角色名称字符串
+        """
+        role_source = self.asktable_client.roles.retrieve(role_id=role_id)
+        return role_source.name
+
+    def get_datasources_info(self):
+        """
+        返回当前 API KEY 下所有数据库 meta data
+        :return: dict, {status, data}
+        """
+        meta_data_list = self.asktable_client.datasources.list()
+        if not meta_data_list.items:
+            return {
+                "status": "failure",
+                "data": "该用户还没有创建任何数据库"
+            }
+        result = [
+            {
+                "datasource_id": ds.id,
+                "数据库引擎": ds.engine,
+                "数据库描述": ds.desc,
+            }
+            for ds in meta_data_list.items
+        ]
+        return {
+            "status": "success",
+            "data": json.dumps(result, ensure_ascii=False, indent=2)
+        }
+
+    def get_datasources_info_by_role(self, role_id):
+        """
+        输入 role_id，返回该角色可访问的数据源的描述、引擎和id
+        :param role_id: 角色ID
+        :return: dict, {status, data}
+        """
+        # 1. 获取所有数据源
+        meta_data_list = self.asktable_client.datasources.list()
+        if not meta_data_list.items:
+            return {
+                "status": "failure",
+                "data": "该用户还没有创建任何数据库"
+            }
+
+        # 2. 获取该角色可访问的数据源ID
+        datasource_ids = set(self.get_datasource_ids_by_role(role_id))
+
+        # 3. 只保留属于该 role 的数据源
+        result = [
+            {
+                "datasource_id": ds.id,
+                "数据库引擎": ds.engine,
+                "数据库描述": ds.desc,
+            }
+            for ds in meta_data_list.items if ds.id in datasource_ids
+        ]
+
+        if not result:
+            return {
+                "status": "failure",
+                "data": "该角色没有可访问的数据库"
+            }
+
+        return {
+            "status": "success",
+            "data": json.dumps(result, ensure_ascii=False, indent=2)
+        }
