suna/agentpress/db_connection.py

"""
Centralized database connection management for AgentPress.
"""

import asyncpg
import logging
from contextlib import asynccontextmanager
import os
import asyncio
import json

class DBConnection:
    """Singleton database connection manager."""
    
    _instance = None
    _initialized = False
    _db_config = {
        'dsn': 'postgresql://agentpress_owner:DkrT4eUS7hpb@ep-empty-bonus-a8nmv66l-pooler.eastus2.azure.neon.tech/agentpress?sslmode=require',
        'min_size': 2,  # Minimum number of connections
        'max_size': 10,  # Maximum number of connections
        'max_inactive_connection_lifetime': 300.0,  # 5 minutes
        'command_timeout': 60.0  # 1 minute command timeout
    }
    _init_lock = asyncio.Lock()
    _initialization_task = None
    _pool = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._initialization_task = asyncio.create_task(cls._instance._initialize())
        return cls._instance

    def __init__(self):
        """No initialization needed in __init__ as it's handled in __new__"""
        pass

    @classmethod
    async def _initialize(cls):
        """Internal initialization method."""
        if cls._initialized:
            return

        async with cls._init_lock:
            if cls._initialized:  # Double-check after acquiring lock
                return
                
            try:
                # Create connection pool
                cls._pool = await asyncpg.create_pool(**cls._db_config)
                
                async with cls._pool.acquire() as conn:
                    # Threads table
                    await conn.execute("""
                        CREATE TABLE IF NOT EXISTS threads (
                            id TEXT PRIMARY KEY,
                            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                        )
                    """)
                    
                    # Messages table - Fixed syntax for PostgreSQL
                    await conn.execute("""
                        CREATE TABLE IF NOT EXISTS messages (
                            id TEXT PRIMARY KEY,
                            thread_id TEXT REFERENCES threads(id),
                            type TEXT,
                            content TEXT,
                            include_in_llm_message_history BOOLEAN DEFAULT TRUE,
                            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                        )
                    """)
                    
                cls._initialized = True
                logging.info("Database schema initialized")
            except Exception as e:
                logging.error(f"Database initialization error: {e}")
                raise

    @classmethod
    def set_db_config(cls, config: dict):
        """Set custom database configuration."""
        if cls._initialized:
            raise RuntimeError("Cannot change database config after initialization")
        cls._db_config.update(config)
        logging.info(f"Updated database configuration")

    @asynccontextmanager
    async def connection(self):
        """Get a database connection."""
        if self._initialization_task and not self._initialized:
            await self._initialization_task
            
        async with self._pool.acquire() as conn:
            try:
                yield conn
            except Exception as e:
                logging.error(f"Database error: {e}")
                raise

    @asynccontextmanager
    async def transaction(self):
        """Execute operations in a transaction."""
        async with self.connection() as conn:
            async with conn.transaction():
                try:
                    yield conn
                except Exception as e:
                    logging.error(f"Transaction error: {e}")
                    raise

    async def execute(self, query: str, params: tuple = ()):
        """Execute a single query."""
        async with self.connection() as conn:
            try:
                # Handle empty params case
                if not params:
                    return await conn.execute(query)
                    
                # Convert params tuple to list and unpack for asyncpg
                params_list = list(params)
                return await conn.execute(query, *params_list)
            except Exception as e:
                logging.error(f"Query execution error: {e}")
                raise

    async def fetch_one(self, query: str, params: tuple = ()):
        """Fetch a single row."""
        async with self.connection() as conn:
            try:
                # Handle empty params case
                if not params:
                    return await conn.fetchrow(query)
                    
                # Convert params tuple to list and unpack for asyncpg
                params_list = list(params)
                return await conn.fetchrow(query, *params_list)
            except Exception as e:
                logging.error(f"Query execution error: {e}")
                raise

    async def fetch_all(self, query: str, params: tuple = ()):
        """Fetch all rows."""
        async with self.connection() as conn:
            try:
                # Handle empty params case
                if not params:
                    return await conn.fetch(query)
                    
                # Convert params tuple to list and unpack for asyncpg
                params_list = list(params)
                return await conn.fetch(query, *params_list)
            except Exception as e:
                logging.error(f"Query execution error: {e}")
                raise