feat: add clickhouse micro-batch storage

src/core/worker/consumer.py

@@ -1,8 +1,19 @@
 import logging
 from typing import Any
 
-from confluent_kafka import Consumer, KafkaException, Message
+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,
@@ -48,6 +59,17 @@ class ProcessingWorker:
 
         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:
@@ -71,6 +93,7 @@ class ProcessingWorker:
                 )
 
                 if message is None:
+                    self._flush_if_due()
                     continue
 
                 if message.error():
@@ -81,6 +104,7 @@ class ProcessingWorker:
                     continue
 
                 self._handle_message(message)
+                self._flush_if_due()
 
         finally:
             self._shutdown()
@@ -106,6 +130,9 @@ class ProcessingWorker:
             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,
@@ -124,6 +151,9 @@ class ProcessingWorker:
                 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,
@@ -132,20 +162,30 @@ class ProcessingWorker:
 
             return
 
-        logger.info(
-            (
-                "Normalized CDR. "
-                "raw_id=%s cdr_id=%s "
-                "service_type=%s rat=%s"
-            ),
-            normalized.raw_id,
-            normalized.cdr_id,
-            normalized.service_type,
-            normalized.radio_access_type,
+        position = KafkaPosition(
+            topic=message.topic(),
+            partition=message.partition(),
+            offset=message.offset(),
         )
 
-        # Temporary until ClickHouse batching is implemented.
-        self._commit(message)
+        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,
@@ -202,6 +242,96 @@ class ProcessingWorker:
             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
 
@@ -209,7 +339,20 @@ class ProcessingWorker:
         logger.info("Worker shutting down.")
 
         try:
-            self._dlq.close()
+            if not self._batch.is_empty:
+                logger.info(
+                    "Flushing remaining records before shutdown."
+                )
+
+                self._flush_batch()
 
         finally:
-            self._consumer.close()
+            try:
+                self._dlq.close()
+
+            finally:
+                try:
+                    self._clickhouse.close()
+
+                finally:
+                    self._consumer.close()