""" PostgreSQL Database Connection Utility Supports multi-user mode with connection pooling and SQLite compatibility layer. """ import os import threading from typing import Optional, Any, List, Dict from contextlib import contextmanager from app.utils.logger import get_logger logger = get_logger(__name__) # Try to import psycopg2 try: import psycopg2 from psycopg2 import pool from psycopg2.extras import RealDictCursor HAS_PSYCOPG2 = True except ImportError: HAS_PSYCOPG2 = False logger.warning("psycopg2 not installed. PostgreSQL support disabled.") # Connection pool (global singleton) _connection_pool: Optional[Any] = None _pool_lock = threading.Lock() def _get_database_url() -> str: """Get database connection URL from environment""" return os.getenv('DATABASE_URL', '').strip() def _parse_database_url(url: str) -> Dict[str, Any]: """ Parse DATABASE_URL format: postgresql://user:password@host:port/dbname """ if not url: return {} # Remove protocol prefix if url.startswith('postgresql://'): url = url[13:] elif url.startswith('postgres://'): url = url[11:] else: return {} result = {} # Split user:password@host:port/dbname if '@' in url: auth, hostpart = url.rsplit('@', 1) if ':' in auth: result['user'], result['password'] = auth.split(':', 1) else: result['user'] = auth else: hostpart = url # Split host:port/dbname if '/' in hostpart: hostport, result['dbname'] = hostpart.split('/', 1) else: hostport = hostpart if ':' in hostport: result['host'], port_str = hostport.split(':', 1) result['port'] = int(port_str) else: result['host'] = hostport result['port'] = 5432 return result def _get_connection_pool(): """Get or create connection pool""" global _connection_pool if _connection_pool is not None: return _connection_pool with _pool_lock: if _connection_pool is not None: return _connection_pool if not HAS_PSYCOPG2: raise RuntimeError("psycopg2 is not installed. Cannot use PostgreSQL.") db_url = _get_database_url() if not db_url: raise RuntimeError("DATABASE_URL environment variable is not set.") params = _parse_database_url(db_url) if not params: raise RuntimeError(f"Invalid DATABASE_URL format: {db_url}") try: _connection_pool = pool.ThreadedConnectionPool( minconn=2, maxconn=20, host=params.get('host', 'localhost'), port=params.get('port', 5432), user=params.get('user', 'quantdinger'), password=params.get('password', ''), dbname=params.get('dbname', 'quantdinger'), connect_timeout=10, ) logger.info(f"PostgreSQL connection pool created: {params.get('host')}:{params.get('port')}/{params.get('dbname')}") except Exception as e: logger.error(f"Failed to create PostgreSQL connection pool: {e}") raise return _connection_pool class PostgresCursor: """PostgreSQL cursor wrapper with SQLite placeholder compatibility""" def __init__(self, cursor): self._cursor = cursor self._last_insert_id = None def _convert_placeholders(self, query: str) -> str: """ Convert SQLite-style ? placeholders to PostgreSQL %s Also handle some SQL syntax differences """ # Replace ? -> %s query = query.replace('?', '%s') # SQLite: INSERT OR IGNORE -> PostgreSQL: INSERT ... ON CONFLICT DO NOTHING query = query.replace('INSERT OR IGNORE', 'INSERT') return query def execute(self, query: str, args: Any = None): """Execute SQL statement""" query = self._convert_placeholders(query) # Check if this is an INSERT and add RETURNING id if not present is_insert = query.strip().upper().startswith('INSERT') if is_insert and 'RETURNING' not in query.upper(): query = query.rstrip(';').rstrip() + ' RETURNING id' if args: if not isinstance(args, (tuple, list)): args = (args,) result = self._cursor.execute(query, args) else: result = self._cursor.execute(query) # Capture last insert id for INSERT statements if is_insert: try: row = self._cursor.fetchone() if row and 'id' in row: self._last_insert_id = row['id'] except Exception: pass return result def fetchone(self) -> Optional[Dict[str, Any]]: """Fetch single row""" row = self._cursor.fetchone() if row is None: return None return dict(row) if row else None def fetchall(self) -> List[Dict[str, Any]]: """Fetch all rows""" rows = self._cursor.fetchall() return [dict(row) for row in rows] if rows else [] def close(self): """Close cursor""" self._cursor.close() @property def lastrowid(self) -> Optional[int]: """Get last inserted row ID""" return self._last_insert_id @property def rowcount(self) -> int: """Get affected row count""" return self._cursor.rowcount class PostgresConnection: """PostgreSQL connection wrapper""" def __init__(self, conn): self._conn = conn self._pool = _get_connection_pool() def cursor(self) -> PostgresCursor: """Create cursor""" return PostgresCursor(self._conn.cursor(cursor_factory=RealDictCursor)) def commit(self): """Commit transaction""" self._conn.commit() def rollback(self): """Rollback transaction""" self._conn.rollback() def close(self): """Return connection to pool""" if self._pool and self._conn: try: self._pool.putconn(self._conn) except Exception as e: logger.warning(f"Failed to return connection to pool: {e}") @contextmanager def get_pg_connection(): """ Get PostgreSQL database connection (Context Manager) """ pool = _get_connection_pool() conn = None try: conn = pool.getconn() pg_conn = PostgresConnection(conn) yield pg_conn except Exception as e: if conn: try: conn.rollback() except Exception: pass logger.error(f"PostgreSQL operation error: {e}") raise finally: if conn: try: pool.putconn(conn) except Exception: pass def get_pg_connection_sync() -> PostgresConnection: """ Get connection synchronously (caller must close) """ pool = _get_connection_pool() conn = pool.getconn() return PostgresConnection(conn) def execute_sql(sql: str, params: tuple = None) -> List[Dict[str, Any]]: """ Execute SQL and return results (convenience function) """ with get_pg_connection() as conn: cursor = conn.cursor() cursor.execute(sql, params) if sql.strip().upper().startswith('SELECT'): return cursor.fetchall() conn.commit() return [] def is_postgres_available() -> bool: """Check if PostgreSQL is available""" if not HAS_PSYCOPG2: return False db_url = _get_database_url() if not db_url: return False try: with get_pg_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT 1") return True except Exception as e: logger.debug(f"PostgreSQL not available: {e}") return False def close_pool(): """Close connection pool (call on app shutdown)""" global _connection_pool if _connection_pool: try: _connection_pool.closeall() _connection_pool = None logger.info("PostgreSQL connection pool closed") except Exception as e: logger.warning(f"Error closing connection pool: {e}")