Kita Trofimov
65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
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() |