handler-base/handler_base/handler.py

"""
Base handler class for building NATS-based services.
"""
import asyncio
import logging
import signal
from abc import ABC, abstractmethod
from typing import Any, Optional

from nats.aio.msg import Msg

from handler_base.config import Settings
from handler_base.health import HealthServer
from handler_base.nats_client import NATSClient
from handler_base.telemetry import setup_telemetry, create_span

logger = logging.getLogger(__name__)


class Handler(ABC):
    """
    Base class for NATS message handlers.
    
    Subclass and implement:
    - setup(): Initialize your service clients
    - handle_message(): Process incoming messages
    - teardown(): Clean up resources (optional)
    
    Example:
        class MyHandler(Handler):
            async def setup(self):
                self.embeddings = EmbeddingsClient()
            
            async def handle_message(self, msg: Msg, data: dict) -> Optional[dict]:
                result = await self.embeddings.embed(data["text"])
                return {"embedding": result}
        
        if __name__ == "__main__":
            MyHandler(subject="my.subject").run()
    """
    
    def __init__(
        self,
        subject: str,
        settings: Optional[Settings] = None,
        queue_group: Optional[str] = None,
    ):
        """
        Initialize the handler.
        
        Args:
            subject: NATS subject to subscribe to
            settings: Configuration settings
            queue_group: Optional queue group for load balancing
        """
        self.subject = subject
        self.settings = settings or Settings()
        self.queue_group = queue_group or self.settings.nats_queue_group
        
        self.nats = NATSClient(self.settings)
        self.health_server = HealthServer(self.settings, self._check_ready)
        
        self._running = False
        self._shutdown_event = asyncio.Event()
    
    @abstractmethod
    async def setup(self) -> None:
        """
        Initialize service clients and resources.
        
        Called once before starting to handle messages.
        Override this to set up your service-specific clients.
        """
        pass
    
    @abstractmethod
    async def handle_message(self, msg: Msg, data: Any) -> Optional[Any]:
        """
        Handle an incoming message.
        
        Args:
            msg: Raw NATS message
            data: Decoded message data (msgpack unpacked)
            
        Returns:
            Optional response data. If returned and msg has a reply subject,
            the response will be sent automatically.
        """
        pass
    
    async def teardown(self) -> None:
        """
        Clean up resources.
        
        Called during graceful shutdown.
        Override to add custom cleanup logic.
        """
        pass
    
    async def _check_ready(self) -> bool:
        """Check if the service is ready to handle requests."""
        return self._running and self.nats._nc is not None
    
    async def _message_handler(self, msg: Msg) -> None:
        """Internal message handler with tracing and error handling."""
        with create_span(f"handle.{self.subject}") as span:
            try:
                # Decode message
                data = NATSClient.decode_msgpack(msg)
                
                if span:
                    span.set_attribute("messaging.destination", msg.subject)
                    if isinstance(data, dict):
                        request_id = data.get("request_id", data.get("id"))
                        if request_id:
                            span.set_attribute("request.id", str(request_id))
                
                # Handle message
                response = await self.handle_message(msg, data)
                
                # Send response if applicable
                if response is not None and msg.reply:
                    await self.nats.publish(msg.reply, response)
                
            except Exception as e:
                logger.exception(f"Error handling message on {msg.subject}")
                if span:
                    span.set_attribute("error", True)
                    span.set_attribute("error.message", str(e))
                
                # Send error response if reply expected
                if msg.reply:
                    error_response = {
                        "error": True,
                        "message": str(e),
                        "type": type(e).__name__,
                    }
                    await self.nats.publish(msg.reply, error_response)
    
    def _setup_signals(self) -> None:
        """Set up signal handlers for graceful shutdown."""
        loop = asyncio.get_event_loop()
        
        for sig in (signal.SIGTERM, signal.SIGINT):
            loop.add_signal_handler(sig, self._handle_signal, sig)
    
    def _handle_signal(self, sig: signal.Signals) -> None:
        """Handle shutdown signal."""
        logger.info(f"Received {sig.name}, initiating graceful shutdown...")
        self._shutdown_event.set()
    
    async def _run(self) -> None:
        """Main async run loop."""
        # Setup telemetry
        setup_telemetry(self.settings)
        
        # Start health server
        self.health_server.start()
        
        try:
            # Connect to NATS
            await self.nats.connect()
            
            # Run user setup
            logger.info("Running service setup...")
            await self.setup()
            
            # Subscribe to subject
            await self.nats.subscribe(
                self.subject,
                self._message_handler,
                queue=self.queue_group,
            )
            
            self._running = True
            logger.info(f"Handler ready, listening on {self.subject}")
            
            # Wait for shutdown signal
            await self._shutdown_event.wait()
            
        except Exception as e:
            logger.exception("Fatal error in handler")
            raise
        finally:
            self._running = False
            
            # Graceful shutdown
            logger.info("Shutting down...")
            
            try:
                await self.teardown()
            except Exception as e:
                logger.warning(f"Error during teardown: {e}")
            
            await self.nats.close()
            self.health_server.stop()
            
            logger.info("Shutdown complete")
    
    def run(self) -> None:
        """
        Run the handler.
        
        This is the main entry point. It sets up signal handlers
        and runs the async event loop.
        """
        # Configure logging
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        )
        
        logger.info(f"Starting {self.settings.service_name} v{self.settings.service_version}")
        
        # Run the async loop
        asyncio.run(self._run_with_signals())
    
    async def _run_with_signals(self) -> None:
        """Run with signal handling."""
        self._setup_signals()
        await self._run()