database/__init__.py

@@ -0,0 +1,3 @@
+from .connection import get_connection, get_db_config
+
+__all__ = ['get_connection', 'get_db_config']

database/connection.py

@@ -0,0 +1,54 @@
+# database/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)),
+        'database': 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