stock_system/src/db/db_connection.py

55 lines
1.5 KiB
Python

import logging
import os
import time
import clickhouse_connect
from dotenv import load_dotenv
from contextlib import contextmanager
# Load environment variables from .env file
load_dotenv()
# Configure logging
logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)
@contextmanager
def create_client():
"""
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.")
client = clickhouse_connect.get_client(
host="clickhouse.abellana.work",
port=443,
username="default",
password=clickhouse_password,
secure=True,
connect_timeout=10,
send_receive_timeout=300,
query_limit=0,
compress=True,
settings={
'enable_http_compression': 1,
'database': 'stock_db',
'max_execution_time': 300,
'mutations_sync': 0
}
)
yield client
except Exception as e:
logger.error(f"Connection error: {str(e)}")
raise
finally:
if client:
try:
client.close()
except Exception as e:
logger.warning(f"Error closing connection: {str(e)}")