bmc_hub/app/core/database.py

"""
Database Module
PostgreSQL connection and helpers using psycopg2
"""

import psycopg2
from psycopg2.extras import RealDictCursor
from psycopg2.pool import SimpleConnectionPool
from typing import Optional
import logging

from app.core.config import settings

logger = logging.getLogger(__name__)

# Connection pool
connection_pool: Optional[SimpleConnectionPool] = None


def init_db():
    """Initialize database connection pool"""
    global connection_pool
    
    try:
        connection_pool = SimpleConnectionPool(
            minconn=1,
            maxconn=10,
            dsn=settings.DATABASE_URL
        )
        logger.info("✅ Database connection pool initialized")
    except Exception as e:
        logger.error(f"❌ Failed to initialize database: {e}")
        raise


def get_db_connection():
    """Get a connection from the pool"""
    if connection_pool:
        return connection_pool.getconn()
    raise Exception("Database pool not initialized")


def release_db_connection(conn):
    """Return a connection to the pool"""
    if connection_pool:
        connection_pool.putconn(conn)


def get_db():
    """Context manager for database connections"""
    conn = get_db_connection()
    try:
        yield conn
    finally:
        release_db_connection(conn)


def execute_query(query: str, params: tuple = None, fetchone: bool = False):
    """
    Execute a SQL query and return results
    
    Args:
        query: SQL query string
        params: Query parameters tuple
        fetchone: If True, return single row dict, otherwise list of dicts
        
    Returns:
        Single dict if fetchone=True, otherwise list of dicts
    """
    conn = get_db_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, params or ())
            if fetchone:
                row = cursor.fetchone()
                return dict(row) if row else None
            else:
                rows = cursor.fetchall()
                return [dict(row) for row in rows]
    except Exception as e:
        conn.rollback()
        logger.error(f"Query error: {e}")
        raise
    finally:
        release_db_connection(conn)


def execute_insert(query: str, params: tuple = ()) -> Optional[int]:
    """
    Execute an INSERT query and return last row id
    
    Args:
        query: SQL INSERT query (will add RETURNING id if not present)
        params: Query parameters tuple
    
    Returns:
        Last inserted row ID or None
    """
    conn = get_db_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cursor:
            # PostgreSQL requires RETURNING clause
            if "RETURNING" not in query.upper():
                query = query.rstrip(";") + " RETURNING id"
            cursor.execute(query, params)
            result = cursor.fetchone()
            conn.commit()
            
            # If result exists, return the first column value (typically ID)
            if result:
                # If it's a dict, get first value
                if isinstance(result, dict):
                    return list(result.values())[0]
                # If it's a tuple/list, get first element
                return result[0]
            return None
    except Exception as e:
        conn.rollback()
        logger.error(f"Insert error: {e}")
        raise
    finally:
        release_db_connection(conn)


def execute_update(query: str, params: tuple = ()) -> int:
    """
    Execute an UPDATE/DELETE query and return affected rows
    
    Args:
        query: SQL UPDATE/DELETE query
        params: Query parameters tuple
    
    Returns:
        Number of affected rows
    """
    conn = get_db_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, params)
            rowcount = cursor.rowcount
            conn.commit()
            return rowcount
    except Exception as e:
        conn.rollback()
        logger.error(f"Update error: {e}")
        raise
    finally:
        release_db_connection(conn)