File size: 3,368 Bytes
e109700
 
48b9b47
e109700
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
db280f4
 
e109700
 
 
 
 
 
 
 
 
dfb1c84
 
e109700
 
 
 
 
 
 
 
 
 
 
 
 
 
0053356
 
48b9b47
0053356
 
 
 
 
 
 
 
 
 
 
 
 
db280f4
 
0053356
 
 
 
 
 
 
 
 
dfb1c84
 
0053356
 
 
 
 
 
c2cd706
0053356
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
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}")