diff --git a/src/dbjavagenix/database/atomic_codegen_tools.py b/src/dbjavagenix/database/atomic_codegen_tools.py index 17615a0..77325c1 100644 --- a/src/dbjavagenix/database/atomic_codegen_tools.py +++ b/src/dbjavagenix/database/atomic_codegen_tools.py @@ -529,7 +529,7 @@ def _compute_file_path( package_name = context.get("package", "com.example") package_path = package_name.replace(".", "/") if file_path.endswith(".java"): - return f"{package_path}/{file_path}" + return "/".join(part for part in f"{package_path}/{file_path}".split("/") if part) return f"resources/{file_path}" diff --git a/src/dbjavagenix/database/connection_manager.py b/src/dbjavagenix/database/connection_manager.py index b1319a0..184449e 100644 --- a/src/dbjavagenix/database/connection_manager.py +++ b/src/dbjavagenix/database/connection_manager.py @@ -3,10 +3,13 @@ """ import uuid import threading +from collections import OrderedDict +from copy import deepcopy from typing import Dict, List, Any, Optional import pymysql import sqlite3 import logging +import re from contextlib import contextmanager from ..core.models import DatabaseConfig, DatabaseType @@ -25,6 +28,42 @@ def __init__(self): self.connection_configs: Dict[str, DatabaseConfig] = {} self._registry_lock = threading.RLock() self._connection_locks: Dict[str, Any] = {} + # Metadata is scoped by connection and schema so DDL cannot leak across sessions. + self._metadata_cache: OrderedDict[tuple[str, str, str | None], Dict[str, Any]] = OrderedDict() + self._metadata_cache_limit = 256 + + def cache_metadata( + self, connection_id: str, table_name: str, schema: str | None, metadata: Dict[str, Any] + ) -> None: + """Store a defensive copy of table metadata in the bounded LRU cache.""" + key = (connection_id, table_name, schema) + with self._registry_lock: + self._metadata_cache[key] = deepcopy(metadata) + self._metadata_cache.move_to_end(key) + while len(self._metadata_cache) > self._metadata_cache_limit: + self._metadata_cache.popitem(last=False) + + def get_cached_metadata( + self, connection_id: str, table_name: str, schema: str | None = None + ) -> Dict[str, Any] | None: + """Return a defensive copy of cached metadata, if present.""" + key = (connection_id, table_name, schema) + with self._registry_lock: + metadata = self._metadata_cache.get(key) + if metadata is None: + return None + self._metadata_cache.move_to_end(key) + return deepcopy(metadata) + + def metadata_cache_size(self) -> int: + with self._registry_lock: + return len(self._metadata_cache) + + def invalidate_metadata_cache(self, connection_id: str) -> None: + """Invalidate all table metadata associated with one connection.""" + with self._registry_lock: + for key in [key for key in self._metadata_cache if key[0] == connection_id]: + self._metadata_cache.pop(key, None) def create_connection(self, config: DatabaseConfig) -> str: """ @@ -152,6 +191,7 @@ def _remove_connection(self, connection_id: str, connection: Any) -> None: self.connections.pop(connection_id, None) self.connection_configs.pop(connection_id, None) self._connection_locks.pop(connection_id, None) + self.invalidate_metadata_cache(connection_id) logger.info("Closed connection %s", connection_id) def close_connection(self, connection_id: str) -> bool: @@ -294,6 +334,14 @@ def execute_query(self, connection_id: str, query: str, params: Optional[tuple] if isinstance(connection, sqlite3.Connection): connection.commit() + statement = re.sub( + r"^\s*(?:(?:/\*.*?\*/)|(?:--[^\n]*(?:\n|$)))\s*", + "", + query, + flags=re.DOTALL, + ).upper() + if re.match(r"(?:ALTER|CREATE|DROP|RENAME|TRUNCATE)\b", statement): + self.invalidate_metadata_cache(connection_id) return result except Exception as e: diff --git a/src/dbjavagenix/database/mcp_tools.py b/src/dbjavagenix/database/mcp_tools.py index a4a229b..a575f29 100644 --- a/src/dbjavagenix/database/mcp_tools.py +++ b/src/dbjavagenix/database/mcp_tools.py @@ -179,6 +179,19 @@ def _tokenize_read_only_sql(query: str) -> List[tuple[str, str]]: tokens.append(("quoted", query[start:index])) continue + if char == "$": + dollar_match = re.match(r"\$([A-Za-z_][A-Za-z0-9_]*)?\$", query[index:]) + if dollar_match: + delimiter = dollar_match.group(0) + start = index + content_start = index + len(delimiter) + end = query.find(delimiter, content_start) + if end < 0: + raise MCPServiceError("Unterminated dollar-quoted literal") + index = end + len(delimiter) + tokens.append(("quoted", query[start:index])) + continue + if char.isalpha() or char == "_": start = index index += 1 @@ -236,6 +249,13 @@ def _validate_read_only_query(query: Any) -> str: if _contains_locking_read_clause(tokens): raise MCPServiceError("Only non-locking read-only SELECT queries are allowed") + for index, token in enumerate(tokens): + if token != ("word", "FETCH"): + continue + fetch_tail = [value for kind, value in tokens[index + 1 : index + 10] if kind == "word"] + if "WITH" in fetch_tail and "TIES" in fetch_tail[fetch_tail.index("WITH") + 1 :]: + raise MCPServiceError("FETCH WITH TIES is not supported for bounded read-only queries") + depth = 0 top_level_select = first_word == "SELECT" for kind, value in tokens: @@ -276,6 +296,11 @@ def _has_top_level_limit_clause(query: str) -> bool: next_token = tokens[index + 1] if index + 1 < len(tokens) else None if next_token and (next_token[0] == "symbol" or next_token[1] == "ALL"): return True + elif kind == "word" and value == "FETCH" and depth == 0: + words = [token[1] for token in tokens[index : index + 5] if token[0] == "word"] + if len(words) >= 4 and words[:2] == ["FETCH", words[1]] and words[1] in {"FIRST", "NEXT"}: + if words[-1] == "ONLY" or (len(words) >= 5 and words[-1] == "ROWS"): + return True return False @@ -289,6 +314,17 @@ def _top_level_limit_span(query: str) -> tuple[int, int, int | None] | None: masked = list(query) index = 0 while index < len(masked): + if masked[index] == "$": + dollar_match = re.match(r"\$([A-Za-z_][A-Za-z0-9_]*)?\$", query[index:]) + if dollar_match: + delimiter = dollar_match.group(0) + end = query.find(delimiter, index + len(delimiter)) + if end < 0: + raise MCPServiceError("Unterminated dollar-quoted literal") + for position in range(index, end + len(delimiter)): + masked[position] = " " + index = end + len(delimiter) + continue if masked[index] not in "'\"`": index += 1 continue @@ -330,6 +366,13 @@ def is_top_level(position: int) -> bool: if is_top_level(match.start()): value = match.group(1).upper() return match.start(1), match.end(1), None if value == "ALL" else int(value) + + fetch_form = re.compile( + r"\bFETCH\s+(FIRST|NEXT)\s+(\d+)\s+ROWS?\s+ONLY\b", re.IGNORECASE + ) + for match in fetch_form.finditer(masked_query): + if is_top_level(match.start()): + return match.start(2), match.end(2), int(match.group(2)) return None