agent-governance/integrations/common/redis_config.py

"""
Redis/DragonflyDB Configuration for Agent Governance System

Provides:
- Connection management with fallback
- Health checks
- Mock mode for testing
"""

import os
from typing import Optional, Dict, Any
from dataclasses import dataclass

from .secrets import get_secret


@dataclass
class RedisConfig:
    """Redis connection configuration"""
    host: str = "localhost"
    port: int = 6379
    password: Optional[str] = None
    db: int = 0
    ssl: bool = False
    socket_timeout: float = 5.0
    retry_on_timeout: bool = True
    max_connections: int = 10

    @classmethod
    def from_env(cls) -> "RedisConfig":
        """Load configuration from environment"""
        url = get_secret("REDIS_URL")
        password = get_secret("REDIS_PASSWORD")

        if url:
            # Parse URL format: redis://[:password@]host:port/db
            return cls._from_url(url, password)

        return cls(
            host=os.environ.get("REDIS_HOST", "localhost"),
            port=int(os.environ.get("REDIS_PORT", "6379")),
            password=password,
            db=int(os.environ.get("REDIS_DB", "0")),
            ssl=os.environ.get("REDIS_SSL", "false").lower() == "true"
        )

    @classmethod
    def _from_url(cls, url: str, password: Optional[str] = None) -> "RedisConfig":
        """Parse Redis URL into config"""
        # Handle redis:// and rediss:// (SSL)
        ssl = url.startswith("rediss://")
        url = url.replace("redis://", "").replace("rediss://", "")

        # Extract password from URL if present
        if "@" in url:
            auth_part, host_part = url.split("@", 1)
            if auth_part.startswith(":"):
                password = password or auth_part[1:]
            url = host_part
        else:
            host_part = url

        # Parse host:port/db
        db = 0
        if "/" in host_part:
            host_port, db_str = host_part.split("/", 1)
            db = int(db_str) if db_str else 0
        else:
            host_port = host_part

        if ":" in host_port:
            host, port_str = host_port.split(":", 1)
            port = int(port_str)
        else:
            host = host_port
            port = 6379

        return cls(
            host=host,
            port=port,
            password=password,
            db=db,
            ssl=ssl
        )

    def to_url(self) -> str:
        """Convert config to URL format"""
        scheme = "rediss" if self.ssl else "redis"
        auth = f":{self.password}@" if self.password else ""
        return f"{scheme}://{auth}{self.host}:{self.port}/{self.db}"

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for client initialization"""
        return {
            "host": self.host,
            "port": self.port,
            "password": self.password,
            "db": self.db,
            "socket_timeout": self.socket_timeout,
            "retry_on_timeout": self.retry_on_timeout,
            "max_connections": self.max_connections
        }


class MockRedisClient:
    """Mock Redis client for testing"""

    def __init__(self):
        self._data: Dict[str, Any] = {}
        self._connected = True

    def ping(self) -> bool:
        return self._connected

    def get(self, key: str) -> Optional[str]:
        return self._data.get(key)

    def set(self, key: str, value: Any, ex: int = None) -> bool:
        self._data[key] = str(value)
        return True

    def delete(self, *keys: str) -> int:
        count = 0
        for key in keys:
            if key in self._data:
                del self._data[key]
                count += 1
        return count

    def exists(self, *keys: str) -> int:
        return sum(1 for k in keys if k in self._data)

    def keys(self, pattern: str = "*") -> list:
        import fnmatch
        return [k for k in self._data.keys() if fnmatch.fnmatch(k, pattern)]

    def flushdb(self):
        self._data.clear()

    def close(self):
        self._connected = False


def get_redis_client(config: RedisConfig = None, mock: bool = None):
    """
    Get a Redis client.

    Args:
        config: Redis configuration (uses env if not provided)
        mock: Force mock mode (auto-detects if not provided)

    Returns:
        Redis client or MockRedisClient
    """
    if mock is None:
        mock = os.environ.get("INTEGRATION_DRY_RUN", "false").lower() == "true"

    if mock:
        return MockRedisClient()

    config = config or RedisConfig.from_env()

    try:
        import redis
        return redis.Redis(**config.to_dict())
    except ImportError:
        # Fallback to mock if redis package not installed
        return MockRedisClient()


def test_redis_connection(config: RedisConfig = None) -> Dict[str, Any]:
    """Test Redis connection and return status"""
    config = config or RedisConfig.from_env()

    result = {
        "host": config.host,
        "port": config.port,
        "connected": False,
        "mock_mode": False,
        "error": None
    }

    try:
        client = get_redis_client(config, mock=False)

        if isinstance(client, MockRedisClient):
            result["mock_mode"] = True
            result["connected"] = True
        else:
            result["connected"] = client.ping()
            client.close()
    except Exception as e:
        result["error"] = str(e)

    return result