Initial commit - cleaned repository

services/pipeline/shared/queue_manager.py

@@ -0,0 +1,176 @@
+"""
+Queue Manager
+Redis 기반 큐 관리 시스템
+"""
+import redis.asyncio as redis
+import json
+import logging
+from typing import Optional, Dict, Any, List
+from datetime import datetime
+
+from .models import PipelineJob, QueueMessage
+
+logger = logging.getLogger(__name__)
+
+class QueueManager:
+    """Redis 기반 큐 매니저"""
+    
+    QUEUES = {
+        "keyword_processing": "queue:keyword_processing",
+        "rss_collection": "queue:rss_collection",
+        "search_enrichment": "queue:search_enrichment",
+        "google_search": "queue:google_search",
+        "ai_article_generation": "queue:ai_article_generation",
+        "image_generation": "queue:image_generation",
+        "translation": "queue:translation",
+        "failed": "queue:failed",
+        "scheduled": "queue:scheduled"
+    }
+    
+    def __init__(self, redis_url: str = "redis://redis:6379"):
+        self.redis_url = redis_url
+        self.redis_client: Optional[redis.Redis] = None
+        
+    async def connect(self):
+        """Redis 연결"""
+        if not self.redis_client:
+            self.redis_client = await redis.from_url(
+                self.redis_url,
+                encoding="utf-8",
+                decode_responses=True
+            )
+            logger.info("Connected to Redis")
+    
+    async def disconnect(self):
+        """Redis 연결 해제"""
+        if self.redis_client:
+            await self.redis_client.close()
+            self.redis_client = None
+            
+    async def enqueue(self, queue_name: str, job: PipelineJob, priority: int = 0) -> str:
+        """작업을 큐에 추가"""
+        try:
+            queue_key = self.QUEUES.get(queue_name, f"queue:{queue_name}")
+            
+            message = QueueMessage(
+                queue_name=queue_name,
+                job=job
+            )
+            
+            # 우선순위에 따라 추가
+            if priority > 0:
+                await self.redis_client.lpush(queue_key, message.json())
+            else:
+                await self.redis_client.rpush(queue_key, message.json())
+            
+            # 통계 업데이트
+            await self.redis_client.hincrby("stats:queues", queue_name, 1)
+            
+            logger.info(f"Job {job.job_id} enqueued to {queue_name}")
+            return job.job_id
+            
+        except Exception as e:
+            logger.error(f"Failed to enqueue job: {e}")
+            raise
+    
+    async def dequeue(self, queue_name: str, timeout: int = 0) -> Optional[PipelineJob]:
+        """큐에서 작업 가져오기"""
+        try:
+            queue_key = self.QUEUES.get(queue_name, f"queue:{queue_name}")
+            logger.info(f"Attempting to dequeue from {queue_key} with timeout={timeout}")
+            
+            if timeout > 0:
+                result = await self.redis_client.blpop(queue_key, timeout)
+                if result:
+                    _, data = result
+                    logger.info(f"Dequeued item from {queue_key}")
+                else:
+                    logger.debug(f"No item available in {queue_key}")
+                    return None
+            else:
+                data = await self.redis_client.lpop(queue_key)
+                
+            if data:
+                message = QueueMessage.parse_raw(data)
+                
+                # 처리 중 목록에 추가
+                processing_key = f"processing:{queue_name}"
+                await self.redis_client.hset(
+                    processing_key,
+                    message.job.job_id,
+                    message.json()
+                )
+                
+                return message.job
+            
+            return None
+            
+        except Exception as e:
+            logger.error(f"Failed to dequeue job: {e}")
+            return None
+    
+    async def mark_completed(self, queue_name: str, job_id: str):
+        """작업 완료 표시"""
+        try:
+            processing_key = f"processing:{queue_name}"
+            await self.redis_client.hdel(processing_key, job_id)
+            
+            # 통계 업데이트
+            await self.redis_client.hincrby("stats:completed", queue_name, 1)
+            
+            logger.info(f"Job {job_id} completed in {queue_name}")
+            
+        except Exception as e:
+            logger.error(f"Failed to mark job as completed: {e}")
+    
+    async def mark_failed(self, queue_name: str, job: PipelineJob, error: str):
+        """작업 실패 처리"""
+        try:
+            processing_key = f"processing:{queue_name}"
+            await self.redis_client.hdel(processing_key, job.job_id)
+            
+            # 재시도 확인
+            if job.retry_count < job.max_retries:
+                job.retry_count += 1
+                await self.enqueue(queue_name, job)
+                logger.info(f"Job {job.job_id} requeued (retry {job.retry_count}/{job.max_retries})")
+            else:
+                # 실패 큐로 이동
+                job.data["error"] = error
+                job.data["failed_stage"] = queue_name
+                await self.enqueue("failed", job)
+                
+                # 통계 업데이트
+                await self.redis_client.hincrby("stats:failed", queue_name, 1)
+                logger.error(f"Job {job.job_id} failed: {error}")
+                
+        except Exception as e:
+            logger.error(f"Failed to mark job as failed: {e}")
+    
+    async def get_queue_stats(self) -> Dict[str, Any]:
+        """큐 통계 조회"""
+        try:
+            stats = {}
+            
+            for name, key in self.QUEUES.items():
+                stats[name] = {
+                    "pending": await self.redis_client.llen(key),
+                    "processing": await self.redis_client.hlen(f"processing:{name}"),
+                }
+            
+            # 완료/실패 통계
+            stats["completed"] = await self.redis_client.hgetall("stats:completed") or {}
+            stats["failed"] = await self.redis_client.hgetall("stats:failed") or {}
+            
+            return stats
+            
+        except Exception as e:
+            logger.error(f"Failed to get queue stats: {e}")
+            return {}
+    
+    async def clear_queue(self, queue_name: str):
+        """큐 초기화 (테스트용)"""
+        queue_key = self.QUEUES.get(queue_name, f"queue:{queue_name}")
+        await self.redis_client.delete(queue_key)
+        await self.redis_client.delete(f"processing:{queue_name}")
+        logger.info(f"Queue {queue_name} cleared")