from prometheus_client import start_http_server, Gauge
import psycopg2
import time

# Configuration de la connexion à la base de données
conn_params = {
    "host": "admin",
    "database": "admin",
    "user": "admin",
    "password": "admin"
}

# Définition des métriques Prometheus
QUERY_CALLS = Gauge('postgresql_query_calls', 'Nombre d\'appels de requêtes PostgreSQL', ['query'])
QUERY_TOTAL_TIME = Gauge('postgresql_query_total_time_ms', 'Temps total des requêtes PostgreSQL en millisecondes', ['query'])

def generate_metrics():
    try:
        # Connexion à la base de données
        conn = psycopg2.connect(**conn_params)
        cur = conn.cursor()

        # Exécution de la requête pour récupérer les informations sur les requêtes
        cur.execute("""
            SELECT
                query,
                calls,
                total_time
            FROM pg_stat_statements
            ORDER BY total_time DESC;
        """)

        # Mise à jour des métriques Prometheus
        for row in cur:
            query = row[0].replace("\\", "\\\\").replace('"', '\\"')  # échappement pour Prometheus
            QUERY_CALLS.labels(query=query).set(row[1])
            QUERY_TOTAL_TIME.labels(query=query).set(row[2] * 1000)  # Convertir le temps en millisecondes

        cur.close()
        conn.close()

    except psycopg2.Error as e:
        print(f"Erreur de connexion à la base de données: {e}")

# Fonction principale pour démarrer l'exporteur
if __name__ == '__main__':
    start_http_server(8000)  # Démarre un serveur HTTP sur le port 8000 pour exposer les métriques
    while True:
        generate_metrics()  # Génére les métriques
        time.sleep(60)  # Intervalle d'exécution, ici toutes les 60 secondes