"""Database operations for Analysis Orchestrator""" import logging from typing import List, Optional, Dict, Any from datetime import datetime from enum import Enum from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy import create_engine, select, update, delete, text from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker, Session from sqlalchemy import Integer, String, Boolean, DateTime, Text, JSON, Enum as SQLEnum from app.config import settings logger = logging.getLogger(__name__) # SQLAlchemy Base class Base(DeclarativeBase): pass # Analysis Status Enum class AnalysisStatus(str, Enum): PENDING = "pending" RUNNING = "running" STOPPED = "stopped" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" INTERRUPTED = "interrupted" # Collection was interrupted (e.g., ingestion-service restart) # Analysis Type Enum class AnalysisType(str, Enum): DEPENDENCY_MAPPING = "dependency_mapping" CHANGE_DETECTION = "change_detection" ANOMALY_DETECTION = "anomaly_detection" BASELINE_CREATION = "baseline_creation" RISK_ASSESSMENT = "risk_assessment" CUSTOM = "custom" # Analysis Model - Must match backend schema exactly! # NOTE: This model reads from the same 'analyses' table created by the backend. # All columns must be defined here even if not used by orchestrator. class Analysis(Base): __tablename__ = "analyses" id: Mapped[int] = mapped_column(Integer, primary_key=True) # Basic info name: Mapped[str] = mapped_column(String(255), nullable=False) description: Mapped[Optional[str]] = mapped_column(Text) # Cluster configuration cluster_id: Mapped[int] = mapped_column(Integer, nullable=False) cluster_ids: Mapped[Optional[str]] = mapped_column(JSON, default=[]) # Multi-cluster: JSONB array is_multi_cluster: Mapped[bool] = mapped_column(Boolean, default=False) # Multi-cluster flag # Scope configuration scope_type: Mapped[str] = mapped_column(String(50), nullable=False, default="cluster") scope_config: Mapped[Optional[dict]] = mapped_column(JSON, default={}) # Gadget configuration gadget_config: Mapped[Optional[dict]] = mapped_column(JSON, default={}) gadget_modules: Mapped[Optional[dict]] = mapped_column(JSON, default=[]) # Legacy column from backend # Time & output configuration time_config: Mapped[Optional[dict]] = mapped_column(JSON, default={}) output_config: Mapped[Optional[dict]] = mapped_column(JSON, default={}) # L7 (Beyla) — must match backend `analyses` table analysis_level: Mapped[str] = mapped_column(String(20), default="l4", nullable=False) l7_config: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True, default=None) # Status fields status: Mapped[str] = mapped_column(String(50), default="draft") is_active: Mapped[bool] = mapped_column(Boolean, default=True) # From backend model # Audit fields created_by: Mapped[Optional[int]] = mapped_column(Integer) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) # Execution timing - for auto-stop monitoring started_at: Mapped[Optional[datetime]] = mapped_column(DateTime) stopped_at: Mapped[Optional[datetime]] = mapped_column(DateTime) # Scheduling support is_scheduled: Mapped[bool] = mapped_column(Boolean, default=False) schedule_expression: Mapped[Optional[str]] = mapped_column(String(100)) schedule_duration_seconds: Mapped[Optional[int]] = mapped_column(Integer) next_run_at: Mapped[Optional[datetime]] = mapped_column(DateTime) last_run_at: Mapped[Optional[datetime]] = mapped_column(DateTime) schedule_run_count: Mapped[int] = mapped_column(Integer, default=0) max_scheduled_runs: Mapped[Optional[int]] = mapped_column(Integer) # Metadata (from backend) - 'metadata' is reserved in SQLAlchemy, use different name analysis_metadata: Mapped[Optional[dict]] = mapped_column("metadata", JSON, default={}) # Analysis Run Model (execution history) # NOTE: Column names MUST match the database schema in 03-migrations-job.yaml (Migration 007) class AnalysisRun(Base): __tablename__ = "analysis_runs" id: Mapped[int] = mapped_column(Integer, primary_key=True) analysis_id: Mapped[int] = mapped_column(Integer, nullable=False) run_number: Mapped[int] = mapped_column(Integer, nullable=False, default=1) status: Mapped[str] = mapped_column(String(50), default="running") start_time: Mapped[Optional[datetime]] = mapped_column(DateTime) # Was: started_at (WRONG) end_time: Mapped[Optional[datetime]] = mapped_column(DateTime) # Was: completed_at (WRONG) duration_seconds: Mapped[Optional[int]] = mapped_column(Integer) events_collected: Mapped[int] = mapped_column(Integer, default=0) workloads_discovered: Mapped[int] = mapped_column(Integer, default=0) communications_discovered: Mapped[int] = mapped_column(Integer, default=0) anomalies_detected: Mapped[int] = mapped_column(Integer, default=0) changes_detected: Mapped[int] = mapped_column(Integer, default=0) error_message: Mapped[Optional[str]] = mapped_column(Text) logs: Mapped[Optional[dict]] = mapped_column(JSON, default=[]) run_metadata: Mapped[Optional[dict]] = mapped_column("metadata", JSON, default={}) # 'metadata' is reserved in SQLAlchemy created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) # Database Manager class DatabaseManager: """Async and Sync database manager""" def __init__(self): # Async engine for async methods self.engine = create_async_engine( settings.database_url, echo=False, pool_pre_ping=True, ) self.async_session = async_sessionmaker( self.engine, class_=AsyncSession, expire_on_commit=False, ) # Sync engine for gRPC sync methods (avoids event loop issues) sync_db_url = settings.database_url.replace("postgresql+asyncpg://", "postgresql://") self.sync_engine = create_engine( sync_db_url, echo=False, pool_pre_ping=True, ) self.sync_session = sessionmaker( self.sync_engine, class_=Session, expire_on_commit=False, ) logger.info(f"Database manager initialized: {settings.postgres_host}") async def create_tables(self): """ Create tables if they don't exist. NOTE: The 'analyses' table is created by the Backend service. We only create 'analysis_runs' table here for execution history. The Analysis model must match the backend's schema exactly. """ async with self.engine.begin() as conn: # Only create analysis_runs table, not analyses (managed by backend) await conn.run_sync( lambda sync_conn: AnalysisRun.__table__.create(sync_conn, checkfirst=True) ) logger.info("Database tables verified (analysis_runs)") # Analysis CRUD async def create_analysis(self, analysis_data: Dict[str, Any]) -> Analysis: """Create new analysis""" async with self.async_session() as session: analysis = Analysis(**analysis_data) session.add(analysis) await session.commit() await session.refresh(analysis) logger.info(f"Created analysis: {analysis.name}") return analysis async def get_analysis(self, analysis_id: int) -> Optional[Analysis]: """Get analysis by ID""" async with self.async_session() as session: result = await session.execute( select(Analysis).where(Analysis.id == analysis_id) ) return result.scalar_one_or_none() async def list_analyses( self, cluster_id: Optional[int] = None, scope_type: Optional[str] = None, status: Optional[str] = None ) -> List[Analysis]: """List analyses with filters""" async with self.async_session() as session: query = select(Analysis) if cluster_id: query = query.where(Analysis.cluster_id == cluster_id) if scope_type: query = query.where(Analysis.scope_type == scope_type) if status: query = query.where(Analysis.status == status) result = await session.execute(query) return list(result.scalars().all()) async def update_analysis( self, analysis_id: int, analysis_data: Dict[str, Any] ) -> Optional[Analysis]: """Update analysis""" async with self.async_session() as session: analysis = await session.get(Analysis, analysis_id) if not analysis: return None now = datetime.utcnow() for key, value in analysis_data.items(): setattr(analysis, key, value) analysis.updated_at = now # Check if status is being updated and set timing fields new_status = analysis_data.get('status') if new_status: if new_status == "running": # Always update started_at on every Start (for auto-stop timing on restart) analysis.started_at = now logger.info(f"Analysis {analysis_id} started at {now}") elif new_status in ("stopped", "completed", "failed"): analysis.stopped_at = now logger.info(f"Analysis {analysis_id} stopped at {now}") await session.commit() await session.refresh(analysis) logger.info(f"Updated analysis: {analysis.name}") return analysis async def delete_analysis(self, analysis_id: int) -> bool: """Delete analysis (async version)""" async with self.async_session() as session: analysis = await session.get(Analysis, analysis_id) if not analysis: return False await session.delete(analysis) await session.commit() logger.info(f"Deleted analysis: {analysis.name}") return True def delete_analysis_sync(self, analysis_id: int) -> bool: """Delete analysis (sync version for gRPC)""" try: with self.sync_session() as session: analysis = session.get(Analysis, analysis_id) if not analysis: return False session.delete(analysis) session.commit() logger.info(f"Deleted analysis: {analysis.name}") return True except Exception as e: logger.error(f"Failed to delete analysis: {e}") return False def update_analysis_sync( self, analysis_id: int, analysis_data: Dict[str, Any] ) -> bool: """Update analysis (sync version for gRPC)""" try: with self.sync_session() as session: analysis = session.get(Analysis, analysis_id) if not analysis: return False now = datetime.utcnow() for key, value in analysis_data.items(): setattr(analysis, key, value) analysis.updated_at = now session.commit() logger.info(f"Updated analysis {analysis_id} (sync)") return True except Exception as e: logger.error(f"Failed to update analysis (sync): {e}") return False async def update_analysis_status( self, analysis_id: int, status: AnalysisStatus ) -> bool: """Update analysis status (async version)""" async with self.async_session() as session: analysis = await session.get(Analysis, analysis_id) if not analysis: return False now = datetime.utcnow() status_value = status.value if hasattr(status, 'value') else status analysis.status = status_value analysis.updated_at = now # Set started_at when analysis starts running if status_value == "running": analysis.started_at = now logger.info(f"Analysis {analysis_id} started at {now}") # Set stopped_at when analysis stops if status_value in ("stopped", "completed", "failed"): analysis.stopped_at = now logger.info(f"Analysis {analysis_id} stopped at {now}") await session.commit() logger.info(f"Updated analysis {analysis_id} status to {status}") return True def update_analysis_status_sync( self, analysis_id: int, status: AnalysisStatus ) -> bool: """Update analysis status (sync version for gRPC)""" try: with self.sync_session() as session: analysis = session.get(Analysis, analysis_id) if not analysis: return False now = datetime.utcnow() status_value = status.value if hasattr(status, 'value') else status analysis.status = status_value analysis.updated_at = now # Set started_at when analysis starts running if status_value == "running": analysis.started_at = now logger.info(f"Analysis {analysis_id} started at {now}") # Set stopped_at when analysis stops if status_value in ("stopped", "completed", "failed"): analysis.stopped_at = now logger.info(f"Analysis {analysis_id} stopped at {now}") session.commit() logger.info(f"Updated analysis {analysis_id} status to {status}") return True except Exception as e: logger.error(f"Failed to update analysis status: {e}") return False def get_analysis_sync(self, analysis_id: int) -> Optional[Analysis]: """Get analysis by ID (sync version for gRPC)""" try: with self.sync_session() as session: analysis = session.get(Analysis, analysis_id) if analysis: logger.info(f"Found analysis {analysis_id}: {analysis.name}, status={analysis.status}") else: # Debug: try raw SQL to see if data exists from sqlalchemy import text result = session.execute(text(f"SELECT id, name, status FROM analyses WHERE id = {analysis_id}")) row = result.fetchone() if row: logger.error(f"Analysis {analysis_id} exists in DB but SQLAlchemy can't load it! Row: {row}") else: logger.warning(f"Analysis {analysis_id} not found in database (raw SQL confirms)") return analysis except Exception as e: logger.error(f"Failed to get analysis {analysis_id}: {e}", exc_info=True) return None def get_running_analyses_sync(self) -> List[Dict[str, Any]]: """Get all analyses with status 'running' (sync version)""" try: with self.sync_session() as session: result = session.execute( select(Analysis).where(Analysis.status == "running") ) analyses = result.scalars().all() return [ { "id": a.id, "name": a.name, "cluster_id": a.cluster_id, "status": a.status } for a in analyses ] except Exception as e: logger.error(f"Failed to get running analyses: {e}") return [] def get_cluster_sync(self, cluster_id: int) -> Optional[Dict[str, Any]]: """Get cluster by ID (sync version for gRPC) - includes credentials for remote clusters""" try: with self.sync_session() as session: result = session.execute( text(""" SELECT id, name, gadget_endpoint, connection_type, api_server_url, token_encrypted, ca_cert_encrypted, kubeconfig_encrypted, skip_tls_verify, gadget_namespace, gadget_version, beyla_namespace, beyla_health_status FROM clusters WHERE id = :id """), {"id": cluster_id} ) row = result.fetchone() if row: return { "id": row[0], "name": row[1], "gadget_endpoint": row[2], "connection_type": row[3], "api_server_url": row[4], "token_encrypted": row[5], "ca_cert_encrypted": row[6], "kubeconfig_encrypted": row[7], "skip_tls_verify": row[8] or False, "gadget_namespace": row[9], "gadget_version": row[10] or '', "beyla_namespace": row[11], "beyla_health_status": row[12], } return None except Exception as e: logger.error(f"Failed to get cluster: {e}") return None # Cluster lookup (read-only, uses raw SQL since clusters table is managed by backend) async def get_cluster(self, cluster_id: int) -> Optional[Dict[str, Any]]: """Get cluster by ID from shared database - includes credentials for remote clusters""" from sqlalchemy import text async with self.async_session() as session: result = await session.execute( text(""" SELECT id, name, gadget_endpoint, connection_type, api_server_url, token_encrypted, ca_cert_encrypted, kubeconfig_encrypted, skip_tls_verify, gadget_namespace, beyla_namespace, beyla_health_status FROM clusters WHERE id = :id """), {"id": cluster_id} ) row = result.fetchone() if row: return { "id": row[0], "name": row[1], "gadget_endpoint": row[2], "connection_type": row[3], "api_server_url": row[4], "token_encrypted": row[5], "ca_cert_encrypted": row[6], "kubeconfig_encrypted": row[7], "skip_tls_verify": row[8] or False, "gadget_namespace": row[9], # From UI, no fallback "beyla_namespace": row[10], "beyla_health_status": row[11], } return None # Analysis Run CRUD async def create_analysis_run(self, run_data: Dict[str, Any]) -> AnalysisRun: """Create new analysis run""" async with self.async_session() as session: run = AnalysisRun(**run_data) session.add(run) await session.commit() await session.refresh(run) return run async def update_analysis_run( self, run_id: int, run_data: Dict[str, Any] ) -> Optional[AnalysisRun]: """Update analysis run""" async with self.async_session() as session: run = await session.get(AnalysisRun, run_id) if not run: return None for key, value in run_data.items(): setattr(run, key, value) await session.commit() await session.refresh(run) return run async def list_analysis_runs( self, analysis_id: int, limit: int = 10 ) -> List[AnalysisRun]: """List analysis runs for an analysis""" async with self.async_session() as session: query = ( select(AnalysisRun) .where(AnalysisRun.analysis_id == analysis_id) .order_by(AnalysisRun.created_at.desc()) .limit(limit) ) result = await session.execute(query) return list(result.scalars().all()) async def get_running_run_for_analysis(self, analysis_id: int) -> Optional[AnalysisRun]: """Get the currently running run for an analysis (if any)""" async with self.async_session() as session: query = ( select(AnalysisRun) .where(AnalysisRun.analysis_id == analysis_id) .where(AnalysisRun.status == "running") .order_by(AnalysisRun.created_at.desc()) .limit(1) ) result = await session.execute(query) return result.scalars().first() def get_scheduled_analyses_sync(self) -> List[Dict[str, Any]]: """Get all analyses with is_scheduled=true (for startup job restoration)""" try: with self.sync_session() as session: result = session.execute( select(Analysis).where(Analysis.is_scheduled == True) ) analyses = result.scalars().all() return [ { "id": a.id, "name": a.name, "schedule_expression": a.schedule_expression, "schedule_duration_seconds": a.schedule_duration_seconds, "max_scheduled_runs": a.max_scheduled_runs, "schedule_run_count": a.schedule_run_count, } for a in analyses ] except Exception as e: logger.error(f"Failed to get scheduled analyses: {e}") return [] async def close(self): """Close database connection""" await self.engine.dispose() logger.info("Database connection closed") # Global database instance db_manager = DatabaseManager()