telecom-cdr-pipeline/src/core/worker/consumer.py

import logging
from typing import Any

from confluent_kafka import (
    Consumer,
    KafkaException,
    Message,
    TopicPartition,
)

from core.worker.batch import (
    BufferedCDR,
    KafkaPosition,
    MicroBatch,
)
from core.worker.clickhouse import ClickHouseWriter
from core.config import Settings
from core.domain.errors import (
    CDRProcessingError,
    FailureReason,
)
from core.worker.dlq import (
    DLQPublishError,
    DLQPublisher,
)
from core.worker.processor import (
    decode_payload,
    process_payload,
)


logger = logging.getLogger(__name__)


class ProcessingWorker:
    def __init__(
        self,
        settings: Settings,
    ) -> None:
        self._settings = settings

        self._consumer = Consumer(
            {
                "bootstrap.servers":
                    settings.kafka_bootstrap_servers,

                "group.id":
                    settings.kafka_consumer_group,

                "client.id":
                    settings.worker_instance_id,

                "auto.offset.reset":
                    settings.kafka_auto_offset_reset,

                "enable.auto.commit": False,
            }
        )

        self._dlq = DLQPublisher(settings)

        self._clickhouse = ClickHouseWriter(
            settings
        )

        self._batch = MicroBatch(
            max_size=settings.batch_size,
            flush_interval_seconds=(
                settings.batch_flush_interval_seconds
            ),
        )

        self._running = False

    def run(self) -> None:
        self._running = True

        self._consumer.subscribe(
            [self._settings.kafka_raw_topic]
        )

        logger.info(
            "Worker started. topic=%s group=%s instance=%s",
            self._settings.kafka_raw_topic,
            self._settings.kafka_consumer_group,
            self._settings.worker_instance_id,
        )

        try:
            while self._running:
                message = self._consumer.poll(
                    self._settings.kafka_consumer_poll_timeout_seconds
                )

                if message is None:
                    self._flush_if_due()
                    continue

                if message.error():
                    logger.error(
                        "Kafka consumer error: %s",
                        message.error(),
                    )
                    continue

                self._handle_message(message)
                self._flush_if_due()

        finally:
            self._shutdown()

    def _handle_message(
        self,
        message: Message,
    ) -> None:
        value = message.value() or b""

        original_payload: Any = {
            "raw": value.decode(
                "utf-8",
                errors="replace",
            )
        }

        try:
            payload = decode_payload(value)

            original_payload = payload

            normalized = process_payload(payload)

        except CDRProcessingError as exc:
            if not self._batch.is_empty:
                self._flush_batch()

            self._route_to_dlq(
                message=message,
                original_payload=original_payload,
                error=exc,
            )

            return

        except Exception as exc:
            logger.exception(
                "Unexpected CDR processing error."
            )

            error = CDRProcessingError(
                reason=FailureReason.NORMALIZATION_ERROR,
                detail=f"Unexpected processing error: {exc}",
            )

            if not self._batch.is_empty:
                self._flush_batch()

            self._route_to_dlq(
                message=message,
                original_payload=original_payload,
                error=error,
            )

            return

        position = KafkaPosition(
            topic=message.topic(),
            partition=message.partition(),
            offset=message.offset(),
        )

        self._batch.add(
            BufferedCDR(
                cdr=normalized,
                position=position,
            )
        )

        logger.debug(
            (
                "Buffered normalized CDR. "
                "raw_id=%s batch_size=%s"
            ),
            normalized.raw_id,
            self._batch.size,
        )

        if self._batch.is_full:
            self._flush_batch()

    def _route_to_dlq(
        self,
        message: Message,
        original_payload: Any,
        error: CDRProcessingError,
    ) -> None:
        logger.warning(
            (
                "Routing CDR to DLQ. "
                "reason=%s field=%s "
                "partition=%s offset=%s"
            ),
            error.reason.value,
            error.failed_field,
            message.partition(),
            message.offset(),
        )

        try:
            self._dlq.publish(
                original_payload=original_payload,
                error=error,
            )

        except DLQPublishError:
            logger.exception(
                "DLQ publishing failed. "
                "Original Kafka offset will not be committed."
            )

            raise

        self._commit(message)

    def _commit(
        self,
        message: Message,
    ) -> None:
        committed = self._consumer.commit(
            message=message,
            asynchronous=False,
        )

        for partition in committed or []:
            if partition.error is not None:
                raise KafkaException(
                    partition.error
                )

        logger.debug(
            "Committed Kafka offset. partition=%s offset=%s",
            message.partition(),
            message.offset(),
        )

    def _flush_if_due(self) -> None:
        if self._batch.is_due():
            self._flush_batch()

    def _flush_batch(self) -> None:
        items = self._batch.snapshot()

        if not items:
            return

        records = [
            item.cdr
            for item in items
        ]

        logger.info(
            "Flushing ClickHouse batch. size=%s",
            len(records),
        )

        self._clickhouse.insert(
            records
        )

        positions = [
            item.position
            for item in items
        ]

        self._commit_positions(
            positions
        )

        self._batch.clear()

        logger.info(
            "ClickHouse batch flushed successfully. size=%s",
            len(records),
        )

    def _commit_positions(
        self,
        positions: list[KafkaPosition],
    ) -> None:
        highest_offsets: dict[
            tuple[str, int],
            int,
        ] = {}

        for position in positions:
            key = (
                position.topic,
                position.partition,
            )

            next_offset = (
                position.offset + 1
            )

            current = highest_offsets.get(key)

            if (
                current is None
                or next_offset > current
            ):
                highest_offsets[key] = next_offset

        offsets = [
            TopicPartition(
                topic,
                partition,
                offset,
            )
            for (
                topic,
                partition,
            ), offset in highest_offsets.items()
        ]

        committed = self._consumer.commit(
            offsets=offsets,
            asynchronous=False,
        )

        for partition in committed or []:
            if partition.error is not None:
                raise KafkaException(
                    partition.error
                )

    def stop(self) -> None:
        self._running = False

    def _shutdown(self) -> None:
        logger.info("Worker shutting down.")

        try:
            if not self._batch.is_empty:
                logger.info(
                    "Flushing remaining records before shutdown."
                )

                self._flush_batch()

        finally:
            try:
                self._dlq.close()

            finally:
                try:
                    self._clickhouse.close()

                finally:
                    self._consumer.close()