telecom-cdr-pipeline-ai/app/api/producer.py

import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, Optional

from app.common.config import config
from app.common.logging import setup_logger

logger = setup_logger("ingestion-api")


class MessageProducer(ABC):
    @abstractmethod
    def produce(self, topic: str, key: Optional[str], value: bytes) -> bool:
        """Produce a message synchronously or enqueue in broker buffer."""
        pass

    @abstractmethod
    def is_healthy(self) -> bool:
        """Check if producer is healthy and connected to broker."""
        pass

    @abstractmethod
    def flush(self, timeout: float = 5.0) -> int:
        """Flush any pending messages."""
        pass

    @abstractmethod
    def close(self) -> None:
        """Close producer connection."""
        pass


class ConfluentKafkaProducer(MessageProducer):
    """
    Production Kafka producer implementation wrapping confluent-kafka.
    Uses 'raw_id' as partition key and LZ4 compression.
    """

    def __init__(self, bootstrap_servers: Optional[str] = None):
        try:
            import confluent_kafka
            self._ck = confluent_kafka
        except ImportError:
            raise ImportError(
                "confluent-kafka library is not installed. Run 'pip install confluent-kafka'."
            )

        self.bootstrap_servers = bootstrap_servers or config.kafka.bootstrap_servers
        self._producer_config = {
            "bootstrap.servers": self.bootstrap_servers,
            "compression.type": "lz4",
            "queue.buffering.max.messages": 100000,
            "queue.buffering.max.ms": 5,
            "acks": "all",
            "retries": 3,
        }
        self._producer = self._ck.Producer(self._producer_config)
        self._healthy = True

    def _delivery_report(self, err: Any, msg: Any) -> None:
        if err is not None:
            self._healthy = False
            raw_id = msg.key().decode("utf-8") if msg and msg.key() else None
            logger.error(
                f"Kafka delivery failed: {err}",
                extra={"raw_id": raw_id, "event": "kafka_produce_failed"},
            )
        else:
            self._healthy = True

    def produce(self, topic: str, key: Optional[str], value: bytes) -> bool:
        try:
            key_bytes = key.encode("utf-8") if key else None
            self._producer.produce(
                topic=topic,
                key=key_bytes,
                value=value,
                on_delivery=self._delivery_report,
            )
            # Serve delivery callback queue non-blockingly
            self._producer.poll(0)
            return True
        except BufferError:
            logger.warning(
                "Kafka local producer buffer full",
                extra={"raw_id": key, "event": "kafka_buffer_full"},
            )
            self._producer.poll(0.1)
            return False
        except Exception as e:
            self._healthy = False
            logger.error(
                f"Kafka produce error: {e}",
                extra={"raw_id": key, "event": "kafka_error"},
            )
            return False

    def is_healthy(self) -> bool:
        return self._healthy

    def flush(self, timeout: float = 5.0) -> int:
        return self._producer.flush(timeout)

    def close(self) -> None:
        self.flush()


class MockMessageProducer(MessageProducer):
    """
    In-memory mock producer for offline testing and verification.
    """

    def __init__(self):
        self.messages = []
        self._healthy = True

    def produce(self, topic: str, key: Optional[str], value: bytes) -> bool:
        if not self._healthy:
            return False
        self.messages.append({"topic": topic, "key": key, "value": value})
        return True

    def is_healthy(self) -> bool:
        return self._healthy

    def set_healthy(self, healthy: bool) -> None:
        self._healthy = healthy

    def flush(self, timeout: float = 5.0) -> int:
        return 0

    def close(self) -> None:
        pass


class AsyncProducerWorker:
    """
    Background worker consuming from the bounded async queue and writing to the Kafka producer.
    """

    def __init__(self, queue: asyncio.Queue, producer: MessageProducer, topic: str):
        self.queue = queue
        self.producer = producer
        self.topic = topic
        self._running = False
        self._task: Optional[asyncio.Task] = None

    async def start(self) -> None:
        self._running = True
        self._task = asyncio.create_task(self._run_loop())

    async def _run_loop(self) -> None:
        while self._running:
            try:
                item = await self.queue.get()
                if item is None:  # Sentinel to stop
                    self.queue.task_done()
                    break

                key, value = item
                success = self.producer.produce(self.topic, key, value)
                if not success:
                    logger.error(
                        "Failed to publish message from async queue to Kafka",
                        extra={"raw_id": key, "event": "queue_dispatch_failed"},
                    )
                self.queue.task_done()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(
                    f"Unexpected error in async producer worker: {e}",
                    extra={"event": "worker_unexpected_error"},
                )

    async def stop(self) -> None:
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self.producer.flush()