clickstream-ch-kafka-superset-demo/generator/src/clickstream_generator/kafka_io.py

"""Kafka-интеграция генератора."""

import json
import logging
import sys
import time
from dataclasses import dataclass
from datetime import datetime

from clickstream_generator.metrics import METRICS_ERRORS_TOTAL, METRICS_EVENTS_TOTAL
from clickstream_generator.state import GeneratorState


logger = logging.getLogger("generator")

_kafka_imported = False
KafkaProducer = None
KafkaError = None


def _import_kafka():
    """Лениво импортирует Kafka-клиент, чтобы тесты могли работать без Kafka."""
    global _kafka_imported, KafkaProducer, KafkaError
    if not _kafka_imported:
        from kafka import KafkaProducer
        from kafka.errors import KafkaError
        _kafka_imported = True
    return KafkaProducer, KafkaError


def _with_retry(operation, max_retries: int = 5, base_delay: float = 1.0, max_delay: float = 30.0):
    """Выполняет операцию с экспоненциальным backoff."""
    last_exception = None
    for attempt in range(max_retries):
        try:
            return operation()
        except Exception as e:
            last_exception = e
            if attempt < max_retries - 1:
                delay = min(base_delay * (2 ** attempt), max_delay)
                logger.warning(
                    f"Operation failed (attempt {attempt + 1}/{max_retries}): "
                    f"{e}. Retrying in {delay:.1f}s..."
                )
                time.sleep(delay)
            else:
                logger.error(f"Operation failed after {max_retries} attempts: {e}")
                raise last_exception


def _facade_attr(name: str, fallback):
    """Берёт совместимый mock из фасада generator, если тест его подменил."""
    facade = sys.modules.get("generator")
    return getattr(facade, name, fallback) if facade else fallback


def _retry(operation, **kwargs):
    return _facade_attr("_with_retry", _with_retry)(operation, **kwargs)


def _kafka_importer():
    return _facade_attr("_import_kafka", _import_kafka)


@dataclass
class BatchRecord:
    """Запись об отправленном батче."""

    batch_id: str
    started_at: datetime
    finished_at: datetime
    sent_total: int
    sent_browser: int
    sent_location: int
    sent_device: int
    sent_geo: int
    status: str
    error_message: str | None = None

    def to_dict(self) -> dict:
        """Конвертирует в словарь для сериализации."""
        return {
            "batch_id": self.batch_id,
            "started_at": self.started_at.isoformat(),
            "finished_at": self.finished_at.isoformat(),
            "sent_total": self.sent_total,
            "sent_browser": self.sent_browser,
            "sent_location": self.sent_location,
            "sent_device": self.sent_device,
            "sent_geo": self.sent_geo,
            "status": self.status,
            "error_message": self.error_message,
        }


class KafkaStateManager:
    """Управление состоянием генератора в Kafka compact topic."""

    STATE_TOPIC = "generator_state"
    STATE_KEY = "default"

    def __init__(self, bootstrap_servers: str):
        self.bootstrap_servers = bootstrap_servers
        KafkaProducerCls, _ = _kafka_importer()()

        logger.info(f"Connecting to Kafka for state management at {self.bootstrap_servers}")
        self.producer = KafkaProducerCls(
            bootstrap_servers=self.bootstrap_servers,
            value_serializer=lambda v: json.dumps(v).encode("utf-8"),
            key_serializer=lambda k: k.encode("utf-8") if k else None,
            retries=3,
            retry_backoff_ms=1000,
        )
        logger.info("Connected to Kafka for state management successfully")

    def save(self, state: GeneratorState) -> None:
        """Сохраняет состояние в топик."""
        value = state.to_dict()

        def _do_send():
            self.producer.send(self.STATE_TOPIC, key=self.STATE_KEY, value=value)

        _retry(_do_send, max_retries=3, base_delay=0.5)

    def flush(self) -> None:
        """Сбрасывает буфер с retry."""
        def _do_flush():
            self.producer.flush()

        _retry(_do_flush, max_retries=3, base_delay=0.5)

    def close(self) -> None:
        """Закрывает соединение."""
        try:
            self.producer.close()
        except Exception as e:
            logger.debug(f"Error closing producer (ignored): {e}")

    def load(self) -> GeneratorState | None:
        """Загружает последнее состояние из топика."""
        from kafka import KafkaConsumer

        logger.info(f"Loading state from topic {self.STATE_TOPIC}")

        def _do_load():
            consumer = KafkaConsumer(
                self.STATE_TOPIC,
                bootstrap_servers=self.bootstrap_servers,
                auto_offset_reset="earliest",
                enable_auto_commit=False,
                consumer_timeout_ms=5000,
                value_deserializer=lambda v: json.loads(v.decode("utf-8")),
            )

            last_state = None
            for message in consumer:
                if message.key and message.key.decode("utf-8") == self.STATE_KEY:
                    last_state = message.value

            consumer.close()
            return last_state

        try:
            last_state = _retry(_do_load, max_retries=3, base_delay=0.5)

            if last_state:
                logger.info(
                    f"Restored state: tick={last_state.get('tick')}, "
                    f"last_batch_id={last_state.get('last_batch_id')}"
                )
                restored = GeneratorState.from_dict_safe(last_state)
                if restored is None:
                    logger.warning("State data was invalid, starting fresh")
                return restored

            logger.info("No previous state found, starting fresh")
            return None

        except Exception as e:
            logger.warning(f"Failed to load state: {e}, starting fresh")
            return None


