joelgilbert/NL2SQL

1"""2Audit logging for query attempts and DBA actions.3"""4 5import logging6import json7import uuid8from datetime import datetime9from pathlib import Path10from typing import Dict, List, Optional11 12logger = logging.getLogger(__name__)13 14 15class AuditLogger:16    """Logs all query attempts and DBA actions for compliance and debugging."""17    18    def __init__(self, log_dir: str = "audit_logs"):19        """20        Initialize audit logger.21        22        Args:23            log_dir: Directory to store audit logs24        """25        self.log_dir = Path(log_dir)26        self.log_dir.mkdir(exist_ok=True)27        self._log_cache: Dict[str, Dict] = {}28    29    def log_query_attempt(30        self,31        user_id: str,32        question: str,33        sql: str,34        mode: str35    ) -> str:36        """37        Log a query attempt.38        39        Args:40            user_id: Session or user identifier41            question: Natural language question42            sql: Generated SQL query43            mode: Execution mode (readonly or dba)44            45        Returns:46            Log entry ID47        """48        log_id = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}"49        50        log_entry = {51            "log_id": log_id,52            "timestamp": datetime.now().isoformat(),53            "session_id": user_id,54            "mode": mode,55            "question": question,56            "generated_sql": sql,57            "validation_passed": None,58            "execution_success": None,59            "error_message": None,60            "execution_time_seconds": None,61            "rows_returned": None,62            "dba_approved": False,63            "approver_id": None64        }65        66        # Cache in memory67        self._log_cache[log_id] = log_entry68        69        # Write to file70        self._write_log_entry(log_id, log_entry)71        72        logger.info(f"Logged query attempt: {log_id}")73        return log_id74    75    def log_query_result(76        self,77        log_id: str,78        success: bool,79        error: Optional[str],80        execution_time: float,81        row_count: int = 082    ):83        """84        Update log entry with query execution results.85        86        Args:87            log_id: Log entry ID from log_query_attempt88            success: Whether execution succeeded89            error: Error message if failed90            execution_time: Execution time in seconds91            row_count: Number of rows returned/affected92        """93        if log_id not in self._log_cache:94            logger.warning(f"Log ID {log_id} not found in cache")95            # Try to load from file96            log_entry = self._read_log_entry(log_id)97            if log_entry:98                self._log_cache[log_id] = log_entry99            else:100                return101        102        log_entry = self._log_cache[log_id]103        log_entry['execution_success'] = success104        log_entry['error_message'] = error105        log_entry['execution_time_seconds'] = execution_time106        log_entry['rows_returned'] = row_count107        108        # Write updated entry109        self._write_log_entry(log_id, log_entry)110        111        logger.info(f"Updated query result for log: {log_id}")112    113    def log_dba_action(114        self,115        action: str,116        sql: str,117        approved: bool,118        approver: str119    ) -> str:120        """121        Log a DBA mode action with approval status.122        123        Args:124            action: Action type (APPROVED, REJECTED, EXECUTED)125            sql: SQL query126            approved: Whether action was approved127            approver: Session ID of approver128            129        Returns:130            Log entry ID131        """132        log_id = f"dba_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}"133        134        log_entry = {135            "log_id": log_id,136            "timestamp": datetime.now().isoformat(),137            "log_type": "DBA_ACTION",138            "action": action,139            "sql": sql,140            "approved": approved,141            "approver_id": approver142        }143        144        # Write to file145        self._write_log_entry(log_id, log_entry)146        147        logger.info(f"Logged DBA action: {log_id} - {action}")148        return log_id149    150    def get_audit_trail(151        self,152        limit: int = 100,153        mode: Optional[str] = None,154        success_only: bool = False155    ) -> List[Dict]:156        """157        Retrieve audit logs with optional filtering.158        159        Args:160            limit: Maximum number of entries to return161            mode: Filter by mode (readonly or dba)162            success_only: Only return successful queries163            164        Returns:165            List of log entry dictionaries166        """167        audit_trail = []168        169        # Read all log files170        log_files = sorted(self.log_dir.glob("*.json"), reverse=True)171        172        for log_file in log_files[:limit]:173            try:174                with open(log_file, 'r') as f:175                    log_entry = json.load(f)176                177                # Apply filters178                if mode and log_entry.get('mode') != mode:179                    continue180                181                if success_only and not log_entry.get('execution_success'):182                    continue183                184                audit_trail.append(log_entry)185                186            except Exception as e:187                logger.error(f"Failed to read log file {log_file}: {e}")188        189        return audit_trail190    191    def _write_log_entry(self, log_id: str, log_entry: Dict):192        """193        Write log entry to file.194        195        Args:196            log_id: Log entry ID197            log_entry: Log entry dictionary198        """199        try:200            log_file = self.log_dir / f"{log_id}.json"201            with open(log_file, 'w') as f:202                json.dump(log_entry, f, indent=2)203        except Exception as e:204            logger.error(f"Failed to write log entry {log_id}: {e}")205    206    def _read_log_entry(self, log_id: str) -> Optional[Dict]:207        """208        Read log entry from file.209        210        Args:211            log_id: Log entry ID212            213        Returns:214            Log entry dictionary or None if not found215        """216        try:217            log_file = self.log_dir / f"{log_id}.json"218            if log_file.exists():219                with open(log_file, 'r') as f:220                    return json.load(f)221        except Exception as e:222            logger.error(f"Failed to read log entry {log_id}: {e}")223        224        return None225    226    def get_statistics(self) -> Dict:227        """228        Get audit log statistics.229        230        Returns:231            Dictionary with statistics232        """233        total_queries = 0234        successful_queries = 0235        failed_queries = 0236        dba_actions = 0237        238        log_files = list(self.log_dir.glob("*.json"))239        240        for log_file in log_files:241            try:242                with open(log_file, 'r') as f:243                    log_entry = json.load(f)244                245                if log_entry.get('log_type') == 'DBA_ACTION':246                    dba_actions += 1247                else:248                    total_queries += 1249                    if log_entry.get('execution_success'):250                        successful_queries += 1251                    elif log_entry.get('execution_success') is False:252                        failed_queries += 1253                        254            except Exception:255                pass256        257        success_rate = (successful_queries / total_queries * 100) if total_queries > 0 else 0258        259        return {260            "total_queries": total_queries,261            "successful_queries": successful_queries,262            "failed_queries": failed_queries,263            "success_rate": success_rate,264            "dba_actions": dba_actions265        }266 267 268# Global audit logger instance269audit_logger = AuditLogger()270