Skip to content

Commit 4dbcc7a

Browse files
committed
🐛 fix(database): 修复元数据缓存与查询安全回归
修复元数据缓存缺失、代码生成路径重复斜杠,以及只读 SQL 对 PostgreSQL dollar-quoted literal 和 FETCH 限制的误判。\n\n通过连接级有界缓存、路径规范化和 tokenizer/limit 回归测试恢复既有契约。\n\n验证:PYTHONPATH=src pytest -q tests/unit/test_connection_manager.py tests/unit/test_database_introspection.py tests/unit/test_introspection_benchmark.py tests/unit/test_atomic_codegen_tools.py tests/unit/test_mcp_query_safety.py(88 passed);ruff check 通过。\n\nCloses #214
1 parent 46730ff commit 4dbcc7a

3 files changed

Lines changed: 91 additions & 1 deletion

File tree

‎src/dbjavagenix/database/atomic_codegen_tools.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -529,7 +529,7 @@ def _compute_file_path(
529529
package_name = context.get("package", "com.example")
530530
package_path = package_name.replace(".", "/")
531531
if file_path.endswith(".java"):
532-
return f"{package_path}/{file_path}"
532+
return "/".join(part for part in f"{package_path}/{file_path}".split("/") if part)
533533
return f"resources/{file_path}"
534534

535535

‎src/dbjavagenix/database/connection_manager.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@
33
"""
44
import uuid
55
import threading
6+
from collections import OrderedDict
7+
from copy import deepcopy
68
from typing import Dict, List, Any, Optional
79
import pymysql
810
import sqlite3
911
import logging
12+
import re
1013
from contextlib import contextmanager
1114

1215
from ..core.models import DatabaseConfig, DatabaseType
@@ -25,6 +28,41 @@ def __init__(self):
2528
self.connection_configs: Dict[str, DatabaseConfig] = {}
2629
self._registry_lock = threading.RLock()
2730
self._connection_locks: Dict[str, Any] = {}
31+
self._metadata_cache: OrderedDict[tuple[str, str, str | None], Dict[str, Any]] = OrderedDict()
32+
self._metadata_cache_limit = 256
33+
34+
def cache_metadata(
35+
self, connection_id: str, table_name: str, schema: str | None, metadata: Dict[str, Any]
36+
) -> None:
37+
"""Store a defensive copy of table metadata in the bounded LRU cache."""
38+
key = (connection_id, table_name, schema)
39+
with self._registry_lock:
40+
self._metadata_cache[key] = deepcopy(metadata)
41+
self._metadata_cache.move_to_end(key)
42+
while len(self._metadata_cache) > self._metadata_cache_limit:
43+
self._metadata_cache.popitem(last=False)
44+
45+
def get_cached_metadata(
46+
self, connection_id: str, table_name: str, schema: str | None = None
47+
) -> Dict[str, Any] | None:
48+
"""Return a defensive copy of cached metadata, if present."""
49+
key = (connection_id, table_name, schema)
50+
with self._registry_lock:
51+
metadata = self._metadata_cache.get(key)
52+
if metadata is None:
53+
return None
54+
self._metadata_cache.move_to_end(key)
55+
return deepcopy(metadata)
56+
57+
def metadata_cache_size(self) -> int:
58+
with self._registry_lock:
59+
return len(self._metadata_cache)
60+
61+
def invalidate_metadata_cache(self, connection_id: str) -> None:
62+
"""Invalidate all table metadata associated with one connection."""
63+
with self._registry_lock:
64+
for key in [key for key in self._metadata_cache if key[0] == connection_id]:
65+
self._metadata_cache.pop(key, None)
2866

2967
def create_connection(self, config: DatabaseConfig) -> str:
3068
"""
@@ -152,6 +190,7 @@ def _remove_connection(self, connection_id: str, connection: Any) -> None:
152190
self.connections.pop(connection_id, None)
153191
self.connection_configs.pop(connection_id, None)
154192
self._connection_locks.pop(connection_id, None)
193+
self.invalidate_metadata_cache(connection_id)
155194
logger.info("Closed connection %s", connection_id)
156195

157196
def close_connection(self, connection_id: str) -> bool:
@@ -294,6 +333,14 @@ def execute_query(self, connection_id: str, query: str, params: Optional[tuple]
294333

295334
if isinstance(connection, sqlite3.Connection):
296335
connection.commit()
336+
statement = re.sub(
337+
r"^\s*(?:(?:/\*.*?\*/)|(?:--[^\n]*(?:\n|$)))\s*",
338+
"",
339+
query,
340+
flags=re.DOTALL,
341+
).upper()
342+
if re.match(r"(?:ALTER|CREATE|DROP|RENAME|TRUNCATE)\b", statement):
343+
self.invalidate_metadata_cache(connection_id)
297344
return result
298345

299346
except Exception as e:

‎src/dbjavagenix/database/mcp_tools.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,19 @@ def _tokenize_read_only_sql(query: str) -> List[tuple[str, str]]:
179179
tokens.append(("quoted", query[start:index]))
180180
continue
181181

182+
if char == "$":
183+
dollar_match = re.match(r"\$([A-Za-z_][A-Za-z0-9_]*)?\$", query[index:])
184+
if dollar_match:
185+
delimiter = dollar_match.group(0)
186+
start = index
187+
content_start = index + len(delimiter)
188+
end = query.find(delimiter, content_start)
189+
if end < 0:
190+
raise MCPServiceError("Unterminated dollar-quoted literal")
191+
index = end + len(delimiter)
192+
tokens.append(("quoted", query[start:index]))
193+
continue
194+
182195
if char.isalpha() or char == "_":
183196
start = index
184197
index += 1
@@ -236,6 +249,13 @@ def _validate_read_only_query(query: Any) -> str:
236249
if _contains_locking_read_clause(tokens):
237250
raise MCPServiceError("Only non-locking read-only SELECT queries are allowed")
238251

252+
for index, token in enumerate(tokens):
253+
if token != ("word", "FETCH"):
254+
continue
255+
fetch_tail = [value for kind, value in tokens[index + 1 : index + 10] if kind == "word"]
256+
if "WITH" in fetch_tail and "TIES" in fetch_tail[fetch_tail.index("WITH") + 1 :]:
257+
raise MCPServiceError("FETCH WITH TIES is not supported for bounded read-only queries")
258+
239259
depth = 0
240260
top_level_select = first_word == "SELECT"
241261
for kind, value in tokens:
@@ -276,6 +296,11 @@ def _has_top_level_limit_clause(query: str) -> bool:
276296
next_token = tokens[index + 1] if index + 1 < len(tokens) else None
277297
if next_token and (next_token[0] == "symbol" or next_token[1] == "ALL"):
278298
return True
299+
elif kind == "word" and value == "FETCH" and depth == 0:
300+
words = [token[1] for token in tokens[index : index + 5] if token[0] == "word"]
301+
if len(words) >= 4 and words[:2] == ["FETCH", words[1]] and words[1] in {"FIRST", "NEXT"}:
302+
if words[-1] == "ONLY" or (len(words) >= 5 and words[-1] == "ROWS"):
303+
return True
279304
return False
280305

281306

@@ -289,6 +314,17 @@ def _top_level_limit_span(query: str) -> tuple[int, int, int | None] | None:
289314
masked = list(query)
290315
index = 0
291316
while index < len(masked):
317+
if masked[index] == "$":
318+
dollar_match = re.match(r"\$([A-Za-z_][A-Za-z0-9_]*)?\$", query[index:])
319+
if dollar_match:
320+
delimiter = dollar_match.group(0)
321+
end = query.find(delimiter, index + len(delimiter))
322+
if end < 0:
323+
raise MCPServiceError("Unterminated dollar-quoted literal")
324+
for position in range(index, end + len(delimiter)):
325+
masked[position] = " "
326+
index = end + len(delimiter)
327+
continue
292328
if masked[index] not in "'\"`":
293329
index += 1
294330
continue
@@ -330,6 +366,13 @@ def is_top_level(position: int) -> bool:
330366
if is_top_level(match.start()):
331367
value = match.group(1).upper()
332368
return match.start(1), match.end(1), None if value == "ALL" else int(value)
369+
370+
fetch_form = re.compile(
371+
r"\bFETCH\s+(FIRST|NEXT)\s+(\d+)\s+ROWS?\s+ONLY\b", re.IGNORECASE
372+
)
373+
for match in fetch_form.finditer(masked_query):
374+
if is_top_level(match.start()):
375+
return match.start(2), match.end(2), int(match.group(2))
333376
return None
334377

335378

0 commit comments

Comments
 (0)