db/models.py

import asyncpg
import logging
from typing import List, Dict, Any, Union

from config.settings import DATABASE_URL

logger = logging.getLogger(__name__)

async def fetch_models_for_group(resource_group: str) -> List[Dict[str, Any]]:
    """Récupérer les détails des modèles depuis la base de données pour un groupe de ressources spécifique.
    
    Args:
        resource_group: Identifiant du groupe de ressources
        
    Returns:
        Liste de dictionnaires contenant les informations des modèles
    
    Raises:
        Exception: Si une erreur se produit lors de la connexion ou de la requête
    """
    conn = None
    try:
        # Désactiver le cache des prepared statements pour compatibilité avec pgbouncer
        conn = await asyncpg.connect(DATABASE_URL, statement_cache_size=0)
        logger.info(f"Successfully connected to database")

        # Récupérer les modèles spécifiques à ce groupe
        query = """
            SELECT 
                model_id, 
                display_name, 
                hf_repo_id, 
                hf_subfolder, 
                hf_filename,
                updated_at
            FROM models
            WHERE hf_resource_group = $1
        """
        rows = await conn.fetch(query, resource_group)
        logger.info(f"Found {len(rows)} models for group '{resource_group}'")
        
        return [dict(row) for row in rows]
    except Exception as e:
        logger.error(f"Database error: {e}", exc_info=True)
        raise
    finally:
        if conn and not conn.is_closed():
            await conn.close()
            logger.debug("Database connection closed")


async def fetch_model_by_id(model_id: str) -> Union[Dict[str, Any], None]:
    """Récupérer les détails d'un modèle spécifique par son ID de base de données.

    Args:
        model_id: L'ID du modèle dans la base de données (peut être int ou str selon le schéma).

    Returns:
        Un dictionnaire contenant les informations du modèle si trouvé, sinon None.

    Raises:
        Exception: Si une erreur se produit lors de la connexion ou de la requête.
    """
    conn = None
    try:
        # Désactiver le cache des prepared statements pour compatibilité avec pgbouncer
        conn = await asyncpg.connect(DATABASE_URL, statement_cache_size=0)
        logger.debug(f"Successfully connected to database to fetch model ID: {model_id}")

        # Récupérer le modèle spécifique par son ID
        query = """
            SELECT 
                model_id, 
                display_name, 
                hf_repo_id, 
                hf_subfolder, 
                hf_filename,
                updated_at
            FROM models
            WHERE model_id = $1
        """
        row = await conn.fetchrow(query, model_id)
        
        if row:
            logger.info(f"Found model with ID '{model_id}': {row['display_name']}")
            return dict(row)
        else:
            logger.warning(f"No model found with ID '{model_id}'")
            return None
            
    except Exception as e:
        logger.error(f"Database error fetching model ID {model_id}: {e}", exc_info=True)
        raise
    finally:
        if conn and not conn.is_closed():
            await conn.close()
            logger.debug(f"Database connection closed after fetching model ID: {model_id}")