database_manager.py

#!/usr/bin/env python3
"""
Database Manager for Catacomb - PostgreSQL connection and operations.

This module provides database connection management and CRUD operations
for all Catacomb data models.
"""

import os
from typing import Optional, List, Dict, Any
from sqlalchemy import create_engine, MetaData
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool
from contextlib import contextmanager
import logging

from database_schema import (
    Base,
    User,
    Asset,
    Evaluation,
    ReviewCard,
    AuditEntry,
    JobDescription,
    Competency,
    CompetencyGraph,
    Outcome,
    UserRole,
    AssetClassification,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DatabaseManager:
    """
    Database manager for Catacomb.
    
    Handles PostgreSQL connections, session management, and CRUD operations.
    """
    
    def __init__(self, database_url: Optional[str] = None):
        """
        Initialize the database manager.
        
        Args:
            database_url: PostgreSQL connection URL. If not provided, uses DATABASE_URL env var.
        """
        self.database_url = database_url or os.getenv("DATABASE_URL")
        
        if not self.database_url:
            raise ValueError(
                "DATABASE_URL environment variable not set. "
                "Please set DATABASE_URL to your PostgreSQL connection string."
            )
        
        # Create engine with connection pooling
        self.engine = create_engine(
            self.database_url,
            poolclass=QueuePool,
            pool_size=10,
            max_overflow=20,
            pool_pre_ping=True,  # Verify connections before using
            echo=False,  # Set to True for SQL query logging
        )
        
        # Create session factory
        self.SessionLocal = sessionmaker(
            autocommit=False,
            autoflush=False,
            bind=self.engine,
        )
        
        logger.info("Database manager initialized")
    
    def create_tables(self) -> None:
        """Create all tables in the database."""
        Base.metadata.create_all(bind=self.engine)
        logger.info("All tables created successfully")
    
    def drop_tables(self) -> None:
        """Drop all tables from the database. Use with caution!"""
        Base.metadata.drop_all(bind=self.engine)
        logger.warning("All tables dropped")
    
    @contextmanager
    def get_session(self) -> Session:
        """
        Context manager for database sessions.
        
        Yields:
            Session: SQLAlchemy session
        """
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"Database error: {e}")
            raise
        finally:
            session.close()
    
    # User operations
    def create_user(
        self,
        email: str,
        password_hash: str,
        role: UserRole = UserRole.PUBLIC,
    ) -> User:
        """Create a new user."""
        with self.get_session() as session:
            user = User(
                email=email,
                password_hash=password_hash,
                role=role,
            )
            session.add(user)
            session.flush()
            session.refresh(user)
            return user
    
    def get_user_by_email(self, email: str) -> Optional[User]:
        """Get a user by email."""
        with self.get_session() as session:
            return session.query(User).filter(User.email == email).first()
    
    def get_user_by_id(self, user_id: int) -> Optional[User]:
        """Get a user by ID."""
        with self.get_session() as session:
            return session.query(User).filter(User.id == user_id).first()
    
    # Asset operations
    def create_asset(
        self,
        asset_id: str,
        asset_name: str,
        classification: AssetClassification,
        owner_id: Optional[int] = None,
        **kwargs,
    ) -> Asset:
        """Create a new asset."""
        with self.get_session() as session:
            asset = Asset(
                asset_id=asset_id,
                asset_name=asset_name,
                classification=classification,
                owner_id=owner_id,
                **kwargs,
            )
            session.add(asset)
            session.flush()
            session.refresh(asset)
            return asset
    
    def get_asset_by_id(self, asset_id: str) -> Optional[Asset]:
        """Get an asset by asset_id."""
        with self.get_session() as session:
            return session.query(Asset).filter(Asset.asset_id == asset_id).first()
    
    def get_assets_by_owner(self, owner_id: int) -> List[Asset]:
        """Get all assets for a specific owner."""
        with self.get_session() as session:
            return session.query(Asset).filter(Asset.owner_id == owner_id).all()
    
    def update_asset(self, asset_id: str, **kwargs) -> Optional[Asset]:
        """Update an asset."""
        with self.get_session() as session:
            asset = session.query(Asset).filter(Asset.asset_id == asset_id).first()
            if asset:
                for key, value in kwargs.items():
                    setattr(asset, key, value)
                session.flush()
                session.refresh(asset)
            return asset
    
    # Evaluation operations
    def create_evaluation(
        self,
        evaluation_id: str,
        asset_id: int,
        user_id: Optional[int] = None,
        **kwargs,
    ) -> Evaluation:
        """Create a new evaluation."""
        with self.get_session() as session:
            evaluation = Evaluation(
                evaluation_id=evaluation_id,
                asset_id=asset_id,
                user_id=user_id,
                **kwargs,
            )
            session.add(evaluation)
            session.flush()
            session.refresh(evaluation)
            return evaluation
    
    def get_evaluation_by_id(self, evaluation_id: str) -> Optional[Evaluation]:
        """Get an evaluation by ID."""
        with self.get_session() as session:
            return session.query(Evaluation).filter(Evaluation.evaluation_id == evaluation_id).first()
    
    def get_evaluations_by_asset(self, asset_id: int) -> List[Evaluation]:
        """Get all evaluations for an asset."""
        with self.get_session() as session:
            return session.query(Evaluation).filter(Evaluation.asset_id == asset_id).all()
    
    def get_evaluations_by_user(self, user_id: int) -> List[Evaluation]:
        """Get all evaluations by a user."""
        with self.get_session() as session:
            return session.query(Evaluation).filter(Evaluation.user_id == user_id).all()
    
    # Review card operations
    def create_review_card(
        self,
        card_id: str,
        asset_id: int,
        **kwargs,
    ) -> ReviewCard:
        """Create a new review card."""
        with self.get_session() as session:
            review_card = ReviewCard(
                card_id=card_id,
                asset_id=asset_id,
                **kwargs,
            )
            session.add(review_card)
            session.flush()
            session.refresh(review_card)
            return review_card
    
    def get_review_card_by_id(self, card_id: str) -> Optional[ReviewCard]:
        """Get a review card by ID."""
        with self.get_session() as session:
            return session.query(ReviewCard).filter(ReviewCard.card_id == card_id).first()
    
    def get_review_cards_by_asset(self, asset_id: int) -> List[ReviewCard]:
        """Get all review cards for an asset."""
        with self.get_session() as session:
            return session.query(ReviewCard).filter(ReviewCard.asset_id == asset_id).all()
    
    # Audit entry operations
    def create_audit_entry(
        self,
        entry_id: str,
        evaluator_name: str,
        decision_type: str,
        decision_value: Any,
        reasoning: str,
        evaluation_id: Optional[int] = None,
        **kwargs,
    ) -> AuditEntry:
        """Create a new audit entry."""
        with self.get_session() as session:
            audit_entry = AuditEntry(
                entry_id=entry_id,
                evaluator_name=evaluator_name,
                decision_type=decision_type,
                decision_value=decision_value,
                reasoning=reasoning,
                evaluation_id=evaluation_id,
                **kwargs,
            )
            session.add(audit_entry)
            session.flush()
            session.refresh(audit_entry)
            return audit_entry
    
    def get_audit_entries_by_evaluation(self, evaluation_id: int) -> List[AuditEntry]:
        """Get all audit entries for an evaluation."""
        with self.get_session() as session:
            return session.query(AuditEntry).filter(AuditEntry.evaluation_id == evaluation_id).all()
    
    def get_audit_entries_by_evaluator(self, evaluator_name: str) -> List[AuditEntry]:
        """Get all audit entries from a specific evaluator."""
        with self.get_session() as session:
            return session.query(AuditEntry).filter(AuditEntry.evaluator_name == evaluator_name).all()
    
    # Job description operations
    def create_job_description(
        self,
        job_id: str,
        role_family: str,
        source: str,
        **kwargs,
    ) -> JobDescription:
        """Create a new job description."""
        with self.get_session() as session:
            job_description = JobDescription(
                job_id=job_id,
                role_family=role_family,
                source=source,
                **kwargs,
            )
            session.add(job_description)
            session.flush()
            session.refresh(job_description)
            return job_description
    
    def get_job_descriptions_by_role_family(self, role_family: str) -> List[JobDescription]:
        """Get all job descriptions for a role family."""
        with self.get_session() as session:
            return session.query(JobDescription).filter(JobDescription.role_family == role_family).all()
    
    def get_job_descriptions_by_source(self, source: str) -> List[JobDescription]:
        """Get all job descriptions from a source."""
        with self.get_session() as session:
            return session.query(JobDescription).filter(JobDescription.source == source).all()
    
    # Competency operations
    def create_competency(
        self,
        competency_id: str,
        name: str,
        domain: str,
        **kwargs,
    ) -> Competency:
        """Create a new competency."""
        with self.get_session() as session:
            competency = Competency(
                competency_id=competency_id,
                name=name,
                domain=domain,
                **kwargs,
            )
            session.add(competency)
            session.flush()
            session.refresh(competency)
            return competency
    
    def get_competency_by_id(self, competency_id: str) -> Optional[Competency]:
        """Get a competency by ID."""
        with self.get_session() as session:
            return session.query(Competency).filter(Competency.competency_id == competency_id).first()
    
    def get_competencies_by_domain(self, domain: str) -> List[Competency]:
        """Get all competencies in a domain."""
        with self.get_session() as session:
            return session.query(Competency).filter(Competency.domain == domain).all()
    
    # Competency graph operations
    def create_competency_graph(
        self,
        graph_id: str,
        name: str,
        **kwargs,
    ) -> CompetencyGraph:
        """Create a new competency graph."""
        with self.get_session() as session:
            competency_graph = CompetencyGraph(
                graph_id=graph_id,
                name=name,
                **kwargs,
            )
            session.add(competency_graph)
            session.flush()
            session.refresh(competency_graph)
            return competency_graph
    
    def get_active_competency_graph(self) -> Optional[CompetencyGraph]:
        """Get the active competency graph."""
        with self.get_session() as session:
            return session.query(CompetencyGraph).filter(CompetencyGraph.is_active == True).first()
    
    def set_active_competency_graph(self, graph_id: str) -> Optional[CompetencyGraph]:
        """Set a competency graph as active."""
        with self.get_session() as session:
            # Deactivate all graphs
            session.query(CompetencyGraph).update({"is_active": False})
            
            # Activate the specified graph
            graph = session.query(CompetencyGraph).filter(CompetencyGraph.graph_id == graph_id).first()
            if graph:
                graph.is_active = True
                session.flush()
                session.refresh(graph)
            return graph
    
    # Outcome operations
    def create_outcome(
        self,
        outcome_id: str,
        intervention_type: str,
        asset_id: str,
        **kwargs,
    ) -> Outcome:
        """Create a new outcome."""
        with self.get_session() as session:
            outcome = Outcome(
                outcome_id=outcome_id,
                intervention_type=intervention_type,
                asset_id=asset_id,
                **kwargs,
            )
            session.add(outcome)
            session.flush()
            session.refresh(outcome)
            return outcome
    
    def get_verified_outcomes(self) -> List[Outcome]:
        """Get all verified outcomes."""
        with self.get_session() as session:
            return session.query(Outcome).filter(Outcome.is_verified == True).all()
    
    def get_outcomes_by_asset(self, asset_id: str) -> List[Outcome]:
        """Get all outcomes for an asset."""
        with self.get_session() as session:
            return session.query(Outcome).filter(Outcome.asset_id == asset_id).all()
    
    # Health check
    def health_check(self) -> bool:
        """Check if database connection is healthy."""
        try:
            with self.get_session() as session:
                session.execute("SELECT 1")
            return True
        except Exception as e:
            logger.error(f"Database health check failed: {e}")
            return False


# Singleton instance
_db_manager: Optional[DatabaseManager] = None


def get_database_manager() -> DatabaseManager:
    """Get the singleton database manager instance."""
    global _db_manager
    if _db_manager is None:
        _db_manager = DatabaseManager()
    return _db_manager


def init_database(database_url: Optional[str] = None) -> DatabaseManager:
    """
    Initialize the database manager and create tables.
    
    Args:
        database_url: Optional database URL. If not provided, uses DATABASE_URL env var.
    
    Returns:
        DatabaseManager instance
    """
    db_manager = DatabaseManager(database_url)
    db_manager.create_tables()
    return db_manager