# SQL query executor with timeout and result formatting import logging import time from typing import Optional, Dict, Any, List from dataclasses import dataclass, field from datetime import datetime from contextlib import contextmanager from sqlalchemy import text from sqlalchemy.exc import SQLAlchemyError, OperationalError, ProgrammingError from .db_adapter import DatabaseAdapter, get_database_adapter logger = logging.getLogger(__name__) @dataclass class QueryResult: # Container for query execution results success: bool query: str columns: List[str] = field(default_factory=list) rows: List[Dict[str, Any]] = field(default_factory=list) row_count: int = 0 execution_time_ms: float = 0.0 error: Optional[str] = None truncated: bool = False original_count: Optional[int] = None timestamp: datetime = field(default_factory=datetime.now) def to_dict(self) -> Dict[str, Any]: # Convert to dictionary for JSON serialization return { "success": self.success, "query": self.query, "columns": self.columns, "rows": self.rows, "row_count": self.row_count, "execution_time_ms": round(self.execution_time_ms, 2), "error": self.error, "truncated": self.truncated, "original_count": self.original_count, "timestamp": self.timestamp.isoformat(), } def get_statistics(self) -> Dict[str, Any]: # Calculate statistics for numeric columns if not self.rows: return {} stats = {} for col in self.columns: values = [row.get(col) for row in self.rows if row.get(col) is not None] numeric_values = [v for v in values if isinstance(v, (int, float))] if numeric_values: stats[col] = { "count": len(numeric_values), "sum": sum(numeric_values), "avg": sum(numeric_values) / len(numeric_values), "min": min(numeric_values), "max": max(numeric_values), } return stats class SQLExecutor: # Executes validated SQL queries with timeout and result formatting def __init__( self, adapter: DatabaseAdapter = None, query_timeout: int = 10, max_rows: int = 1000, truncate_at: int = 100, ): self._adapter = adapter self._query_timeout = query_timeout self._max_rows = max_rows self._truncate_at = truncate_at @property def adapter(self) -> DatabaseAdapter: # Get database adapter if self._adapter is None: self._adapter = get_database_adapter() if self._adapter is None: raise RuntimeError("Database adapter not configured") return self._adapter def execute( self, sql: str, validate: bool = True, params: Dict[str, Any] = None ) -> QueryResult: start_time = time.time() if not self.adapter.is_connected: return QueryResult( success=False, query=sql, error="Database not connected", execution_time_ms=(time.time() - start_time) * 1000, ) try: rows = self.adapter.execute_query(sql, params, timeout=self._query_timeout) execution_time_ms = (time.time() - start_time) * 1000 if not rows: return QueryResult( success=True, query=sql, columns=[], rows=[], row_count=0, execution_time_ms=execution_time_ms, ) columns = list(rows[0].keys()) if rows else [] original_count = len(rows) truncated = False if len(rows) > self._truncate_at: rows = rows[:self._truncate_at] truncated = True return QueryResult( success=True, query=sql, columns=columns, rows=rows, row_count=len(rows), execution_time_ms=execution_time_ms, truncated=truncated, original_count=original_count if truncated else None, ) except OperationalError as e: error_msg = str(e) if "timeout" in error_msg.lower(): error_msg = "Query timed out, please narrow your search criteria" else: error_msg = self._sanitize_error(error_msg) return QueryResult( success=False, query=sql, error=error_msg, execution_time_ms=(time.time() - start_time) * 1000, ) except ProgrammingError as e: return QueryResult( success=False, query=sql, error=self._sanitize_error(str(e)), execution_time_ms=(time.time() - start_time) * 1000, ) except SQLAlchemyError as e: logger.error(f"Query execution failed: {e}") return QueryResult( success=False, query=sql, error=self._sanitize_error(str(e)), execution_time_ms=(time.time() - start_time) * 1000, ) except Exception as e: logger.error(f"Unexpected error during query execution: {e}") return QueryResult( success=False, query=sql, error="Query execution failed", execution_time_ms=(time.time() - start_time) * 1000, ) def _sanitize_error(self, error: str) -> str: # Remove sensitive information from error messages patterns_to_remove = [ r"password\s*=\s*['\"][^'\"]*['\"]", r"user\s*=\s*['\"][^'\"]*['\"]", r"host\s*=\s*['\"][^'\"]*['\"]", r"at\s+0x[0-9a-fA-F]+", r"\/[^\s]+\/[^\s]+\.py", ] import re for pattern in patterns_to_remove: error = re.sub(pattern, "[REDACTED]", error, flags=re.IGNORECASE) if len(error) > 200: error = error[:200] + "..." return error def execute_with_retry( self, sql: str, max_retries: int = 2, params: Dict[str, Any] = None ) -> QueryResult: # Execute with automatic retry on transient failures last_result = None for attempt in range(max_retries + 1): result = self.execute(sql, validate=True, params=params) if result.success: return result last_result = result if not self._is_retriable_error(result.error): break logger.warning(f"Query attempt {attempt + 1} failed, retrying...") time.sleep(0.5 * (attempt + 1)) return last_result def _is_retriable_error(self, error: str) -> bool: # Check if error is retriable if not error: return False retriable_patterns = [ "connection", "timeout", "deadlock", "lock wait", "too many connections", ] error_lower = error.lower() return any(pattern in error_lower for pattern in retriable_patterns) def get_row_count(self, table_name: str) -> Optional[int]: result = self.execute( f"SELECT COUNT(*) as count FROM {table_name}", validate=False ) if result.success and result.rows: return result.rows[0].get("count", 0) return None def test_connection(self) -> bool: # Test database connection try: result = self.execute("SELECT 1 as test", validate=False) return result.success except Exception: return False # Module-level singleton _executor_instance: Optional[SQLExecutor] = None def get_sql_executor(adapter: DatabaseAdapter = None) -> SQLExecutor: global _executor_instance if _executor_instance is None: _executor_instance = SQLExecutor(adapter=adapter) return _executor_instance def reset_sql_executor() -> None: # Reset singleton instance for testing global _executor_instance _executor_instance = None