Kita Trofimov
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import psycopg2
|
|
import os
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from typing import Generator
|
|
|
|
from utils.loadDotEnv import initializeENV
|
|
|
|
initializeENV()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
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"Ошибка подключения к БД: {e}")
|
|
raise
|
|
except psycopg2.Error as e:
|
|
logger.error(f"Ошибка PostgreSQL: {e}")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Неожиданная ошибка при работе с БД: {e}")
|
|
raise
|
|
finally:
|
|
if conn:
|
|
try:
|
|
conn.close()
|
|
logger.info("Соединение с БД закрыто")
|
|
except Exception as e:
|
|
logger.warning(f"Ошибка при закрытии соединения: {e}")
|
|
|
|
def test_connection() -> bool:
|
|
try:
|
|
with get_connection() as conn:
|
|
cur = PSQLCursor(conn)
|
|
try:
|
|
cur.execute("SELECT version();")
|
|
version = cur.fetchone()
|
|
logger.info(f"Версия PostgreSQL: {version[0]}")
|
|
return True
|
|
finally:
|
|
cur.close()
|
|
logger.debug("Курсор закрыт")
|
|
except Exception as e:
|
|
logger.error(f"Тест подключения к БД провален: {e}")
|
|
return False
|
|
print(test_connection()) |