65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
|
"""NL2SQL Supervisor 的意图和会话并发控制。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from uuid import uuid4
|
||
|
|
|
||
|
|
|
||
|
|
class UnsupportedIntent(ValueError):
|
||
|
|
"""当前版本不支持的查询意图。"""
|
||
|
|
|
||
|
|
|
||
|
|
class SessionBusyError(RuntimeError):
|
||
|
|
"""同一会话已有查询运行。"""
|
||
|
|
|
||
|
|
|
||
|
|
def detect_intent(question: str) -> str:
|
||
|
|
"""根据有限关键词识别当前版本的任务类型。"""
|
||
|
|
text = question.lower()
|
||
|
|
if "etl" in text or "数据任务" in text:
|
||
|
|
return "etl_build"
|
||
|
|
if "口径" in text or "caliber" in text:
|
||
|
|
return "caliber_maintain"
|
||
|
|
if "血缘" in text or "lineage" in text:
|
||
|
|
return "lineage_query"
|
||
|
|
return "query"
|
||
|
|
|
||
|
|
|
||
|
|
def ensure_query_intent(question: str) -> None:
|
||
|
|
"""M1 只允许 query 意图进入 SQL 主链路。"""
|
||
|
|
intent = detect_intent(question)
|
||
|
|
if intent != "query":
|
||
|
|
raise UnsupportedIntent(f"当前不支持 {intent} 意图")
|
||
|
|
|
||
|
|
|
||
|
|
class SessionLock:
|
||
|
|
"""基于 Redis NX 的同会话互斥锁。"""
|
||
|
|
|
||
|
|
def __init__(self, redis, session_id: str, *, ttl: int = 120):
|
||
|
|
self.redis = redis
|
||
|
|
self.key = f"nl2sql:session:{session_id}:lock"
|
||
|
|
self.token = uuid4().hex
|
||
|
|
self.ttl = ttl
|
||
|
|
self.acquired = False
|
||
|
|
|
||
|
|
async def acquire(self) -> None:
|
||
|
|
if not await self.redis.set(self.key, self.token, nx=True, ex=self.ttl):
|
||
|
|
raise SessionBusyError("同一会话已有查询运行")
|
||
|
|
self.acquired = True
|
||
|
|
|
||
|
|
async def release(self) -> None:
|
||
|
|
if not self.acquired:
|
||
|
|
return
|
||
|
|
release_script = """
|
||
|
|
if redis.call('get', KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call('del', KEYS[1])
|
||
|
|
end
|
||
|
|
return 0
|
||
|
|
"""
|
||
|
|
if hasattr(self.redis, "eval"):
|
||
|
|
await self.redis.eval(release_script, 1, self.key, self.token)
|
||
|
|
else:
|
||
|
|
current = await self.redis.get(self.key)
|
||
|
|
if current in {self.token, self.token.encode()}:
|
||
|
|
await self.redis.delete(self.key)
|
||
|
|
self.acquired = False
|