Kairo-Jinkendo/backend/db.py

"""PostgreSQL connection helpers."""

from __future__ import annotations

import os
import time

import psycopg2
from psycopg2.extensions import connection


def db_params() -> dict[str, str]:
    return {
        "host": os.getenv("DB_HOST", "localhost"),
        "port": os.getenv("DB_PORT", "5432"),
        "dbname": os.getenv("DB_NAME", "kairo_dev"),
        "user": os.getenv("DB_USER", "kairo_dev"),
        "password": os.getenv("DB_PASSWORD", "dev_password"),
    }


def get_connection() -> connection:
    p = db_params()
    return psycopg2.connect(
        host=p["host"],
        port=p["port"],
        database=p["dbname"],
        user=p["user"],
        password=p["password"],
    )


def is_retryable_db_error(exc: Exception) -> bool:
    """Nur temporäre Verbindungsprobleme erneut versuchen (nicht Auth/Config)."""
    msg = str(exc).lower()
    if "password authentication failed" in msg:
        return False
    if "does not exist" in msg and any(token in msg for token in ("role", "database")):
        return False
    if "no pg_hba.conf entry" in msg:
        return False
    return True


def connect_with_retry(max_retries: int = 30, sleep_seconds: float = 2.0) -> connection:
    p = db_params()
    last_exc: Exception | None = None
    for attempt in range(max_retries):
        try:
            conn = get_connection()
            conn.autocommit = False
            print(f"[OK] Connected to database: {p['dbname']}")
            return conn
        except psycopg2.OperationalError as exc:
            last_exc = exc
            if not is_retryable_db_error(exc):
                raise
            if attempt >= max_retries - 1:
                raise
            print(f"Waiting for database... ({attempt + 1}/{max_retries})")
            time.sleep(sleep_seconds)
    if last_exc:
        raise last_exc
    raise RuntimeError("database connection failed")


def check_db() -> bool:
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT 1")
            cur.fetchone()
        return True
    finally:
        conn.close()