cob integration scaffold

COBY/storage/connection_pool.py

@@ -0,0 +1,140 @@
+"""
+Database connection pool management for TimescaleDB.
+"""
+
+import asyncio
+import asyncpg
+from typing import Optional, Dict, Any
+from contextlib import asynccontextmanager
+from ..config import config
+from ..utils.logging import get_logger
+from ..utils.exceptions import StorageError
+
+logger = get_logger(__name__)
+
+
+class DatabaseConnectionPool:
+    """Manages database connection pool for TimescaleDB"""
+    
+    def __init__(self):
+        self._pool: Optional[asyncpg.Pool] = None
+        self._is_initialized = False
+    
+    async def initialize(self) -> None:
+        """Initialize the connection pool"""
+        if self._is_initialized:
+            return
+        
+        try:
+            # Build connection string
+            dsn = (
+                f"postgresql://{config.database.user}:{config.database.password}"
+                f"@{config.database.host}:{config.database.port}/{config.database.name}"
+            )
+            
+            # Create connection pool
+            self._pool = await asyncpg.create_pool(
+                dsn,
+                min_size=5,
+                max_size=config.database.pool_size,
+                max_queries=50000,
+                max_inactive_connection_lifetime=300,
+                command_timeout=config.database.pool_timeout,
+                server_settings={
+                    'search_path': config.database.schema,
+                    'timezone': 'UTC'
+                }
+            )
+            
+            self._is_initialized = True
+            logger.info(f"Database connection pool initialized with {config.database.pool_size} connections")
+            
+            # Test connection
+            await self.health_check()
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize database connection pool: {e}")
+            raise StorageError(f"Database connection failed: {e}", "DB_INIT_ERROR")
+    
+    async def close(self) -> None:
+        """Close the connection pool"""
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+            self._is_initialized = False
+            logger.info("Database connection pool closed")
+    
+    @asynccontextmanager
+    async def get_connection(self):
+        """Get a database connection from the pool"""
+        if not self._is_initialized:
+            await self.initialize()
+        
+        if not self._pool:
+            raise StorageError("Connection pool not initialized", "POOL_NOT_READY")
+        
+        async with self._pool.acquire() as connection:
+            try:
+                yield connection
+            except Exception as e:
+                logger.error(f"Database operation failed: {e}")
+                raise
+    
+    @asynccontextmanager
+    async def get_transaction(self):
+        """Get a database transaction"""
+        async with self.get_connection() as conn:
+            async with conn.transaction():
+                yield conn
+    
+    async def execute_query(self, query: str, *args) -> Any:
+        """Execute a query and return results"""
+        async with self.get_connection() as conn:
+            return await conn.fetch(query, *args)
+    
+    async def execute_command(self, command: str, *args) -> str:
+        """Execute a command and return status"""
+        async with self.get_connection() as conn:
+            return await conn.execute(command, *args)
+    
+    async def execute_many(self, command: str, args_list) -> None:
+        """Execute a command multiple times with different arguments"""
+        async with self.get_connection() as conn:
+            await conn.executemany(command, args_list)
+    
+    async def health_check(self) -> bool:
+        """Check database health"""
+        try:
+            async with self.get_connection() as conn:
+                result = await conn.fetchval("SELECT 1")
+                if result == 1:
+                    logger.debug("Database health check passed")
+                    return True
+                else:
+                    logger.warning("Database health check returned unexpected result")
+                    return False
+        except Exception as e:
+            logger.error(f"Database health check failed: {e}")
+            return False
+    
+    async def get_pool_stats(self) -> Dict[str, Any]:
+        """Get connection pool statistics"""
+        if not self._pool:
+            return {}
+        
+        return {
+            'size': self._pool.get_size(),
+            'min_size': self._pool.get_min_size(),
+            'max_size': self._pool.get_max_size(),
+            'idle_size': self._pool.get_idle_size(),
+            'is_closing': self._pool.is_closing()
+        }
+    
+    @property
+    def is_initialized(self) -> bool:
+        """Check if pool is initialized"""
+        return self._is_initialized
+
+
+# Global connection pool instance
+db_pool = DatabaseConnectionPool()