llm_api/db/postgres.py - master - gnomdeployer/VUZPlus

llm_api/db/postgres.py
79 строк3 KB
daniil
i hate python
16 апр 2025, 23:22
16 апр 2025, 23:2252bdc6e
import psycopg2
from psycopg2.extras import RealDictCursor
from contextlib import contextmanager
import logging
from config import settings
from models.tables import create_all_tables

logger = logging.getLogger(__name__)

def get_connection_string():
    """Generate PostgreSQL connection string from settings"""
    return f"dbname={settings.DB_NAME} user={settings.DB_USER} password={settings.DB_PASSWORD} host={settings.DB_HOST} port={settings.DB_PORT}"

@contextmanager
def get_db_connection():
    """Context manager for database connections to ensure proper closing"""
    conn = None
    try:
        conn = psycopg2.connect(get_connection_string())
        yield conn
    except Exception as e:
        logger.error(f"Database connection error: {e}")
        raise
    finally:
        if conn:
            conn.close()

@contextmanager
def get_db_cursor(commit=False):
    """Context manager for database cursors"""
    with get_db_connection() as conn:
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        try:
            yield cursor
            if commit:
                conn.commit()
        except Exception as e:
            conn.rollback()
            logger.error(f"Database operation error: {e}")
            raise
        finally:
            cursor.close()

def initialize_db():
    """Initialize database with required tables"""
    try:
        with get_db_cursor(commit=True) as cursor:
            for table_sql in create_all_tables:
                cursor.execute(table_sql)
        logger.info("Database tables initialized successfully")
    except Exception as e:
        logger.error(f"Error initializing database tables: {e}")
        raise

def execute_query(query, params=None, fetch=True, commit=True):
    """Execute a query with error handling"""
    try:
        with get_db_cursor(commit=commit) as cursor:
            cursor.execute(query, params or {})
            if fetch:
                return cursor.fetchall()
            return None
    except Exception as e:
        logger.error(f"Query execution error: {e}")
        logger.error(f"Query: {query}")
        logger.error(f"Params: {params}")
        raise

def execute_batch(query, params_list, commit=True):
    """Execute a batch of queries with error handling"""
    try:
        with get_db_connection() as conn:
            with conn.cursor() as cursor:
                psycopg2.extras.execute_batch(cursor, query, params_list)
            if commit:
                conn.commit()
    except Exception as e:
        logger.error(f"Batch execution error: {e}")
        raise