chartbastan/backend/app/queues/rabbitmq_client.py

"""
RabbitMQ client module.

This module provides a RabbitMQ client with connection management
and queue declaration functionality.
"""

import json
import logging
from datetime import datetime
from typing import Dict, Optional, Callable
import pika
from pika.exceptions import AMQPConnectionError, AMQPChannelError

logger = logging.getLogger(__name__)


class RabbitMQClient:
    """
    RabbitMQ client with connection management and queue declaration.
    
    Features:
    - Connection management with reconnection logic
    - Queue declaration with durability
    - Message publishing with standard event format
    - Task consumption with error handling
    - Automatic acknowledgment management
    """
    
    def __init__(
        self,
        rabbitmq_url: str = "amqp://guest:guest@localhost:5672",
        prefetch_count: int = 1
    ):
        """
        Initialize RabbitMQ client.
        
        Args:
            rabbitmq_url: RabbitMQ connection URL
            prefetch_count: Number of unacknowledged messages to prefetch
        """
        self.rabbitmq_url = rabbitmq_url
        self.prefetch_count = prefetch_count
        self.connection: Optional[pika.BlockingConnection] = None
        self.channel: Optional[pika.adapters.blocking_connection.BlockingChannel] = None
        
        # Queue names
        self.queues = {
            'scraping_tasks': 'scraping_tasks',
            'sentiment_analysis_tasks': 'sentiment_analysis_tasks',
            'energy_calculation_tasks': 'energy_calculation_tasks',
            'results': 'results'
        }
    
    def connect(self) -> None:
        """
        Establish connection to RabbitMQ server.
        
        Raises:
            AMQPConnectionError: If connection fails
        """
        try:
            logger.info(f"🔗 Connecting to RabbitMQ at {self.rabbitmq_url}")
            
            # Create connection
            self.connection = pika.BlockingConnection(
                pika.URLParameters(self.rabbitmq_url)
            )
            self.channel = self.connection.channel()
            
            # Set prefetch count for fair dispatch
            self.channel.basic_qos(prefetch_count=self.prefetch_count)
            
            # Declare queues
            self._declare_queues()
            
            logger.info("✅ Connected to RabbitMQ successfully")
            
        except AMQPConnectionError as e:
            logger.error(f"❌ Failed to connect to RabbitMQ: {e}")
            raise
        except AMQPChannelError as e:
            logger.error(f"❌ Failed to create RabbitMQ channel: {e}")
            raise
    
    def _declare_queues(self) -> None:
        """
        Declare all required queues with durability.
        """
        for queue_name in self.queues.values():
            self.channel.queue_declare(
                queue=queue_name,
                durable=True
            )
            logger.debug(f"✅ Declared queue: {queue_name}")
    
    def close(self) -> None:
        """Close connection to RabbitMQ server."""
        if self.channel:
            self.channel.close()
        if self.connection:
            self.connection.close()
        logger.info("🔌 Closed RabbitMQ connection")
    
    def publish_message(
        self,
        queue_name: str,
        data: Dict,
        event_type: str = "task.created",
        version: str = "1.0"
    ) -> None:
        """
        Publish a message to a queue with standard event format.
        
        Args:
            queue_name: Target queue name
            data: Message payload (will be JSON serialized)
            event_type: Event type for message header
            version: Event version
        """
        if not self.channel:
            raise RuntimeError("RabbitMQ channel not initialized. Call connect() first.")
        
        # Create standard event message format
        message = {
            "event": event_type,
            "version": version,
            "timestamp": datetime.utcnow().isoformat(),
            "data": data,
            "metadata": {
                "source": "api",
                "queue": queue_name
            }
        }
        
        try:
            self.channel.basic_publish(
                exchange='',
                routing_key=queue_name,
                body=json.dumps(message, ensure_ascii=False, default=str),
                properties=pika.BasicProperties(
                    delivery_mode=2,  # Make message persistent
                    content_type='application/json'
                )
            )
            logger.debug(f"📤 Published message to {queue_name}: {event_type}")
            
        except Exception as e:
            logger.error(f"❌ Failed to publish message to {queue_name}: {e}")
            raise
    
    def consume_messages(
        self,
        queue_name: str,
        callback: Callable
    ) -> None:
        """
        Start consuming messages from a queue.
        
        Args:
            queue_name: Queue to consume from
            callback: Callback function to process messages
        """
        if not self.channel:
            raise RuntimeError("RabbitMQ channel not initialized. Call connect() first.")
        
        try:
            logger.info(f"👂 Starting to consume from queue: {queue_name}")
            
            self.channel.basic_consume(
                queue=queue_name,
                on_message_callback=callback,
                auto_ack=False
            )
            
            self.channel.start_consuming()
            
        except KeyboardInterrupt:
            logger.info("⏹️ Stopping consumer...")
            self.channel.stop_consuming()
        except Exception as e:
            logger.error(f"❌ Error consuming from {queue_name}: {e}")
            raise
    
    def ack_message(self, delivery_tag: int) -> None:
        """
        Acknowledge a message as processed.
        
        Args:
            delivery_tag: Message delivery tag
        """
        if not self.channel:
            raise RuntimeError("RabbitMQ channel not initialized.")
        
        self.channel.basic_ack(delivery_tag=delivery_tag)
    
    def reject_message(
        self,
        delivery_tag: int,
        requeue: bool = False
    ) -> None:
        """
        Reject a message (e.g., on processing failure).
        
        Args:
            delivery_tag: Message delivery tag
            requeue: Whether to requeue the message
        """
        if not self.channel:
            raise RuntimeError("RabbitMQ channel not initialized.")
        
        self.channel.basic_reject(delivery_tag=delivery_tag, requeue=requeue)


def create_rabbitmq_client(
    rabbitmq_url: Optional[str] = None,
    prefetch_count: int = 1
) -> RabbitMQClient:
    """
    Factory function to create a RabbitMQ client.
    
    Args:
        rabbitmq_url: Optional RabbitMQ connection URL (defaults to localhost)
        prefetch_count: Number of unacknowledged messages to prefetch
        
    Returns:
        Configured RabbitMQClient instance
    """
    if rabbitmq_url is None:
        # TODO: Load from environment variables or config file
        rabbitmq_url = "amqp://guest:guest@localhost:5672"
    
    client = RabbitMQClient(
        rabbitmq_url=rabbitmq_url,
        prefetch_count=prefetch_count
    )
    
    return client