Admin
9b72e6127f
Some checks failed
continuous-integration/drone/push Build is failing
- 新增重试和失败告警功能（支持自动重试N次，失败后钉钉/企微通知） - 新增密钥管理（安全存储API Key等敏感信息） - 新增脚本模板库（预置常用脚本模板） - 新增脚本版本管理（自动保存历史版本，支持回滚） - 新增执行统计（成功率、平均耗时、7日趋势） - SDK 新增多租户遍历能力（get_tenants/get_tenant_config/get_all_tenant_configs） - SDK 新增密钥读取方法（get_secret）
612 lines
20 KiB
Python
612 lines
20 KiB
Python
"""脚本执行 SDK - 提供给 Python 脚本使用的内置能力"""
|
||
import json
|
||
import logging
|
||
from datetime import datetime
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import httpx
|
||
from sqlalchemy import text
|
||
from sqlalchemy.orm import Session
|
||
|
||
from ..database import SessionLocal
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ScriptSDK:
|
||
"""
|
||
脚本执行 SDK
|
||
|
||
提供以下能力：
|
||
- AI 大模型调用
|
||
- 钉钉/企微通知
|
||
- 数据库查询（只读）
|
||
- HTTP 请求
|
||
- 变量存储（跨执行持久化）
|
||
- 日志记录
|
||
- 多租户遍历
|
||
- 密钥管理
|
||
"""
|
||
|
||
def __init__(self, tenant_id: str, task_id: int, trace_id: str = None):
|
||
self.tenant_id = tenant_id
|
||
self.task_id = task_id
|
||
self.trace_id = trace_id or f"script_{task_id}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
|
||
self._logs: List[str] = []
|
||
self._db: Optional[Session] = None
|
||
self._tenants_cache: Optional[List[Dict]] = None # 租户列表缓存
|
||
|
||
def _get_db(self) -> Session:
|
||
"""获取数据库会话"""
|
||
if self._db is None:
|
||
self._db = SessionLocal()
|
||
return self._db
|
||
|
||
def _close_db(self):
|
||
"""关闭数据库会话"""
|
||
if self._db:
|
||
self._db.close()
|
||
self._db = None
|
||
|
||
# ============ AI 服务 ============
|
||
|
||
async def ai_chat(
|
||
self,
|
||
prompt: str,
|
||
system: str = None,
|
||
model: str = "gemini-2.5-flash",
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 2000
|
||
) -> str:
|
||
"""
|
||
调用大模型
|
||
|
||
Args:
|
||
prompt: 用户提示词
|
||
system: 系统提示词（可选）
|
||
model: 模型名称，默认 gemini-2.5-flash
|
||
temperature: 温度，默认 0.7
|
||
max_tokens: 最大 token 数，默认 2000
|
||
|
||
Returns:
|
||
AI 生成的文本
|
||
"""
|
||
# 使用 4sapi 作为 AI 服务
|
||
api_key = "sk-9yMCXjRGANbacz20kJY8doSNy6Rf446aYwmgGIuIXQ7DAyBw"
|
||
base_url = "https://4sapi.com/v1"
|
||
|
||
messages = []
|
||
if system:
|
||
messages.append({"role": "system", "content": system})
|
||
messages.append({"role": "user", "content": prompt})
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
response = await client.post(
|
||
f"{base_url}/chat/completions",
|
||
headers={
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json"
|
||
},
|
||
json={
|
||
"model": model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens
|
||
}
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
content = data["choices"][0]["message"]["content"]
|
||
self.log(f"AI 调用成功，模型: {model}，响应长度: {len(content)}")
|
||
return content
|
||
except Exception as e:
|
||
self.log(f"AI 调用失败: {str(e)}", level="ERROR")
|
||
raise
|
||
|
||
# ============ 通知服务 ============
|
||
|
||
async def send_dingtalk(
|
||
self,
|
||
webhook: str,
|
||
content: str,
|
||
msg_type: str = "text",
|
||
at_mobiles: List[str] = None,
|
||
at_all: bool = False
|
||
) -> bool:
|
||
"""
|
||
发送钉钉群消息
|
||
|
||
Args:
|
||
webhook: 钉钉机器人 Webhook URL
|
||
content: 消息内容
|
||
msg_type: 消息类型，text 或 markdown
|
||
at_mobiles: @的手机号列表
|
||
at_all: 是否@所有人
|
||
|
||
Returns:
|
||
是否发送成功
|
||
"""
|
||
try:
|
||
if msg_type == "text":
|
||
data = {
|
||
"msgtype": "text",
|
||
"text": {"content": content},
|
||
"at": {
|
||
"atMobiles": at_mobiles or [],
|
||
"isAtAll": at_all
|
||
}
|
||
}
|
||
else:
|
||
data = {
|
||
"msgtype": "markdown",
|
||
"markdown": {
|
||
"title": content[:20] if len(content) > 20 else content,
|
||
"text": content
|
||
},
|
||
"at": {
|
||
"atMobiles": at_mobiles or [],
|
||
"isAtAll": at_all
|
||
}
|
||
}
|
||
|
||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||
response = await client.post(webhook, json=data)
|
||
result = response.json()
|
||
|
||
if result.get("errcode") == 0:
|
||
self.log(f"钉钉消息发送成功")
|
||
return True
|
||
else:
|
||
self.log(f"钉钉消息发送失败: {result.get('errmsg')}", level="ERROR")
|
||
return False
|
||
except Exception as e:
|
||
self.log(f"钉钉消息发送异常: {str(e)}", level="ERROR")
|
||
return False
|
||
|
||
async def send_wecom(
|
||
self,
|
||
webhook: str,
|
||
content: str,
|
||
msg_type: str = "text"
|
||
) -> bool:
|
||
"""
|
||
发送企业微信群消息
|
||
|
||
Args:
|
||
webhook: 企微机器人 Webhook URL
|
||
content: 消息内容
|
||
msg_type: 消息类型，text 或 markdown
|
||
|
||
Returns:
|
||
是否发送成功
|
||
"""
|
||
try:
|
||
if msg_type == "text":
|
||
data = {
|
||
"msgtype": "text",
|
||
"text": {"content": content}
|
||
}
|
||
else:
|
||
data = {
|
||
"msgtype": "markdown",
|
||
"markdown": {"content": content}
|
||
}
|
||
|
||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||
response = await client.post(webhook, json=data)
|
||
result = response.json()
|
||
|
||
if result.get("errcode") == 0:
|
||
self.log(f"企微消息发送成功")
|
||
return True
|
||
else:
|
||
self.log(f"企微消息发送失败: {result.get('errmsg')}", level="ERROR")
|
||
return False
|
||
except Exception as e:
|
||
self.log(f"企微消息发送异常: {str(e)}", level="ERROR")
|
||
return False
|
||
|
||
# ============ 数据库查询 ============
|
||
|
||
def db_query(self, sql: str, params: Dict[str, Any] = None) -> List[Dict]:
|
||
"""
|
||
执行 SQL 查询（只读）
|
||
|
||
Args:
|
||
sql: SQL 语句（仅支持 SELECT）
|
||
params: 查询参数
|
||
|
||
Returns:
|
||
查询结果列表
|
||
"""
|
||
# 安全检查：只允许 SELECT
|
||
sql_upper = sql.strip().upper()
|
||
if not sql_upper.startswith("SELECT"):
|
||
raise ValueError("只允许 SELECT 查询")
|
||
|
||
forbidden = ["INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER", "TRUNCATE"]
|
||
for word in forbidden:
|
||
if word in sql_upper:
|
||
raise ValueError(f"禁止使用 {word} 语句")
|
||
|
||
db = self._get_db()
|
||
try:
|
||
result = db.execute(text(sql), params or {})
|
||
rows = [dict(row) for row in result.mappings().all()]
|
||
self.log(f"SQL 查询成功，返回 {len(rows)} 条记录")
|
||
return rows
|
||
except Exception as e:
|
||
self.log(f"SQL 查询失败: {str(e)}", level="ERROR")
|
||
raise
|
||
|
||
# ============ HTTP 请求 ============
|
||
|
||
async def http_get(
|
||
self,
|
||
url: str,
|
||
headers: Dict[str, str] = None,
|
||
params: Dict[str, Any] = None,
|
||
timeout: int = 30
|
||
) -> Dict:
|
||
"""
|
||
发送 HTTP GET 请求
|
||
|
||
Args:
|
||
url: 请求 URL
|
||
headers: 请求头
|
||
params: 查询参数
|
||
timeout: 超时时间（秒）
|
||
|
||
Returns:
|
||
响应数据（JSON 解析后）
|
||
"""
|
||
try:
|
||
async with httpx.AsyncClient(timeout=float(timeout)) as client:
|
||
response = await client.get(url, headers=headers, params=params)
|
||
self.log(f"HTTP GET {url} -> {response.status_code}")
|
||
|
||
try:
|
||
return response.json()
|
||
except:
|
||
return {"text": response.text, "status_code": response.status_code}
|
||
except Exception as e:
|
||
self.log(f"HTTP GET 失败: {str(e)}", level="ERROR")
|
||
raise
|
||
|
||
async def http_post(
|
||
self,
|
||
url: str,
|
||
data: Dict[str, Any] = None,
|
||
json_data: Dict[str, Any] = None,
|
||
headers: Dict[str, str] = None,
|
||
timeout: int = 30
|
||
) -> Dict:
|
||
"""
|
||
发送 HTTP POST 请求
|
||
|
||
Args:
|
||
url: 请求 URL
|
||
data: 表单数据
|
||
json_data: JSON 数据
|
||
headers: 请求头
|
||
timeout: 超时时间（秒）
|
||
|
||
Returns:
|
||
响应数据（JSON 解析后）
|
||
"""
|
||
try:
|
||
async with httpx.AsyncClient(timeout=float(timeout)) as client:
|
||
response = await client.post(
|
||
url,
|
||
data=data,
|
||
json=json_data,
|
||
headers=headers
|
||
)
|
||
self.log(f"HTTP POST {url} -> {response.status_code}")
|
||
|
||
try:
|
||
return response.json()
|
||
except:
|
||
return {"text": response.text, "status_code": response.status_code}
|
||
except Exception as e:
|
||
self.log(f"HTTP POST 失败: {str(e)}", level="ERROR")
|
||
raise
|
||
|
||
# ============ 变量存储 ============
|
||
|
||
def get_var(self, key: str, default: Any = None) -> Any:
|
||
"""
|
||
获取存储的变量
|
||
|
||
Args:
|
||
key: 变量名
|
||
default: 默认值
|
||
|
||
Returns:
|
||
变量值（JSON 解析后）
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
result = db.execute(
|
||
text("SELECT var_value FROM platform_script_vars WHERE tenant_id = :tid AND var_key = :key"),
|
||
{"tid": self.tenant_id, "key": key}
|
||
)
|
||
row = result.first()
|
||
if row:
|
||
try:
|
||
return json.loads(row[0])
|
||
except:
|
||
return row[0]
|
||
return default
|
||
except Exception as e:
|
||
self.log(f"获取变量失败: {str(e)}", level="ERROR")
|
||
return default
|
||
|
||
def set_var(self, key: str, value: Any) -> bool:
|
||
"""
|
||
存储变量（跨执行持久化）
|
||
|
||
Args:
|
||
key: 变量名
|
||
value: 变量值（会 JSON 序列化）
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
value_str = json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
|
||
|
||
# 使用 REPLACE INTO 实现 upsert
|
||
db.execute(
|
||
text("""
|
||
REPLACE INTO platform_script_vars (tenant_id, var_key, var_value, updated_at)
|
||
VALUES (:tid, :key, :value, NOW())
|
||
"""),
|
||
{"tid": self.tenant_id, "key": key, "value": value_str}
|
||
)
|
||
db.commit()
|
||
self.log(f"变量已存储: {key}")
|
||
return True
|
||
except Exception as e:
|
||
self.log(f"存储变量失败: {str(e)}", level="ERROR")
|
||
return False
|
||
|
||
def delete_var(self, key: str) -> bool:
|
||
"""
|
||
删除变量
|
||
|
||
Args:
|
||
key: 变量名
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
db.execute(
|
||
text("DELETE FROM platform_script_vars WHERE tenant_id = :tid AND var_key = :key"),
|
||
{"tid": self.tenant_id, "key": key}
|
||
)
|
||
db.commit()
|
||
self.log(f"变量已删除: {key}")
|
||
return True
|
||
except Exception as e:
|
||
self.log(f"删除变量失败: {str(e)}", level="ERROR")
|
||
return False
|
||
|
||
# ============ 多租户遍历 ============
|
||
|
||
def get_tenants(self, app_code: str = None) -> List[Dict]:
|
||
"""
|
||
获取租户列表（用于多租户任务遍历）
|
||
|
||
Args:
|
||
app_code: 应用代码（可选），筛选订阅了该应用的租户
|
||
|
||
Returns:
|
||
租户列表 [{"code": "xxx", "name": "租户名", "custom_configs": {...}}]
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
if app_code:
|
||
# 筛选订阅了指定应用的租户
|
||
result = db.execute(
|
||
text("""
|
||
SELECT DISTINCT t.code, t.name, ta.custom_configs
|
||
FROM platform_tenants t
|
||
INNER JOIN platform_tenant_apps ta ON t.code = ta.tenant_id
|
||
WHERE ta.app_code = :app_code AND t.status = 1
|
||
"""),
|
||
{"app_code": app_code}
|
||
)
|
||
else:
|
||
# 获取所有启用的租户
|
||
result = db.execute(
|
||
text("SELECT code, name FROM platform_tenants WHERE status = 1")
|
||
)
|
||
|
||
tenants = []
|
||
for row in result.mappings().all():
|
||
tenant = dict(row)
|
||
# 解析 custom_configs
|
||
if "custom_configs" in tenant and tenant["custom_configs"]:
|
||
try:
|
||
tenant["custom_configs"] = json.loads(tenant["custom_configs"])
|
||
except:
|
||
pass
|
||
tenants.append(tenant)
|
||
|
||
self.log(f"获取租户列表成功，共 {len(tenants)} 个租户")
|
||
return tenants
|
||
except Exception as e:
|
||
self.log(f"获取租户列表失败: {str(e)}", level="ERROR")
|
||
return []
|
||
|
||
def get_tenant_config(self, tenant_id: str, app_code: str, key: str = None) -> Any:
|
||
"""
|
||
获取指定租户的应用配置
|
||
|
||
Args:
|
||
tenant_id: 租户ID
|
||
app_code: 应用代码
|
||
key: 配置项键名（可选，不传返回全部配置）
|
||
|
||
Returns:
|
||
配置值或配置字典
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
result = db.execute(
|
||
text("""
|
||
SELECT custom_configs FROM platform_tenant_apps
|
||
WHERE tenant_id = :tenant_id AND app_code = :app_code
|
||
"""),
|
||
{"tenant_id": tenant_id, "app_code": app_code}
|
||
)
|
||
row = result.first()
|
||
|
||
if not row or not row[0]:
|
||
return None if key else {}
|
||
|
||
try:
|
||
configs = json.loads(row[0])
|
||
except:
|
||
configs = {}
|
||
|
||
if key:
|
||
return configs.get(key)
|
||
return configs
|
||
except Exception as e:
|
||
self.log(f"获取租户配置失败: {str(e)}", level="ERROR")
|
||
return None if key else {}
|
||
|
||
def get_all_tenant_configs(self, app_code: str) -> List[Dict]:
|
||
"""
|
||
获取所有租户的应用配置（便捷方法，用于批量操作）
|
||
|
||
Args:
|
||
app_code: 应用代码
|
||
|
||
Returns:
|
||
[{"tenant_id": "xxx", "tenant_name": "租户名", "configs": {...}}]
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
result = db.execute(
|
||
text("""
|
||
SELECT t.code as tenant_id, t.name as tenant_name, ta.custom_configs
|
||
FROM platform_tenants t
|
||
INNER JOIN platform_tenant_apps ta ON t.code = ta.tenant_id
|
||
WHERE ta.app_code = :app_code AND t.status = 1
|
||
"""),
|
||
{"app_code": app_code}
|
||
)
|
||
|
||
tenants = []
|
||
for row in result.mappings().all():
|
||
configs = {}
|
||
if row["custom_configs"]:
|
||
try:
|
||
configs = json.loads(row["custom_configs"])
|
||
except:
|
||
pass
|
||
tenants.append({
|
||
"tenant_id": row["tenant_id"],
|
||
"tenant_name": row["tenant_name"],
|
||
"configs": configs
|
||
})
|
||
|
||
self.log(f"获取 {app_code} 应用的租户配置，共 {len(tenants)} 个")
|
||
return tenants
|
||
except Exception as e:
|
||
self.log(f"获取租户配置失败: {str(e)}", level="ERROR")
|
||
return []
|
||
|
||
# ============ 密钥管理 ============
|
||
|
||
def get_secret(self, key: str) -> Optional[str]:
|
||
"""
|
||
获取密钥（优先读取租户级密钥，其次读取全局密钥）
|
||
|
||
Args:
|
||
key: 密钥名称
|
||
|
||
Returns:
|
||
密钥值（如不存在返回 None）
|
||
"""
|
||
db = self._get_db()
|
||
try:
|
||
# 优先查询租户级密钥
|
||
result = db.execute(
|
||
text("""
|
||
SELECT secret_value FROM platform_secrets
|
||
WHERE (tenant_id = :tenant_id OR tenant_id IS NULL)
|
||
AND secret_key = :key
|
||
ORDER BY tenant_id DESC
|
||
LIMIT 1
|
||
"""),
|
||
{"tenant_id": self.tenant_id, "key": key}
|
||
)
|
||
row = result.first()
|
||
|
||
if row:
|
||
self.log(f"获取密钥成功: {key}")
|
||
return row[0]
|
||
|
||
self.log(f"密钥不存在: {key}", level="WARN")
|
||
return None
|
||
except Exception as e:
|
||
self.log(f"获取密钥失败: {str(e)}", level="ERROR")
|
||
return None
|
||
|
||
# ============ 日志 ============
|
||
|
||
def log(self, message: str, level: str = "INFO"):
|
||
"""
|
||
记录日志
|
||
|
||
Args:
|
||
message: 日志内容
|
||
level: 日志级别（INFO, WARN, ERROR）
|
||
"""
|
||
timestamp = datetime.now().strftime("%H:%M:%S")
|
||
log_entry = f"[{timestamp}] [{level}] {message}"
|
||
self._logs.append(log_entry)
|
||
|
||
# 同时输出到标准日志
|
||
if level == "ERROR":
|
||
logger.error(f"[Script {self.task_id}] {message}")
|
||
else:
|
||
logger.info(f"[Script {self.task_id}] {message}")
|
||
|
||
# 写入 platform_logs
|
||
try:
|
||
db = self._get_db()
|
||
db.execute(
|
||
text("""
|
||
INSERT INTO platform_logs
|
||
(trace_id, app_code, module, level, message, created_at)
|
||
VALUES (:trace_id, :app_code, :module, :level, :message, NOW())
|
||
"""),
|
||
{
|
||
"trace_id": self.trace_id,
|
||
"app_code": "000-platform",
|
||
"module": "script",
|
||
"level": level,
|
||
"message": message[:2000] # 限制长度
|
||
}
|
||
)
|
||
db.commit()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to write log to database: {e}")
|
||
|
||
def get_logs(self) -> List[str]:
|
||
"""获取所有日志"""
|
||
return self._logs.copy()
|
||
|
||
def cleanup(self):
|
||
"""清理资源"""
|
||
self._close_db()
|