class KafkaStartupHistoryManifest:
    """Хранение манифеста стартовой истории в Kafka compact topic."""

    MANIFEST_TOPIC = "generator_startup_history_manifest"
    MANIFEST_KEY = "default"

    def __init__(self, bootstrap_servers: str):
        self.bootstrap_servers = bootstrap_servers
        KafkaProducerCls, _ = _kafka_importer()()

        logger.info(
            f"Connecting to Kafka for startup history manifest at {self.bootstrap_servers}"
        )
        self.producer = KafkaProducerCls(
            bootstrap_servers=self.bootstrap_servers,
            value_serializer=lambda v: json.dumps(v).encode("utf-8"),
            key_serializer=lambda k: k.encode("utf-8") if k else None,
            retries=3,
            retry_backoff_ms=1000,
        )
        logger.info("Connected to Kafka for startup history manifest successfully")

    def save(self, manifest: dict) -> None:
        """Сохраняет манифест стартовой истории."""
        def _do_send():
            self.producer.send(
                self.MANIFEST_TOPIC,
                key=self.MANIFEST_KEY,
                value=manifest,
            )

        _retry(_do_send, max_retries=3, base_delay=0.5)

    def flush(self) -> None:
        """Сбрасывает буфер с retry."""
        def _do_flush():
            self.producer.flush()

        _retry(_do_flush, max_retries=3, base_delay=0.5)

    def close(self) -> None:
        """Закрывает соединение."""
        try:
            self.producer.close()
        except Exception as e:
            logger.debug(f"Error closing manifest producer (ignored): {e}")

    def load(self) -> dict | None:
        """Загружает последний манифест стартовой истории."""
        from kafka import KafkaConsumer

        logger.info(f"Loading startup history manifest from topic {self.MANIFEST_TOPIC}")

        def _do_load():
            consumer = KafkaConsumer(
                self.MANIFEST_TOPIC,
                bootstrap_servers=self.bootstrap_servers,
                auto_offset_reset="earliest",
                enable_auto_commit=False,
                consumer_timeout_ms=5000,
                value_deserializer=lambda v: json.loads(v.decode("utf-8")),
            )

            last_manifest = None
            for message in consumer:
                if message.key and message.key.decode("utf-8") == self.MANIFEST_KEY:
                    last_manifest = message.value

            consumer.close()
            return last_manifest

        try:
            return _retry(_do_load, max_retries=3, base_delay=0.5)
        except Exception as e:
            logger.warning(f"Failed to load startup history manifest: {e}")
            return None


def ensure_topics(bootstrap_servers: str) -> None:
    """Создаёт служебные топики, если их ещё нет."""
    from kafka import KafkaAdminClient
    from kafka.admin import NewTopic
    from kafka.errors import TopicAlreadyExistsError

    def _create_topics():
        admin_client = KafkaAdminClient(bootstrap_servers=bootstrap_servers)
        try:
            history_topic = NewTopic(
                name=KafkaBatchHistory.HISTORY_TOPIC,
                num_partitions=1,
                replication_factor=1,
            )
            state_topic = NewTopic(
                name=KafkaStateManager.STATE_TOPIC,
                num_partitions=1,
                replication_factor=1,
                topic_configs={
                    "cleanup.policy": "compact",
                    "min.cleanable.dirty.ratio": "0.1",
                    "delete.retention.ms": "100",
                },
            )
            manifest_topic = NewTopic(
                name=KafkaStartupHistoryManifest.MANIFEST_TOPIC,
                num_partitions=1,
                replication_factor=1,
                topic_configs={
                    "cleanup.policy": "compact",
                    "min.cleanable.dirty.ratio": "0.1",
                    "delete.retention.ms": "100",
                },
            )

            for topic in [history_topic, state_topic, manifest_topic]:
                try:
                    admin_client.create_topics([topic])
                    logger.info(f"Created topic: {topic.name}")
                except TopicAlreadyExistsError:
                    logger.debug(f"Topic already exists: {topic.name}")
        finally:
            admin_client.close()

    _retry(_create_topics, max_retries=5, base_delay=1.0)


