Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions config_example.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
15 changes: 11 additions & 4 deletions src/meeseeql/database_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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:
Expand Down
29 changes: 22 additions & 7 deletions src/meeseeql/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -72,7 +71,6 @@ def show_database_config() -> ToolResult:
)


@mcp.tool()
async def execute_query(
database: str,
query: str,
Expand All @@ -98,7 +96,6 @@ async def execute_query(
)


@mcp.tool()
async def table_summary(
database: str,
table_name: str,
Expand All @@ -124,7 +121,6 @@ async def table_summary(
)


@mcp.tool()
async def search(
database: str,
search_term: str,
Expand All @@ -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)
Expand All @@ -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()
Expand All @@ -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()


Expand Down
6 changes: 2 additions & 4 deletions src/meeseeql/tools/execute_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/meeseeql/tools/search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
4 changes: 1 addition & 3 deletions src/meeseeql/tools/table_summary.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
53 changes: 53 additions & 0 deletions tests/database_manager/test_available_tools.py
Original file line number Diff line number Diff line change
@@ -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 == []
33 changes: 31 additions & 2 deletions tests/tools/test_show_database_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")