Backend/db/connection.py

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

from utils.loadDotEnv import initializeENV

initializeENV()

def PSQLConnect():
    conn = psycopg2.connect(os.getenv('POSTDRESS_CONNECTION'))
    return conn

def PSQLCursor(conn):
    cur =  conn.cursor()
    return cur


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

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

logger.debug(test_connection())