# websocket_server/config.py import os import redis import logging from logging import StreamHandler, FileHandler from colorlog import ColoredFormatter from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker, declarative_base import pymysql # Redis configuration REDIS_URL = "redis://172.17.0.1:6379/2" # Configurable intervals PING_INTERVAL_SECONDS = 15 # How often to send PING requests REDIS_PING_EXPIRY_SECONDS = 120 # Redis key expiry value PING_EXPIRY_SECONDS = 5 # How long between PING and PONG before we raise the alarm - TODO - Impliment this ASSIGN_INTERVAL_SECONDS = 30 LIVENESS_CHECK_INTERVAL_SECONDS = 5 # How often to cross check connected_workers with ping results from Redis ENABLE_WEBSOCKET_PING_DEBUG = os.getenv('ENABLE_WEBSOCKET_PING_DEBUG', 'false').lower() == 'true' ENABLE_TASK_ASSIGNMENT_DEBUG = os.getenv('ENABLE_TASK_ASSIGNMENT_DEBUG', 'false').lower() == 'true' # Database configuration DATABASE_URL = "mysql://root:password@172.17.0.1:3306/theapi" # Connection pool for Redis redis_connection_pool = redis.ConnectionPool.from_url( REDIS_URL, decode_responses=True, max_connections=10000, socket_timeout=5 ) def get_redis_client(): return redis.StrictRedis(connection_pool=redis_connection_pool) # SQLAlchemy setup pymysql.install_as_MySQLdb() engine = create_engine( DATABASE_URL, isolation_level='SERIALIZABLE', pool_size=5, max_overflow=10, pool_timeout=30, pool_recycle=3600, pool_pre_ping=True ) SessionLocal = sessionmaker( autocommit=False, autoflush=False, bind=engine, expire_on_commit=False ) Base = declarative_base() # Logging setup def setup_logging(): class FunctionNameFilter(logging.Filter): def filter(self, record): record.funcName = getattr(record, 'funcName', '') return True log_format = ( "%(log_color)s%(asctime)s - %(levelname)s - %(funcName)s - %(message)s" ) date_format = "%Y-%m-%d %H:%M:%S" formatter = ColoredFormatter( log_format, datefmt=date_format, log_colors={ "DEBUG": "cyan", "INFO": "green", "WARNING": "yellow", "ERROR": "red", "CRITICAL": "bold_red", }, ) logger = logging.getLogger("websocket_server") logger.setLevel(logging.DEBUG) logger.addFilter(FunctionNameFilter()) stream_handler = StreamHandler() stream_handler.setFormatter(formatter) logger.addHandler(stream_handler) # Ensure logs directory exists log_dir = "logs" os.makedirs(log_dir, exist_ok=True) file_handler = FileHandler("logs/app.log") file_handler.setFormatter(formatter) file_handler.setLevel(logging.WARNING) logger.addHandler(file_handler) return logger # Context manager for DB session from contextlib import contextmanager @contextmanager def get_db_session(): session = SessionLocal() try: yield session session.commit() except Exception: session.rollback() raise finally: session.close()