api/routers/maintenance/redis_admin.py

import os
from typing import Any, Dict, List, Optional

from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel
import redis

from modules.logger_tool import initialise_logger

logger = initialise_logger(__name__, os.getenv("LOG_LEVEL"), os.getenv("LOG_PATH"), 'default', True)

router = APIRouter()


def _redis_client() -> redis.Redis:
    host = os.getenv('REDIS_HOST', '127.0.0.1')
    port = int(os.getenv('REDIS_PORT', '6379'))
    return redis.Redis(host=host, port=port, decode_responses=True)


@router.get("/ping")
def ping() -> Dict[str, Any]:
    try:
        r = _redis_client()
        pong = r.ping()
        return {"ok": True, "pong": pong}
    except Exception as e:
        logger.error(f"Redis ping failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis ping failed: {e}")


@router.get("/info")
def info(section: Optional[str] = Query(None, description="Optional INFO section, e.g. memory, server, clients, keyspace")) -> Dict[str, Any]:
    try:
        r = _redis_client()
        data = r.info(section) if section else r.info()
        return {"ok": True, "info": data}
    except Exception as e:
        logger.error(f"Redis info failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis info failed: {e}")


@router.get("/scan")
def scan(cursor: int = 0, pattern: Optional[str] = None, count: int = Query(100, ge=1, le=10000)) -> Dict[str, Any]:
    try:
        r = _redis_client()
        next_cursor, keys = r.scan(cursor=cursor, match=pattern, count=count)
        return {"ok": True, "cursor": next_cursor, "keys": keys, "count": len(keys)}
    except Exception as e:
        logger.error(f"Redis scan failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis scan failed: {e}")


@router.get("/keys")
def list_keys(pattern: str = Query("*", description="Glob-style pattern to match keys"), limit: int = Query(200, ge=1, le=5000)) -> Dict[str, Any]:
    try:
        r = _redis_client()
        keys: List[str] = []
        cursor = 0
        while True and len(keys) < limit:
            cursor, batch = r.scan(cursor=cursor, match=pattern, count=min(1000, limit - len(keys)))
            keys.extend(batch)
            if cursor == 0:
                break
        return {"ok": True, "keys": keys[:limit], "count": len(keys[:limit])}
    except Exception as e:
        logger.error(f"Redis keys failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis keys failed: {e}")


@router.get("/key")
def get_key(key: str) -> Dict[str, Any]:
    try:
        r = _redis_client()
        t = r.type(key)
        value: Any = None
        if t == 'string':
            value = r.get(key)
        elif t == 'hash':
            value = r.hgetall(key)
        elif t == 'list':
            value = r.lrange(key, 0, 99)
        elif t == 'set':
            value = list(r.smembers(key))
        elif t == 'zset':
            value = r.zrange(key, 0, 99, withscores=True)
        ttl = r.ttl(key)
        return {"ok": True, "type": t, "ttl": ttl, "value": value}
    except Exception as e:
        logger.error(f"Redis get key failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis get key failed: {e}")


class DeleteKeysBody(BaseModel):
    keys: Optional[List[str]] = None
    pattern: Optional[str] = None


@router.post("/delete")
def delete_keys(body: DeleteKeysBody) -> Dict[str, Any]:
    if not body.keys and not body.pattern:
        raise HTTPException(status_code=400, detail="Provide 'keys' or 'pattern'")
    try:
        r = _redis_client()
        to_delete: List[str] = body.keys or []
        if body.pattern:
            # Resolve pattern to keys via SCAN to avoid blocking
            keys: List[str] = []
            cursor = 0
            while True:
                cursor, batch = r.scan(cursor=cursor, match=body.pattern, count=1000)
                keys.extend(batch)
                if cursor == 0:
                    break
            to_delete.extend(keys)
        deleted = 0
        for k in set(to_delete):
            try:
                deleted += r.delete(k)
            except Exception:
                pass
        return {"ok": True, "deleted": deleted}
    except Exception as e:
        logger.error(f"Redis delete failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis delete failed: {e}")


@router.post("/flushdb")
def flush_db(confirm: bool = Query(False, description="Must be true to flush the current DB")) -> Dict[str, Any]:
    if not confirm:
        raise HTTPException(status_code=400, detail="Set confirm=true to flush the current Redis DB")
    try:
        r = _redis_client()
        r.flushdb()
        return {"ok": True, "message": "Flushed current Redis DB"}
    except Exception as e:
        logger.error(f"Redis FLUSHDB failed: {e}")
        raise HTTPException(status_code=500, detail=f"Redis FLUSHDB failed: {e}")


@router.get("/queues")
def queue_stats() -> Dict[str, Any]:
    """Report queue sizes and basic counters used by DocumentProcessingQueue."""
    try:
        r = _redis_client()
        keys = {
            "high": "queue:high",
            "normal": "queue:normal",
            "low": "queue:low",
            "processing": "processing",
            "dead_letter": "dead_letter",
            "metrics": "metrics",
        }
        data = {
            "queues": {
                "high": r.llen(keys["high"]),
                "normal": r.llen(keys["normal"]),
                "low": r.llen(keys["low"]),
            },
            "processing": r.hgetall(keys["processing"]),
            "dead_letter": r.llen(keys["dead_letter"]),
            "metrics": r.hgetall(keys["metrics"]),
        }
        return {"ok": True, **data}
    except Exception as e:
        logger.error(f"Queue stats failed: {e}")
        raise HTTPException(status_code=500, detail=f"Queue stats failed: {e}")