Backend/db/connection.py

import psycopg2
import os
from contextlib import contextmanager
from typing import Generator
from loguru import logger

from utils.loadDotEnv import initializeENV

initializeENV()


def PSQLConnect():
    conn_str = os.getenv('POSTGRES_CONNECTION')

    if not conn_str:
        logger.error("POSTGRES_CONNECTION не найден в .env файле")
        raise ValueError("POSTGRES_CONNECTION не найден в .env файле")

    conn = psycopg2.connect(conn_str)
    logger.debug("Подключение к БД установлено")
    return conn


def PSQLCursor(conn):
    cur = conn.cursor()
    logger.debug("Курсор БД создан")
    return cur


@contextmanager
def get_connection() -> Generator[psycopg2.extensions.connection, None, None]:
    conn = None
    try:
        conn = PSQLConnect()
        logger.debug("Контекст подключения к БД открыт")
        yield conn
    except Exception as e:
        logger.error(f"Ошибка в контексте подключения: {e}")
        if conn:
            conn.rollback()
            logger.debug("Откат транзакции выполнен")
        raise
    finally:
        if conn:
            conn.close()
            logger.debug("Подключение к БД закрыто")


def test_connection() -> bool:
    try:
        with get_connection() as conn:
            cur = PSQLCursor(conn)
            cur.execute("SELECT version();")
            version = cur.fetchone()
            logger.info(f"Подключение к БД успешно: {version[0]}")
            cur.close()
            logger.debug("Курсор БД закрыт")
            return True
    except Exception as e:
        logger.error(f"Ошибка подключения к БД: {e}")
        return False


if __name__ == "__main__":
    test_connection()