From fd3cab2fbf4509dbc880ab09d275ae988bae0467 Mon Sep 17 00:00:00 2001 From: Thomas Smith Date: Wed, 6 Aug 2025 21:49:08 +0200 Subject: [PATCH] Add config option for enabling only a specific set of tools --- config_example.yaml | 8 ++- src/meeseeql/database_manager.py | 15 ++++-- src/meeseeql/main.py | 29 +++++++--- src/meeseeql/tools/execute_query.py | 6 +-- src/meeseeql/tools/search.py | 2 +- src/meeseeql/tools/table_summary.py | 4 +- .../database_manager/test_available_tools.py | 53 +++++++++++++++++++ tests/tools/test_show_database_config.py | 33 +++++++++++- 8 files changed, 127 insertions(+), 23 deletions(-) create mode 100644 tests/database_manager/test_available_tools.py diff --git a/config_example.yaml b/config_example.yaml index 9961c8e..5bc59d2 100644 --- a/config_example.yaml +++ b/config_example.yaml @@ -31,5 +31,9 @@ databases: # Global MCP server settings settings: - max_query_timeout: 30 # Maximum time in seconds for any single query - max_rows_per_query: 1000 # Maximum rows returned per query (pagination will apply) + max_query_timeout: 30 # Maximum time in seconds for any single query + max_rows_per_query: 1000 # Maximum rows returned per query, this is a hard limit on top of any LLM-configurable pagination + + # Optional: restrict which tools are available (omit to enable all tools) + # Note: show_database_config is always available + available_tools: ["execute_query", "table_summary", "search", "test_connection", "reload_config"] diff --git a/src/meeseeql/database_manager.py b/src/meeseeql/database_manager.py index 46ce3eb..58142bb 100644 --- a/src/meeseeql/database_manager.py +++ b/src/meeseeql/database_manager.py @@ -87,9 +87,15 @@ def dialect(self) -> str: return dialect_map[self.type] +class GlobalSettings(BaseModel): + max_query_timeout: int = 30 + max_rows_per_query: int = 500 + available_tools: List[str] | None = None + + class AppConfig(BaseModel): databases: Dict[str, DatabaseConfig] - settings: Dict[str, Any] + settings: GlobalSettings = GlobalSettings() config_path: str | None = None @model_validator(mode="after") @@ -165,9 +171,7 @@ def _get_engine(self, db_label: str) -> AsyncEngine | Engine: # Only add pool_timeout for non-SQLite databases if not str(url).startswith("sqlite"): - engine_kwargs["pool_timeout"] = self.config.settings.get( - "max_query_timeout", 30 - ) + engine_kwargs["pool_timeout"] = self.config.settings.max_query_timeout # Snowflake doesn't have native async support, use sync engine with async wrapper if db_config.type == "snowflake": @@ -273,6 +277,9 @@ def get_table_filter_type(self, db_name: str) -> str | None: else: return None + def get_available_tools(self) -> List[str] | None: + return self.config.settings.available_tools + def reload_config(self, new_config: AppConfig, changed_db_names: set[str]): for db_name in changed_db_names: if db_name in self.engines: diff --git a/src/meeseeql/main.py b/src/meeseeql/main.py index 159bbdb..e232d0c 100644 --- a/src/meeseeql/main.py +++ b/src/meeseeql/main.py @@ -58,7 +58,6 @@ def get_or_init_db_manager(): return db_manager -@mcp.tool() def show_database_config() -> ToolResult: """List all configured databases and their settings. This tool also gives you the full path to the config file @@ -72,7 +71,6 @@ def show_database_config() -> ToolResult: ) -@mcp.tool() async def execute_query( database: str, query: str, @@ -98,7 +96,6 @@ async def execute_query( ) -@mcp.tool() async def table_summary( database: str, table_name: str, @@ -124,7 +121,6 @@ async def table_summary( ) -@mcp.tool() async def search( database: str, search_term: str, @@ -147,7 +143,6 @@ async def search( ) -@mcp.tool() async def test_connection(database: str) -> ToolResult: """Test database connection, useful for debugging issues""" result = await tools.test_connection(get_or_init_db_manager(), database) @@ -157,7 +152,6 @@ async def test_connection(database: str) -> ToolResult: ) -@mcp.tool() def reload_config() -> ToolResult: """Reload configuration file and report what changed""" config_path = find_config_file() @@ -168,8 +162,29 @@ def reload_config() -> ToolResult: ) +def register_tools(): + available_tools = get_or_init_db_manager().get_available_tools() + + mcp.tool()(show_database_config) + + if available_tools is None or "execute_query" in available_tools: + mcp.tool()(execute_query) + + if available_tools is None or "table_summary" in available_tools: + mcp.tool()(table_summary) + + if available_tools is None or "search" in available_tools: + mcp.tool()(search) + + if available_tools is None or "test_connection" in available_tools: + mcp.tool()(test_connection) + + if available_tools is None or "reload_config" in available_tools: + mcp.tool()(reload_config) + + def main(): - get_or_init_db_manager() + register_tools() mcp.run() diff --git a/src/meeseeql/tools/execute_query.py b/src/meeseeql/tools/execute_query.py index b449990..6388187 100644 --- a/src/meeseeql/tools/execute_query.py +++ b/src/meeseeql/tools/execute_query.py @@ -2,7 +2,7 @@ from typing import Dict, Any, List from pydantic import BaseModel from meeseeql.database_manager import DatabaseManager -from meeseeql.sql_transformer import SqlQueryTransformer, TableAccessError +from meeseeql.sql_transformer import SqlQueryTransformer class QueryResponse(BaseModel): @@ -78,9 +78,7 @@ async def execute_query( if page < 1: raise ValueError("Page number must be greater than 0") - max_rows = db_manager.config.settings.get("max_rows_per_query", 1000) - if limit > max_rows: - limit = max_rows + limit = min(limit, db_manager.config.settings.max_rows_per_query) dialect = db_manager.get_dialect_name(database) transformer = SqlQueryTransformer(query.strip(), dialect) diff --git a/src/meeseeql/tools/search.py b/src/meeseeql/tools/search.py index 02a8161..d0ea567 100644 --- a/src/meeseeql/tools/search.py +++ b/src/meeseeql/tools/search.py @@ -99,7 +99,7 @@ async def search( sql_template = load_sql_query(dialect, "search") - limit = min(250, db_manager.config.settings.get("max_rows_per_query", 250)) + limit = min(250, db_manager.config.settings.max_rows_per_query) sql_query = sql_template.replace("{{search_term}}", search_term) diff --git a/src/meeseeql/tools/table_summary.py b/src/meeseeql/tools/table_summary.py index 7709d3f..a2ded5e 100644 --- a/src/meeseeql/tools/table_summary.py +++ b/src/meeseeql/tools/table_summary.py @@ -404,9 +404,7 @@ async def table_summary( db_manager, database, table_name, schema_value, limit, page ) - max_rows = db_manager.config.settings.get("max_rows_per_query", 1000) - if limit > max_rows: - limit = max_rows + limit = min(limit, db_manager.config.settings.max_rows_per_query) dialect = db_manager.get_dialect_name(database) diff --git a/tests/database_manager/test_available_tools.py b/tests/database_manager/test_available_tools.py new file mode 100644 index 0000000..b89e76d --- /dev/null +++ b/tests/database_manager/test_available_tools.py @@ -0,0 +1,53 @@ +from meeseeql.database_manager import ( + AppConfig, + DatabaseConfig, + DatabaseManager, + GlobalSettings, +) + + +def test_get_available_tools_returns_none_when_not_configured(): + """Test that get_available_tools returns None when available_tools is not configured""" + config = AppConfig( + databases={ + "test_db": DatabaseConfig( + type="sqlite", description="Test", database=":memory:" + ) + }, + settings=GlobalSettings(), + ) + db_manager = DatabaseManager(config) + + assert db_manager.get_available_tools() is None + + +def test_get_available_tools_returns_configured_tools(): + """Test that get_available_tools returns the configured tools""" + config = AppConfig( + databases={ + "test_db": DatabaseConfig( + type="sqlite", description="Test", database=":memory:" + ) + }, + settings=GlobalSettings(available_tools=["execute_query", "table_summary"]), + ) + db_manager = DatabaseManager(config) + + tools = db_manager.get_available_tools() + assert tools == ["execute_query", "table_summary"] + + +def test_get_available_tools_returns_empty_list(): + """Test that get_available_tools can return an empty list""" + config = AppConfig( + databases={ + "test_db": DatabaseConfig( + type="sqlite", description="Test", database=":memory:" + ) + }, + settings=GlobalSettings(available_tools=[]), + ) + db_manager = DatabaseManager(config) + + tools = db_manager.get_available_tools() + assert tools == [] diff --git a/tests/tools/test_show_database_config.py b/tests/tools/test_show_database_config.py index 6c9eb3a..808bb5d 100644 --- a/tests/tools/test_show_database_config.py +++ b/tests/tools/test_show_database_config.py @@ -49,12 +49,41 @@ def test_show_database_config_includes_correct_database_info(db_manager): def test_show_database_config_with_empty_config(): """Test show_database_config with empty database config""" - from meeseeql.database_manager import AppConfig + from meeseeql.database_manager import AppConfig, GlobalSettings - empty_config = AppConfig(databases={}, settings={}) + empty_config = AppConfig(databases={}, settings=GlobalSettings()) empty_db_manager = DatabaseManager(empty_config) result = show_database_config(empty_db_manager) assert isinstance(result, DatabaseList) assert result.total_count == 0 + + +def test_show_database_config_does_not_include_passwords(): + """Test that passwords are not included in the output""" + from meeseeql.database_manager import AppConfig, DatabaseConfig, GlobalSettings + + config_with_password = AppConfig( + databases={ + "test_db": DatabaseConfig( + type="postgresql", + description="Test DB with credentials", + host="localhost", + database="test", + username="user", + password="secret_password", + ) + }, + settings=GlobalSettings(), + ) + + db_manager = DatabaseManager(config_with_password) + result = show_database_config(db_manager) + + output_str = str(result) + assert "secret_password" not in output_str + assert "password:" not in output_str + + db_info = result.databases[0] + assert not hasattr(db_info, "password")