diff --git a/src/db/db_connection.py b/src/db/db_connection.py index f828e5a..ea070e4 100644 --- a/src/db/db_connection.py +++ b/src/db/db_connection.py @@ -18,26 +18,50 @@ def create_client(): if not clickhouse_password: raise ValueError("CLICKHOUSE_PASSWORD environment variable not set.") - max_retries = 3 - retry_delay = 2 # seconds - + max_retries = 5 # Increased from 3 to 5 + retry_delay = 1 # Reduced initial delay to 1 second + + last_exception = None + for attempt in range(max_retries): try: + # Generate a new UUID for each attempt + session_id = str(uuid.uuid4()) + client = clickhouse_connect.get_client( host="clickhouse.abellana.work", port=443, username="default", password=clickhouse_password, secure=True, - session_id=str(uuid.uuid4()), - settings={'session_timeout': 60} + session_id=session_id, + settings={ + 'session_timeout': 60, + 'session_check': 1 + } ) - # Test the connection. + + # Test the connection client.query('SELECT 1') return client + except Exception as e: - if attempt < max_retries - 1: + last_exception = e + + # Check specifically for session locked error + if "SESSION_IS_LOCKED" in str(e): + if attempt < max_retries - 1: + # Use exponential backoff for session locks + wait_time = retry_delay * (2 ** attempt) + logging.warning(f"Session locked, retrying in {wait_time} seconds...") + time.sleep(wait_time) + continue + + # For other errors, use normal retry logic + elif attempt < max_retries - 1: logging.warning(f"Connection attempt {attempt + 1} failed: {str(e)}") time.sleep(retry_delay) - else: - raise + continue + + # If we've exhausted all retries, raise the last exception + raise last_exception