giglez/config/database.py

"""
Database configuration for GigLez

Manages PostgreSQL connection with PostGIS support
"""

import os
from typing import Optional
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool
from loguru import logger


class DatabaseConfig:
    """Database configuration and connection management"""

    def __init__(
        self,
        host: str = "localhost",
        port: int = 5432,
        database: str = "giglez",
        user: str = "giglez_user",
        password: str = "giglez_secure_password_2026",
        pool_size: int = 10,
        max_overflow: int = 20,
        echo: bool = False
    ):
        """
        Initialize database configuration

        Args:
            host: Database host
            port: Database port
            database: Database name
            user: Database user
            password: Database password
            pool_size: Connection pool size (Wigle pattern: moderate pooling)
            max_overflow: Max overflow connections
            echo: Echo SQL queries (debug mode)
        """
        self.host = host
        self.port = port
        self.database = database
        self.user = user
        self.password = password
        self.pool_size = pool_size
        self.max_overflow = max_overflow
        self.echo = echo

        self._engine: Optional = None
        self._session_factory: Optional[sessionmaker] = None

    @property
    def connection_string(self) -> str:
        """Generate PostgreSQL connection string"""
        return f"postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}"

    @property
    def connection_string_safe(self) -> str:
        """Generate connection string without password (for logging)"""
        return f"postgresql://{self.user}:****@{self.host}:{self.port}/{self.database}"

    def get_engine(self):
        """
        Get or create SQLAlchemy engine

        Uses connection pooling for performance (Wigle pattern)
        """
        if self._engine is None:
            logger.info(f"Creating database engine: {self.connection_string_safe}")

            self._engine = create_engine(
                self.connection_string,
                poolclass=QueuePool,
                pool_size=self.pool_size,
                max_overflow=self.max_overflow,
                pool_pre_ping=True,  # Verify connections before using
                echo=self.echo,
                connect_args={
                    "options": "-c timezone=utc"  # Always use UTC
                }
            )

            logger.success("Database engine created successfully")

        return self._engine

    def get_session_factory(self) -> sessionmaker:
        """Get or create session factory"""
        if self._session_factory is None:
            engine = self.get_engine()
            self._session_factory = sessionmaker(
                bind=engine,
                autocommit=False,
                autoflush=False
            )

        return self._session_factory

    def get_session(self) -> Session:
        """Create a new database session"""
        factory = self.get_session_factory()
        return factory()

    def test_connection(self) -> bool:
        """
        Test database connection

        Returns:
            True if connection successful, False otherwise
        """
        try:
            engine = self.get_engine()
            with engine.connect() as conn:
                result = conn.execute("SELECT 1")
                assert result.fetchone()[0] == 1

            logger.success("Database connection test successful")
            return True

        except Exception as e:
            logger.error(f"Database connection test failed: {e}")
            return False

    def test_postgis(self) -> bool:
        """
        Test PostGIS extension

        Returns:
            True if PostGIS is available, False otherwise
        """
        try:
            engine = self.get_engine()
            with engine.connect() as conn:
                result = conn.execute("SELECT PostGIS_Version()")
                version = result.fetchone()[0]

            logger.success(f"PostGIS available: {version}")
            return True

        except Exception as e:
            logger.error(f"PostGIS test failed: {e}")
            return False

    def close(self):
        """Close database connections"""
        if self._engine:
            self._engine.dispose()
            logger.info("Database connections closed")


# =============================================================================
# GLOBAL CONFIGURATION
# =============================================================================

# Load from environment variables (production pattern)
_db_config = DatabaseConfig(
    host=os.getenv("GIGLEZ_DB_HOST", "localhost"),
    port=int(os.getenv("GIGLEZ_DB_PORT", "5432")),
    database=os.getenv("GIGLEZ_DB_NAME", "giglez"),
    user=os.getenv("GIGLEZ_DB_USER", "giglez_user"),
    password=os.getenv("GIGLEZ_DB_PASSWORD", "giglez_secure_password_2026"),
    pool_size=int(os.getenv("GIGLEZ_DB_POOL_SIZE", "10")),
    max_overflow=int(os.getenv("GIGLEZ_DB_MAX_OVERFLOW", "20")),
    echo=os.getenv("GIGLEZ_DB_ECHO", "false").lower() == "true"
)


def get_db_config() -> DatabaseConfig:
    """Get global database configuration"""
    return _db_config


def get_db_session() -> Session:
    """
    Get a new database session

    Usage:
        with get_db_session() as session:
            # Use session
            pass
    """
    return _db_config.get_session()


def get_db_engine():
    """Get global database engine"""
    return _db_config.get_engine()


# =============================================================================
# CONTEXT MANAGER
# =============================================================================

class DatabaseSession:
    """
    Context manager for database sessions

    Usage:
        with DatabaseSession() as session:
            captures = session.query(Capture).all()
    """

    def __init__(self):
        self.session: Optional[Session] = None

    def __enter__(self) -> Session:
        self.session = get_db_session()
        return self.session

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            # Rollback on exception
            self.session.rollback()
            logger.warning(f"Database session rolled back due to: {exc_val}")
        else:
            # Commit on success
            self.session.commit()

        self.session.close()


# =============================================================================
# TESTING UTILITIES
# =============================================================================

def test_database_connection():
    """Test database connection and PostGIS"""
    config = get_db_config()

    logger.info("Testing database connection...")
    conn_ok = config.test_connection()

    logger.info("Testing PostGIS extension...")
    postgis_ok = config.test_postgis()

    if conn_ok and postgis_ok:
        logger.success("✅ Database fully operational")
        return True
    else:
        logger.error("❌ Database tests failed")
        return False


if __name__ == "__main__":
    # Run tests when executed directly
    test_database_connection()