This is an automated email from the ASF dual-hosted git repository.
FreeOnePlus pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris-mcp-server.git
The following commit(s) were added to refs/heads/master by this push:
new 45e39a6 feat: enforce bounded query results (#150)
45e39a6 is described below
commit 45e39a6b56af5d7063fea8f16fcd03451d39fc6f
Author: Yijia Su <[email protected]>
AuthorDate: Thu Jul 30 14:10:37 2026 +0800
feat: enforce bounded query results (#150)
---
.env.example | 7 +-
CHANGELOG.md | 3 +
README.md | 12 +-
docker-compose.yml | 2 +
docs/tool-registry.md | 4 +-
doris_mcp_server/result_limits.py | 150 +++++++++++
doris_mcp_server/tools/tools_manager.py | 64 ++++-
doris_mcp_server/utils/adbc_query_tools.py | 330 +++++++++++++++++--------
doris_mcp_server/utils/config.py | 51 +++-
doris_mcp_server/utils/db.py | 127 +++++++++-
doris_mcp_server/utils/query_executor.py | 236 +++++++++++++-----
doris_mcp_server/utils/schema_extractor.py | 12 +-
test/integration/test_real_doris_transports.py | 139 ++++++++++-
test/tools/test_tools_operation_guard.py | 25 +-
test/utils/test_adbc_query_tools.py | 188 ++++++++++++++
test/utils/test_db.py | 104 ++++++++
test/utils/test_doris_user_pool_manager.py | 11 +-
test/utils/test_result_limits.py | 168 +++++++++++++
18 files changed, 1431 insertions(+), 202 deletions(-)
diff --git a/.env.example b/.env.example
index b878b74..4f80861 100644
--- a/.env.example
+++ b/.env.example
@@ -366,6 +366,7 @@
BLOCKED_KEYWORDS=DROP,CREATE,ALTER,TRUNCATE,DELETE,INSERT,UPDATE,GRANT,REVOKE,EX
# Query limits
MAX_QUERY_COMPLEXITY=100
+# Deployment ceiling; absolute hard cap: 100000
MAX_RESULT_ROWS=10000
# Data masking
@@ -382,7 +383,10 @@ MAX_CACHE_SIZE=1000
# Concurrency control
MAX_CONCURRENT_QUERIES=50
+# Deployment ceiling; absolute hard cap: 300 seconds
QUERY_TIMEOUT=300
+# UTF-8 JSON row-data budget; range: 256-16777216 bytes
+MAX_RESULT_BYTES=1048576
# Response content size limit (characters)
MAX_RESPONSE_CONTENT_SIZE=4096
@@ -394,7 +398,8 @@ MAX_RESPONSE_CONTENT_SIZE=4096
ADBC_ENABLED=true
# Default ADBC query parameters
-ADBC_DEFAULT_MAX_ROWS=100000
+# Must not exceed MAX_RESULT_ROWS
+ADBC_DEFAULT_MAX_ROWS=10000
ADBC_DEFAULT_TIMEOUT=60
# Format: "arrow", "pandas", "dict"
ADBC_DEFAULT_RETURN_FORMAT=arrow
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b8d6721..7026e3a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -42,6 +42,9 @@ under **Unreleased** until a new version is selected and
published.
- W3C `traceparent`, `tracestate`, and `baggage` propagation from request
`_meta`, with value-safe validation, credential-like baggage redaction,
per-request isolation, and no trace metadata in model-facing results.
+- Bounded query result streaming with deployment and absolute ceilings for
+ rows, serialized bytes, and execution time, plus cancellation-safe database
+ connection disposal.
- Real Doris process tests covering Streamable HTTP and stdio.
### Changed
diff --git a/README.md b/README.md
index 006da86..b83df17 100644
--- a/README.md
+++ b/README.md
@@ -330,9 +330,11 @@ cp .env.example .env
* `ENABLE_SECURITY_CHECK`: Enable/disable SQL security validation
(default: true)
* `BLOCKED_KEYWORDS`: Comma-separated list of blocked SQL keywords
* `ENABLE_MASKING`: Enable data masking (default: true)
- * `MAX_RESULT_ROWS`: Maximum result rows (default: 10000)
+ * `MAX_RESULT_ROWS`: Deployment ceiling for returned query rows
+ (default: 10000; absolute hard cap: 100000)
* **ADBC Configuration (New in v0.5.0)**:
- * `ADBC_DEFAULT_MAX_ROWS`: Default maximum rows for ADBC queries
(default: 100000)
+ * `ADBC_DEFAULT_MAX_ROWS`: Default maximum rows for ADBC queries
+ (default: 10000; cannot exceed `MAX_RESULT_ROWS`)
* `ADBC_DEFAULT_TIMEOUT`: Default ADBC query timeout in seconds
(default: 60)
* `ADBC_DEFAULT_RETURN_FORMAT`: Default return format -
arrow/pandas/dict (default: arrow)
* `ADBC_CONNECTION_TIMEOUT`: ADBC connection timeout in seconds
(default: 30)
@@ -341,6 +343,10 @@ cp .env.example .env
* `ENABLE_QUERY_CACHE`: Enable query caching (default: true)
* `CACHE_TTL`: Cache time-to-live in seconds (default: 300)
* `MAX_CONCURRENT_QUERIES`: Maximum concurrent queries (default: 50)
+ * `QUERY_TIMEOUT`: Deployment ceiling for query execution time
+ (default and absolute hard cap: 300 seconds)
+ * `MAX_RESULT_BYTES`: Deployment ceiling for UTF-8 JSON row data
+ (default: 1048576; allowed range: 256-16777216 bytes)
* `MAX_RESPONSE_CONTENT_SIZE`: Maximum response content size for LLM
compatibility (default: 4096, New in v0.4.0)
* **Enhanced Logging Configuration (Improved in v0.5.0)**:
* `LOG_LEVEL`: Log level (DEBUG/INFO/WARNING/ERROR, default: INFO)
@@ -1847,7 +1853,7 @@ the request metadata, HTTP headers, migration steps, and
deployment limits.
3. **Optional ADBC Customization**:
```bash
# Customize ADBC behavior (optional)
- ADBC_DEFAULT_MAX_ROWS=200000
+ ADBC_DEFAULT_MAX_ROWS=10000
ADBC_DEFAULT_TIMEOUT=120
ADBC_DEFAULT_RETURN_FORMAT=pandas # arrow/pandas/dict
```
diff --git a/docker-compose.yml b/docker-compose.yml
index 1fb1f17..87a6706 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -42,6 +42,8 @@ services:
- ENABLE_TOKEN_AUTH=true
- TOKEN_ADMIN_FILE=/run/secrets/mcp_static_token
- MAX_RESULT_ROWS=10000
+ - MAX_RESULT_BYTES=1048576
+ - QUERY_TIMEOUT=300
# Performance configuration
- ENABLE_QUERY_CACHE=true
diff --git a/docs/tool-registry.md b/docs/tool-registry.md
index 2001621..3d3e09a 100644
--- a/docs/tool-registry.md
+++ b/docs/tool-registry.md
@@ -4,7 +4,7 @@
| Tool | Policy | Risk | Handler | Audit event | Parameters |
|---|---|---|---|---|---|
-| `exec_query` | `query` | `query` | `_exec_query_tool` |
`mcp.tool.call.exec_query` | `catalog_name`, `db_name`, `max_rows`, `sql`*,
`timeout` |
+| `exec_query` | `query` | `query` | `_exec_query_tool` |
`mcp.tool.call.exec_query` | `catalog_name`, `db_name`, `max_bytes`,
`max_rows`, `sql`*, `timeout` |
| `get_table_schema` | `metadata` | `metadata` | `_get_table_schema_tool` |
`mcp.tool.call.get_table_schema` | `catalog_name`, `db_name`, `table_name`* |
| `get_db_table_list` | `metadata` | `metadata` | `_get_db_table_list_tool` |
`mcp.tool.call.get_db_table_list` | `catalog_name`, `db_name` |
| `get_db_list` | `metadata` | `metadata` | `_get_db_list_tool` |
`mcp.tool.call.get_db_list` | `catalog_name` |
@@ -27,7 +27,7 @@
| `analyze_data_flow_dependencies` | `restricted` | `high` |
`_analyze_data_flow_dependencies_tool` |
`mcp.tool.call.analyze_data_flow_dependencies` | `analysis_depth`,
`catalog_name`, `db_name`, `include_views`, `target_table` |
| `analyze_slow_queries_topn` | `restricted` | `high` |
`_analyze_slow_queries_topn_tool` | `mcp.tool.call.analyze_slow_queries_topn` |
`days`, `include_patterns`, `min_execution_time_ms`, `top_n` |
| `analyze_resource_growth_curves` | `restricted` | `high` |
`_analyze_resource_growth_curves_tool` |
`mcp.tool.call.analyze_resource_growth_curves` | `days`, `detailed_response`,
`include_predictions`, `resource_types` |
-| `exec_adbc_query` | `restricted` | `high` | `_exec_adbc_query_tool` |
`mcp.tool.call.exec_adbc_query` | `max_rows`, `return_format`, `sql`*,
`timeout` |
+| `exec_adbc_query` | `restricted` | `high` | `_exec_adbc_query_tool` |
`mcp.tool.call.exec_adbc_query` | `max_bytes`, `max_rows`, `return_format`,
`sql`*, `timeout` |
| `get_adbc_connection_info` | `restricted` | `high` |
`_get_adbc_connection_info_tool` | `mcp.tool.call.get_adbc_connection_info` |
None |
Required parameters are marked with `*`. Tool descriptions and JSON Schemas
are exposed directly by MCP `tools/list` from the same registry entries.
diff --git a/doris_mcp_server/result_limits.py
b/doris_mcp_server/result_limits.py
new file mode 100644
index 0000000..003d17a
--- /dev/null
+++ b/doris_mcp_server/result_limits.py
@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Central query-result and execution limits.
+
+Configuration limits are deployment ceilings. Tool arguments may request a
+smaller budget, but cannot raise those ceilings. Absolute constants remain a
+last line of defence if configuration objects are constructed without running
+``DorisConfig.validate()``.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+ABSOLUTE_MAX_RESULT_ROWS = 100_000
+ABSOLUTE_MAX_RESULT_BYTES = 16 * 1024 * 1024
+ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS = 300
+MIN_RESULT_BYTES = 256
+
+DEFAULT_MAX_RESULT_ROWS = 10_000
+DEFAULT_MAX_RESULT_BYTES = 1024 * 1024
+DEFAULT_QUERY_TIMEOUT_SECONDS = 300
+
+
+class ResultLimitError(ValueError):
+ """Raised when a caller asks for an invalid or excessive result budget."""
+
+
+@dataclass(frozen=True)
+class ResultLimits:
+ """Effective per-call query result and execution budgets."""
+
+ max_rows: int
+ max_bytes: int
+ timeout_seconds: int
+
+
+def _configured_positive_int(
+ value: object,
+ *,
+ default: int,
+ hard_maximum: int,
+) -> int:
+ if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+ return default
+ return min(value, hard_maximum)
+
+
+def configured_result_limits(config: object | None) -> ResultLimits:
+ """Read bounded deployment ceilings from a Doris configuration object."""
+ security = getattr(config, "security", None)
+ performance = getattr(config, "performance", None)
+ return ResultLimits(
+ max_rows=_configured_positive_int(
+ getattr(security, "max_result_rows", None),
+ default=DEFAULT_MAX_RESULT_ROWS,
+ hard_maximum=ABSOLUTE_MAX_RESULT_ROWS,
+ ),
+ max_bytes=_configured_positive_int(
+ getattr(performance, "max_result_bytes", None),
+ default=DEFAULT_MAX_RESULT_BYTES,
+ hard_maximum=ABSOLUTE_MAX_RESULT_BYTES,
+ ),
+ timeout_seconds=_configured_positive_int(
+ getattr(performance, "query_timeout", None),
+ default=DEFAULT_QUERY_TIMEOUT_SECONDS,
+ hard_maximum=ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS,
+ ),
+ )
+
+
+def _requested_limit(
+ value: object | None,
+ *,
+ name: str,
+ configured_maximum: int,
+ minimum: int = 1,
+) -> int:
+ if value is None:
+ return configured_maximum
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise ResultLimitError(f"{name} must be an integer")
+ if value < minimum:
+ raise ResultLimitError(f"{name} must be at least {minimum}")
+ if value > configured_maximum:
+ raise ResultLimitError(
+ f"{name} exceeds the configured maximum of {configured_maximum}"
+ )
+ return value
+
+
+def resolve_result_limits(
+ config: object | None,
+ *,
+ max_rows: object | None,
+ max_bytes: object | None,
+ timeout_seconds: object | None,
+) -> ResultLimits:
+ """Resolve request budgets without allowing configuration escalation."""
+ ceilings = configured_result_limits(config)
+ return ResultLimits(
+ max_rows=_requested_limit(
+ max_rows,
+ name="max_rows",
+ configured_maximum=ceilings.max_rows,
+ ),
+ max_bytes=_requested_limit(
+ max_bytes,
+ name="max_bytes",
+ configured_maximum=ceilings.max_bytes,
+ minimum=MIN_RESULT_BYTES,
+ ),
+ timeout_seconds=_requested_limit(
+ timeout_seconds,
+ name="timeout",
+ configured_maximum=ceilings.timeout_seconds,
+ ),
+ )
+
+
+def json_array_row_size(row: dict[str, Any], *, first: bool) -> int:
+ """Return a conservative UTF-8 JSON contribution for one result row.
+
+ ``ensure_ascii=True`` intentionally counts escaped non-ASCII data, which is
+ never smaller than the UTF-8 representation emitted by the MCP stack.
+ Array brackets are accounted for separately by the caller.
+ """
+ encoded = json.dumps(
+ row,
+ ensure_ascii=True,
+ separators=(",", ":"),
+ default=str,
+ ).encode("utf-8")
+ return len(encoded) + (0 if first else 1)
diff --git a/doris_mcp_server/tools/tools_manager.py
b/doris_mcp_server/tools/tools_manager.py
index b103fe2..ab556fd 100644
--- a/doris_mcp_server/tools/tools_manager.py
+++ b/doris_mcp_server/tools/tools_manager.py
@@ -31,6 +31,7 @@ from ..auth.operation_policy import (
authorize_operation,
filter_tools_for_auth_context,
)
+from ..result_limits import configured_result_limits
from ..utils.adbc_query_tools import DorisADBCQueryTools
from ..utils.analysis_tools import MemoryTracker, SQLAnalyzer, TableAnalyzer
from ..utils.config import ADBCConfig
@@ -120,6 +121,15 @@ class DorisToolsManager:
connection_manager = getattr(self, "connection_manager", None)
config = getattr(connection_manager, "config", None)
adbc_config = getattr(config, "adbc", None) or ADBCConfig()
+ result_limits = configured_result_limits(config)
+ adbc_default_max_rows = min(
+ adbc_config.default_max_rows,
+ result_limits.max_rows,
+ )
+ adbc_default_timeout = min(
+ adbc_config.default_timeout,
+ result_limits.timeout_seconds,
+ )
tools = [
Tool(
@@ -136,6 +146,8 @@ class DorisToolsManager:
- max_rows (integer) [Optional] - Maximum number of rows to return, default 100
+- max_bytes (integer) [Optional] - Maximum UTF-8 JSON bytes for returned row
data
+
- timeout (integer) [Optional] - Query timeout in seconds, default 30
""",
input_schema={
@@ -157,11 +169,22 @@ class DorisToolsManager:
"type": "integer",
"description": "Maximum number of rows to return",
"default": 100,
+ "minimum": 1,
+ "maximum": result_limits.max_rows,
+ },
+ "max_bytes": {
+ "type": "integer",
+ "description": "Maximum UTF-8 JSON bytes for
returned row data",
+ "default": result_limits.max_bytes,
+ "minimum": 256,
+ "maximum": result_limits.max_bytes,
},
"timeout": {
"type": "integer",
"description": "Timeout in seconds",
"default": 30,
+ "minimum": 1,
+ "maximum": result_limits.timeout_seconds,
},
},
"required": ["sql"],
@@ -936,8 +959,9 @@ class DorisToolsManager:
[Parameter Content]:
- sql (string) [Required] - SQL statement to execute
-- max_rows (integer) [Optional] - Maximum number of rows to return, default is
{adbc_config.default_max_rows}
-- timeout (integer) [Optional] - Query timeout in seconds, default is
{adbc_config.default_timeout}
+- max_rows (integer) [Optional] - Maximum number of rows to return, default is
{adbc_default_max_rows}
+- max_bytes (integer) [Optional] - Maximum UTF-8 JSON bytes for returned row
data, default is {result_limits.max_bytes}
+- timeout (integer) [Optional] - Query timeout in seconds, default is
{adbc_default_timeout}
- return_format (string) [Optional] - Format for returned data, default is
"{adbc_config.default_return_format}"
* "arrow": Return Arrow format with metadata
* "pandas": Return Pandas DataFrame format
@@ -958,12 +982,23 @@ class DorisToolsManager:
"max_rows": {
"type": "integer",
"description": "Maximum number of rows to return",
- "default": adbc_config.default_max_rows,
+ "minimum": 1,
+ "maximum": result_limits.max_rows,
+ "default": adbc_default_max_rows,
+ },
+ "max_bytes": {
+ "type": "integer",
+ "description": "Maximum UTF-8 JSON result bytes",
+ "minimum": 256,
+ "maximum": result_limits.max_bytes,
+ "default": result_limits.max_bytes,
},
"timeout": {
"type": "integer",
"description": "Query timeout in seconds",
- "default": adbc_config.default_timeout,
+ "minimum": 1,
+ "maximum": result_limits.timeout_seconds,
+ "default": adbc_default_timeout,
},
"return_format": {
"type": "string",
@@ -1092,11 +1127,17 @@ No parameters required. Returns connection status,
configuration, and diagnostic
db_name = arguments.get("db_name")
catalog_name = arguments.get("catalog_name")
max_rows = arguments.get("max_rows", 100)
+ max_bytes = arguments.get("max_bytes")
timeout = arguments.get("timeout", 30)
# Delegate to metadata extractor for processing
return await self.metadata_extractor.exec_query_for_mcp(
- sql, db_name, catalog_name, max_rows, timeout
+ sql,
+ db_name,
+ catalog_name,
+ max_rows,
+ timeout,
+ max_bytes=max_bytes,
)
async def _get_table_schema_tool(self, arguments: dict[str, Any]) ->
dict[str, Any]:
@@ -1588,13 +1629,18 @@ No parameters required. Returns connection status,
configuration, and diagnostic
async def _exec_adbc_query_tool(self, arguments: dict[str, Any]) ->
dict[str, Any]:
"""ADBC query execution tool routing"""
sql = self._required_string(arguments, "sql")
- max_rows = arguments.get("max_rows", 100000)
- timeout = arguments.get("timeout", 60)
- return_format = arguments.get("return_format", "arrow")
+ max_rows = arguments.get("max_rows")
+ max_bytes = arguments.get("max_bytes")
+ timeout = arguments.get("timeout")
+ return_format = arguments.get("return_format")
# Delegate to ADBC query tools for processing
return await self.adbc_query_tools.exec_adbc_query(
- sql, max_rows, timeout, return_format
+ sql,
+ max_rows=max_rows,
+ timeout=timeout,
+ return_format=return_format,
+ max_bytes=max_bytes,
)
async def _get_adbc_connection_info_tool(
diff --git a/doris_mcp_server/utils/adbc_query_tools.py
b/doris_mcp_server/utils/adbc_query_tools.py
index 478532b..7de8b44 100644
--- a/doris_mcp_server/utils/adbc_query_tools.py
+++ b/doris_mcp_server/utils/adbc_query_tools.py
@@ -20,13 +20,19 @@ Apache Doris ADBC Query Tools
High-performance data querying using Apache Arrow Flight SQL protocol
"""
+import asyncio
import os
import socket
import time
from datetime import datetime
-from importlib.util import find_spec
-from typing import Any, cast
-
+from typing import Any
+
+from ..result_limits import (
+ ResultLimitError,
+ ResultLimits,
+ json_array_row_size,
+ resolve_result_limits,
+)
from ..utils.db import DorisConnectionManager
from ..utils.logger import get_logger
from ..utils.security import AuthContext
@@ -61,28 +67,6 @@ def _convert_numpy_types(obj: Any) -> Any:
return obj
-def _convert_dataframe_to_json_serializable(
- df: Any,
-) -> list[dict[str, Any]]:
- """Convert DataFrame to JSON serializable format"""
- if find_spec("numpy") is None or find_spec("pandas") is None:
- # Fallback to basic dict conversion
- return cast(list[dict[str, Any]], df.to_dict("records"))
-
- # Convert DataFrame to records
- records = df.to_dict('records')
-
- # Convert each record's values
- converted_records = []
- for record in records:
- converted_record = {}
- for key, value in record.items():
- converted_record[key] = _convert_numpy_types(value)
- converted_records.append(converted_record)
-
- return converted_records
-
-
class DorisADBCQueryTools:
"""ADBC Query Tools for high-performance data transfer using Arrow Flight
SQL"""
@@ -91,13 +75,16 @@ class DorisADBCQueryTools:
self.adbc_client: Any | None = None
self.flight_sql_module: Any | None = None
self.adbc_manager_module: Any | None = None
+ # ADBC connections and cursors are not safe to replace concurrently.
+ self._query_lock = asyncio.Lock()
async def exec_adbc_query(
self,
sql: str,
max_rows: int | None = None,
timeout: int | None = None,
- return_format: str | None = None
+ return_format: str | None = None,
+ max_bytes: int | None = None,
) -> dict[str, Any]:
"""
Execute SQL query using ADBC (Arrow Flight SQL) protocol
@@ -107,18 +94,42 @@ class DorisADBCQueryTools:
max_rows: Maximum number of rows to return (uses config default if
None)
timeout: Query timeout in seconds (uses config default if None)
return_format: Format for returned data ("arrow", "pandas",
"dict", uses config default if None)
+ max_bytes: Maximum UTF-8 JSON bytes for returned row data
Returns:
Query results in specified format with metadata
"""
try:
- start_time = time.time()
-
# Use configuration defaults if parameters not specified
adbc_config = self.connection_manager.config.adbc
- max_rows = max_rows if max_rows is not None else
adbc_config.default_max_rows
- timeout = timeout if timeout is not None else
adbc_config.default_timeout
- return_format = return_format if return_format is not None else
adbc_config.default_return_format
+ try:
+ limits = resolve_result_limits(
+ self.connection_manager.config,
+ max_rows=(
+ max_rows
+ if max_rows is not None
+ else adbc_config.default_max_rows
+ ),
+ max_bytes=max_bytes,
+ timeout_seconds=(
+ timeout
+ if timeout is not None
+ else adbc_config.default_timeout
+ ),
+ )
+ except ResultLimitError as exc:
+ return {
+ "success": False,
+ "error": str(exc),
+ "error_type": "invalid_result_limits",
+ "timestamp": datetime.now().isoformat(),
+ }
+ return_format = (
+ return_format
+ if return_format is not None
+ else adbc_config.default_return_format
+ )
+ start_time = time.time()
# Step 1: Check environment variables and port availability
port_check_result = await self._check_arrow_flight_ports()
@@ -130,15 +141,21 @@ class DorisADBCQueryTools:
if not import_result["success"]:
return import_result
- # Step 3: Create ADBC connection
- connection_result = await self._create_adbc_connection()
- if not connection_result["success"]:
- return connection_result
-
- # Step 4: Execute query using ADBC
- query_result = await self._execute_query_with_adbc(
- sql, max_rows, timeout, return_format
- )
+ async with self._query_lock:
+ # Step 3: Create a per-call ADBC connection.
+ connection_result = await self._create_adbc_connection()
+ if not connection_result["success"]:
+ return connection_result
+
+ try:
+ # Step 4: Execute query using ADBC.
+ query_result = await self._execute_query_with_adbc(
+ sql,
+ limits,
+ return_format,
+ )
+ finally:
+ await self._close_adbc_client()
execution_time = time.time() - start_time
@@ -158,6 +175,16 @@ class DorisADBCQueryTools:
"timestamp": datetime.now().isoformat()
}
+ async def _close_adbc_client(self) -> None:
+ """Close and clear the current ADBC connection without blocking MCP."""
+ client, self.adbc_client = self.adbc_client, None
+ if client is None:
+ return
+ try:
+ await asyncio.to_thread(client.close)
+ except Exception as exc:
+ logger.debug(f"Failed to close ADBC client: {exc}")
+
async def _check_arrow_flight_ports(self) -> dict[str, Any]:
"""Check Arrow Flight SQL port configuration and availability"""
try:
@@ -397,11 +424,10 @@ class DorisADBCQueryTools:
async def _execute_query_with_adbc(
self,
sql: str,
- max_rows: int,
- timeout: int,
- return_format: str
+ limits: ResultLimits,
+ return_format: str,
) -> dict[str, Any]:
- """Execute query using ADBC"""
+ """Execute an ADBC query with bounded output and real cancellation."""
try:
if not self.adbc_client:
return {
@@ -425,88 +451,174 @@ class DorisADBCQueryTools:
}
cursor = self.adbc_client.cursor()
- start_time = time.time()
+ worker = asyncio.create_task(
+ asyncio.to_thread(
+ self._execute_query_with_adbc_sync,
+ cursor,
+ sql,
+ limits,
+ return_format,
+ )
+ )
+ try:
+ return await asyncio.wait_for(
+ asyncio.shield(worker),
+ timeout=limits.timeout_seconds,
+ )
+ except TimeoutError:
+ await self._cancel_adbc_worker(cursor, worker)
+ return {
+ "success": False,
+ "error": (
+ "ADBC query execution timed out after "
+ f"{limits.timeout_seconds} seconds"
+ ),
+ "error_type": "timeout",
+ "sql": sql,
+ }
+ except asyncio.CancelledError:
+ await self._cancel_adbc_worker(cursor, worker)
+ raise
- # Execute query
- cursor.execute(sql)
+ except Exception as e:
+ logger.error(f"ADBC query execution failed: {str(e)}")
+ return {
+ "success": False,
+ "error": f"ADBC query execution failed: {str(e)}",
+ "error_type": "query_execution_error",
+ "sql": sql
+ }
- # Get results based on return format
- if return_format == "arrow":
- # Return Arrow format
- arrow_data = cursor.fetchallarrow()
+ async def _cancel_adbc_worker(
+ self,
+ cursor: Any,
+ worker: asyncio.Task[dict[str, Any]],
+ ) -> None:
+ """Request driver cancellation and reap the worker thread."""
+ try:
+ await asyncio.wait_for(
+ asyncio.to_thread(cursor.adbc_cancel),
+ timeout=2,
+ )
+ except Exception as exc:
+ logger.warning(f"Failed to cancel ADBC query: {exc}")
- # Limit rows
- if len(arrow_data) > max_rows:
- arrow_data = arrow_data.slice(0, max_rows)
+ try:
+ await asyncio.wait_for(asyncio.shield(worker), timeout=5)
+ except (asyncio.CancelledError, Exception) as exc:
+ # The driver owns the worker thread. Keep the result consumed even
+ # if a non-cooperative implementation takes longer to unwind.
+ logger.debug(f"ADBC query worker did not stop promptly: {exc}")
+ worker.add_done_callback(self._consume_adbc_worker_result)
+
+ @staticmethod
+ def _consume_adbc_worker_result(worker: asyncio.Task[dict[str, Any]]) ->
None:
+ try:
+ worker.result()
+ except (asyncio.CancelledError, Exception):
+ pass
+
+ @staticmethod
+ def _execute_query_with_adbc_sync(
+ cursor: Any,
+ sql: str,
+ limits: ResultLimits,
+ return_format: str,
+ ) -> dict[str, Any]:
+ """Run blocking ADBC calls in a worker and stream bounded batches."""
+ start_time = time.time()
+ rows: list[dict[str, Any]] = []
+ result_bytes = 2 # JSON array brackets
+ truncated = False
+ truncation_reason: str | None = None
- # Convert Arrow data to serializable format
- preview_df = arrow_data.to_pandas().head(10) if
len(arrow_data) > 0 else None
+ try:
+ cursor.execute(sql)
+ description = cursor.description or []
+ column_names = [str(column[0]) for column in description]
+ column_types = [str(column[1]) for column in description]
+
+ while not truncated:
+ fetch_size = min(256, limits.max_rows + 1 - len(rows))
+ batch = cursor.fetchmany(max(fetch_size, 1))
+ if not batch:
+ break
+
+ for raw_row in batch:
+ if len(rows) >= limits.max_rows:
+ truncated = True
+ truncation_reason = "max_rows"
+ break
+
+ if isinstance(raw_row, dict):
+ row = {
+ str(key): _convert_numpy_types(value)
+ for key, value in raw_row.items()
+ }
+ else:
+ row = {
+ name: _convert_numpy_types(value)
+ for name, value in zip(
+ column_names,
+ raw_row,
+ strict=False,
+ )
+ }
+
+ contribution = json_array_row_size(
+ row,
+ first=not rows,
+ )
+ if result_bytes + contribution > limits.max_bytes:
+ truncated = True
+ truncation_reason = "max_bytes"
+ break
+
+ rows.append(row)
+ result_bytes += contribution
+
+ common_result = {
+ "format": return_format,
+ "num_rows": len(rows),
+ "num_columns": len(column_names),
+ "column_names": column_names,
+ "column_types": column_types,
+ }
+ if return_format == "arrow":
result_data = {
- "format": "arrow",
- "num_rows": len(arrow_data),
- "num_columns": len(arrow_data.schema),
- "column_names": arrow_data.schema.names,
- "column_types": [str(field.type) for field in
arrow_data.schema],
- "data_preview":
_convert_dataframe_to_json_serializable(preview_df) if preview_df is not None
else [],
- "total_bytes": arrow_data.nbytes if hasattr(arrow_data,
'nbytes') else 0
+ **common_result,
+ "data_preview": rows[:10],
+ "total_bytes": result_bytes,
}
-
elif return_format == "pandas":
- # Return Pandas DataFrame
- df = cursor.fetch_df()
-
- # Limit rows
- if len(df) > max_rows:
- df = df.head(max_rows)
-
result_data = {
- "format": "pandas",
- "num_rows": len(df),
- "num_columns": len(df.columns),
- "column_names": df.columns.tolist(),
- "column_types": df.dtypes.astype(str).tolist(),
- "data": _convert_dataframe_to_json_serializable(df),
- "memory_usage": int(df.memory_usage(deep=True).sum())
+ **common_result,
+ "data": rows,
+ "memory_usage": result_bytes,
}
-
- else: # return_format == "dict"
- # Return dictionary format
- arrow_data = cursor.fetchallarrow()
- df = arrow_data.to_pandas()
-
- # Limit rows
- if len(df) > max_rows:
- df = df.head(max_rows)
-
+ else:
result_data = {
- "format": "dict",
- "num_rows": len(df),
- "num_columns": len(df.columns),
- "column_names": df.columns.tolist(),
- "column_types": df.dtypes.astype(str).tolist(),
- "data": _convert_dataframe_to_json_serializable(df)
+ **common_result,
+ "data": rows,
}
- execution_time = time.time() - start_time
-
- cursor.close()
-
return {
"success": True,
"result": result_data,
- "execution_time": round(execution_time, 3),
+ "execution_time": round(time.time() - start_time, 3),
"sql": sql,
- "max_rows_applied": len(result_data.get("data", [])) >=
max_rows
- }
-
- except Exception as e:
- logger.error(f"ADBC query execution failed: {str(e)}")
- return {
- "success": False,
- "error": f"ADBC query execution failed: {str(e)}",
- "error_type": "query_execution_error",
- "sql": sql
+ "max_rows_applied": truncation_reason == "max_rows",
+ "result_bytes": result_bytes,
+ "truncated": truncated,
+ "truncation_reason": truncation_reason,
+ "limits": {
+ "max_rows": limits.max_rows,
+ "max_bytes": limits.max_bytes,
+ "timeout_seconds": limits.timeout_seconds,
+ },
}
+ finally:
+ cursor.close()
async def get_adbc_connection_info(self) -> dict[str, Any]:
"""Get ADBC connection information and status"""
diff --git a/doris_mcp_server/utils/config.py b/doris_mcp_server/utils/config.py
index d41359d..ba061b7 100644
--- a/doris_mcp_server/utils/config.py
+++ b/doris_mcp_server/utils/config.py
@@ -34,6 +34,14 @@ from urllib.parse import urlparse
from dotenv import load_dotenv
from .._version import __version__
+from ..result_limits import (
+ ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS,
+ ABSOLUTE_MAX_RESULT_BYTES,
+ ABSOLUTE_MAX_RESULT_ROWS,
+ DEFAULT_MAX_RESULT_BYTES,
+ DEFAULT_MAX_RESULT_ROWS,
+ MIN_RESULT_BYTES,
+)
from ..tools.tool_registry import (
DORIS_OAUTH_EXPLAIN_TOOL_SET,
DORIS_OAUTH_METADATA_TOOL_NAMES,
@@ -499,6 +507,7 @@ class PerformanceConfig:
# Concurrency control configuration
max_concurrent_queries: int = 50
query_timeout: int = 300
+ max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES
# Connection pool optimization configuration
connection_pool_size: int = 20
@@ -540,7 +549,7 @@ class ADBCConfig:
"""ADBC (Arrow Flight SQL) configuration"""
# Default query parameters
- default_max_rows: int = 100000
+ default_max_rows: int = DEFAULT_MAX_RESULT_ROWS
default_timeout: int = 60
default_return_format: str = "arrow" # "arrow", "pandas", "dict"
@@ -1036,6 +1045,12 @@ class DorisConfig:
config.performance.query_timeout = int(
os.getenv("QUERY_TIMEOUT", str(config.performance.query_timeout))
)
+ config.performance.max_result_bytes = int(
+ os.getenv(
+ "MAX_RESULT_BYTES",
+ str(config.performance.max_result_bytes),
+ )
+ )
config.performance.max_response_content_size = int(
os.getenv("MAX_RESPONSE_CONTENT_SIZE",
str(config.performance.max_response_content_size))
)
@@ -1348,6 +1363,7 @@ class DorisConfig:
"max_cache_size": self.performance.max_cache_size,
"max_concurrent_queries":
self.performance.max_concurrent_queries,
"query_timeout": self.performance.query_timeout,
+ "max_result_bytes": self.performance.max_result_bytes,
"connection_pool_size": self.performance.connection_pool_size,
"idle_timeout": self.performance.idle_timeout,
"max_response_content_size":
self.performance.max_response_content_size,
@@ -1480,6 +1496,11 @@ class DorisConfig:
if self.security.max_result_rows <= 0:
errors.append("Maximum result rows must be greater than 0")
+ elif self.security.max_result_rows > ABSOLUTE_MAX_RESULT_ROWS:
+ errors.append(
+ "Maximum result rows must not exceed "
+ f"{ABSOLUTE_MAX_RESULT_ROWS}"
+ )
# Validate performance configuration
if self.performance.cache_ttl <= 0:
@@ -1490,6 +1511,24 @@ class DorisConfig:
if self.performance.query_timeout <= 0:
errors.append("Query timeout must be greater than 0")
+ elif (
+ self.performance.query_timeout
+ > ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS
+ ):
+ errors.append(
+ "Query timeout must not exceed "
+ f"{ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS} seconds"
+ )
+
+ if not (
+ MIN_RESULT_BYTES
+ <= self.performance.max_result_bytes
+ <= ABSOLUTE_MAX_RESULT_BYTES
+ ):
+ errors.append(
+ "Maximum result bytes must be in the range "
+ f"{MIN_RESULT_BYTES}-{ABSOLUTE_MAX_RESULT_BYTES}"
+ )
# Validate data quality configuration
if self.data_quality.max_columns_per_batch <= 0:
@@ -1542,9 +1581,19 @@ class DorisConfig:
# Validate ADBC configuration
if self.adbc.default_max_rows <= 0:
errors.append("ADBC default max rows must be greater than 0")
+ elif self.adbc.default_max_rows > self.security.max_result_rows:
+ errors.append(
+ "ADBC default max rows must not exceed the configured "
+ f"maximum result rows ({self.security.max_result_rows})"
+ )
if self.adbc.default_timeout <= 0:
errors.append("ADBC default timeout must be greater than 0")
+ elif self.adbc.default_timeout > self.performance.query_timeout:
+ errors.append(
+ "ADBC default timeout must not exceed the configured query "
+ f"timeout ({self.performance.query_timeout} seconds)"
+ )
if self.adbc.default_return_format not in ["arrow", "pandas", "dict"]:
errors.append("ADBC default return format must be one of arrow,
pandas, or dict")
diff --git a/doris_mcp_server/utils/db.py b/doris_mcp_server/utils/db.py
index 0285d4d..71760c6 100644
--- a/doris_mcp_server/utils/db.py
+++ b/doris_mcp_server/utils/db.py
@@ -41,6 +41,7 @@ from typing import TYPE_CHECKING, Any, TypedDict
import aiomysql
from aiomysql import Connection, Pool
+from ..result_limits import json_array_row_size
from .config import DorisConfig
from .datetime_utils import utc_now
from .logger import get_logger
@@ -174,6 +175,8 @@ class DorisConnection:
auth_context: AuthContext | None = None,
*,
mask_result: bool = True,
+ max_rows: int | None = None,
+ max_bytes: int | None = None,
) -> QueryResult:
"""Execute SQL after validation, with optional result masking.
@@ -199,14 +202,37 @@ class DorisConnection:
"blocked_operations": validation_result.blocked_operations,
}
- async with self.connection.cursor(aiomysql.DictCursor) as cursor:
+ bounded_result = max_rows is not None or max_bytes is not None
+ cursor_type = (
+ aiomysql.SSDictCursor if bounded_result else
aiomysql.DictCursor
+ )
+ async with self.connection.cursor(cursor_type) as cursor:
await cursor.execute(sql, params)
# cursor.description is set by the DB driver for any statement
that returns rows,
# avoiding a brittle hardcoded keyword list (e.g. missing
WITH/CTE, comments before keywords).
+ result_bytes: int | None = None
+ truncated = False
+ truncation_reason: str | None = None
if cursor.description:
- data = await cursor.fetchall()
- row_count = len(data)
+ if bounded_result:
+ (
+ final_data,
+ result_bytes,
+ truncated,
+ truncation_reason,
+ ) = await self._fetch_bounded_rows(
+ cursor,
+ auth_context=auth_context,
+ mask_result=mask_result,
+ max_rows=max_rows,
+ max_bytes=max_bytes,
+ )
+ data = final_data
+ row_count = len(final_data)
+ else:
+ data = await cursor.fetchall()
+ row_count = len(data)
else:
data = []
row_count = cursor.rowcount
@@ -223,6 +249,8 @@ class DorisConnection:
# If security manager exists and has auth context, apply data
masking
final_data = list(data) if data else []
if (
+ not bounded_result
+ and
self.security_manager
and auth_context
and final_data
@@ -240,6 +268,21 @@ class DorisConnection:
}
if security_result:
metadata["security_check"] = security_result
+ if bounded_result:
+ metadata.update(
+ {
+ "result_bytes": result_bytes or 2,
+ "truncated": truncated,
+ "truncation_reason": truncation_reason,
+ }
+ )
+
+ if truncated:
+ # An unbuffered cursor still has unread rows. Closing the
+ # physical connection avoids draining an attacker-sized
+ # result and prevents it from returning to the pool.
+ self.is_healthy = False
+ await self.connection.ensure_closed()
return QueryResult(
data=final_data,
@@ -249,11 +292,78 @@ class DorisConnection:
sql=sql,
)
+ except asyncio.CancelledError:
+ # MCP cancellation and asyncio timeouts both arrive here as task
+ # cancellation. A running MySQL command cannot safely be returned
+ # to the pool, so terminate its physical connection first.
+ self.is_healthy = False
+ try:
+ await self.connection.ensure_closed()
+ finally:
+ raise
except Exception as e:
self.is_healthy = False
logging.error(f"Query execution failed: {e}")
raise
+ async def _fetch_bounded_rows(
+ self,
+ cursor: Any,
+ *,
+ auth_context: AuthContext | None,
+ mask_result: bool,
+ max_rows: int | None,
+ max_bytes: int | None,
+ ) -> tuple[list[dict[str, Any]], int, bool, str | None]:
+ """Fetch a result incrementally under row and serialized-byte
budgets."""
+ row_budget = max_rows if max_rows is not None else 100_000
+ byte_budget = max_bytes if max_bytes is not None else 16 * 1024 * 1024
+ final_data: list[dict[str, Any]] = []
+ result_bytes = 2 # JSON array brackets.
+ truncation_reason: str | None = None
+
+ while truncation_reason is None:
+ remaining_rows = row_budget - len(final_data)
+ fetch_size = min(128, remaining_rows + 1)
+ rows = list(await cursor.fetchmany(fetch_size))
+ if not rows:
+ break
+
+ processed_rows = [dict(row) for row in rows]
+ if (
+ self.security_manager
+ and auth_context
+ and processed_rows
+ and mask_result
+ ):
+ processed_rows = list(
+ await self.security_manager.apply_data_masking(
+ processed_rows,
+ auth_context,
+ )
+ )
+
+ for row in processed_rows:
+ if len(final_data) >= row_budget:
+ truncation_reason = "row_limit"
+ break
+ row_size = json_array_row_size(row, first=not final_data)
+ if result_bytes + row_size > byte_budget:
+ truncation_reason = "byte_limit"
+ break
+ final_data.append(row)
+ result_bytes += row_size
+
+ if len(rows) < fetch_size:
+ break
+
+ return (
+ final_data,
+ result_bytes,
+ truncation_reason is not None,
+ truncation_reason,
+ )
+
async def ping(self) -> bool:
"""Check connection health status with enhanced at_eof error
detection"""
try:
@@ -2460,6 +2570,9 @@ class DorisConnectionManager:
sql: str,
params: Mapping[str, Any] | tuple[Any, ...] | None = None,
auth_context: AuthContext | None = None,
+ *,
+ max_rows: int | None = None,
+ max_bytes: int | None = None,
) -> QueryResult:
"""Execute query using the same routed acquire/release contract as
get_connection()."""
connection = None
@@ -2471,7 +2584,13 @@ class DorisConnectionManager:
)
# Execute query
- result = await connection.execute(sql, params,
effective_auth_context)
+ result = await connection.execute(
+ sql,
+ params,
+ effective_auth_context,
+ max_rows=max_rows,
+ max_bytes=max_bytes,
+ )
return result
diff --git a/doris_mcp_server/utils/query_executor.py
b/doris_mcp_server/utils/query_executor.py
index 77ead62..e30c7b6 100644
--- a/doris_mcp_server/utils/query_executor.py
+++ b/doris_mcp_server/utils/query_executor.py
@@ -33,6 +33,12 @@ from typing import TYPE_CHECKING, Any, cast
import sqlparse
+from ..result_limits import (
+ ResultLimitError,
+ ResultLimits,
+ configured_result_limits,
+ resolve_result_limits,
+)
from .auth_credentials import EMPTY_CREDENTIAL
from .datetime_utils import utc_now
from .db import (
@@ -62,6 +68,8 @@ class QueryRequest:
user_id: str
parameters: dict[str, Any] | None = None
timeout: int | None = None
+ max_rows: int | None = None
+ max_bytes: int | None = None
cache_enabled: bool = True
@@ -338,7 +346,11 @@ class DorisQueryExecutor:
config: Any | None = None,
) -> None:
self.connection_manager = connection_manager
- self.config = config or self._create_default_config()
+ self.config = (
+ config
+ or getattr(connection_manager, "config", None)
+ or self._create_default_config()
+ )
self.logger = get_logger(__name__)
# Initialize components
@@ -359,6 +371,7 @@ class DorisQueryExecutor:
self.max_concurrent_queries = getattr(
getattr(self.config, 'performance', None),
'max_concurrent_queries', 50
) if hasattr(self.config, 'performance') else 50
+ self.configured_result_limits = configured_result_limits(self.config)
# Background tasks
self._background_tasks: list[asyncio.Task[None]] = []
@@ -368,12 +381,19 @@ class DorisQueryExecutor:
class DefaultConfig:
def __init__(self) -> None:
self.performance = DefaultPerformanceConfig()
+ self.security = DefaultSecurityConfig()
class DefaultPerformanceConfig:
def __init__(self) -> None:
self.max_cache_size = 1000
self.cache_ttl = 300
self.max_concurrent_queries = 50
+ self.query_timeout = 300
+ self.max_result_bytes = 1024 * 1024
+
+ class DefaultSecurityConfig:
+ def __init__(self) -> None:
+ self.max_result_rows = 10_000
return DefaultConfig()
@@ -475,13 +495,27 @@ class DorisQueryExecutor:
if query_request.timeout:
try:
result = await asyncio.wait_for(
-
self.connection_manager.execute_query(query_request.session_id, optimized_sql,
query_request.parameters, auth_context),
+ self.connection_manager.execute_query(
+ query_request.session_id,
+ optimized_sql,
+ query_request.parameters,
+ auth_context,
+ max_rows=query_request.max_rows,
+ max_bytes=query_request.max_bytes,
+ ),
timeout=query_request.timeout
)
except TimeoutError:
raise Exception(f"Query timeout after {query_request.timeout}
seconds")
else:
- result = await
self.connection_manager.execute_query(query_request.session_id, optimized_sql,
query_request.parameters, auth_context)
+ result = await self.connection_manager.execute_query(
+ query_request.session_id,
+ optimized_sql,
+ query_request.parameters,
+ auth_context,
+ max_rows=query_request.max_rows,
+ max_bytes=query_request.max_bytes,
+ )
return result
@@ -505,6 +539,8 @@ class DorisQueryExecutor:
async def execute_batch_sqls_for_mcp(
self,
sqls: list[str],
+ max_rows: int = 1000,
+ max_bytes: int | None = None,
timeout: int = 30,
session_id: str = "mcp_session",
user_id: str = "mcp_user",
@@ -517,35 +553,53 @@ class DorisQueryExecutor:
"error": "SQL query is required",
"data": None
}
- query_requests = [
- QueryRequest(
- sql=sql,
- session_id=session_id,
- user_id=user_id,
- timeout=timeout,
- cache_enabled=False
- )
- for sql in sqls
- ]
- query_results = await self.execute_batch_queries(query_requests,
auth_context)
+ limits = resolve_result_limits(
+ self.config,
+ max_rows=max_rows,
+ max_bytes=max_bytes,
+ timeout_seconds=timeout,
+ )
+ query_results: list[QueryResult] = []
+ remaining_rows = limits.max_rows
+ remaining_bytes = limits.max_bytes
+ async with asyncio.timeout(limits.timeout_seconds):
+ for sql in sqls:
+ if remaining_rows <= 0 or remaining_bytes < 256:
+ break
+ result = await self.execute_query(
+ QueryRequest(
+ sql=sql,
+ session_id=session_id,
+ user_id=user_id,
+ timeout=None,
+ max_rows=remaining_rows,
+ max_bytes=remaining_bytes,
+ cache_enabled=False,
+ ),
+ auth_context,
+ )
+ query_results.append(result)
+ remaining_rows -= result.row_count
+ remaining_bytes -= int(result.metadata.get("result_bytes", 2))
+
# Serialize data for JSON response
results = [
- {
- "data": [self._serialize_row_data(data) for data in
result.data],
- "row_count": result.row_count,
- "execution_time": result.execution_time,
- "metadata": {
- "columns": result.metadata.get("columns", []),
- "query": result.sql
- }
- }
+ self._query_result_payload(result, limits)
for result in query_results
]
return {
"success": True,
"multiple_results": True,
- "results": results
+ "results": results,
+ "metadata": {
+ "truncated": len(query_results) < len(sqls),
+ "limits": {
+ "max_rows": limits.max_rows,
+ "max_bytes": limits.max_bytes,
+ "timeout_seconds": limits.timeout_seconds,
+ },
+ },
}
async def execute_batch_queries(
@@ -630,7 +684,16 @@ class DorisQueryExecutor:
if callable(release_connection):
await release_connection(session_id, connection)
- def _query_result_payload(self, result: QueryResult) -> dict[str, Any]:
+ def _query_result_payload(
+ self,
+ result: QueryResult,
+ limits: ResultLimits,
+ ) -> dict[str, Any]:
+ boundary_metadata = {
+ key: result.metadata.get(key)
+ for key in ("result_bytes", "truncated", "truncation_reason")
+ if key in result.metadata
+ }
return {
"data": [self._serialize_row_data(data) for data in result.data],
"row_count": result.row_count,
@@ -638,6 +701,12 @@ class DorisQueryExecutor:
"metadata": {
"columns": result.metadata.get("columns", []),
"query": result.sql,
+ **boundary_metadata,
+ "limits": {
+ "max_rows": limits.max_rows,
+ "max_bytes": limits.max_bytes,
+ "timeout_seconds": limits.timeout_seconds,
+ },
},
}
@@ -648,12 +717,19 @@ class DorisQueryExecutor:
db_name: str | None = None,
catalog_name: str | None = None,
limit: int = 1000,
+ max_bytes: int | None = None,
timeout: int = 30,
session_id: str = "mcp_session",
user_id: str = "mcp_user",
auth_context: AuthContext | None = None,
) -> dict[str, Any]:
"""Execute optional catalog/db context and target SQL on one routed
connection."""
+ limits = resolve_result_limits(
+ self.config,
+ max_rows=limit,
+ max_bytes=max_bytes,
+ timeout_seconds=timeout,
+ )
try:
context_statements = self._build_context_statements(db_name,
catalog_name)
except SQLSecurityError as exc:
@@ -681,29 +757,42 @@ class DorisQueryExecutor:
connection = None
try:
- connection = await self._acquire_routed_connection(session_id,
auth_context)
- for context_sql in context_statements:
- if timeout:
- await asyncio.wait_for(
- connection.execute(context_sql,
auth_context=auth_context),
- timeout=timeout,
+ async with asyncio.timeout(limits.timeout_seconds):
+ connection = await self._acquire_routed_connection(
+ session_id,
+ auth_context,
+ )
+ query_results: list[QueryResult] = []
+ for context_sql in context_statements:
+ await connection.execute(
+ context_sql,
+ auth_context=auth_context,
)
- else:
- await connection.execute(context_sql,
auth_context=auth_context)
-
- query_results: list[QueryResult] = []
- for statement in target_statements:
- if timeout:
- result = await asyncio.wait_for(
- connection.execute(statement,
auth_context=auth_context),
- timeout=timeout,
+
+ remaining_rows = limits.max_rows
+ remaining_bytes = limits.max_bytes
+ for statement in target_statements:
+ statement_limits = ResultLimits(
+ max_rows=max(1, remaining_rows),
+ max_bytes=max(256, remaining_bytes),
+ timeout_seconds=limits.timeout_seconds,
)
- else:
- result = await connection.execute(statement,
auth_context=auth_context)
- query_results.append(result)
+ result = await connection.execute(
+ statement,
+ auth_context=auth_context,
+ max_rows=statement_limits.max_rows,
+ max_bytes=statement_limits.max_bytes,
+ )
+ query_results.append(result)
+ remaining_rows -= result.row_count
+ remaining_bytes -= int(
+ result.metadata.get("result_bytes", 2)
+ )
+ if remaining_rows <= 0 or remaining_bytes < 256:
+ break
if len(query_results) == 1:
- payload = self._query_result_payload(query_results[0])
+ payload = self._query_result_payload(query_results[0], limits)
return {
"success": True,
**payload,
@@ -712,7 +801,10 @@ class DorisQueryExecutor:
return {
"success": True,
"multiple_results": True,
- "results": [self._query_result_payload(result) for result in
query_results],
+ "results": [
+ self._query_result_payload(result, limits)
+ for result in query_results
+ ],
}
finally:
if connection is not None:
@@ -776,6 +868,7 @@ class DorisQueryExecutor:
self,
sql: str,
limit: int = 1000,
+ max_bytes: int | None = None,
timeout: int = 30,
session_id: str = "mcp_session",
user_id: str = "mcp_user",
@@ -787,6 +880,21 @@ class DorisQueryExecutor:
FIX for Issue #62 Bug 1: Now accepts auth_context parameter to support
token-bound database configuration
"""
+ try:
+ limits = resolve_result_limits(
+ self.config,
+ max_rows=limit,
+ max_bytes=max_bytes,
+ timeout_seconds=timeout,
+ )
+ except ResultLimitError as exc:
+ return {
+ "success": False,
+ "error": str(exc),
+ "error_type": "invalid_result_limits",
+ "data": None,
+ }
+
max_retries = 2
retry_count = 0
@@ -876,8 +984,9 @@ class DorisQueryExecutor:
sql,
db_name=db_name,
catalog_name=catalog_name,
- limit=limit,
- timeout=timeout,
+ limit=limits.max_rows,
+ max_bytes=limits.max_bytes,
+ timeout=limits.timeout_seconds,
session_id=session_id,
user_id=user_id,
auth_context=auth_context,
@@ -889,9 +998,15 @@ class DorisQueryExecutor:
if s.strip()
]
if len(all_statements) > 1:
- return await
self.execute_batch_sqls_for_mcp(sqls=all_statements, timeout=timeout,
-
session_id=session_id, user_id=user_id,
-
auth_context=auth_context)
+ return await self.execute_batch_sqls_for_mcp(
+ sqls=all_statements,
+ max_rows=limits.max_rows,
+ max_bytes=limits.max_bytes,
+ timeout=limits.timeout_seconds,
+ session_id=session_id,
+ user_id=user_id,
+ auth_context=auth_context,
+ )
# Add LIMIT if not present and it's a single SELECT query.
# Split first so a multi-statement SQL block is not turned into
@@ -900,34 +1015,25 @@ class DorisQueryExecutor:
if get_first_sql_keyword(sql) == "SELECT" and "LIMIT" not in
sql_upper:
if sql.endswith(";"):
sql = sql[:-1]
- sql = f"{sql} LIMIT {limit}"
+ sql = f"{sql} LIMIT {limits.max_rows}"
# Create query request
query_request = QueryRequest(
sql=sql,
session_id=session_id,
user_id=user_id,
- timeout=timeout,
+ timeout=limits.timeout_seconds,
+ max_rows=limits.max_rows,
+ max_bytes=limits.max_bytes,
cache_enabled=False # Disable cache for MCP calls to
ensure fresh data
)
# Execute query with retry logic
result = await self.execute_query(query_request, auth_context)
- # Serialize data for JSON response
- serialized_data = []
- for row in result.data:
- serialized_data.append(self._serialize_row_data(row))
-
return {
"success": True,
- "data": serialized_data,
- "row_count": result.row_count,
- "execution_time": result.execution_time,
- "metadata": {
- "columns": result.metadata.get("columns", []),
- "query": sql
- }
+ **self._query_result_payload(result, limits),
}
except Exception as e:
@@ -1170,6 +1276,7 @@ async def execute_sql_query(
# Extract parameters from kwargs or use defaults
limit = kwargs.get("limit", 1000)
+ max_bytes = kwargs.get("max_bytes")
timeout = kwargs.get("timeout", 30)
session_id = kwargs.get("session_id", "mcp_session")
user_id = kwargs.get("user_id", "mcp_user")
@@ -1181,6 +1288,7 @@ async def execute_sql_query(
result = await executor.execute_sql_for_mcp(
sql=sql,
limit=limit,
+ max_bytes=max_bytes,
timeout=timeout,
session_id=session_id,
user_id=user_id,
diff --git a/doris_mcp_server/utils/schema_extractor.py
b/doris_mcp_server/utils/schema_extractor.py
index 0d99b0c..b43f735 100644
--- a/doris_mcp_server/utils/schema_extractor.py
+++ b/doris_mcp_server/utils/schema_extractor.py
@@ -2102,6 +2102,8 @@ class MetadataExtractor:
catalog_name: str | None = None,
max_rows: int = 100,
timeout: int = 30,
+ *,
+ max_bytes: int | None = None,
) -> dict[str, Any]:
"""
Execute SQL query and return results, supports catalog federation
queries
@@ -2193,6 +2195,7 @@ class MetadataExtractor:
sql=sql,
connection_manager=self.connection_manager,
limit=max_rows,
+ max_bytes=max_bytes,
timeout=timeout,
db_name=db_name,
catalog_name=catalog_name,
@@ -2598,10 +2601,17 @@ class MetadataManager:
catalog_name: str | None = None,
max_rows: int = 100,
timeout: int = 30,
+ *,
+ max_bytes: int | None = None,
) -> dict[str, Any]:
"""Execute SQL query and return results, supports catalog federation
queries"""
return await self.extractor.exec_query_for_mcp(
- sql, db_name, catalog_name, max_rows, timeout
+ sql,
+ db_name,
+ catalog_name,
+ max_rows,
+ timeout,
+ max_bytes=max_bytes,
)
async def get_table_schema(
diff --git a/test/integration/test_real_doris_transports.py
b/test/integration/test_real_doris_transports.py
index b59c1a9..3c1fcc9 100644
--- a/test/integration/test_real_doris_transports.py
+++ b/test/integration/test_real_doris_transports.py
@@ -376,19 +376,148 @@ async def _exec_query(
sql: str,
*,
timeout: int = 5,
+ max_rows: int = 20,
+ max_bytes: int | None = None,
) -> tuple[Any, dict[str, Any]]:
+ arguments: dict[str, Any] = {
+ "sql": sql,
+ "max_rows": max_rows,
+ "timeout": timeout,
+ }
+ if max_bytes is not None:
+ arguments["max_bytes"] = max_bytes
result = await client.call_tool(
"exec_query",
- {
- "sql": sql,
- "max_rows": 20,
- "timeout": timeout,
- },
+ arguments,
)
assert isinstance(result.structured_content, dict)
return result, result.structured_content
[email protected]("transport", ["http", "stdio"])
+async def test_real_doris_result_boundaries_and_cancellation(
+ transport: str,
+ doris_sandbox: DorisSandbox,
+) -> None:
+ environment = _server_environment(
+ doris_sandbox.settings,
+ user=doris_sandbox.settings.user,
+ password=doris_sandbox.settings.password,
+ )
+ environment.update(
+ {
+ "MAX_RESULT_ROWS": "5",
+ "MAX_RESULT_BYTES": "256",
+ "QUERY_TIMEOUT": "5",
+ }
+ )
+ with doris_sandbox.admin_connection.cursor() as cursor:
+ cursor.executemany(
+ f"INSERT INTO {doris_sandbox.qualified_table} VALUES (%s, %s)",
+ [
+ (index, f"{doris_sandbox.marker}-{index}-" + ("x" * 32))
+ for index in range(1, 7)
+ ],
+ )
+
+ async with _transport_client(
+ transport,
+ environment,
+ read_timeout_seconds=10,
+ ) as client:
+ tools = {
+ tool.name: tool
+ for tool in (await client.list_tools(cache_mode="bypass")).tools
+ }
+ query_schema = tools["exec_query"].input_schema["properties"]
+ assert query_schema["max_rows"]["maximum"] == 5
+ assert query_schema["max_bytes"]["maximum"] == 256
+ assert query_schema["timeout"]["maximum"] == 5
+
+ with pytest.raises(MCPError) as excessive_rows:
+ await client.call_tool(
+ "exec_query",
+ {
+ "sql": f"SELECT * FROM {doris_sandbox.qualified_table}",
+ "max_rows": 6,
+ "max_bytes": 256,
+ "timeout": 5,
+ },
+ )
+ assert excessive_rows.value.code == -32602
+
+ row_result, row_payload = await _exec_query(
+ client,
+ f"SELECT id FROM {doris_sandbox.qualified_table} ORDER BY id",
+ max_rows=2,
+ max_bytes=256,
+ timeout=5,
+ )
+ assert row_result.is_error is False
+ assert len(row_payload["data"]) == 2
+ assert row_payload["metadata"]["result_bytes"] <= 256
+ assert row_payload["metadata"]["limits"] == {
+ "max_rows": 2,
+ "max_bytes": 256,
+ "timeout_seconds": 5,
+ }
+
+ byte_result, byte_payload = await _exec_query(
+ client,
+ (
+ "SELECT id, marker "
+ f"FROM {doris_sandbox.qualified_table} ORDER BY id"
+ ),
+ max_rows=5,
+ max_bytes=256,
+ timeout=5,
+ )
+ assert byte_result.is_error is False
+ assert byte_payload["metadata"]["result_bytes"] <= 256
+ assert byte_payload["metadata"]["truncated"] is True
+ assert byte_payload["metadata"]["truncation_reason"] == "byte_limit"
+ assert 0 < len(byte_payload["data"]) < 5
+
+ if transport == "stdio":
+ cancelled_query = asyncio.create_task(
+ client.call_tool(
+ "exec_query",
+ {
+ "sql": "SELECT SLEEP(5) AS slept",
+ "max_rows": 1,
+ "max_bytes": 256,
+ "timeout": 5,
+ },
+ )
+ )
+ await asyncio.sleep(0.2)
+ cancelled_query.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await cancelled_query
+ else:
+ timed_out, timed_out_payload = await _exec_query(
+ client,
+ "SELECT SLEEP(5) AS slept",
+ max_rows=1,
+ max_bytes=256,
+ timeout=1,
+ )
+ assert timed_out.is_error is True
+ assert timed_out_payload["error_type"] == "timeout"
+
+ recovery_started = time.monotonic()
+ recovered_result, recovered_payload = await _exec_query(
+ client,
+ "SELECT 1 AS recovered",
+ max_rows=1,
+ max_bytes=256,
+ timeout=5,
+ )
+ assert time.monotonic() - recovery_started < 3
+ assert recovered_result.is_error is False
+ assert recovered_payload["data"] == [{"recovered": 1}]
+
+
async def _collect_list_pages(
list_method: Any,
*,
diff --git a/test/tools/test_tools_operation_guard.py
b/test/tools/test_tools_operation_guard.py
index ceda039..8a679c9 100644
--- a/test/tools/test_tools_operation_guard.py
+++ b/test/tools/test_tools_operation_guard.py
@@ -32,12 +32,22 @@ class FakeRoutedConnection:
self.session_id = session_id
self.auth_context = auth_context
- async def execute(self, sql, params=None, auth_context=None):
+ async def execute(
+ self,
+ sql,
+ params=None,
+ auth_context=None,
+ *,
+ max_rows=None,
+ max_bytes=None,
+ ):
return await self.manager.execute_query(
self.session_id,
sql,
params,
auth_context or self.auth_context,
+ max_rows=max_rows,
+ max_bytes=max_bytes,
)
@@ -81,7 +91,16 @@ class FakeRoutedConnectionManager:
async def release_connection(self, session_id, connection):
self.connection_releases += 1
- async def execute_query(self, session_id, sql, params=None,
auth_context=None):
+ async def execute_query(
+ self,
+ session_id,
+ sql,
+ params=None,
+ auth_context=None,
+ *,
+ max_rows=None,
+ max_bytes=None,
+ ):
if getattr(auth_context, "auth_method", "") != "doris_oauth":
raise AssertionError("Doris OAuth tool path did not pass
AuthContext")
if getattr(auth_context, "doris_user", "") != "alice":
@@ -92,6 +111,8 @@ class FakeRoutedConnectionManager:
"sql": sql,
"params": params,
"doris_user": auth_context.doris_user,
+ "max_rows": max_rows,
+ "max_bytes": max_bytes,
}
)
if sql.strip().upper().startswith("EXPLAIN"):
diff --git a/test/utils/test_adbc_query_tools.py
b/test/utils/test_adbc_query_tools.py
new file mode 100644
index 0000000..e7f2bf0
--- /dev/null
+++ b/test/utils/test_adbc_query_tools.py
@@ -0,0 +1,188 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bounded streaming and cancellation tests for Arrow Flight SQL queries."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from doris_mcp_server.result_limits import ResultLimits
+from doris_mcp_server.utils.adbc_query_tools import DorisADBCQueryTools
+
+
+def _manager() -> SimpleNamespace:
+ return SimpleNamespace(
+ config=SimpleNamespace(
+ security=SimpleNamespace(max_result_rows=50),
+ performance=SimpleNamespace(
+ max_result_bytes=4096,
+ query_timeout=20,
+ ),
+ adbc=SimpleNamespace(
+ default_max_rows=20,
+ default_timeout=10,
+ default_return_format="dict",
+ ),
+ ),
+ security_manager=None,
+ )
+
+
+class _StreamingCursor:
+ description = (("id", "BIGINT"), ("payload", "VARCHAR"))
+
+ def __init__(self, rows: list[tuple[Any, ...]]) -> None:
+ self._rows = rows
+ self._offset = 0
+ self.executed_sql: str | None = None
+ self.fetch_sizes: list[int] = []
+ self.closed = False
+
+ def execute(self, sql: str) -> None:
+ self.executed_sql = sql
+
+ def fetchmany(self, size: int) -> list[tuple[Any, ...]]:
+ self.fetch_sizes.append(size)
+ batch = self._rows[self._offset : self._offset + size]
+ self._offset += len(batch)
+ return batch
+
+ def close(self) -> None:
+ self.closed = True
+
+ def adbc_cancel(self) -> None:
+ raise AssertionError("completed queries must not be cancelled")
+
+
+class _BlockingCursor:
+ description: tuple[()] = ()
+
+ def __init__(self) -> None:
+ self.started = threading.Event()
+ self.cancelled = threading.Event()
+ self.closed = False
+ self.cancel_calls = 0
+
+ def execute(self, sql: str) -> None:
+ del sql
+ self.started.set()
+ self.cancelled.wait(timeout=5)
+
+ def fetchmany(self, size: int) -> list[tuple[Any, ...]]:
+ del size
+ return []
+
+ def adbc_cancel(self) -> None:
+ self.cancel_calls += 1
+ self.cancelled.set()
+
+ def close(self) -> None:
+ self.closed = True
+
+
[email protected]
+async def test_adbc_streams_to_row_limit_without_fetchall() -> None:
+ cursor = _StreamingCursor(
+ [(1, "one"), (2, "two"), (3, "three")]
+ )
+ tools = DorisADBCQueryTools(_manager())
+ tools.adbc_client = SimpleNamespace(cursor=lambda: cursor)
+
+ result = await tools._execute_query_with_adbc(
+ "SELECT id, payload FROM bounded",
+ ResultLimits(max_rows=2, max_bytes=4096, timeout_seconds=5),
+ "dict",
+ )
+
+ assert result["success"] is True
+ assert result["result"]["data"] == [
+ {"id": 1, "payload": "one"},
+ {"id": 2, "payload": "two"},
+ ]
+ assert result["truncated"] is True
+ assert result["truncation_reason"] == "max_rows"
+ assert result["result_bytes"] <= 4096
+ assert cursor.closed is True
+ assert cursor.fetch_sizes == [3]
+
+
[email protected]
+async def test_adbc_streams_to_byte_limit() -> None:
+ cursor = _StreamingCursor(
+ [(1, "x" * 180), (2, "y" * 180)]
+ )
+ tools = DorisADBCQueryTools(_manager())
+ tools.adbc_client = SimpleNamespace(cursor=lambda: cursor)
+
+ result = await tools._execute_query_with_adbc(
+ "SELECT id, payload FROM bounded",
+ ResultLimits(max_rows=10, max_bytes=256, timeout_seconds=5),
+ "pandas",
+ )
+
+ assert result["success"] is True
+ assert result["result"]["num_rows"] == 1
+ assert result["truncated"] is True
+ assert result["truncation_reason"] == "max_bytes"
+ assert result["result_bytes"] <= 256
+ assert cursor.closed is True
+
+
[email protected]
+async def test_adbc_cancellation_calls_driver_and_reaps_worker() -> None:
+ cursor = _BlockingCursor()
+ tools = DorisADBCQueryTools(_manager())
+ tools.adbc_client = SimpleNamespace(cursor=lambda: cursor)
+
+ task = asyncio.create_task(
+ tools._execute_query_with_adbc(
+ "SELECT SLEEP(5)",
+ ResultLimits(max_rows=1, max_bytes=256, timeout_seconds=10),
+ "dict",
+ )
+ )
+ assert await asyncio.to_thread(cursor.started.wait, 1)
+
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+
+ assert cursor.cancel_calls == 1
+ assert cursor.closed is True
+
+
[email protected]
+async def test_adbc_rejects_limit_escalation_before_network_checks() -> None:
+ tools = DorisADBCQueryTools(_manager())
+
+ result = await tools.exec_adbc_query(
+ "SELECT 1",
+ max_rows=51,
+ max_bytes=1024,
+ timeout=5,
+ )
+
+ assert result["success"] is False
+ assert result["error_type"] == "invalid_result_limits"
+ assert result["error"] == (
+ "max_rows exceeds the configured maximum of 50"
+ )
diff --git a/test/utils/test_db.py b/test/utils/test_db.py
index 158e4e9..36d9e9d 100644
--- a/test/utils/test_db.py
+++ b/test/utils/test_db.py
@@ -1,3 +1,4 @@
+import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest
@@ -153,6 +154,28 @@ def _make_doris_connection(cursor_description,
fetchall_rows, rowcount=0):
return DorisConnection(connection=raw_connection, session_id="test")
+def _make_bounded_doris_connection(cursor_description, fetchmany_rows):
+ cursor = MagicMock()
+ cursor.execute = AsyncMock(return_value=None)
+ cursor.fetchmany = AsyncMock(side_effect=fetchmany_rows)
+ cursor.description = cursor_description
+ cursor.rowcount = 0
+
+ cursor_ctx = MagicMock()
+ cursor_ctx.__aenter__ = AsyncMock(return_value=cursor)
+ cursor_ctx.__aexit__ = AsyncMock(return_value=None)
+
+ raw_connection = MagicMock()
+ raw_connection.cursor = MagicMock(return_value=cursor_ctx)
+ raw_connection.ensure_closed = AsyncMock()
+
+ return (
+ DorisConnection(connection=raw_connection, session_id="bounded"),
+ cursor,
+ raw_connection,
+ )
+
+
class TestExecuteResultSetDetection:
"""Behavior contract for DorisConnection.execute().
@@ -270,3 +293,84 @@ class TestExecuteResultSetDetection:
assert result.data == []
assert result.row_count == affected
+
+
+class TestBoundedResultFetch:
+ async def
test_row_limit_stops_streaming_and_closes_physical_connection(self):
+ conn, cursor, raw_connection = _make_bounded_doris_connection(
+ [("id", None, None, None, None, None, None)],
+ [[{"id": 1}, {"id": 2}, {"id": 3}]],
+ )
+
+ result = await conn.execute(
+ "WITH data AS (...) SELECT * FROM data",
+ max_rows=2,
+ max_bytes=4096,
+ )
+
+ assert result.data == [{"id": 1}, {"id": 2}]
+ assert result.row_count == 2
+ assert result.metadata["truncated"] is True
+ assert result.metadata["truncation_reason"] == "row_limit"
+ assert result.metadata["result_bytes"] <= 4096
+ cursor.fetchall.assert_not_called()
+ raw_connection.ensure_closed.assert_awaited_once()
+ assert conn.is_healthy is False
+
+ async def test_byte_limit_is_measured_after_masking(self):
+ conn, _, raw_connection = _make_bounded_doris_connection(
+ [("secret", None, None, None, None, None, None)],
+ [[{"secret": "x"}], []],
+ )
+ security_manager = MagicMock()
+ security_manager.validate_sql_security = AsyncMock(
+ return_value=MagicMock(
+ is_valid=True,
+ risk_level="low",
+ blocked_operations=[],
+ )
+ )
+ security_manager.apply_data_masking = AsyncMock(
+ return_value=[{"secret": "masked-value-that-is-too-large"}]
+ )
+ conn.security_manager = security_manager
+
+ result = await conn.execute(
+ "SELECT secret FROM t",
+ auth_context=object(),
+ max_rows=10,
+ max_bytes=24,
+ )
+
+ assert result.data == []
+ assert result.metadata["truncated"] is True
+ assert result.metadata["truncation_reason"] == "byte_limit"
+ assert result.metadata["result_bytes"] == 2
+ security_manager.apply_data_masking.assert_awaited_once()
+ raw_connection.ensure_closed.assert_awaited_once()
+
+ async def test_task_cancellation_closes_connection_and_propagates(self):
+ conn, cursor, raw_connection = _make_bounded_doris_connection(
+ [("id", None, None, None, None, None, None)],
+ [],
+ )
+ async def never_returns(*args, **kwargs):
+ del args, kwargs
+ await asyncio.Event().wait()
+
+ cursor.execute.side_effect = never_returns
+ task = asyncio.create_task(
+ conn.execute(
+ "SELECT SLEEP(30)",
+ max_rows=1,
+ max_bytes=4096,
+ )
+ )
+ await asyncio.sleep(0)
+ task.cancel()
+
+ with pytest.raises(asyncio.CancelledError):
+ await task
+
+ raw_connection.ensure_closed.assert_awaited_once()
+ assert conn.is_healthy is False
diff --git a/test/utils/test_doris_user_pool_manager.py
b/test/utils/test_doris_user_pool_manager.py
index 9d64b42..e8285f0 100644
--- a/test/utils/test_doris_user_pool_manager.py
+++ b/test/utils/test_doris_user_pool_manager.py
@@ -271,7 +271,16 @@ async def
test_execute_query_releases_to_captured_doris_user_owner(manager, monk
await manager.create_or_replace_doris_user_pool("alice", "pw1")
- async def fake_execute(self, sql, params=None, auth_context=None):
+ async def fake_execute(
+ self,
+ sql,
+ params=None,
+ auth_context=None,
+ *,
+ max_rows=None,
+ max_bytes=None,
+ ):
+ del max_rows, max_bytes
return QueryResult(data=[{"ok": 1}], metadata={}, execution_time=0.0,
row_count=1, sql=sql)
monkeypatch.setattr(DorisConnection, "execute", fake_execute)
diff --git a/test/utils/test_result_limits.py b/test/utils/test_result_limits.py
new file mode 100644
index 0000000..0475b22
--- /dev/null
+++ b/test/utils/test_result_limits.py
@@ -0,0 +1,168 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for non-escalating query result and timeout budgets."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from doris_mcp_server.result_limits import (
+ ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS,
+ ABSOLUTE_MAX_RESULT_BYTES,
+ ABSOLUTE_MAX_RESULT_ROWS,
+ ResultLimitError,
+ configured_result_limits,
+ resolve_result_limits,
+)
+from doris_mcp_server.tools.tools_manager import DorisToolsManager
+from doris_mcp_server.utils.config import DorisConfig
+from doris_mcp_server.utils.query_executor import DorisQueryExecutor
+
+
+def _config(
+ *,
+ rows: int = 50,
+ result_bytes: int = 4096,
+ timeout: int = 20,
+) -> SimpleNamespace:
+ return SimpleNamespace(
+ security=SimpleNamespace(
+ max_result_rows=rows,
+ enable_security_check=False,
+ ),
+ performance=SimpleNamespace(
+ max_result_bytes=result_bytes,
+ query_timeout=timeout,
+ max_cache_size=10,
+ cache_ttl=10,
+ max_concurrent_queries=2,
+ ),
+ adbc=SimpleNamespace(
+ default_max_rows=rows,
+ default_timeout=timeout,
+ default_return_format="dict",
+ ),
+ )
+
+
+def test_configured_limits_cannot_exceed_absolute_hard_caps() -> None:
+ limits = configured_result_limits(
+ _config(
+ rows=ABSOLUTE_MAX_RESULT_ROWS + 1,
+ result_bytes=ABSOLUTE_MAX_RESULT_BYTES + 1,
+ timeout=ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS + 1,
+ )
+ )
+
+ assert limits.max_rows == ABSOLUTE_MAX_RESULT_ROWS
+ assert limits.max_bytes == ABSOLUTE_MAX_RESULT_BYTES
+ assert limits.timeout_seconds == ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS
+
+
[email protected](
+ ("field", "value", "message"),
+ [
+ ("max_rows", True, "max_rows must be an integer"),
+ ("max_rows", 51, "max_rows exceeds the configured maximum of 50"),
+ ("max_bytes", 255, "max_bytes must be at least 256"),
+ ("max_bytes", 4097, "max_bytes exceeds the configured maximum of
4096"),
+ ("timeout_seconds", 21, "timeout exceeds the configured maximum of
20"),
+ ],
+)
+def test_request_cannot_raise_deployment_ceiling(
+ field: str,
+ value: object,
+ message: str,
+) -> None:
+ kwargs = {
+ "max_rows": 10,
+ "max_bytes": 1024,
+ "timeout_seconds": 10,
+ }
+ kwargs[field] = value
+
+ with pytest.raises(ResultLimitError, match=message):
+ resolve_result_limits(_config(), **kwargs)
+
+
+def test_doris_config_validation_rejects_values_above_hard_caps() -> None:
+ config = DorisConfig()
+ config.security.max_result_rows = ABSOLUTE_MAX_RESULT_ROWS + 1
+ config.performance.max_result_bytes = ABSOLUTE_MAX_RESULT_BYTES + 1
+ config.performance.query_timeout = ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS + 1
+
+ errors = config.validate()
+
+ assert f"Maximum result rows must not exceed {ABSOLUTE_MAX_RESULT_ROWS}"
in errors
+ assert (
+ "Maximum result bytes must be in the range "
+ f"256-{ABSOLUTE_MAX_RESULT_BYTES}"
+ ) in errors
+ assert (
+ "Query timeout must not exceed "
+ f"{ABSOLUTE_MAX_QUERY_TIMEOUT_SECONDS} seconds"
+ ) in errors
+
+
+def test_exec_query_schema_advertises_effective_ceilings() -> None:
+ connection_manager = Mock()
+ connection_manager.config = _config()
+ manager = DorisToolsManager(connection_manager)
+ tool = manager.tool_registry.resolve("exec_query").tool
+ assert tool is not None
+ properties = tool.input_schema["properties"]
+
+ assert properties["max_rows"]["maximum"] == 50
+ assert properties["max_bytes"]["maximum"] == 4096
+ assert properties["timeout"]["maximum"] == 20
+
+
+def test_exec_adbc_query_schema_advertises_effective_ceilings() -> None:
+ connection_manager = Mock()
+ connection_manager.config = _config()
+ manager = DorisToolsManager(connection_manager)
+ tool = manager.tool_registry.resolve("exec_adbc_query").tool
+ assert tool is not None
+ properties = tool.input_schema["properties"]
+
+ assert properties["max_rows"]["maximum"] == 50
+ assert properties["max_bytes"]["maximum"] == 4096
+ assert properties["timeout"]["maximum"] == 20
+
+
[email protected]
+async def test_executor_returns_typed_error_before_database_dispatch() -> None:
+ connection_manager = Mock()
+ connection_manager.config = _config()
+ connection_manager.execute_query = Mock()
+ executor = DorisQueryExecutor(connection_manager)
+
+ result = await executor.execute_sql_for_mcp(
+ "SELECT 1",
+ limit=51,
+ max_bytes=1024,
+ timeout=10,
+ )
+
+ assert result == {
+ "success": False,
+ "error": "max_rows exceeds the configured maximum of 50",
+ "error_type": "invalid_result_limits",
+ "data": None,
+ }
+ connection_manager.execute_query.assert_not_called()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
