telecom-cdr-pipeline-ai/app/worker/clickhouse.py

import time
from abc import ABC, abstractmethod
from typing import Any, List, Optional

from app.common.config import config
from app.common.logging import setup_logger
from app.worker.normalizer import NormalizedRecord

logger = setup_logger("processing-worker")

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",
]


class BaseClickHouseClient(ABC):
    @abstractmethod
    def insert_batch(self, records: List[NormalizedRecord]) -> bool:
        pass

    @abstractmethod
    def is_healthy(self) -> bool:
        pass

    @abstractmethod
    def close(self) -> None:
        pass


class ClickHouseService(BaseClickHouseClient):
    """
    ClickHouse client using clickhouse-connect with exponential retry backoff.
    """

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        database: Optional[str] = None,
        table: Optional[str] = None,
        max_attempts: Optional[int] = None,
        initial_delay_ms: Optional[int] = None,
        backoff_factor: Optional[float] = None,
    ):
        try:
            import clickhouse_connect
            self._cc = clickhouse_connect
        except ImportError:
            raise ImportError(
                "clickhouse-connect is not installed. Run 'pip install clickhouse-connect'."
            )

        self.host = host or config.clickhouse.host
        self.port = port or config.clickhouse.port
        self.user = user or config.clickhouse.user
        self.password = password or config.clickhouse.password
        self.database = database or config.clickhouse.database
        self.table = table or config.clickhouse.table

        self.max_attempts = max_attempts or config.retry.max_attempts
        self.initial_delay_ms = initial_delay_ms or config.retry.initial_delay_ms
        self.backoff_factor = backoff_factor or config.retry.backoff_factor

        self._client = None
        self._connect()

    def _connect(self) -> None:
        try:
            self._client = self._cc.get_client(
                host=self.host,
                port=self.port,
                username=self.user,
                password=self.password,
                database=self.database,
                connect_timeout=5,
                send_receive_timeout=10,
            )
        except Exception as e:
            logger.error(
                f"ClickHouse connection failed: {e}",
                extra={"event": "clickhouse_connection_failed"},
            )
            self._client = None

    def is_healthy(self) -> bool:
        if self._client is None:
            self._connect()
        if self._client is None:
            return False
        try:
            return self._client.ping()
        except Exception:
            return False

    def insert_batch(self, records: List[NormalizedRecord]) -> bool:
        """
        Inserts a micro-batch of NormalizedRecord into ClickHouse with exponential backoff retry.
        Returns True if successful, raises exception if all retries fail.
        """
        if not records:
            return True

        rows = [record.to_row() for record in records]
        last_exception = None

        for attempt in range(1, self.max_attempts + 1):
            try:
                if self._client is None:
                    self._connect()
                if self._client is None:
                    raise ConnectionError("ClickHouse client is not connected")

                self._client.insert(
                    table=self.table,
                    data=rows,
                    column_names=COLUMN_NAMES,
                )
                return True
            except Exception as e:
                last_exception = e
                delay = (self.initial_delay_ms / 1000.0) * (
                    self.backoff_factor ** (attempt - 1)
                )
                logger.warning(
                    f"ClickHouse insert attempt {attempt}/{self.max_attempts} failed: {e}. "
                    f"Retrying in {delay:.2f}s...",
                    extra={"event": "clickhouse_insert_retry"},
                )
                time.sleep(delay)

        logger.error(
            f"All {self.max_attempts} ClickHouse insert attempts failed: {last_exception}",
            extra={"event": "clickhouse_batch_failed"},
        )
        raise last_exception or RuntimeError("ClickHouse insert failed")

    def close(self) -> None:
        if self._client:
            self._client.close()
            self._client = None


class MockClickHouseClient(BaseClickHouseClient):
    """
    In-memory ClickHouse mock client for offline tests.
    """

    def __init__(self, fail_attempts: int = 0):
        self.rows: List[Any] = []
        self.fail_attempts = fail_attempts
        self.attempt_count = 0
        self._healthy = True

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

    def insert_batch(self, records: List[NormalizedRecord]) -> bool:
        self.attempt_count += 1
        if self.attempt_count <= self.fail_attempts:
            raise TimeoutError("ClickHouse insert timed out (simulated)")

        for record in records:
            self.rows.append(record.to_row())
        return True

    def close(self) -> None:
        pass