187 lines
7.6 KiB
Python
187 lines
7.6 KiB
Python
"""内存版 DbPort 实现 + 受限 SQL 解释器（仅供离线单测，不进生产路径）。
|
||
|
||
支持 assoc.py 生成的全部 SQL 形态：
|
||
- SELECT <cols|*> FROM t [WHERE a=%s AND b=%s ...] ORDER BY col LIMIT n OFFSET m
|
||
- SELECT COUNT(*) AS cnt FROM t [WHERE ...]
|
||
- INSERT INTO t (c1, c2, ...) VALUES (%s, %s, ...)
|
||
- UPDATE t SET c=%s, ..., updated_at=NOW() WHERE a=%s AND b=%s
|
||
- DELETE FROM t WHERE a=%s AND b=%s
|
||
|
||
任何无法解析的语句直接抛错——测试里出现意外 SQL 形态必须暴露，不能静默通过。
|
||
同时充当「不改基表」铁律的守卫：对 world/scene/entity 的任何写操作立即 AssertionError。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||
|
||
from pbl_domain_ext.assoc import DbPort
|
||
|
||
_SELECT_RE = re.compile(
|
||
r"^SELECT\s+(?P<cols>.+?)\s+FROM\s+(?P<table>\w+)(?P<where>\s+WHERE\s+.+?)?"
|
||
r"(?:\s+ORDER\s+BY\s+(?P<order>\w+))?(?:\s+LIMIT\s+(?P<limit>\d+))?"
|
||
r"(?:\s+OFFSET\s+(?P<offset>\d+))?\s*$",
|
||
re.IGNORECASE | re.DOTALL,
|
||
)
|
||
_INSERT_RE = re.compile(
|
||
r"^INSERT\s+INTO\s+(?P<table>\w+)\s*\((?P<cols>[^)]*)\)\s*VALUES\s*\((?P<vals>[^)]*)\)\s*$",
|
||
re.IGNORECASE | re.DOTALL,
|
||
)
|
||
_UPDATE_RE = re.compile(
|
||
r"^UPDATE\s+(?P<table>\w+)\s+SET\s+(?P<assigns>.+?)\s+WHERE\s+(?P<where>.+?)\s*$",
|
||
re.IGNORECASE | re.DOTALL,
|
||
)
|
||
_DELETE_RE = re.compile(
|
||
r"^DELETE\s+FROM\s+(?P<table>\w+)\s+WHERE\s+(?P<where>.+?)\s*$",
|
||
re.IGNORECASE | re.DOTALL,
|
||
)
|
||
_COND_RE = re.compile(r"(\w+)\s*=\s*%s")
|
||
_BASE_TABLES = ("world", "scene", "entity")
|
||
|
||
|
||
def _split_top_level(text: str) -> List[str]:
|
||
"""按逗号切分（本模块 SQL 无嵌套函数参数含逗号，直接切即可）。"""
|
||
return [part.strip() for part in text.split(",") if part.strip()]
|
||
|
||
|
||
class FakeDb(DbPort):
|
||
"""表名 -> 行列表；行是 dict，含自增 id。"""
|
||
|
||
def __init__(self, tables: Optional[Sequence[str]] = None) -> None:
|
||
self.tables: Dict[str, List[Dict[str, Any]]] = {name: [] for name in (tables or [])}
|
||
self._seq: Dict[str, int] = {name: 0 for name in (tables or [])}
|
||
self.executed: List[Tuple[str, Tuple[Any, ...]]] = []
|
||
self.base_table_writes: List[str] = []
|
||
|
||
# ------------------------------------------------------------ 内部工具
|
||
def _ensure(self, table: str) -> List[Dict[str, Any]]:
|
||
if table not in self.tables:
|
||
self.tables[table] = []
|
||
self._seq[table] = 0
|
||
return self.tables[table]
|
||
|
||
def _guard_base(self, table: str, kind: str) -> None:
|
||
if table in _BASE_TABLES:
|
||
self.base_table_writes.append("%s:%s" % (kind, table))
|
||
raise AssertionError("铁律违规：薄扩展不得写基表 %s（%s）" % (table, kind))
|
||
|
||
@staticmethod
|
||
def _row_matches(row: Dict[str, Any], cols: Sequence[str], params: Sequence[Any]) -> bool:
|
||
for col, expected in zip(cols, params):
|
||
if str(row.get(col)) != str(expected):
|
||
return False
|
||
return True
|
||
|
||
def _where_rows(self, table: str, where_sql: str, params: Sequence[Any]) -> List[Dict[str, Any]]:
|
||
cols = _COND_RE.findall(where_sql or "")
|
||
return [row for row in self._ensure(table) if self._row_matches(row, cols, params)]
|
||
|
||
# ------------------------------------------------------------ DbPort
|
||
def select(self, sql: str, params: Sequence[Any] = ()) -> List[Dict[str, Any]]:
|
||
self.executed.append((sql, tuple(params)))
|
||
text = " ".join(sql.split())
|
||
m = _SELECT_RE.match(text)
|
||
if not m:
|
||
raise AssertionError("FakeDb 无法解析 SELECT: %s" % sql)
|
||
table = m.group("table")
|
||
cols = m.group("cols").strip()
|
||
where = m.group("where") or ""
|
||
rows = [dict(r) for r in self._where_rows(table, where, tuple(params))]
|
||
|
||
if cols.upper().startswith("COUNT("):
|
||
return [{"cnt": len(rows)}]
|
||
|
||
order = m.group("order")
|
||
if order:
|
||
rows.sort(key=lambda r: (r.get(order) is None, r.get(order)))
|
||
offset = int(m.group("offset") or 0)
|
||
limit = m.group("limit")
|
||
rows = rows[offset:offset + int(limit)] if limit is not None else rows[offset:]
|
||
|
||
if cols != "*":
|
||
wanted = _split_top_level(cols)
|
||
rows = [{c: r.get(c) for c in wanted} for r in rows]
|
||
return rows
|
||
|
||
def insert(self, sql: str, params: Sequence[Any] = ()) -> int:
|
||
self.executed.append((sql, tuple(params)))
|
||
text = " ".join(sql.split())
|
||
m = _INSERT_RE.match(text)
|
||
if not m:
|
||
raise AssertionError("FakeDb 无法解析 INSERT: %s" % sql)
|
||
table = m.group("table")
|
||
self._guard_base(table, "INSERT")
|
||
cols = _split_top_level(m.group("cols"))
|
||
placeholders = _split_top_level(m.group("vals"))
|
||
if len(cols) != len(placeholders) or len(cols) != len(params):
|
||
raise AssertionError(
|
||
"INSERT 列数与参数数不一致: cols=%d ph=%d params=%d"
|
||
% (len(cols), len(placeholders), len(params))
|
||
)
|
||
rows = self._ensure(table)
|
||
self._seq[table] = self._seq.get(table, 0) + 1
|
||
row: Dict[str, Any] = {"id": self._seq[table]}
|
||
for col, value in zip(cols, params):
|
||
row[col] = value
|
||
uniq_cols = [c for c in ("tenant_id", "world_id", "scene_id", "entity_id", "team_id", "role_code")
|
||
if c in row]
|
||
for exist in rows:
|
||
if self._row_matches(exist, uniq_cols, [row[c] for c in uniq_cols]):
|
||
raise AssertionError("唯一键冲突（幂等 upsert 失效）: %s %s" % (table, row))
|
||
rows.append(row)
|
||
return int(row["id"])
|
||
|
||
def execute(self, sql: str, params: Sequence[Any] = ()) -> int:
|
||
self.executed.append((sql, tuple(params)))
|
||
text = " ".join(sql.split())
|
||
|
||
m = _UPDATE_RE.match(text)
|
||
if m:
|
||
table = m.group("table")
|
||
self._guard_base(table, "UPDATE")
|
||
assigns = _split_top_level(m.group("assigns"))
|
||
set_cols: List[str] = []
|
||
set_values: List[Any] = []
|
||
cursor = 0
|
||
for assign in assigns:
|
||
col, _, expr = assign.partition("=")
|
||
col = col.strip()
|
||
expr = expr.strip()
|
||
if expr == "%s":
|
||
set_cols.append(col)
|
||
set_values.append(params[cursor])
|
||
cursor += 1
|
||
elif expr.upper() == "NOW()":
|
||
set_cols.append(col)
|
||
set_values.append("NOW()")
|
||
else:
|
||
raise AssertionError("FakeDb 不支持的 SET 表达式: %s" % assign)
|
||
where_cols = _COND_RE.findall(m.group("where"))
|
||
where_values = list(params[cursor:cursor + len(where_cols)])
|
||
affected = 0
|
||
for row in self._ensure(table):
|
||
if not self._row_matches(row, where_cols, where_values):
|
||
continue
|
||
for col, value in zip(set_cols, set_values):
|
||
row[col] = value
|
||
affected += 1
|
||
return affected
|
||
|
||
m = _DELETE_RE.match(text)
|
||
if m:
|
||
table = m.group("table")
|
||
self._guard_base(table, "DELETE")
|
||
where_cols = _COND_RE.findall(m.group("where"))
|
||
keep: List[Dict[str, Any]] = []
|
||
affected = 0
|
||
for row in self._ensure(table):
|
||
if self._row_matches(row, where_cols, tuple(params)):
|
||
affected += 1
|
||
else:
|
||
keep.append(row)
|
||
self.tables[table] = keep
|
||
return affected
|
||
|
||
raise AssertionError("FakeDb 无法解析语句: %s" % sql)
|