57 lines
1.4 KiB
Python
57 lines
1.4 KiB
Python
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
import os
|
|
from contextlib import contextmanager
|
|
|
|
# Database configuration from environment variables
|
|
DB_HOST = os.getenv("DB_HOST", "localhost")
|
|
DB_PORT = os.getenv("DB_PORT", "5432")
|
|
DB_NAME = os.getenv("DB_NAME", "navix")
|
|
DB_USER = os.getenv("DB_USER", "postgres")
|
|
DB_PASSWORD = os.getenv("DB_PASSWORD", "postgres")
|
|
|
|
|
|
def get_connection():
|
|
"""Create a new database connection"""
|
|
return psycopg2.connect(
|
|
host=DB_HOST,
|
|
port=DB_PORT,
|
|
database=DB_NAME,
|
|
user=DB_USER,
|
|
password=DB_PASSWORD,
|
|
cursor_factory=RealDictCursor
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def get_db_cursor(commit=False):
|
|
"""Context manager for database operations"""
|
|
conn = get_connection()
|
|
cursor = conn.cursor()
|
|
try:
|
|
yield cursor
|
|
if commit:
|
|
conn.commit()
|
|
except Exception as e:
|
|
conn.rollback()
|
|
raise e
|
|
finally:
|
|
cursor.close()
|
|
conn.close()
|
|
|
|
|
|
def init_db():
|
|
"""Test database connection"""
|
|
try:
|
|
conn = get_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT version();")
|
|
version = cursor.fetchone()
|
|
print(f"Connected to PostgreSQL: {version}")
|
|
cursor.close()
|
|
conn.close()
|
|
return True
|
|
except Exception as e:
|
|
print(f"Database connection failed: {e}")
|
|
return False
|