From e68b0b87f17cbf22aab6b673979dd6f262bcf100 Mon Sep 17 00:00:00 2001 From: FreeOnePlus Date: Thu, 30 Jul 2026 18:15:42 +0800 Subject: [PATCH] fix: make default query row limit configurable --- .env.example | 2 + README.md | 2 + docker-compose.yml | 1 + doris_mcp_server/result_limits.py | 12 +++++ doris_mcp_server/tools/tool_catalog.py | 7 +-- doris_mcp_server/tools/tools_manager.py | 6 ++- doris_mcp_server/utils/config.py | 17 ++++++ .../integration/test_real_doris_transports.py | 19 +++++++ test/utils/test_result_limits.py | 52 ++++++++++++++++++- 9 files changed, 113 insertions(+), 5 deletions(-) diff --git a/.env.example b/.env.example index 4f80861..676e8b9 100644 --- a/.env.example +++ b/.env.example @@ -368,6 +368,8 @@ BLOCKED_KEYWORDS=DROP,CREATE,ALTER,TRUNCATE,DELETE,INSERT,UPDATE,GRANT,REVOKE,EX MAX_QUERY_COMPLEXITY=100 # Deployment ceiling; absolute hard cap: 100000 MAX_RESULT_ROWS=10000 +# Per-query default when exec_query omits max_rows; must not exceed MAX_RESULT_ROWS +DEFAULT_RESULT_ROWS=100 # Data masking ENABLE_MASKING=true diff --git a/README.md b/README.md index 5fbbbd8..5bcd34a 100644 --- a/README.md +++ b/README.md @@ -346,6 +346,8 @@ cp .env.example .env * `ENABLE_MASKING`: Enable data masking (default: true) * `MAX_RESULT_ROWS`: Deployment ceiling for returned query rows (default: 10000; absolute hard cap: 100000) + * `DEFAULT_RESULT_ROWS`: Default row budget when `exec_query.max_rows` + is omitted (default: 100; cannot exceed `MAX_RESULT_ROWS`) * **ADBC Configuration (New in v0.5.0)**: * `ADBC_DEFAULT_MAX_ROWS`: Default maximum rows for ADBC queries (default: 10000; cannot exceed `MAX_RESULT_ROWS`) diff --git a/docker-compose.yml b/docker-compose.yml index 87a6706..072d2b3 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -42,6 +42,7 @@ services: - ENABLE_TOKEN_AUTH=true - TOKEN_ADMIN_FILE=/run/secrets/mcp_static_token - MAX_RESULT_ROWS=10000 + - DEFAULT_RESULT_ROWS=${DEFAULT_RESULT_ROWS:-100} - MAX_RESULT_BYTES=1048576 - QUERY_TIMEOUT=300 diff --git a/doris_mcp_server/result_limits.py b/doris_mcp_server/result_limits.py index 003d17a..8c071fd 100644 --- a/doris_mcp_server/result_limits.py +++ b/doris_mcp_server/result_limits.py @@ -34,6 +34,7 @@ MIN_RESULT_BYTES = 256 DEFAULT_MAX_RESULT_ROWS = 10_000 +DEFAULT_RESULT_ROWS = 100 DEFAULT_MAX_RESULT_BYTES = 1024 * 1024 DEFAULT_QUERY_TIMEOUT_SECONDS = 300 @@ -85,6 +86,17 @@ def configured_result_limits(config: object | None) -> ResultLimits: ) +def configured_default_result_rows(config: object | None) -> int: + """Return the configured default row budget bounded by the deployment ceiling.""" + ceilings = configured_result_limits(config) + performance = getattr(config, "performance", None) + return _configured_positive_int( + getattr(performance, "default_result_rows", None), + default=min(DEFAULT_RESULT_ROWS, ceilings.max_rows), + hard_maximum=ceilings.max_rows, + ) + + def _requested_limit( value: object | None, *, diff --git a/doris_mcp_server/tools/tool_catalog.py b/doris_mcp_server/tools/tool_catalog.py index a7e7b07..8c26208 100644 --- a/doris_mcp_server/tools/tool_catalog.py +++ b/doris_mcp_server/tools/tool_catalog.py @@ -20,7 +20,7 @@ from mcp.types import Tool -from ..result_limits import configured_result_limits +from ..result_limits import configured_default_result_rows, configured_result_limits from ..utils.config import ADBCConfig from .tool_registry import ToolDefinitionRegistry @@ -32,6 +32,7 @@ def build_tool_registry( """Build immutable tool metadata and bind it to a handler owner.""" adbc_config = getattr(config, "adbc", None) or ADBCConfig() result_limits = configured_result_limits(config) + default_result_rows = configured_default_result_rows(config) adbc_default_max_rows = min( adbc_config.default_max_rows, result_limits.max_rows, @@ -54,7 +55,7 @@ def build_tool_registry( - catalog_name (string) [Optional] - Reference catalog name for context, defaults to current catalog -- max_rows (integer) [Optional] - Maximum number of rows to return, default 100 +- max_rows (integer) [Optional] - Maximum number of rows to return, defaults to the configured DEFAULT_RESULT_ROWS value - max_bytes (integer) [Optional] - Maximum UTF-8 JSON bytes for returned row data @@ -78,7 +79,7 @@ def build_tool_registry( "max_rows": { "type": "integer", "description": "Maximum number of rows to return", - "default": 100, + "default": default_result_rows, "minimum": 1, "maximum": result_limits.max_rows, }, diff --git a/doris_mcp_server/tools/tools_manager.py b/doris_mcp_server/tools/tools_manager.py index 5ca292b..0242197 100644 --- a/doris_mcp_server/tools/tools_manager.py +++ b/doris_mcp_server/tools/tools_manager.py @@ -31,6 +31,7 @@ authorize_operation, filter_tools_for_auth_context, ) +from ..result_limits import configured_default_result_rows from ..utils.adbc_query_tools import DorisADBCQueryTools from ..utils.analysis_tools import MemoryTracker, SQLAnalyzer, TableAnalyzer from ..utils.data_exploration_tools import DataExplorationTools @@ -224,7 +225,10 @@ async def _exec_query_tool(self, arguments: dict[str, Any]) -> dict[str, Any]: sql = self._required_string(arguments, "sql") db_name = arguments.get("db_name") catalog_name = arguments.get("catalog_name") - max_rows = arguments.get("max_rows", 100) + max_rows = arguments.get( + "max_rows", + configured_default_result_rows(self.connection_manager.config), + ) max_bytes = arguments.get("max_bytes") timeout = arguments.get("timeout", 30) diff --git a/doris_mcp_server/utils/config.py b/doris_mcp_server/utils/config.py index ba061b7..dcbd7c5 100644 --- a/doris_mcp_server/utils/config.py +++ b/doris_mcp_server/utils/config.py @@ -40,6 +40,7 @@ ABSOLUTE_MAX_RESULT_ROWS, DEFAULT_MAX_RESULT_BYTES, DEFAULT_MAX_RESULT_ROWS, + DEFAULT_RESULT_ROWS, MIN_RESULT_BYTES, ) from ..tools.tool_registry import ( @@ -507,6 +508,7 @@ class PerformanceConfig: # Concurrency control configuration max_concurrent_queries: int = 50 query_timeout: int = 300 + default_result_rows: int = DEFAULT_RESULT_ROWS max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES # Connection pool optimization configuration @@ -1045,6 +1047,12 @@ def from_env(cls, env_file: str | None = None) -> "DorisConfig": config.performance.query_timeout = int( os.getenv("QUERY_TIMEOUT", str(config.performance.query_timeout)) ) + config.performance.default_result_rows = int( + os.getenv( + "DEFAULT_RESULT_ROWS", + str(config.performance.default_result_rows), + ) + ) config.performance.max_result_bytes = int( os.getenv( "MAX_RESULT_BYTES", @@ -1363,6 +1371,7 @@ def to_dict(self) -> dict[str, Any]: "max_cache_size": self.performance.max_cache_size, "max_concurrent_queries": self.performance.max_concurrent_queries, "query_timeout": self.performance.query_timeout, + "default_result_rows": self.performance.default_result_rows, "max_result_bytes": self.performance.max_result_bytes, "connection_pool_size": self.performance.connection_pool_size, "idle_timeout": self.performance.idle_timeout, @@ -1520,6 +1529,14 @@ def validate(self) -> list[str]: f"{ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS} seconds" ) + if self.performance.default_result_rows <= 0: + errors.append("Default result rows must be greater than 0") + elif self.performance.default_result_rows > self.security.max_result_rows: + errors.append( + "Default result rows must not exceed the configured " + f"maximum result rows ({self.security.max_result_rows})" + ) + if not ( MIN_RESULT_BYTES <= self.performance.max_result_bytes diff --git a/test/integration/test_real_doris_transports.py b/test/integration/test_real_doris_transports.py index 3c1fcc9..2a61678 100644 --- a/test/integration/test_real_doris_transports.py +++ b/test/integration/test_real_doris_transports.py @@ -407,6 +407,7 @@ async def test_real_doris_result_boundaries_and_cancellation( environment.update( { "MAX_RESULT_ROWS": "5", + "DEFAULT_RESULT_ROWS": "2", "MAX_RESULT_BYTES": "256", "QUERY_TIMEOUT": "5", } @@ -431,9 +432,27 @@ async def test_real_doris_result_boundaries_and_cancellation( } query_schema = tools["exec_query"].input_schema["properties"] assert query_schema["max_rows"]["maximum"] == 5 + assert query_schema["max_rows"]["default"] == 2 assert query_schema["max_bytes"]["maximum"] == 256 assert query_schema["timeout"]["maximum"] == 5 + default_result = await client.call_tool( + "exec_query", + { + "sql": ( + "SELECT id " + f"FROM {doris_sandbox.qualified_table} ORDER BY id" + ), + "max_bytes": 256, + "timeout": 5, + }, + ) + assert isinstance(default_result.structured_content, dict) + default_payload = default_result.structured_content + assert default_result.is_error is False + assert len(default_payload["data"]) == 2 + assert default_payload["metadata"]["limits"]["max_rows"] == 2 + with pytest.raises(MCPError) as excessive_rows: await client.call_tool( "exec_query", diff --git a/test/utils/test_result_limits.py b/test/utils/test_result_limits.py index 0475b22..7e95073 100644 --- a/test/utils/test_result_limits.py +++ b/test/utils/test_result_limits.py @@ -17,7 +17,7 @@ """Tests for non-escalating query result and timeout budgets.""" from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest @@ -37,6 +37,7 @@ def _config( *, rows: int = 50, + default_rows: int = 10, result_bytes: int = 4096, timeout: int = 20, ) -> SimpleNamespace: @@ -46,6 +47,7 @@ def _config( enable_security_check=False, ), performance=SimpleNamespace( + default_result_rows=default_rows, max_result_bytes=result_bytes, query_timeout=timeout, max_cache_size=10, @@ -128,10 +130,58 @@ def test_exec_query_schema_advertises_effective_ceilings() -> None: properties = tool.input_schema["properties"] assert properties["max_rows"]["maximum"] == 50 + assert properties["max_rows"]["default"] == 10 assert properties["max_bytes"]["maximum"] == 4096 assert properties["timeout"]["maximum"] == 20 +def test_config_validation_rejects_default_rows_above_ceiling() -> None: + config = DorisConfig() + config.security.max_result_rows = 50 + config.performance.default_result_rows = 51 + + assert ( + "Default result rows must not exceed the configured maximum result rows (50)" + in config.validate() + ) + + +def test_environment_configures_default_rows_independently_from_ceiling( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + monkeypatch.setenv("MAX_RESULT_ROWS", "500") + monkeypatch.setenv("DEFAULT_RESULT_ROWS", "321") + monkeypatch.setenv("ADBC_DEFAULT_MAX_ROWS", "500") + + config = DorisConfig.from_env(str(tmp_path / "missing.env")) + + assert config.security.max_result_rows == 500 + assert config.performance.default_result_rows == 321 + assert config.validate() == [] + + +@pytest.mark.asyncio +async def test_exec_query_uses_configured_default_rows_when_argument_is_omitted() -> None: + connection_manager = Mock() + connection_manager.config = _config(default_rows=17) + manager = DorisToolsManager(connection_manager) + manager.metadata_extractor.exec_query_for_mcp = AsyncMock() + manager.metadata_extractor.exec_query_for_mcp.return_value = {"success": True} + + result = await manager._exec_query_tool({"sql": "SELECT 1"}) + + assert result == {"success": True} + manager.metadata_extractor.exec_query_for_mcp.assert_called_once_with( + "SELECT 1", + None, + None, + 17, + 30, + max_bytes=None, + ) + + def test_exec_adbc_query_schema_advertises_effective_ceilings() -> None: connection_manager = Mock() connection_manager.config = _config()