diff --git a/src/dbjavagenix/database/atomic_codegen_tools.py b/src/dbjavagenix/database/atomic_codegen_tools.py index 4c28cdd..2b55c2b 100644 --- a/src/dbjavagenix/database/atomic_codegen_tools.py +++ b/src/dbjavagenix/database/atomic_codegen_tools.py @@ -23,6 +23,7 @@ """ import json +import asyncio import logging from typing import Any, Dict, List @@ -34,6 +35,20 @@ logger = logging.getLogger(__name__) +async def _run_db_call(callable_obj, *args): + """Run blocking database work outside the MCP event loop.""" + return await asyncio.to_thread(callable_obj, *args) + + +async def _run_async_db_call(callable_obj, *args, **kwargs): + """Run the async analyzer in a worker because its introspection is synchronous.""" + + def run(): + return asyncio.run(callable_obj(*args, **kwargs)) + + return await asyncio.to_thread(run) + + # ============================================================ # Tool 定义 - 7 个原子工具 # ============================================================ @@ -208,10 +223,11 @@ async def handle_codegen_build_context(arguments: Dict[str, Any]) -> List[TextCo database = config.database or "information_schema" # 收集所有表名用于前缀分析(沿用旧逻辑) - all_table_names = _collect_all_table_names(connection_id, config) + all_table_names = await _run_db_call(_collect_all_table_names, connection_id, config) analyzer = CodegenAnalyzer(connection_manager) - analysis = await analyzer.analyze_table_for_codegen( + analysis = await _run_async_db_call( + analyzer.analyze_table_for_codegen, connection_id, table_name, all_table_names=all_table_names, @@ -508,9 +524,8 @@ def _compute_file_path( def _collect_all_table_names(connection_id: str, config) -> List[str]: """收集库内所有表名(用于前缀分析)。失败时返回空列表。""" try: - conn = connection_manager.get_connection(connection_id) - cursor = conn.cursor() - try: + + def collect(cursor) -> List[str]: if config.type.name == "MYSQL": cursor.execute("SHOW TABLES") return [row[0] for row in cursor.fetchall()] @@ -529,6 +544,16 @@ def _collect_all_table_names(connection_id: str, config) -> List[str]: ) return [row[0] for row in cursor.fetchall()] return [] + + if hasattr(connection_manager, "get_cursor"): + with connection_manager.get_cursor(connection_id) as cursor: + return collect(cursor) + + # Preserve the small test-double contract used by older integrations. + connection = connection_manager.get_connection(connection_id) + cursor = connection.cursor() + try: + return collect(cursor) finally: cursor.close() except Exception as e: # noqa: BLE001 diff --git a/src/dbjavagenix/database/connection_manager.py b/src/dbjavagenix/database/connection_manager.py index a218492..73ccf6f 100644 --- a/src/dbjavagenix/database/connection_manager.py +++ b/src/dbjavagenix/database/connection_manager.py @@ -6,6 +6,7 @@ import re from threading import RLock import uuid +import threading from typing import Dict, List, Any, Optional import pymysql import sqlite3 @@ -52,6 +53,8 @@ class ConnectionManager: def __init__(self): self.connections: Dict[str, Any] = {} self.connection_configs: Dict[str, DatabaseConfig] = {} + self._registry_lock = threading.RLock() + self._connection_locks: Dict[str, Any] = {} self._metadata_cache: OrderedDict[ tuple[str, str, Optional[str]], Dict[str, Any] ] = OrderedDict() @@ -152,16 +155,21 @@ def create_connection(self, config: DatabaseConfig) -> str: ) connection.autocommit = True elif config.type == DatabaseType.SQLITE: - connection = sqlite3.connect(config.database) + # MCP database work runs in worker threads. SQLite's default + # thread affinity would reject a connection created elsewhere. + connection = sqlite3.connect(config.database, check_same_thread=False) connection.row_factory = sqlite3.Row # Enable dict-like access else: raise DatabaseConnectionError(f"Unsupported database type: {config.type}") - self.connections[connection_id] = connection - # Store config without sensitive data for reference - safe_config = config.model_copy() - safe_config.password = "***" # Mask password - self.connection_configs[connection_id] = safe_config + # Store the connection and its lock atomically with the registry. + with self._registry_lock: + self.connections[connection_id] = connection + self._connection_locks[connection_id] = threading.RLock() + # Store config without sensitive data for reference + safe_config = config.model_copy() + safe_config.password = "***" # Mask password + self.connection_configs[connection_id] = safe_config logger.info(f"Created connection {connection_id} to {config.type}://{config.host}:{config.port}") return connection_id @@ -184,23 +192,54 @@ def get_connection(self, connection_id: str) -> Any: Raises: DatabaseConnectionError: If connection not found """ - if connection_id not in self.connections: + with self._connection_lock(connection_id): + connection = self._get_connection_unlocked(connection_id) + try: + if getattr(connection, "closed", 0): + raise DatabaseConnectionError("connection is closed") + if hasattr(connection, "ping"): + connection.ping(reconnect=True) + except Exception as exc: + logger.warning("Connection %s is dead, removing: %s", connection_id, exc) + self._close_connection_unlocked(connection_id, connection) + raise DatabaseConnectionError( + f"Connection {connection_id} is no longer valid" + ) from exc + return connection + + @contextmanager + def _connection_lock(self, connection_id: str): + """Serialize operations for one connection without blocking the registry.""" + with self._registry_lock: + lock = self._connection_locks.get(connection_id) + if lock is None or connection_id not in self.connections: + raise DatabaseConnectionError(f"Connection {connection_id} not found") + with lock: + yield + + def _get_connection_unlocked(self, connection_id: str) -> Any: + with self._registry_lock: + connection = self.connections.get(connection_id) + if connection is None: raise DatabaseConnectionError(f"Connection {connection_id} not found") - - connection = self.connections[connection_id] - - # Test connection is still alive - try: - if getattr(connection, "closed", 0): - raise DatabaseConnectionError("connection is closed") - if hasattr(connection, 'ping'): - connection.ping(reconnect=True) - except Exception as e: - logger.warning(f"Connection {connection_id} is dead, removing: {e}") - self.close_connection(connection_id) - raise DatabaseConnectionError(f"Connection {connection_id} is no longer valid") - return connection + + def _close_connection_unlocked(self, connection_id: str, connection: Any) -> bool: + """Close and remove a connection while its per-connection lock is held.""" + try: + connection.close() + closed = True + except Exception as exc: + logger.error("Error closing connection %s: %s", connection_id, exc) + closed = True + finally: + with self._registry_lock: + 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) + return closed def close_connection(self, connection_id: str) -> bool: """ @@ -212,25 +251,19 @@ def close_connection(self, connection_id: str) -> bool: Returns: True if connection was closed, False if not found """ - if connection_id not in self.connections: - self.invalidate_metadata_cache(connection_id) - return False - - try: - connection = self.connections[connection_id] - connection.close() - self.connections.pop(connection_id, None) - self.connection_configs.pop(connection_id, None) - self.invalidate_metadata_cache(connection_id) - logger.info(f"Closed connection {connection_id}") - return True - except Exception as e: - logger.error(f"Error closing connection {connection_id}: {e}") - # Remove from dict anyway - self.connections.pop(connection_id, None) - self.connection_configs.pop(connection_id, None) - self.invalidate_metadata_cache(connection_id) - return True + with self._registry_lock: + if connection_id not in self.connections: + self.invalidate_metadata_cache(connection_id) + return False + lock = self._connection_locks.setdefault(connection_id, threading.RLock()) + with lock: + # A preceding close may have removed the connection while this + # caller was waiting for the per-connection lock. + with self._registry_lock: + current = self.connections.get(connection_id) + if current is None: + return False + return self._close_connection_unlocked(connection_id, current) def get_connection_info(self, connection_id: str) -> Optional[DatabaseConfig]: """ @@ -242,7 +275,8 @@ def get_connection_info(self, connection_id: str) -> Optional[DatabaseConfig]: Returns: Database configuration or None if not found """ - config = self.connection_configs.get(connection_id) + with self._registry_lock: + config = self.connection_configs.get(connection_id) return config.model_copy(deep=True) if config else None def list_connections(self) -> Dict[str, Dict[str, Any]]: @@ -252,15 +286,18 @@ def list_connections(self) -> Dict[str, Dict[str, Any]]: Returns: Dict of connection_id -> connection_info """ + with self._registry_lock: + configs = list(self.connection_configs.items()) + active_ids = set(self.connections) result = {} - for conn_id, config in self.connection_configs.items(): + for conn_id, config in configs: result[conn_id] = { "type": config.type, "host": config.host, "port": config.port, "database": config.database, "username": config.username, - "status": "active" if conn_id in self.connections else "closed" + "status": "active" if conn_id in active_ids else "closed" } return result @@ -275,12 +312,33 @@ def get_cursor(self, connection_id: str): Yields: Database cursor """ - connection = self.get_connection(connection_id) - cursor = connection.cursor() - try: - yield cursor - finally: - cursor.close() + with self._connection_lock(connection_id): + connection = self._get_connection_unlocked(connection_id) + try: + if getattr(connection, "closed", 0): + raise DatabaseConnectionError("connection is closed") + if hasattr(connection, "ping"): + connection.ping(reconnect=True) + except Exception as exc: + logger.warning("Connection %s is dead, removing: %s", connection_id, exc) + self._close_connection_unlocked(connection_id, connection) + raise DatabaseConnectionError( + f"Connection {connection_id} is no longer valid" + ) from exc + try: + cursor = connection.cursor() + except Exception as exc: + logger.warning("Connection %s could not create a cursor: %s", connection_id, exc) + self._close_connection_unlocked(connection_id, connection) + raise DatabaseConnectionError( + f"Connection {connection_id} is no longer valid" + ) from exc + try: + # Exceptions from the caller's SQL body must propagate without + # evicting a healthy connection from the registry. + yield cursor + finally: + cursor.close() def execute_query(self, connection_id: str, query: str, params: Optional[tuple] = None) -> List[Dict[str, Any]]: """ diff --git a/src/dbjavagenix/database/mcp_tools.py b/src/dbjavagenix/database/mcp_tools.py index 7947d1f..008d240 100644 --- a/src/dbjavagenix/database/mcp_tools.py +++ b/src/dbjavagenix/database/mcp_tools.py @@ -2,6 +2,7 @@ MCP tools for database connection and basic query operations """ from base64 import b64encode +import asyncio from datetime import date, datetime, time, timedelta from decimal import Decimal import json @@ -348,6 +349,66 @@ def _quote_mysql_identifier(identifier: Any) -> str: return quote_mysql_identifier(identifier) except ValueError as exc: raise MCPServiceError(str(exc)) from exc + + +async def _run_db_call(callable_obj, *args): + """Run one blocking database operation outside the MCP event loop.""" + return await asyncio.to_thread(callable_obj, *args) + + +async def _run_async_db_call(callable_obj, *args, **kwargs): + """Run an async analyzer whose internals contain blocking DB calls in a worker.""" + def run() -> Any: + return asyncio.run(callable_obj(*args, **kwargs)) + + return await asyncio.to_thread(run) + + +def _connect_and_probe(config: DatabaseConfig) -> tuple[str, str]: + """Create a connection and perform the initial blocking connectivity probe.""" + connection_id = connection_manager.create_connection(config) + connection_manager.get_connection(connection_id) + server_info = "" + try: + if config.type == DatabaseType.MYSQL: + with connection_manager.get_cursor(connection_id) as cursor: + cursor.execute("SELECT VERSION() as version") + result = cursor.fetchone() + if result: + server_info = f"MySQL {result[0] if isinstance(result, tuple) else result['version']}" + elif config.type == DatabaseType.POSTGRESQL: + with connection_manager.get_cursor(connection_id) as cursor: + cursor.execute("SELECT version() AS version") + result = cursor.fetchone() + if result: + server_info = ( + f"PostgreSQL {result[0] if isinstance(result, tuple) else result['version']}" + ) + elif config.type == DatabaseType.SQLITE: + server_info = "SQLite" + except Exception as exc: + logger.warning("Could not get server info: %s", exc) + server_info = f"{config.type.value} (version unknown)" + return connection_id, server_info + + +def _collect_all_table_names(connection_id: str, config: DatabaseConfig) -> List[str]: + """Collect table names for codegen prefix analysis inside a DB worker.""" + try: + with connection_manager.get_cursor(connection_id) as cursor: + if config.type == DatabaseType.MYSQL: + cursor.execute("SHOW TABLES") + elif config.type == DatabaseType.SQLITE: + cursor.execute( + "SELECT name FROM sqlite_master " + "WHERE type='table' AND name NOT LIKE 'sqlite_%'" + ) + else: + return [] + return [row[0] for row in cursor.fetchall()] + except Exception as exc: + logger.warning("Failed to get all table names for prefix analysis: %s", exc) + return [] def get_connection_tools() -> List[Tool]: @@ -663,37 +724,8 @@ async def handle_db_connect_test(arguments: Dict[str, Any]) -> List[TextContent] charset=arguments.get("charset", "utf8mb4") ) - # Create connection - connection_id = connection_manager.create_connection(config) - - # Test basic connectivity - connection = connection_manager.get_connection(connection_id) - - # Get server information - server_info = "" - try: - if config.type == DatabaseType.MYSQL: - with connection_manager.get_cursor(connection_id) as cursor: - cursor.execute("SELECT VERSION() as version") - result = cursor.fetchone() - if result: - server_info = f"MySQL {result[0] if isinstance(result, tuple) else result['version']}" - - elif config.type == DatabaseType.POSTGRESQL: - with connection_manager.get_cursor(connection_id) as cursor: - cursor.execute("SELECT version() AS version") - result = cursor.fetchone() - if result: - server_info = ( - f"PostgreSQL {result[0] if isinstance(result, tuple) else result['version']}" - ) - - elif config.type == DatabaseType.SQLITE: - server_info = "SQLite" - - except Exception as e: - logger.warning(f"Could not get server info: {e}") - server_info = f"{config.type.value} (version unknown)" + # Connection setup and the initial probe are blocking driver calls. + connection_id, server_info = await _run_db_call(_connect_and_probe, config) response = { "success": True, @@ -785,7 +817,7 @@ async def handle_db_query_databases(arguments: Dict[str, Any]) -> List[TextConte else: raise MCPServiceError(f"Listing databases not implemented for {config.type}") - results = connection_manager.execute_query(connection_id, query) + results = await _run_db_call(connection_manager.execute_query, connection_id, query) # Extract database names databases = [] @@ -876,9 +908,9 @@ async def handle_db_query_tables(arguments: Dict[str, Any]) -> List[TextContent] raise MCPServiceError(f"Listing tables not implemented for {config.type}") if params is None: - results = connection_manager.execute_query(connection_id, query) + results = await _run_db_call(connection_manager.execute_query, connection_id, query) else: - results = connection_manager.execute_query(connection_id, query, params) + results = await _run_db_call(connection_manager.execute_query, connection_id, query, params) # Extract table names tables = [] @@ -964,7 +996,9 @@ async def handle_db_query_table_exists(arguments: Dict[str, Any]) -> List[TextCo FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s """ - results = connection_manager.execute_query(connection_id, query, (database, table)) + results = await _run_db_call( + connection_manager.execute_query, connection_id, query, (database, table) + ) elif config.type == DatabaseType.POSTGRESQL: schema_filter = "AND table_schema = %s" if schema else "" @@ -977,7 +1011,7 @@ async def handle_db_query_table_exists(arguments: Dict[str, Any]) -> List[TextCo {schema_filter} """.format(schema_filter=schema_filter) params = (database, table, schema) if schema else (database, table) - results = connection_manager.execute_query(connection_id, query, params) + results = await _run_db_call(connection_manager.execute_query, connection_id, query, params) elif config.type == DatabaseType.SQLITE: query = """ @@ -985,7 +1019,7 @@ async def handle_db_query_table_exists(arguments: Dict[str, Any]) -> List[TextCo FROM sqlite_master WHERE type='table' AND name = ? """ - results = connection_manager.execute_query(connection_id, query, (table,)) + results = await _run_db_call(connection_manager.execute_query, connection_id, query, (table,)) else: raise MCPServiceError(f"Table existence check not implemented for {config.type}") @@ -1053,7 +1087,7 @@ async def handle_db_query_execute(arguments: Dict[str, Any]) -> List[TextContent query = _apply_query_limit(query, limit) - results = connection_manager.execute_query(connection_id, query) + results = await _run_db_call(connection_manager.execute_query, connection_id, query) response = { "success": True, @@ -1204,7 +1238,7 @@ async def handle_db_table_describe(arguments: Dict[str, Any]) -> List[TextConten include_java_types = arguments.get("include_java_types", True) introspector = DatabaseIntrospector(connection_manager) config = introspector.get_config(connection_id) - metadata = introspector.describe_table(connection_id, table, schema) + metadata = await _run_db_call(introspector.describe_table, connection_id, table, schema) columns = [] java_imports = set() @@ -1341,11 +1375,16 @@ async def handle_db_table_columns(arguments: Dict[str, Any]) -> List[TextContent WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s ORDER BY ORDINAL_POSITION """ - results = connection_manager.execute_query(connection_id, query, (database, table)) + results = await _run_db_call( + connection_manager.execute_query, connection_id, query, (database, table) + ) elif config.type == DatabaseType.POSTGRESQL: - columns = DatabaseIntrospector(connection_manager).get_columns( - connection_id, table, schema + columns = await _run_db_call( + DatabaseIntrospector(connection_manager).get_columns, + connection_id, + table, + schema, ) results = [ { @@ -1364,8 +1403,11 @@ async def handle_db_table_columns(arguments: Dict[str, Any]) -> List[TextContent ] elif config.type == DatabaseType.SQLITE: - columns = DatabaseIntrospector(connection_manager).get_columns( - connection_id, table, schema + columns = await _run_db_call( + DatabaseIntrospector(connection_manager).get_columns, + connection_id, + table, + schema, ) results = [ { @@ -1471,11 +1513,16 @@ async def handle_db_table_primary_keys(arguments: Dict[str, Any]) -> List[TextCo AND CONSTRAINT_NAME = 'PRIMARY' ORDER BY ORDINAL_POSITION """ - results = connection_manager.execute_query(connection_id, query, (database, table)) + results = await _run_db_call( + connection_manager.execute_query, connection_id, query, (database, table) + ) elif config.type == DatabaseType.POSTGRESQL: - primary_keys = DatabaseIntrospector(connection_manager).get_primary_keys( - connection_id, table, schema + primary_keys = await _run_db_call( + DatabaseIntrospector(connection_manager).get_primary_keys, + connection_id, + table, + schema, ) results = [ {"COLUMN_NAME": column_name, "ORDINAL_POSITION": position} @@ -1483,8 +1530,11 @@ async def handle_db_table_primary_keys(arguments: Dict[str, Any]) -> List[TextCo ] elif config.type == DatabaseType.SQLITE: - primary_keys = DatabaseIntrospector(connection_manager).get_primary_keys( - connection_id, table, schema + primary_keys = await _run_db_call( + DatabaseIntrospector(connection_manager).get_primary_keys, + connection_id, + table, + schema, ) results = [ {"COLUMN_NAME": column_name, "ORDINAL_POSITION": position} @@ -1581,11 +1631,16 @@ async def handle_db_table_foreign_keys(arguments: Dict[str, Any]) -> List[TextCo AND kcu.REFERENCED_TABLE_NAME IS NOT NULL ORDER BY kcu.ORDINAL_POSITION """ - results = connection_manager.execute_query(connection_id, query, (database, table)) + results = await _run_db_call( + connection_manager.execute_query, connection_id, query, (database, table) + ) elif config.type == DatabaseType.POSTGRESQL: - foreign_keys = DatabaseIntrospector(connection_manager).get_foreign_keys( - connection_id, table, schema + foreign_keys = await _run_db_call( + DatabaseIntrospector(connection_manager).get_foreign_keys, + connection_id, + table, + schema, ) results = [ { @@ -1601,8 +1656,11 @@ async def handle_db_table_foreign_keys(arguments: Dict[str, Any]) -> List[TextCo ] elif config.type == DatabaseType.SQLITE: - foreign_keys = DatabaseIntrospector(connection_manager).get_foreign_keys( - connection_id, table, schema + foreign_keys = await _run_db_call( + DatabaseIntrospector(connection_manager).get_foreign_keys, + connection_id, + table, + schema, ) results = [ { @@ -1719,11 +1777,16 @@ async def handle_db_table_indexes(arguments: Dict[str, Any]) -> List[TextContent AND TABLE_NAME = %s ORDER BY INDEX_NAME, SEQ_IN_INDEX """ - results = connection_manager.execute_query(connection_id, query, (database, table)) + results = await _run_db_call( + connection_manager.execute_query, connection_id, query, (database, table) + ) elif config.type == DatabaseType.POSTGRESQL: - indexes = DatabaseIntrospector(connection_manager).get_indexes( - connection_id, table, schema + indexes = await _run_db_call( + DatabaseIntrospector(connection_manager).get_indexes, + connection_id, + table, + schema, ) results = [ { @@ -1739,8 +1802,11 @@ async def handle_db_table_indexes(arguments: Dict[str, Any]) -> List[TextContent ] elif config.type == DatabaseType.SQLITE: - indexes = DatabaseIntrospector(connection_manager).get_indexes( - connection_id, table, schema + indexes = await _run_db_call( + DatabaseIntrospector(connection_manager).get_indexes, + connection_id, + table, + schema, ) results = [ { @@ -2007,7 +2073,8 @@ async def handle_db_codegen_analyze(arguments: Dict[str, Any]) -> List[TextConte project_path = arguments.get("project_path") proj_struct = _detect_project_structure(project_path) project_root = str(proj_struct["project_root"]) if proj_struct.get("project_root") else None - analysis_result = await analyzer.analyze_table_for_codegen( + analysis_result = await _run_async_db_call( + analyzer.analyze_table_for_codegen, connection_id, table_name, template_category=template_category, @@ -2209,26 +2276,11 @@ async def handle_db_codegen_generate(arguments: Dict[str, Any]) -> List[TextCont logger.info("🔍 Getting all table names for package structure optimization...") # 获取数据库中的所有表名用于前缀分析 - config = connection_manager.get_connection_info(connection_id) - connection = connection_manager.get_connection(connection_id) - cursor = connection.cursor() - - all_table_names = [] - try: - if config.type.name == "MYSQL": - cursor.execute("SHOW TABLES") - all_table_names = [row[0] for row in cursor.fetchall()] - elif config.type.name == "SQLITE": - cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'") - all_table_names = [row[0] for row in cursor.fetchall()] - - logger.info(f"Found {len(all_table_names)} tables for prefix analysis: {all_table_names}") - - except Exception as e: - logger.warning(f"Failed to get all table names for prefix analysis: {e}") - all_table_names = [table_name] # 至少包含当前表 - finally: - cursor.close() + config = connection_manager.get_connection_info(connection_id) + all_table_names = await _run_db_call(_collect_all_table_names, connection_id, config) + if not all_table_names: + all_table_names = [table_name] + logger.info("Found %s tables for prefix analysis: %s", len(all_table_names), all_table_names) # ===== STEP 2: 分析表结构(包含前缀优化) ===== # Initialize analyzer and generator @@ -2237,7 +2289,8 @@ async def handle_db_codegen_generate(arguments: Dict[str, Any]) -> List[TextCont # Step 1: Analyze table structure with all table names for prefix optimization _ps = _detect_project_structure(project_path) - analysis_result = await analyzer.analyze_table_for_codegen( + analysis_result = await _run_async_db_call( + analyzer.analyze_table_for_codegen, connection_id, table_name, all_table_names=all_table_names, # 传递所有表名用于前缀分析 diff --git a/src/dbjavagenix/database/visualization_tools.py b/src/dbjavagenix/database/visualization_tools.py index 06801e1..64deb73 100644 --- a/src/dbjavagenix/database/visualization_tools.py +++ b/src/dbjavagenix/database/visualization_tools.py @@ -5,6 +5,7 @@ - db_render_er_diagram: 给定多个表名,生成 Mermaid erDiagram (附加 mcp-apps/mermaid meta) """ +import asyncio import logging from typing import Any, Dict, List @@ -81,8 +82,13 @@ async def handle_db_render_er_diagram(arguments: Dict[str, Any]) -> List[TextCon all_fks: List[ERForeignKey] = [] for table_name in tables: - columns, fks = _collect_table_for_er( - connection_id, database, table_name, config.type, include_non_pk + columns, fks = await asyncio.to_thread( + _collect_table_for_er, + connection_id, + database, + table_name, + config.type, + include_non_pk, ) er_tables.append(ERTable(name=table_name, columns=columns)) all_fks.extend(fks) diff --git a/tests/unit/test_async_db_boundary.py b/tests/unit/test_async_db_boundary.py new file mode 100644 index 0000000..5d17c1f --- /dev/null +++ b/tests/unit/test_async_db_boundary.py @@ -0,0 +1,183 @@ +"""Regression tests for the async MCP/database boundary.""" + +import asyncio +import threading +import time + +import pytest + +from dbjavagenix.core.models import DatabaseConfig, DatabaseType +from dbjavagenix.database.connection_manager import ConnectionManager +from dbjavagenix.database import mcp_tools + + +@pytest.fixture +def sqlite_config(): + return DatabaseConfig( + type=DatabaseType.SQLITE, + host="", + port=0, + database=":memory:", + username="", + password="", + ) + + +class _BlockingCursor: + description = None + + def __init__(self, state): + self.state = state + + def execute(self, _query, _params=()): + with self.state["lock"]: + self.state["active"] += 1 + self.state["max_active"] = max(self.state["max_active"], self.state["active"]) + self.state["entered"].set() + self.state["release"].wait(timeout=2) + with self.state["lock"]: + self.state["active"] -= 1 + + def close(self): + self.state["closed_cursors"] += 1 + + +class _BlockingConnection: + closed = 0 + + def __init__(self, state): + self.state = state + self.closed_event = threading.Event() + + def cursor(self): + return _BlockingCursor(self.state) + + def close(self): + self.closed = 1 + self.closed_event.set() + + +def _state(): + return { + "lock": threading.Lock(), + "active": 0, + "max_active": 0, + "entered": threading.Event(), + "release": threading.Event(), + "closed_cursors": 0, + } + + +def _manager_with_fake_connection(config, state): + manager = ConnectionManager() + connection_id = manager.create_connection(config) + manager.connections[connection_id] = _BlockingConnection(state) + return manager, connection_id + + +@pytest.mark.asyncio +async def test_mcp_query_keeps_event_loop_running_during_blocking_driver(monkeypatch): + started = threading.Event() + + def blocking_execute(connection_id, query): + assert connection_id == "db-1" + assert query.endswith("LIMIT 100") + started.set() + time.sleep(0.06) + return [] + + monkeypatch.setattr(mcp_tools.connection_manager, "execute_query", blocking_execute) + + ticks = 0 + + async def heartbeat(): + nonlocal ticks + deadline = asyncio.get_running_loop().time() + 0.04 + while asyncio.get_running_loop().time() < deadline: + ticks += 1 + await asyncio.sleep(0.002) + + await asyncio.gather( + mcp_tools.handle_db_query_execute({"connection_id": "db-1", "query": "SELECT 1"}), + heartbeat(), + ) + + assert started.is_set() + assert ticks >= 2 + + +@pytest.mark.asyncio +async def test_sqlite_connection_can_be_used_by_worker_thread(sqlite_config): + manager = ConnectionManager() + connection_id = manager.create_connection(sqlite_config) + try: + rows = await asyncio.to_thread(manager.execute_query, connection_id, "SELECT 1 AS value") + assert rows == [{"value": 1}] + finally: + manager.close_connection(connection_id) + + +@pytest.mark.asyncio +async def test_same_connection_queries_are_serialized(sqlite_config): + first_state = _state() + manager, connection_id = _manager_with_fake_connection(sqlite_config, first_state) + + first = asyncio.create_task(asyncio.to_thread(manager.execute_query, connection_id, "SELECT 1")) + assert await asyncio.to_thread(first_state["entered"].wait, 1) + second = asyncio.create_task( + asyncio.to_thread(manager.execute_query, connection_id, "SELECT 2") + ) + await asyncio.sleep(0.02) + assert first_state["max_active"] == 1 + + first_state["release"].set() + await asyncio.gather(first, second) + assert first_state["closed_cursors"] == 2 + manager.close_connection(connection_id) + + +@pytest.mark.asyncio +async def test_different_connections_can_execute_in_parallel(sqlite_config): + first_state = _state() + second_state = _state() + manager, first_id = _manager_with_fake_connection(sqlite_config, first_state) + second_id = manager.create_connection(sqlite_config) + manager.connections[second_id] = _BlockingConnection(second_state) + + first = asyncio.create_task(asyncio.to_thread(manager.execute_query, first_id, "SELECT 1")) + second = asyncio.create_task(asyncio.to_thread(manager.execute_query, second_id, "SELECT 2")) + await asyncio.wait_for( + asyncio.gather( + asyncio.to_thread(first_state["entered"].wait, 1), + asyncio.to_thread(second_state["entered"].wait, 1), + ), + timeout=1, + ) + assert first_state["max_active"] == second_state["max_active"] == 1 + + first_state["release"].set() + second_state["release"].set() + await asyncio.gather(first, second) + manager.close_connection(first_id) + manager.close_connection(second_id) + + +@pytest.mark.asyncio +async def test_close_waits_for_query_and_cleans_connection(sqlite_config): + state = _state() + manager, connection_id = _manager_with_fake_connection(sqlite_config, state) + + query = asyncio.create_task(asyncio.to_thread(manager.execute_query, connection_id, "SELECT 1")) + assert await asyncio.to_thread(state["entered"].wait, 1) + close = asyncio.create_task(asyncio.to_thread(manager.close_connection, connection_id)) + await asyncio.sleep(0.02) + assert not close.done() + + state["release"].set() + result, closed = await asyncio.gather(query, close) + assert result == [] + assert closed is True + assert connection_id not in manager.connections + assert connection_id not in manager.connection_configs + assert connection_id not in manager._connection_locks + assert state["closed_cursors"] == 1 diff --git a/tests/unit/test_connection_manager.py b/tests/unit/test_connection_manager.py index 7263edc..d233bf4 100644 --- a/tests/unit/test_connection_manager.py +++ b/tests/unit/test_connection_manager.py @@ -153,6 +153,8 @@ def test_bad_sql_raises(self, manager_with_conn): mgr, cid = manager_with_conn with pytest.raises(DatabaseQueryError): mgr.execute_query(cid, "INVALID SQL STATEMENT") + # A statement error must not evict an otherwise healthy connection. + assert cid in mgr.connections def test_empty_result_table(self, manager_with_conn): mgr, cid = manager_with_conn