87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
|
|
import clickhouse_connect
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables from .env file
|
|
load_dotenv()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def create_client():
|
|
"""
|
|
Create a ClickHouse client with enhanced session management and retry logic
|
|
"""
|
|
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:
|
|
# 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=session_id,
|
|
settings={
|
|
'session_timeout': 3600, # Increased to 1 hour
|
|
'session_check': 1
|
|
},
|
|
connect_timeout=10, # Added explicit connect timeout
|
|
send_receive_timeout=300 # Added send/receive timeout
|
|
)
|
|
|
|
# Test the connection with a simple query
|
|
try:
|
|
client.query('SELECT 1')
|
|
logger.info(f"Successfully established connection with session {session_id}")
|
|
return client
|
|
except Exception as e:
|
|
if "SESSION_NOT_FOUND" in str(e) or "SESSION_IS_LOCKED" in str(e):
|
|
raise # Re-raise these specific errors to trigger retry
|
|
logger.error(f"Connection test failed: {str(e)}")
|
|
raise
|
|
|
|
except Exception as e:
|
|
last_exception = e
|
|
error_message = str(e)
|
|
|
|
# Handle specific error codes
|
|
if "SESSION_NOT_FOUND" in error_message:
|
|
if attempt < max_retries - 1:
|
|
wait_time = retry_delay * (2 ** attempt)
|
|
logger.warning(f"Session not found, retrying in {wait_time} seconds...")
|
|
time.sleep(wait_time)
|
|
continue
|
|
elif "SESSION_IS_LOCKED" in error_message:
|
|
if attempt < max_retries - 1:
|
|
wait_time = retry_delay * (2 ** attempt)
|
|
logger.warning(f"Session locked, retrying in {wait_time} seconds...")
|
|
time.sleep(wait_time)
|
|
continue
|
|
else:
|
|
# For other errors, use normal retry logic
|
|
if attempt < max_retries - 1:
|
|
logger.warning(f"Connection attempt {attempt + 1} failed: {error_message}")
|
|
time.sleep(retry_delay)
|
|
continue
|
|
|
|
# If we've exhausted all retries, log the final error and raise
|
|
logger.error(f"Failed to establish connection after {max_retries} attempts: {error_message}")
|
|
raise last_exception
|