3cloud-backend/websocket_server/config.py

# 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', '<unknown>')
            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()