Spaces:
Sleeping
Sleeping
| # 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__) | |
| 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 | |
| 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 | |