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

37
src/db/db_connection.py
View File

@ -1,7 +1,6 @@
import logging import logging
import os import os
import time import time
import uuid
import clickhouse_connect import clickhouse_connect
from dotenv import load_dotenv from dotenv import load_dotenv
@ -15,7 +14,7 @@ logger = logging.getLogger(__name__)
def create_client(): 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") clickhouse_password = os.getenv("CLICKHOUSE_PASSWORD")
if not clickhouse_password: if not clickhouse_password:
@ -28,32 +27,22 @@ def create_client():
for attempt in range(max_retries): for attempt in range(max_retries):
try: try:
# Generate a new UUID for each attempt
session_id = str(uuid.uuid4())
client = clickhouse_connect.get_client( client = clickhouse_connect.get_client(
host="clickhouse.abellana.work", host="clickhouse.abellana.work",
port=443, port=443,
username="default", username="default",
password=clickhouse_password, password=clickhouse_password,
secure=True, secure=True,
session_id=session_id, connect_timeout=10,
settings={ send_receive_timeout=300
'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 # Test the connection with a simple query
try: try:
client.query('SELECT 1') client.query('SELECT 1')
logger.info(f"Successfully established connection with session {session_id}") logger.info(f"Successfully established connection")
return client return client
except Exception as e: 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)}") logger.error(f"Connection test failed: {str(e)}")
raise raise
@ -61,24 +50,12 @@ def create_client():
last_exception = e last_exception = e
error_message = str(e) error_message = str(e)
# Handle specific error codes # For any errors, use normal retry logic
if "SESSION_NOT_FOUND" in error_message:
if attempt < max_retries - 1: if attempt < max_retries - 1:
wait_time = retry_delay * (2 ** attempt) 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}") logger.warning(f"Connection attempt {attempt + 1} failed: {error_message}")
time.sleep(retry_delay) logger.warning(f"Retrying in {wait_time} seconds...")
time.sleep(wait_time)
continue continue
# If we've exhausted all retries, log the final error and raise # If we've exhausted all retries, log the final error and raise