class KafkaBatchHistory:
    """Хранение истории batch в Kafka."""

    HISTORY_TOPIC = "generator_batch_history"

    def __init__(self, bootstrap_servers: str):
        self.bootstrap_servers = bootstrap_servers
        self.producer = None
        self._connect()

    def _connect(self):
        """Устанавливает соединение с Kafka с retry."""
        KafkaProducerCls, _ = _kafka_importer()()

        def _do_connect():
            logger.info(f"Connecting to Kafka for history at {self.bootstrap_servers}")
            self.producer = KafkaProducerCls(
                bootstrap_servers=self.bootstrap_servers,
                value_serializer=lambda v: json.dumps(v).encode("utf-8"),
                key_serializer=lambda k: k.encode("utf-8") if k else None,
                retries=3,
                retry_backoff_ms=1000,
            )
            logger.info("Connected to Kafka for history successfully")

        _retry(_do_connect, max_retries=5, base_delay=1.0)

    def add(self, record: BatchRecord):
        """Добавляет запись в историю."""
        key = record.batch_id
        value = record.to_dict()

        def _do_send():
            self.producer.send(self.HISTORY_TOPIC, key=key, value=value)

        try:
            _retry(_do_send, max_retries=3, base_delay=0.5)
        except Exception as e:
            logger.warning(f"Failed to send history record after retries: {e}, attempting reconnect")
            self._connect()
            _retry(_do_send, max_retries=2, base_delay=0.5)

    def flush(self):
        """Сбрасывает буфер с retry."""
        def _do_flush():
            self.producer.flush()

        _retry(_do_flush, max_retries=3, base_delay=0.5)

    def close(self):
        """Закрывает соединение."""
        try:
            if self.producer:
                self.producer.close()
        except Exception as e:
            logger.debug(f"Error closing history producer (ignored): {e}")


class KafkaPublisher:
    """Публикация событий в Kafka с retry и реконнектом."""

    def __init__(self, bootstrap_servers: str):
        self.bootstrap_servers = bootstrap_servers
        self.producer = None
        self._connect()

    def _connect(self):
        """Устанавливает соединение с Kafka с retry."""
        KafkaProducerCls, _ = _kafka_importer()()

        def _do_connect():
            logger.info(f"Connecting to Kafka at {self.bootstrap_servers}")
            self.producer = KafkaProducerCls(
                bootstrap_servers=self.bootstrap_servers,
                value_serializer=lambda v: json.dumps(v).encode("utf-8"),
                key_serializer=lambda k: k.encode("utf-8") if k else None,
                batch_size=16384,
                linger_ms=100,
                retries=3,
                retry_backoff_ms=1000,
            )
            logger.info("Connected to Kafka successfully")

        _retry(_do_connect, max_retries=5, base_delay=1.0)

    def _publish_with_retry(self, topic: str, events: list[dict]) -> tuple[int, int]:
        """Внутренняя функция публикации с retry на уровне batch."""
        if not self.producer:
            raise RuntimeError("Producer not connected")

        sent = 0
        errors = 0
        futures = []

        for event in events:
            key = event.get("event_id") or event.get("click_id")
            try:
                future = self.producer.send(topic, key=key, value=event)
                futures.append(future)
            except Exception as e:
                logger.error(f"Failed to send message to {topic}: {e}")
                errors += 1
                METRICS_ERRORS_TOTAL.labels(topic=topic).inc()

        for future in futures:
            try:
                future.get(timeout=10)
                sent += 1
                METRICS_EVENTS_TOTAL.labels(topic=topic).inc()
            except Exception as e:
                logger.error(f"Failed to confirm message delivery: {e}")
                errors += 1
                METRICS_ERRORS_TOTAL.labels(topic=topic).inc()

        return sent, errors

    def publish(self, topic: str, events: list[dict]) -> tuple[int, int]:
        """Публикует события в топик с retry и автоматическим реконнектом."""
        def _do_publish():
            return self._publish_with_retry(topic, events)

        try:
            return _retry(_do_publish, max_retries=3, base_delay=0.5)
        except Exception as e:
            logger.warning(f"Publish failed after retries: {e}, attempting reconnect")
            self._connect()
            return _retry(_do_publish, max_retries=2, base_delay=0.5)

    def flush(self):
        """Сбрасывает буфер с retry."""
        def _do_flush():
            if self.producer:
                self.producer.flush()

        _retry(_do_flush, max_retries=3, base_delay=0.5)

    def close(self):
        """Закрывает соединение."""
        try:
            if self.producer:
                self.producer.close()
        except Exception as e:
            logger.debug(f"Error closing producer (ignored): {e}")