Kita Trofimov
57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
# db/connection.py
|
|
import psycopg2
|
|
import os
|
|
from contextlib import contextmanager
|
|
from typing import Generator
|
|
|
|
from utils.loadDotEnv import initializeENV
|
|
|
|
initializeENV()
|
|
|
|
|
|
def get_db_config() -> dict:
|
|
return {
|
|
'host': os.getenv('DB_HOST', 'localhost'),
|
|
'port': int(os.getenv('DB_PORT', 5432)),
|
|
'db': os.getenv('DB_NAME', 'warehouse_db'),
|
|
'user': os.getenv('DB_USER', 'postgres'),
|
|
'password': os.getenv('DB_PASSWORD', '')
|
|
}
|
|
|
|
|
|
@contextmanager
|
|
def get_connection() -> Generator[psycopg2.extensions.connection, None, None]:
|
|
conn = None
|
|
try:
|
|
config = get_db_config()
|
|
conn = psycopg2.connect(**config)
|
|
print("Подключение к БД установлено")
|
|
yield conn
|
|
except psycopg2.OperationalError as e:
|
|
print(f"Ошибка подключения к БД: {e}")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Неожиданная ошибка: {e}")
|
|
if conn:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
print("БД закрыта")
|
|
|
|
|
|
def test_connection() -> bool:
|
|
try:
|
|
with get_connection() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT version();")
|
|
version = cur.fetchone()
|
|
print(f" Версия PostgreSQL: {version[0]}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"Тест подключения бд провален: {e}")
|
|
return False
|
|
|
|
|
|
print(test_connection()) |