Bobby/stock_system
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Remove session management from db_connection.py

Browse Source
This commit is contained in:
Bobby (aider) 2025-02-08 17:51:43 -08:00
parent cd4691a4cb
commit 0d0c4944d8
1 changed files with 11 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
src/db/db_connection.py
View File

@ -1,7 +1,6 @@
import logging
import os
import time
import uuid
import clickhouse_connect
from dotenv import load_dotenv
@ -15,7 +14,7 @@ logger = logging.getLogger(__name__)
def create_client():
"""
Create a ClickHouse client with enhanced session management and retry logic
Create a ClickHouse client with retry logic but without session management
"""
clickhouse_password = os.getenv("CLICKHOUSE_PASSWORD")
if not clickhouse_password:
@ -28,32 +27,22 @@ def create_client():
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
connect_timeout=10,
send_receive_timeout=300
)
# Test the connection with a simple query
try:
client.query('SELECT 1')
logger.info(f"Successfully established connection with session {session_id}")
logger.info(f"Successfully established connection")
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
@ -61,25 +50,13 @@ def create_client():
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
# For any errors, use normal retry logic
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
# 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}")