""" 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()