| #!/usr/bin/env python3 |
| # 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. |
| """ |
| Doris Query Execution Module |
| Implements query optimization, cache management and performance monitoring functionality |
| """ |
| |
| from __future__ import annotations |
| |
| import asyncio |
| import hashlib |
| import json |
| import time |
| from dataclasses import dataclass, replace |
| from datetime import date, datetime, timedelta |
| from decimal import Decimal |
| 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 ( |
| DorisConnection, |
| DorisConnectionManager, |
| QueryResult, |
| get_first_sql_keyword, |
| ) |
| from .logger import get_logger |
| from .sql_security_utils import ( |
| SQLSecurityError, |
| get_auth_context, |
| quote_identifier, |
| validate_identifier, |
| ) |
| |
| if TYPE_CHECKING: |
| from .security import AuthContext |
| |
| |
| @dataclass |
| class QueryRequest: |
| """Query request wrapper""" |
| |
| sql: str |
| session_id: str |
| 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 |
| |
| |
| @dataclass |
| class CachedQuery: |
| """Cached query result""" |
| |
| result: QueryResult |
| created_at: datetime |
| ttl: int |
| access_count: int = 0 |
| last_accessed: datetime | None = None |
| |
| def is_expired(self) -> bool: |
| """Check if cache is expired""" |
| if self.ttl <= 0: |
| return False |
| return (utc_now() - self.created_at).total_seconds() > self.ttl |
| |
| def access(self) -> None: |
| """Record access""" |
| self.access_count += 1 |
| self.last_accessed = utc_now() |
| |
| |
| @dataclass |
| class QueryMetrics: |
| """Query performance metrics""" |
| |
| total_queries: int = 0 |
| successful_queries: int = 0 |
| failed_queries: int = 0 |
| cache_hits: int = 0 |
| cache_misses: int = 0 |
| avg_execution_time: float = 0.0 |
| total_execution_time: float = 0.0 |
| slow_queries: int = 0 |
| concurrent_queries: int = 0 |
| |
| |
| class QueryCache: |
| """Query result cache manager""" |
| |
| def __init__(self, max_size: int = 1000, default_ttl: int = 300): |
| self.max_size = max_size |
| self.default_ttl = default_ttl |
| self.cache: dict[str, CachedQuery] = {} |
| self.logger = get_logger(__name__) |
| |
| def _generate_cache_key( |
| self, sql: str, parameters: dict[str, Any] | None = None |
| ) -> str: |
| """Generate cache key""" |
| cache_data = {"sql": sql.strip().lower(), "parameters": parameters or {}} |
| cache_string = json.dumps(cache_data, sort_keys=True) |
| return hashlib.sha256(cache_string.encode()).hexdigest() |
| |
| async def get( |
| self, sql: str, parameters: dict[str, Any] | None = None |
| ) -> CachedQuery | None: |
| """Get cached query result""" |
| cache_key = self._generate_cache_key(sql, parameters) |
| |
| if cache_key in self.cache: |
| cached_query = self.cache[cache_key] |
| |
| if not cached_query.is_expired(): |
| cached_query.access() |
| self.logger.debug(f"Cache hit: {cache_key}") |
| return cached_query |
| else: |
| # Clean up expired cache |
| del self.cache[cache_key] |
| self.logger.debug(f"Cache expired, cleaned up: {cache_key}") |
| |
| return None |
| |
| async def set( |
| self, |
| sql: str, |
| result: QueryResult, |
| parameters: dict[str, Any] | None = None, |
| ttl: int | None = None, |
| ) -> str: |
| """Set query result cache""" |
| cache_key = self._generate_cache_key(sql, parameters) |
| |
| # Check cache size limit |
| if len(self.cache) >= self.max_size: |
| await self._evict_oldest() |
| |
| cached_query = CachedQuery( |
| result=result, created_at=utc_now(), ttl=ttl or self.default_ttl |
| ) |
| |
| self.cache[cache_key] = cached_query |
| self.logger.debug(f"Cache set: {cache_key}") |
| |
| return cache_key |
| |
| async def _evict_oldest(self) -> None: |
| """Clean up oldest cache item""" |
| if not self.cache: |
| return |
| |
| # Find oldest cache item |
| oldest_key = min(self.cache.keys(), key=lambda k: self.cache[k].created_at) |
| |
| del self.cache[oldest_key] |
| self.logger.debug(f"Cleaned up oldest cache: {oldest_key}") |
| |
| async def clear_expired(self) -> None: |
| """Clean up all expired cache""" |
| expired_keys = [ |
| key for key, cached_query in self.cache.items() if cached_query.is_expired() |
| ] |
| |
| for key in expired_keys: |
| del self.cache[key] |
| |
| if expired_keys: |
| self.logger.info(f"Cleaned up {len(expired_keys)} expired cache items") |
| |
| async def clear_all(self) -> None: |
| """Clean up all cache""" |
| cache_count = len(self.cache) |
| self.cache.clear() |
| self.logger.info(f"Cleaned up all cache, total {cache_count} items") |
| |
| def get_stats(self) -> dict[str, Any]: |
| """Get cache statistics""" |
| total_access = sum(cached.access_count for cached in self.cache.values()) |
| |
| return { |
| "cache_size": len(self.cache), |
| "max_size": self.max_size, |
| "total_access": total_access, |
| "hit_rate": 0.0 |
| if total_access == 0 |
| else sum(cached.access_count for cached in self.cache.values()) |
| / total_access, |
| } |
| |
| |
| class QueryOptimizer: |
| """Query optimizer""" |
| |
| def __init__(self, config: Any) -> None: |
| self.config = config |
| self.logger = get_logger(__name__) |
| self.optimization_rules = self._load_optimization_rules() |
| |
| def _load_optimization_rules(self) -> list[dict[str, Any]]: |
| """Load query optimization rules""" |
| return [ |
| { |
| "name": "add_limit_clause", |
| "description": "Add default limit for SELECT queries without LIMIT", |
| "pattern": r"^select\s+.*(?!.*limit\s+\d+)", |
| "action": "add_limit", |
| "params": {"default_limit": 1000}, |
| }, |
| { |
| "name": "optimize_count_query", |
| "description": "Optimize COUNT queries", |
| "pattern": r"select\s+count\(\*\)\s+from\s+(\w+)", |
| "action": "optimize_count", |
| "params": {}, |
| }, |
| ] |
| |
| async def optimize_query(self, sql: str, context: dict[str, Any]) -> str: |
| """Apply query optimization""" |
| optimized_sql = sql |
| |
| for rule in self.optimization_rules: |
| if self._should_apply_rule(rule, optimized_sql, context): |
| optimized_sql = await self._apply_optimization_rule( |
| optimized_sql, rule, context |
| ) |
| self.logger.debug(f"Applied optimization rule: {rule['name']}") |
| |
| return optimized_sql |
| |
| def _should_apply_rule( |
| self, rule: dict[str, Any], sql: str, context: dict[str, Any] |
| ) -> bool: |
| """Check if optimization rule should be applied""" |
| import re |
| |
| # Check pattern match |
| if "pattern" in rule: |
| if not re.search(rule["pattern"], sql, re.IGNORECASE): |
| return False |
| |
| # Check conditions |
| if "conditions" in rule: |
| for condition in rule["conditions"]: |
| if not self._check_condition(condition, context): |
| return False |
| |
| return True |
| |
| def _check_condition( |
| self, condition: dict[str, Any], context: dict[str, Any] |
| ) -> bool: |
| """Check optimization condition""" |
| condition_type = condition.get("type") |
| |
| if condition_type == "user_role": |
| required_roles = condition.get("roles", []) |
| user_roles = context.get("user_roles", []) |
| return any(role in user_roles for role in required_roles) |
| |
| elif condition_type == "query_size": |
| max_size = condition.get("max_size", 1000) |
| return len(str(context.get("sql", ""))) <= int(max_size) |
| |
| return True |
| |
| async def _apply_optimization_rule( |
| self, sql: str, rule: dict[str, Any], context: dict[str, Any] |
| ) -> str: |
| """Apply optimization rule""" |
| action = rule.get("action") |
| params = rule.get("params", {}) |
| |
| if action == "add_limit": |
| return await self._add_limit_clause(sql, params) |
| elif action == "optimize_count": |
| return await self._optimize_count_query(sql, params) |
| elif action == "add_hints": |
| return await self._add_query_hints(sql, params) |
| |
| return sql |
| |
| async def _add_limit_clause(self, sql: str, params: dict[str, Any]) -> str: |
| """Add LIMIT clause to query""" |
| import re |
| |
| default_limit = params.get("default_limit", 1000) |
| |
| # Check if LIMIT already exists |
| if re.search(r"\blimit\s+\d+", sql, re.IGNORECASE): |
| return sql |
| |
| # Add LIMIT clause |
| if sql.strip().endswith(";"): |
| sql = sql.strip()[:-1] |
| |
| return f"{sql} LIMIT {default_limit}" |
| |
| async def _optimize_count_query(self, sql: str, params: dict[str, Any]) -> str: |
| """Optimize COUNT query""" |
| # For COUNT queries, we can add optimization hints |
| return sql.replace("COUNT(*)", "COUNT(1)") |
| |
| async def _add_query_hints(self, sql: str, params: dict[str, Any]) -> str: |
| """Add query hints""" |
| hints = params.get("hints", []) |
| if not hints: |
| return sql |
| |
| hint_string = "/*+ " + " ".join(hints) + " */" |
| return f"{hint_string} {sql}" |
| |
| |
| class DorisQueryExecutor: |
| """Doris query executor with caching and optimization""" |
| |
| def __init__( |
| self, |
| connection_manager: DorisConnectionManager, |
| config: Any | None = None, |
| ) -> None: |
| self.connection_manager = connection_manager |
| self.config = ( |
| config |
| or getattr(connection_manager, "config", None) |
| or self._create_default_config() |
| ) |
| self.logger = get_logger(__name__) |
| |
| # Initialize components |
| cache_config = getattr(self.config, 'performance', None) |
| if cache_config: |
| cache_size = getattr(cache_config, 'max_cache_size', 1000) |
| cache_ttl = getattr(cache_config, 'cache_ttl', 300) |
| else: |
| cache_size = 1000 |
| cache_ttl = 300 |
| |
| self.query_cache = QueryCache(max_size=cache_size, default_ttl=cache_ttl) |
| self.query_optimizer = QueryOptimizer(self.config) |
| self.metrics = QueryMetrics() |
| |
| # Performance monitoring |
| self.slow_query_threshold = 5.0 # seconds |
| 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]] = [] |
| |
| def _create_default_config(self) -> Any: |
| """Create default configuration""" |
| 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() |
| |
| async def start(self) -> None: |
| """Start owned background tasks within an explicit runtime lifecycle.""" |
| if any(not task.done() for task in self._background_tasks): |
| return |
| self._background_tasks = [ |
| task for task in self._background_tasks if not task.done() |
| ] |
| cleanup_task = asyncio.create_task(self._cache_cleanup_loop()) |
| self._background_tasks.append(cleanup_task) |
| |
| async def _cache_cleanup_loop(self) -> None: |
| """Background cache cleanup loop""" |
| while True: |
| try: |
| await asyncio.sleep(300) # Run every 5 minutes |
| await self.query_cache.clear_expired() |
| except asyncio.CancelledError: |
| break |
| except Exception as e: |
| self.logger.error(f"Cache cleanup error: {e}") |
| |
| async def execute_query( |
| self, |
| query_request: QueryRequest, |
| auth_context: AuthContext | None = None, |
| ) -> QueryResult: |
| """Execute query with caching and optimization""" |
| start_time = time.time() |
| self.metrics.total_queries += 1 |
| self.metrics.concurrent_queries += 1 |
| |
| try: |
| effective_auth_context = auth_context or get_auth_context() |
| if ( |
| query_request.cache_enabled |
| and getattr(effective_auth_context, "auth_method", "") == "doris_oauth" |
| ): |
| self.logger.warning( |
| "Doris OAuth query cache disabled for session %s", |
| query_request.session_id, |
| ) |
| query_request = replace(query_request, cache_enabled=False) |
| auth_context = effective_auth_context |
| |
| # Check cache first |
| if query_request.cache_enabled: |
| cached_result = await self.query_cache.get( |
| query_request.sql, query_request.parameters |
| ) |
| if cached_result: |
| self.metrics.cache_hits += 1 |
| self.logger.debug("Query cache hit") |
| return cached_result.result |
| |
| self.metrics.cache_misses += 1 |
| |
| # Execute query |
| result = await self._execute_query_internal(query_request, auth_context) |
| |
| # Cache result if enabled |
| if query_request.cache_enabled and result.row_count > 0: |
| await self.query_cache.set( |
| query_request.sql, result, query_request.parameters |
| ) |
| |
| self.metrics.successful_queries += 1 |
| return result |
| |
| except Exception as e: |
| self.metrics.failed_queries += 1 |
| self.logger.error(f"Query execution failed: {e}") |
| raise |
| |
| finally: |
| execution_time = time.time() - start_time |
| self.metrics.concurrent_queries -= 1 |
| self._update_execution_metrics(execution_time) |
| |
| async def _execute_query_internal( |
| self, |
| query_request: QueryRequest, |
| auth_context: AuthContext | None, |
| ) -> QueryResult: |
| """Internal query execution""" |
| |
| # Database configuration should already be handled during authentication |
| # No need to configure again during query execution |
| |
| # Optimize query |
| optimized_sql = await self.query_optimizer.optimize_query( |
| query_request.sql, {"user_roles": getattr(auth_context, 'roles', [])} |
| ) |
| |
| # Execute query |
| # Set timeout if specified |
| 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, |
| 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, |
| max_rows=query_request.max_rows, |
| max_bytes=query_request.max_bytes, |
| ) |
| |
| return result |
| |
| def _update_execution_metrics(self, execution_time: float) -> None: |
| """Update execution metrics""" |
| self.metrics.total_execution_time += execution_time |
| |
| # Update average execution time |
| if self.metrics.successful_queries > 0: |
| self.metrics.avg_execution_time = ( |
| self.metrics.total_execution_time / self.metrics.successful_queries |
| ) |
| |
| # Check for slow queries |
| if execution_time > self.slow_query_threshold: |
| self.metrics.slow_queries += 1 |
| self.logger.warning( |
| f"Slow query detected: {execution_time:.2f}s (threshold: {self.slow_query_threshold}s)" |
| ) |
| |
| 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", |
| auth_context: AuthContext | None = None, |
| ) -> dict[str, Any]: |
| """Execute multiple sqls in batch""" |
| if not sqls: |
| return { |
| "success": False, |
| "error": "SQL query is required", |
| "data": None |
| } |
| 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 = [ |
| self._query_result_payload(result, limits) |
| for result in query_results |
| ] |
| |
| return { |
| "success": True, |
| "multiple_results": True, |
| "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( |
| self, |
| query_requests: list[QueryRequest], |
| auth_context: AuthContext | None = None, |
| ) -> list[QueryResult]: |
| """Execute multiple queries in batch""" |
| # Check concurrent query limit |
| if len(query_requests) > self.max_concurrent_queries: |
| raise Exception( |
| f"Batch size {len(query_requests)} exceeds maximum concurrent queries {self.max_concurrent_queries}" |
| ) |
| |
| # Execute queries concurrently |
| tasks = [ |
| self.execute_query(request, auth_context) for request in query_requests |
| ] |
| |
| query_results: list[QueryResult] = [] |
| results = await asyncio.gather(*tasks, return_exceptions=True) |
| for result in results: |
| if isinstance(result, BaseException): |
| self.logger.error(f"Batch query execution failed: {result}") |
| raise result |
| query_results.append(result) |
| |
| return query_results |
| |
| def _build_context_statements( |
| self, db_name: str | None = None, catalog_name: str | None = None |
| ) -> list[str]: |
| """Build validated Doris context switching statements.""" |
| statements: list[str] = [] |
| if catalog_name: |
| validate_identifier(catalog_name, "catalog name") |
| safe_catalog = quote_identifier(catalog_name, "catalog name") |
| statements.append(f"USE CATALOG {safe_catalog}") |
| if db_name: |
| validate_identifier(db_name, "database name") |
| safe_db = quote_identifier(db_name, "database name") |
| statements.append(f"USE {safe_db}") |
| return statements |
| |
| async def _acquire_routed_connection( |
| self, |
| session_id: str, |
| auth_context: AuthContext | None = None, |
| ) -> DorisConnection: |
| """Acquire a routed connection, preserving an explicit auth context.""" |
| get_connection_for_auth_context = getattr( |
| self.connection_manager, |
| "_get_connection_for_auth_context", |
| None, |
| ) |
| if callable(get_connection_for_auth_context): |
| get_effective_auth_context = getattr( |
| self.connection_manager, |
| "_get_effective_auth_context", |
| None, |
| ) |
| effective_auth_context = ( |
| get_effective_auth_context(auth_context) |
| if callable(get_effective_auth_context) |
| else auth_context |
| ) |
| return cast( |
| DorisConnection, |
| await get_connection_for_auth_context( |
| session_id, |
| effective_auth_context, |
| ), |
| ) |
| return await self.connection_manager.get_connection(session_id) |
| |
| async def _release_routed_connection( |
| self, |
| session_id: str, |
| connection: DorisConnection, |
| ) -> None: |
| release_connection = getattr(self.connection_manager, "release_connection", None) |
| if callable(release_connection): |
| await release_connection(session_id, connection) |
| |
| 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, |
| "execution_time": result.execution_time, |
| "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, |
| }, |
| }, |
| } |
| |
| async def _execute_sql_with_context_for_mcp( |
| self, |
| sql: str, |
| *, |
| 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: |
| return { |
| "success": False, |
| "error": str(exc), |
| "error_type": "invalid_context", |
| "data": None, |
| } |
| |
| target_statements = [s.strip() for s in sqlparse.split(sql) if s.strip()] |
| if not target_statements: |
| return { |
| "success": False, |
| "error": "SQL query is required", |
| "data": None, |
| } |
| |
| if len(target_statements) == 1: |
| target_sql = target_statements[0] |
| if get_first_sql_keyword(target_sql) == "SELECT" and "LIMIT" not in target_sql.upper(): |
| target_sql = target_sql.rstrip(";") |
| target_sql = f"{target_sql} LIMIT {limit}" |
| target_statements = [target_sql] |
| |
| connection = None |
| try: |
| 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, |
| internal_session_control=True, |
| ) |
| |
| 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, |
| ) |
| 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], limits) |
| return { |
| "success": True, |
| **payload, |
| } |
| |
| return { |
| "success": True, |
| "multiple_results": True, |
| "results": [ |
| self._query_result_payload(result, limits) |
| for result in query_results |
| ], |
| } |
| finally: |
| if connection is not None: |
| await self._release_routed_connection(session_id, connection) |
| |
| async def explain_query(self, sql: str, session_id: str) -> dict[str, Any]: |
| """Get query execution plan""" |
| explain_sql = f"EXPLAIN {sql}" |
| |
| async with self.connection_manager.get_connection_context( |
| session_id |
| ) as connection: |
| auth_context = get_auth_context() |
| result = await connection.execute( |
| explain_sql, |
| auth_context=auth_context, |
| ) |
| |
| return { |
| "query": sql, |
| "execution_plan": result.data, |
| "estimated_cost": "N/A", # Doris doesn't provide cost estimates |
| } |
| |
| async def get_query_stats(self) -> dict[str, Any]: |
| """Get query execution statistics""" |
| cache_stats = self.query_cache.get_stats() |
| |
| return { |
| "query_metrics": { |
| "total_queries": self.metrics.total_queries, |
| "successful_queries": self.metrics.successful_queries, |
| "failed_queries": self.metrics.failed_queries, |
| "success_rate": ( |
| self.metrics.successful_queries / self.metrics.total_queries |
| if self.metrics.total_queries > 0 |
| else 0.0 |
| ), |
| "avg_execution_time": self.metrics.avg_execution_time, |
| "slow_queries": self.metrics.slow_queries, |
| "concurrent_queries": self.metrics.concurrent_queries, |
| }, |
| "cache_metrics": { |
| "cache_hits": self.metrics.cache_hits, |
| "cache_misses": self.metrics.cache_misses, |
| "hit_rate": ( |
| self.metrics.cache_hits |
| / (self.metrics.cache_hits + self.metrics.cache_misses) |
| if (self.metrics.cache_hits + self.metrics.cache_misses) > 0 |
| else 0.0 |
| ), |
| **cache_stats, |
| }, |
| } |
| |
| async def clear_cache(self) -> None: |
| """Clear query cache""" |
| await self.query_cache.clear_all() |
| |
| async def execute_sql_for_mcp( |
| self, |
| sql: str, |
| limit: int = 1000, |
| max_bytes: int | None = None, |
| timeout: int = 30, |
| session_id: str = "mcp_session", |
| user_id: str = "mcp_user", |
| db_name: str | None = None, |
| catalog_name: str | None = None, |
| auth_context: AuthContext | None = None, |
| ) -> dict[str, Any]: |
| """Execute SQL query for MCP interface - unified method |
| |
| 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 |
| |
| while retry_count <= max_retries: |
| try: |
| if not sql: |
| return { |
| "success": False, |
| "error": "SQL query is required", |
| "data": None |
| } |
| |
| # Import required security modules |
| from .security import AuthContext, DorisSecurityManager, SecurityLevel |
| |
| # FIX: Use provided auth_context if available (contains token for DB config) |
| # Otherwise create default auth context for backward compatibility |
| if auth_context is None: |
| auth_context = AuthContext( |
| user_id=user_id, |
| roles=["read_only_user"], # Restrictive role for MCP interface |
| permissions=["read_data"], # Only read permissions |
| session_id=session_id, |
| security_level=SecurityLevel.INTERNAL, |
| token=EMPTY_CREDENTIAL, |
| ) |
| else: |
| # Use provided auth_context (may contain token for database configuration) |
| self.logger.debug(f"Using provided auth_context with token: {bool(hasattr(auth_context, 'token') and auth_context.token)}") |
| |
| # Perform SQL security validation if enabled |
| if hasattr(self.connection_manager, 'config') and hasattr(self.connection_manager.config, 'security'): |
| if self.connection_manager.config.security.enable_security_check: |
| try: |
| # 🔧 FIX: Use existing security_manager to avoid creating multiple TokenManager instances |
| # Creating new DorisSecurityManager each time causes multiple hot reload monitors |
| security_manager = getattr(self.connection_manager, 'security_manager', None) |
| if not security_manager: |
| # Fallback: create new one only if not available (should rarely happen) |
| self.logger.warning("No existing security_manager, creating new instance") |
| security_manager = DorisSecurityManager(self.connection_manager.config) |
| validation_result = await security_manager.validate_sql_security(sql, auth_context) |
| |
| if not validation_result.is_valid: |
| self.logger.warning( |
| "SQL security validation rejected a query" |
| ) |
| return { |
| "success": False, |
| "error": f"SQL security validation failed: {validation_result.error_message}", |
| "error_type": "security_violation", |
| "blocked_operations": validation_result.blocked_operations, |
| "risk_level": validation_result.risk_level, |
| "data": None, |
| "metadata": { |
| "query": sql, |
| "validation_details": { |
| "blocked_operations": validation_result.blocked_operations, |
| "risk_level": validation_result.risk_level |
| } |
| } |
| } |
| else: |
| self.logger.debug( |
| "SQL security validation passed" |
| ) |
| except Exception as security_error: |
| self.logger.error(f"Security validation error: {str(security_error)}") |
| # In case of security validation error, fail safe |
| return { |
| "success": False, |
| "error": f"Security validation system error: {str(security_error)}", |
| "error_type": "security_system_error", |
| "data": None, |
| "metadata": { |
| "query": sql, |
| "security_error": str(security_error) |
| } |
| } |
| else: |
| self.logger.info("SQL security check is disabled in configuration") |
| else: |
| self.logger.warning("Security configuration not found, proceeding without validation") |
| |
| if db_name or catalog_name: |
| return await self._execute_sql_with_context_for_mcp( |
| sql, |
| db_name=db_name, |
| catalog_name=catalog_name, |
| 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, |
| ) |
| |
| all_statements = [ |
| s.strip() |
| for s in sqlparse.split(sql) |
| if s.strip() |
| ] |
| if len(all_statements) > 1: |
| 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 |
| # an extra synthetic statement such as a trailing "LIMIT 100". |
| sql_upper = sql.upper() |
| if get_first_sql_keyword(sql) == "SELECT" and "LIMIT" not in sql_upper: |
| if sql.endswith(";"): |
| sql = sql[:-1] |
| sql = f"{sql} LIMIT {limits.max_rows}" |
| |
| # Create query request |
| query_request = QueryRequest( |
| sql=sql, |
| session_id=session_id, |
| user_id=user_id, |
| 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) |
| |
| return { |
| "success": True, |
| **self._query_result_payload(result, limits), |
| } |
| |
| except Exception as e: |
| error_msg = str(e) |
| error_str = error_msg.lower() |
| |
| # Check if it's a connection-related error that we should retry |
| connection_errors = [ |
| "at_eof", "connection", "closed", "nonetype", |
| "transport", "reader", "broken pipe", "connection reset" |
| ] |
| |
| is_connection_error = any(err in error_str for err in connection_errors) |
| |
| if is_connection_error and retry_count < max_retries: |
| retry_count += 1 |
| self.logger.warning(f"Connection error detected, retrying ({retry_count}/{max_retries}): {e}") |
| |
| # Wait a bit before retry |
| await asyncio.sleep(0.5 * retry_count) |
| continue |
| else: |
| # If we've exhausted retries or it's not a connection error, return error |
| error_analysis = self._analyze_error(error_msg) |
| |
| return { |
| "success": False, |
| "error": error_analysis.get("user_message", error_msg), |
| "error_type": error_analysis.get("error_type", "general_error"), |
| "data": None, |
| "metadata": { |
| "query": sql, |
| "error_details": error_msg, |
| "retry_count": retry_count |
| } |
| } |
| |
| # This should never be reached, but just in case |
| return { |
| "success": False, |
| "error": "Maximum retries exceeded", |
| "data": None, |
| "metadata": { |
| "query": sql, |
| "retry_count": retry_count |
| } |
| } |
| |
| def _serialize_row_data(self, row_data: dict[str, Any]) -> dict[str, Any]: |
| """Serialize row data for JSON response""" |
| serialized: dict[str, Any] = {} |
| |
| for key, value in row_data.items(): |
| if value is None: |
| serialized[key] = None |
| elif isinstance(value, str | int | float | bool): |
| serialized[key] = value |
| elif isinstance(value, Decimal): |
| serialized[key] = float(value) |
| elif isinstance(value, datetime | date): |
| serialized[key] = value.isoformat() |
| elif isinstance(value, bytes): |
| try: |
| serialized[key] = value.decode('utf-8') |
| except UnicodeDecodeError: |
| serialized[key] = str(value) |
| else: |
| serialized[key] = str(value) |
| |
| return serialized |
| |
| def _analyze_error(self, error_message: str) -> dict[str, str]: |
| """Analyze error message and provide user-friendly feedback""" |
| error_msg_lower = error_message.lower() |
| |
| if "at_eof" in error_msg_lower or "nonetype" in error_msg_lower and "at_eof" in error_msg_lower: |
| return { |
| "error_type": "connection_lost", |
| "user_message": "Database connection was lost. The query has been automatically retried. If this persists, please restart the server." |
| } |
| elif "table" in error_msg_lower and "doesn't exist" in error_msg_lower: |
| return { |
| "error_type": "table_not_found", |
| "user_message": "The specified table does not exist. Please check the table name and database." |
| } |
| elif "column" in error_msg_lower and ("unknown" in error_msg_lower or "doesn't exist" in error_msg_lower): |
| return { |
| "error_type": "column_not_found", |
| "user_message": "One or more columns in the query do not exist. Please check column names." |
| } |
| elif "syntax error" in error_msg_lower or "sql syntax" in error_msg_lower: |
| return { |
| "error_type": "syntax_error", |
| "user_message": "SQL syntax error. Please check your query syntax." |
| } |
| elif ( |
| "access denied" in error_msg_lower |
| or "command denied" in error_msg_lower |
| or "permission" in error_msg_lower |
| ): |
| return { |
| "error_type": "permission_denied", |
| "user_message": "Access denied. You don't have permission to execute this query." |
| } |
| elif "timeout" in error_msg_lower: |
| return { |
| "error_type": "timeout", |
| "user_message": "Query execution timed out. Try simplifying your query or adding more specific filters." |
| } |
| elif "connection" in error_msg_lower and ("closed" in error_msg_lower or "reset" in error_msg_lower): |
| return { |
| "error_type": "connection_error", |
| "user_message": "Database connection was interrupted. The query has been automatically retried." |
| } |
| else: |
| return { |
| "error_type": "general_error", |
| "user_message": f"Query execution failed: {error_message}" |
| } |
| |
| async def close(self) -> None: |
| """Close query executor and cleanup resources""" |
| # Cancel background tasks |
| for task in self._background_tasks: |
| task.cancel() |
| if self._background_tasks: |
| await asyncio.gather(*self._background_tasks, return_exceptions=True) |
| self._background_tasks.clear() |
| |
| # Clear cache |
| await self.query_cache.clear_all() |
| |
| self.logger.info("Query executor closed") |
| |
| |
| class QueryPerformanceMonitor: |
| """Query performance monitor""" |
| |
| def __init__(self, query_executor: DorisQueryExecutor) -> None: |
| self.query_executor = query_executor |
| self.logger = get_logger(__name__) |
| self.performance_records: list[dict[str, Any]] = [] |
| |
| async def record_query_performance( |
| self, query_request: QueryRequest, result: QueryResult, execution_time: float |
| ) -> None: |
| """Record query performance""" |
| record = { |
| "timestamp": utc_now(), |
| "sql": query_request.sql, |
| "user_id": query_request.user_id, |
| "session_id": query_request.session_id, |
| "execution_time": execution_time, |
| "row_count": result.row_count, |
| "cache_hit": False, # This would need to be passed from executor |
| } |
| |
| self.performance_records.append(record) |
| |
| # Keep only recent records (last 1000) |
| if len(self.performance_records) > 1000: |
| self.performance_records = self.performance_records[-1000:] |
| |
| async def get_performance_report( |
| self, time_range_minutes: int = 60 |
| ) -> dict[str, Any]: |
| """Get performance report""" |
| cutoff_time = utc_now() - timedelta(minutes=time_range_minutes) |
| |
| recent_records = [ |
| record |
| for record in self.performance_records |
| if record["timestamp"] >= cutoff_time |
| ] |
| |
| if not recent_records: |
| return {"message": "No performance data available for the specified time range"} |
| |
| # Calculate statistics |
| execution_times = [record["execution_time"] for record in recent_records] |
| row_counts = [record["row_count"] for record in recent_records] |
| |
| return { |
| "time_range_minutes": time_range_minutes, |
| "total_queries": len(recent_records), |
| "avg_execution_time": sum(execution_times) / len(execution_times), |
| "max_execution_time": max(execution_times), |
| "min_execution_time": min(execution_times), |
| "avg_row_count": sum(row_counts) / len(row_counts), |
| "query_distribution": self._analyze_query_distribution(recent_records), |
| } |
| |
| def _analyze_query_distribution( |
| self, records: list[dict[str, Any]] |
| ) -> dict[str, Any]: |
| """Analyze query distribution""" |
| query_types: dict[str, int] = {} |
| user_distribution: dict[str, int] = {} |
| |
| for record in records: |
| # Analyze query type |
| sql_upper = record["sql"].strip().upper() |
| if sql_upper.startswith("SELECT"): |
| query_type = "SELECT" |
| elif sql_upper.startswith("INSERT"): |
| query_type = "INSERT" |
| elif sql_upper.startswith("UPDATE"): |
| query_type = "UPDATE" |
| elif sql_upper.startswith("DELETE"): |
| query_type = "DELETE" |
| else: |
| query_type = "OTHER" |
| |
| query_types[query_type] = query_types.get(query_type, 0) + 1 |
| |
| # Analyze user distribution |
| user_id = record["user_id"] |
| user_distribution[user_id] = user_distribution.get(user_id, 0) + 1 |
| |
| return {"query_types": query_types, "user_distribution": user_distribution} |
| |
| |
| # Unified convenience function for MCP integration |
| async def execute_sql_query( |
| sql: str, |
| connection_manager: DorisConnectionManager, |
| **kwargs: Any, |
| ) -> dict[str, Any]: |
| """Execute SQL query - unified convenience function for MCP tools |
| |
| This function now includes security validation to ensure safe query execution. |
| All queries are validated against the configured security policies before execution. |
| |
| FIX for Issue #62 Bug 1: Now supports auth_context parameter for token-bound database configuration |
| FIX for Issue #58 Problem 2: Removed executor.close() to prevent ClosedResourceError in multi-worker mode |
| """ |
| try: |
| # Create query executor with the connection manager's configuration |
| executor = DorisQueryExecutor(connection_manager) |
| |
| # 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") |
| auth_context = kwargs.get("auth_context", None) # FIX: Extract auth_context |
| db_name = kwargs.get("db_name", None) |
| catalog_name = kwargs.get("catalog_name", None) |
| |
| # The execute_sql_for_mcp method now includes security validation |
| 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, |
| db_name=db_name, |
| catalog_name=catalog_name, |
| auth_context=auth_context # FIX: Pass auth_context with token |
| ) |
| |
| # FIX for Issue #58 Problem 2: Do NOT close executor here |
| # In multi-worker mode, closing here causes ClosedResourceError |
| # The executor's resources (cache, background tasks) will be managed |
| # by the connection_manager lifecycle and Python's garbage collection |
| # This prevents premature cleanup while MCP session manager is still processing |
| |
| return result |
| |
| except Exception as e: |
| return { |
| "success": False, |
| "error": f"Query execution failed: {str(e)}", |
| "error_type": "execution_error", |
| "data": None, |
| "metadata": { |
| "query": sql, |
| "execution_error": str(e) |
| } |
| } |
