From 1a2f492992e7176fefcc18f7857b84b32b616d74 Mon Sep 17 00:00:00 2001 From: Amirhossein Khalili Date: Sun, 20 Sep 2026 15:33:44 +0330 Subject: [PATCH] feat: add clickhouse micro-batch storage --- .env.sample | 15 ++- README.md | 16 +++- pyproject.toml | 1 + src/core/config.py | 10 ++ src/core/worker/batch.py | 73 ++++++++++++++ src/core/worker/clickhouse.py | 97 +++++++++++++++++++ src/core/worker/consumer.py | 173 +++++++++++++++++++++++++++++++--- tests/unit/test_batch.py | 101 ++++++++++++++++++++ 8 files changed, 468 insertions(+), 18 deletions(-) create mode 100644 src/core/worker/batch.py create mode 100644 src/core/worker/clickhouse.py create mode 100644 tests/unit/test_batch.py diff --git a/.env.sample b/.env.sample index 4860c01..462a8be 100644 --- a/.env.sample +++ b/.env.sample @@ -58,10 +58,23 @@ KAFKA_PRODUCER_MAX_MESSAGES=100000 CLICKHOUSE_VERSION=26.8.7.19 -CLICKHOUSE_HOST=clickhouse +# Used by Python applications running directly on the host. +CLICKHOUSE_HOST=localhost + +# Used later by Dockerized API/worker services. +CLICKHOUSE_INTERNAL_HOST=clickhouse + CLICKHOUSE_PORT=8123 CLICKHOUSE_HTTP_HOST_PORT=8123 CLICKHOUSE_DATABASE=telecom CLICKHOUSE_USER=telecom CLICKHOUSE_PASSWORD=telecom_dev_password + + +# ========================= +# Worker batching +# ========================= + +BATCH_SIZE=1000 +BATCH_FLUSH_INTERVAL_SECONDS=2 diff --git a/README.md b/README.md index 558d142..c6703e7 100644 --- a/README.md +++ b/README.md @@ -28,8 +28,12 @@ Redpanda / cdr.raw Processing Worker ↓ Validation + Normalization - ├── valid → processed by worker - └── invalid → cdr.raw.dlq + ├── invalid → cdr.raw.dlq + └── valid + ↓ + Micro-Batch + ↓ + ClickHouse ``` Infrastructure already available: @@ -72,6 +76,11 @@ worker processing is not connected yet. * DLQ routing for malformed and unprocessable Kafka records * DLQ records with failure reason, failed field, timestamp, and worker instance * Kafka offsets committed only after the current processing outcome succeeds +* ClickHouse bulk storage through the official Python client +* Configurable CDR micro-batching +* Flush on 1,000 records or 2-second timeout +* Kafka offsets committed only after successful ClickHouse batch writes +* Remaining worker buffer flushed during graceful shutdown ## Project Structure @@ -106,6 +115,8 @@ telecom-cdr-pipeline/ │ │ │ └── worker/ │ ├── __init__.py +│ ├── batch.py +│ ├── clickhouse.py │ ├── consumer.py │ ├── dlq.py │ ├── main.py @@ -116,6 +127,7 @@ telecom-cdr-pipeline/ │ │ └── test_ingestion_api.py │ │ │ └── unit/ +│ ├── test_batch.py │ ├── test_dlq.py │ ├── test_normalization.py │ ├── test_project_setup.py diff --git a/pyproject.toml b/pyproject.toml index 8103e65..5d217b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "uvicorn[standard]>=0.30,<1.0", "confluent-kafka>=2.15,<3.0", "pydantic-settings>=2.0,<3.0", + "clickhouse-connect>=1.8,<2.0", ] [project.optional-dependencies] diff --git a/src/core/config.py b/src/core/config.py index db82f95..39582ea 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -31,6 +31,16 @@ class Settings(BaseSettings): worker_instance_id: str = "worker-1" + clickhouse_host: str = "localhost" + clickhouse_port: int = 8123 + clickhouse_database: str = "telecom" + clickhouse_user: str = "telecom" + clickhouse_password: str = "" + + batch_size: int = 1_000 + batch_flush_interval_seconds: float = 2.0 + + @lru_cache def get_settings() -> Settings: return Settings() diff --git a/src/core/worker/batch.py b/src/core/worker/batch.py new file mode 100644 index 0000000..694efa1 --- /dev/null +++ b/src/core/worker/batch.py @@ -0,0 +1,73 @@ +import time +from dataclasses import dataclass + +from core.domain.models import NormalizedCDR + + +@dataclass(frozen=True, slots=True) +class KafkaPosition: + topic: str + partition: int + offset: int + + +@dataclass(frozen=True, slots=True) +class BufferedCDR: + cdr: NormalizedCDR + position: KafkaPosition + + +class MicroBatch: + def __init__( + self, + max_size: int, + flush_interval_seconds: float, + ) -> None: + self._max_size = max_size + self._flush_interval_seconds = flush_interval_seconds + + self._items: list[BufferedCDR] = [] + + self._started_at: float | None = None + + def add( + self, + item: BufferedCDR, + ) -> None: + if not self._items: + self._started_at = time.monotonic() + + self._items.append(item) + + @property + def size(self) -> int: + return len(self._items) + + @property + def is_empty(self) -> bool: + return not self._items + + @property + def is_full(self) -> bool: + return len(self._items) >= self._max_size + + def is_due(self) -> bool: + if not self._items: + return False + + if self._started_at is None: + return False + + elapsed = ( + time.monotonic() + - self._started_at + ) + + return elapsed >= self._flush_interval_seconds + + def snapshot(self) -> list[BufferedCDR]: + return list(self._items) + + def clear(self) -> None: + self._items.clear() + self._started_at = None diff --git a/src/core/worker/clickhouse.py b/src/core/worker/clickhouse.py new file mode 100644 index 0000000..96a753a --- /dev/null +++ b/src/core/worker/clickhouse.py @@ -0,0 +1,97 @@ +from collections.abc import Sequence + +import clickhouse_connect + +from core.config import Settings +from core.domain.models import NormalizedCDR + + +class ClickHouseWriteError(Exception): + pass + + +class ClickHouseWriter: + COLUMN_NAMES = [ + "cdr_id", + "raw_id", + "msisdn", + "imsi", + "service_type", + "radio_access_type", + "cell_id", + "bytes_uploaded", + "bytes_downloaded", + "total_bytes", + "duration_seconds", + "session_start_time", + "session_end_time", + "source_ip", + ] + + def __init__( + self, + settings: Settings, + ) -> None: + self._settings = settings + + self._client = clickhouse_connect.get_client( + host=settings.clickhouse_host, + port=settings.clickhouse_port, + username=settings.clickhouse_user, + password=settings.clickhouse_password, + database=settings.clickhouse_database, + ) + + @staticmethod + def _to_row( + cdr: NormalizedCDR, + ) -> list: + return [ + cdr.cdr_id, + cdr.raw_id, + + cdr.msisdn, + cdr.imsi, + + cdr.service_type, + cdr.radio_access_type, + cdr.cell_id, + + cdr.bytes_uploaded, + cdr.bytes_downloaded, + cdr.total_bytes, + cdr.duration_seconds, + + cdr.session_start_time, + cdr.session_end_time, + + str(cdr.source_ip), + ] + + def insert( + self, + records: Sequence[NormalizedCDR], + ) -> None: + if not records: + return + + rows = [ + self._to_row(record) + for record in records + ] + + try: + self._client.insert( + table="cdr_records", + data=rows, + column_names=self.COLUMN_NAMES, + database=self._settings.clickhouse_database, + ) + + except Exception as exc: + raise ClickHouseWriteError( + f"ClickHouse batch insert failed: {exc}" + ) from exc + + def close(self) -> None: + self._client.close() diff --git a/src/core/worker/consumer.py b/src/core/worker/consumer.py index b00b82c..4058dfc 100644 --- a/src/core/worker/consumer.py +++ b/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() diff --git a/tests/unit/test_batch.py b/tests/unit/test_batch.py new file mode 100644 index 0000000..53017d6 --- /dev/null +++ b/tests/unit/test_batch.py @@ -0,0 +1,101 @@ +from datetime import datetime, timezone +from ipaddress import IPv4Address +from uuid import uuid4 + +from core.domain.models import NormalizedCDR +from core.worker.batch import ( + BufferedCDR, + KafkaPosition, + MicroBatch, +) + + +def normalized_cdr( + raw_id: str, +) -> NormalizedCDR: + now = datetime.now( + timezone.utc + ) + + return NormalizedCDR( + cdr_id=uuid4(), + raw_id=raw_id, + + msisdn="+19875551234", + imsi="310410998877665", + + service_type="4G_DATA", + radio_access_type="LTE_4G", + cell_id="4021_8812", + + bytes_uploaded=100, + bytes_downloaded=200, + total_bytes=300, + duration_seconds=300, + + session_start_time=now, + session_end_time=now, + + source_ip=IPv4Address( + "10.240.12.1" + ), + ) + + +def buffered_item( + raw_id: str, + offset: int, +) -> BufferedCDR: + return BufferedCDR( + cdr=normalized_cdr(raw_id), + position=KafkaPosition( + topic="cdr.raw", + partition=0, + offset=offset, + ), + ) + + +def test_batch_flushes_when_full() -> None: + batch = MicroBatch( + max_size=2, + flush_interval_seconds=10, + ) + + batch.add( + buffered_item( + raw_id="raw-1", + offset=1, + ) + ) + + assert batch.is_full is False + + batch.add( + buffered_item( + raw_id="raw-2", + offset=2, + ) + ) + + assert batch.is_full is True + + +def test_batch_clear_resets_state() -> None: + batch = MicroBatch( + max_size=2, + flush_interval_seconds=10, + ) + + batch.add( + buffered_item( + raw_id="raw-1", + offset=1, + ) + ) + + batch.clear() + + assert batch.size == 0 + assert batch.is_empty is True + assert batch.is_due() is False