joelgilbert/NL2SQL

1"""2PostgreSQL database connection management for Neon Serverless.3"""4 5import psycopg26from psycopg2 import pool, OperationalError, DatabaseError7from typing import Optional8import logging9from contextlib import contextmanager10 11from config.settings import settings12 13logger = logging.getLogger(__name__)14 15 16class DatabaseConnection:17    """Manages PostgreSQL connections with pooling and automatic retry."""18    19    def __init__(self):20        """Initialize connection pools for read-only and DBA modes."""21        self._readonly_pool: Optional[pool.SimpleConnectionPool] = None22        self._dba_pool: Optional[pool.SimpleConnectionPool] = None23        24    def get_readonly_pool(self) -> pool.SimpleConnectionPool:25        """26        Get or create read-only connection pool.27        28        Returns:29            SimpleConnectionPool: Read-only connection pool30            31        Raises:32            OperationalError: If connection fails33        """34        if self._readonly_pool is None:35            try:36                self._readonly_pool = psycopg2.pool.SimpleConnectionPool(37                    minconn=1,38                    maxconn=10,39                    dsn=settings.database.readonly_connection_string40                )41                logger.info("Read-only connection pool created successfully")42            except OperationalError as e:43                logger.error(f"Failed to create read-only connection pool: {e}")44                raise45        46        return self._readonly_pool47    48    def get_dba_pool(self) -> pool.SimpleConnectionPool:49        """50        Get or create DBA connection pool.51        52        Returns:53            SimpleConnectionPool: DBA connection pool54            55        Raises:56            OperationalError: If connection fails57        """58        if self._dba_pool is None:59            try:60                self._dba_pool = psycopg2.pool.SimpleConnectionPool(61                    minconn=1,62                    maxconn=5,63                    dsn=settings.database.dba_connection_string64                )65                logger.info("DBA connection pool created successfully")66            except OperationalError as e:67                logger.error(f"Failed to create DBA connection pool: {e}")68                raise69        70        return self._dba_pool71    72    @contextmanager73    def get_readonly_connection(self):74        """75        Context manager for read-only database connections.76        77        Yields:78            psycopg2.connection: Database connection79            80        Example:81            with db.get_readonly_connection() as conn:82                cursor = conn.cursor()83                cursor.execute("SELECT * FROM users")84        """85        pool = self.get_readonly_pool()86        conn = pool.getconn()87        88        try:89            yield conn90            conn.commit()91        except Exception as e:92            conn.rollback()93            logger.error(f"Read-only connection error: {e}")94            raise95        finally:96            pool.putconn(conn)97    98    @contextmanager99    def get_dba_connection(self):100        """101        Context manager for DBA database connections.102        103        Yields:104            psycopg2.connection: Database connection105            106        Example:107            with db.get_dba_connection() as conn:108                cursor = conn.cursor()109                cursor.execute("UPDATE users SET status = 'active'")110        """111        pool = self.get_dba_pool()112        conn = pool.getconn()113        114        try:115            yield conn116            conn.commit()117        except Exception as e:118            conn.rollback()119            logger.error(f"DBA connection error: {e}")120            raise121        finally:122            pool.putconn(conn)123    124    def test_connection(self, mode: str = "readonly") -> tuple[bool, str]:125        """126        Test database connectivity.127        128        Args:129            mode: Connection mode ("readonly" or "dba")130            131        Returns:132            Tuple of (success: bool, message: str)133        """134        try:135            if mode == "readonly":136                with self.get_readonly_connection() as conn:137                    cursor = conn.cursor()138                    cursor.execute("SELECT 1;")139                    cursor.close()140            else:141                with self.get_dba_connection() as conn:142                    cursor = conn.cursor()143                    cursor.execute("SELECT 1;")144                    cursor.close()145            146            return True, f"{mode.capitalize()} connection successful"147            148        except OperationalError as e:149            error_msg = f"{mode.capitalize()} connection failed: Authentication or network error - {str(e)}"150            logger.error(error_msg)151            return False, error_msg152        except DatabaseError as e:153            error_msg = f"{mode.capitalize()} connection failed: Database error - {str(e)}"154            logger.error(error_msg)155            return False, error_msg156        except Exception as e:157            error_msg = f"{mode.capitalize()} connection failed: {str(e)}"158            logger.error(error_msg)159            return False, error_msg160    161    def close_all_connections(self):162        """Close all connection pools gracefully."""163        if self._readonly_pool:164            self._readonly_pool.closeall()165            logger.info("Read-only connection pool closed")166            self._readonly_pool = None167        168        if self._dba_pool:169            self._dba_pool.closeall()170            logger.info("DBA connection pool closed")171            self._dba_pool = None172 173 174# Global database connection instance175db = DatabaseConnection()176