refactor: Replace connection pool with context manager for ClickHouse connections

src/db/db_connection.py

@@ -1,10 +1,9 @@
 import logging
 import os
 import time
-from queue import Queue
-from threading import Lock
 import clickhouse_connect
 from dotenv import load_dotenv
+from contextlib import contextmanager
 
 # Load environment variables from .env file
 load_dotenv()
@@ -13,26 +12,17 @@ load_dotenv()
 logging.basicConfig(level=logging.WARNING)
 logger = logging.getLogger(__name__)
 
-# Connection pool settings
+@contextmanager
-MAX_POOL_SIZE = 5
+def create_client():
-connection_pool = Queue(maxsize=MAX_POOL_SIZE)
-pool_lock = Lock()
-
-def create_base_client():
     """
-    Create a base ClickHouse client without session management
+    Context manager for creating and properly closing ClickHouse connections
     """
+    client = None
+    try:
         clickhouse_password = os.getenv("CLICKHOUSE_PASSWORD")
         if not clickhouse_password:
             raise ValueError("CLICKHOUSE_PASSWORD environment variable not set.")
 
-    max_retries = 5
-    retry_delay = 1
-    
-    last_exception = None
-    
-    for attempt in range(max_retries):
-        try:
         client = clickhouse_connect.get_client(
             host="clickhouse.abellana.work",
             port=443,
@@ -51,79 +41,14 @@ def create_base_client():
             }
         )
         
-            # Test the connection
+        yield client
-            try:
+        
-                client.query('SELECT 1')
-                logger.debug(f"Successfully established connection")
-                return client
     except Exception as e:
-                logger.error(f"Connection test failed: {str(e)}")
+        logger.error(f"Connection error: {str(e)}")
         raise
-            
+    finally:
-        except Exception as e:
+        if client:
-            last_exception = e
-            error_message = str(e)
-            
-            if attempt < max_retries - 1:
-                wait_time = retry_delay * (2 ** attempt)
-                logger.warning(f"Connection attempt {attempt + 1} failed: {error_message}")
-                logger.warning(f"Retrying in {wait_time} seconds...")
-                time.sleep(wait_time)
-                continue
-            
-            logger.error(f"Failed to establish connection after {max_retries} attempts: {error_message}")
-            raise last_exception
-
-def initialize_pool():
-    """Initialize the connection pool"""
-    while not connection_pool.full():
             try:
-            client = create_base_client()
-            connection_pool.put(client)
-        except Exception as e:
-            logger.error(f"Error initializing pool connection: {str(e)}")
-            break
-
-def create_client():
-    """
-    Get a client from the pool or create a new one if needed.
-    Maintains the original function name for compatibility.
-    """
-    # Initialize pool if empty
-    with pool_lock:
-        if connection_pool.empty():
-            initialize_pool()
-    
-    try:
-        # Try to get a connection from the pool
-        client = connection_pool.get_nowait()
-        
-        # Test the connection
-        try:
-            client.query('SELECT 1')
-            return client
-        except Exception:
-            # Connection is dead, create a new one
-            logger.warning("Replacing dead connection in pool")
-            client = create_base_client()
-            return client
-            
-    except Exception:
-        # Pool is empty or other error, create new connection
-        logger.warning("Creating new connection (pool empty or error)")
-        return create_base_client()
-
-def return_client(client):
-    """
-    Return a client to the pool if possible
-    """
-    try:
-        if not connection_pool.full():
-            connection_pool.put(client)
-        else:
                 client.close()
             except Exception as e:
-        logger.warning(f"Error returning client to pool: {str(e)}")
+                logger.warning(f"Error closing connection: {str(e)}")
-
-# Export the create_client function
-__all__ = ['create_client']

src/utils/data_utils.py

@@ -73,8 +73,7 @@ def get_stock_data(ticker: str, start_date: datetime, end_date: datetime, interv
         pd.DataFrame: Resampled DataFrame with OHLCV data
     """
     try:
-        client = create_client()
+        with create_client() as client:
-        
             # Expand window to get enough data for calculations
             start_date = start_date - timedelta(days=90)