buz 2.14.0__tar.gz → 2.14.1__tar.gz

{buz-2.14.0 → buz-2.14.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: buz
-Version: 2.14.0
+Version: 2.14.1
 Summary: Buz is a set of light, simple and extensible implementations of event, command and query buses.
 License: MIT
 Author: Luis Pintado Lozano

{buz-2.14.0 → buz-2.14.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "buz"
-version = "2.14.0"
+version = "2.14.1"
 description = "Buz is a set of light, simple and extensible implementations of event, command and query buses."
 readme = "README.md"
 authors = ["Luis Pintado Lozano <luis.pintado.lozano@gmail.com>", "Gerardo Parra <gprauxiliar@gmail.com>"]

buz-2.14.1/src/buz/kafka/domain/exceptions/not_valid_partition_number_exception.py

@@ -0,0 +1,10 @@
+class NotValidPartitionNumberException(Exception):
+    def __init__(
+        self,
+        *,
+        partition_number: int,
+        min_partition_number: int,
+    ) -> None:
+        super().__init__(
+            f'"{partition_number}" is not a valid partition number, the minimum partition number is "{min_partition_number}"'
+        )

{buz-2.14.0 → buz-2.14.1}/src/buz/kafka/domain/services/kafka_admin_client.py

@@ -40,6 +40,10 @@ class KafkaAdminClient(ConnectionManager, ABC):
     ) -> set[str]:
         pass
 
+    @abstractmethod
+    def get_number_of_partitions(self, topic: str) -> int:
+        pass
+
     # This function moves the following offset from the provided date
     # if there are no messages with a date greater than the provided offset
     # the offset will be moved to the end
@@ -52,3 +56,12 @@ class KafkaAdminClient(ConnectionManager, ABC):
         target_datetime: datetime,
     ) -> None:
         pass
+
+    @abstractmethod
+    def increase_topic_partitions_and_set_offset_of_related_consumer_groups_to_the_beginning_of_the_new_ones(
+        self,
+        *,
+        topic: str,
+        new_number_of_partitions: int,
+    ) -> None:
+        pass

buz-2.14.1/src/buz/kafka/infrastructure/kafka_python/kafka_python_admin_client.py

@@ -0,0 +1,399 @@
+from __future__ import annotations
+
+from datetime import datetime
+from logging import Logger
+import re
+from typing import Any, Callable, Optional, Sequence, cast
+
+from cachetools import TTLCache
+from kafka import KafkaClient, KafkaConsumer
+from kafka.admin import KafkaAdminClient as KafkaPythonLibraryAdminClient, NewTopic
+from kafka.admin.new_partitions import NewPartitions
+from kafka.errors import TopicAlreadyExistsError
+from kafka.structs import TopicPartition, OffsetAndTimestamp
+
+from buz.kafka.domain.exceptions.not_all_partition_assigned_exception import NotAllPartitionAssignedException
+from buz.kafka.domain.exceptions.not_valid_partition_number_exception import NotValidPartitionNumberException
+from buz.kafka.domain.exceptions.topic_already_created_exception import KafkaTopicsAlreadyCreatedException
+from buz.kafka.domain.exceptions.topic_not_found_exception import TopicNotFoundException
+from buz.kafka.domain.models.consumer_initial_offset_position import ConsumerInitialOffsetPosition
+from buz.kafka.domain.models.create_kafka_topic import CreateKafkaTopic
+from buz.kafka.domain.models.kafka_connection_config import KafkaConnectionConfig
+from buz.kafka.domain.services.kafka_admin_client import KafkaAdminClient
+
+from buz.kafka.infrastructure.kafka_python.translators.consumer_initial_offset_position_translator import (
+    KafkaPythonConsumerInitialOffsetPositionTranslator,
+)
+
+INTERNAL_KAFKA_TOPICS = {"__consumer_offsets", "_schema"}
+TOPIC_CACHE_KEY = "topics"
+
+
+class KafkaPythonAdminClient(KafkaAdminClient):
+    __PYTHON_KAFKA_DUPLICATED_TOPIC_ERROR_CODE = 36
+
+    _kafka_admin: Optional[KafkaPythonLibraryAdminClient] = None
+    _kafka_client: Optional[KafkaClient] = None
+
+    def __init__(
+        self,
+        *,
+        logger: Logger,
+        connection_config: KafkaConnectionConfig,
+        cache_ttl_seconds: int = 0,
+    ):
+        self._logger = logger
+        self.__connection_config = connection_config
+        self._config_in_library_format = self.__get_kafka_config_in_library_format(self.__connection_config)
+        self.__ttl_cache: TTLCache[str, Any] = TTLCache(maxsize=1, ttl=cache_ttl_seconds)
+
+    def __get_kafka_config_in_library_format(self, config: KafkaConnectionConfig) -> dict:
+        return {
+            "client_id": config.client_id,
+            "bootstrap_servers": config.bootstrap_servers,
+            "security_protocol": config.credentials.security_protocol.value,
+            "sasl_mechanism": config.credentials.sasl_mechanism.value if config.credentials.sasl_mechanism else None,
+            "sasl_plain_username": config.credentials.user,
+            "sasl_plain_password": config.credentials.password,
+        }
+
+    def connect(self):
+        self._get_kafka_admin()
+        self._get_kafka_client()
+
+    def disconnect(self):
+        if self._kafka_admin is not None:
+            self._kafka_admin.close()
+            self._kafka_admin = None
+        if self._kafka_client is not None:
+            self._kafka_client.close()
+            self._kafka_client = None
+
+    def _get_kafka_admin(self) -> KafkaPythonLibraryAdminClient:
+        if not self._kafka_admin:
+            self._kafka_admin = KafkaPythonLibraryAdminClient(**self._config_in_library_format)
+        return self._kafka_admin
+
+    def _get_kafka_client(self) -> KafkaClient:
+        if not self._kafka_client:
+            self._kafka_client = KafkaClient(**self._config_in_library_format)
+        return self._kafka_client
+
+    def create_topics(
+        self,
+        *,
+        topics: Sequence[CreateKafkaTopic],
+    ) -> None:
+        new_topics = [
+            NewTopic(
+                name=topic.name,
+                num_partitions=topic.partitions,
+                replication_factor=topic.replication_factor,
+                topic_configs=topic.configs,
+            )
+            for topic in topics
+        ]
+
+        try:
+            self._get_kafka_admin().create_topics(new_topics=new_topics)
+        except TopicAlreadyExistsError as error:
+            topic_names = self.__get_list_of_kafka_topics_from_topic_already_exists_error(error)
+            raise KafkaTopicsAlreadyCreatedException(topic_names=topic_names)
+
+    def __get_list_of_kafka_topics_from_topic_already_exists_error(self, error: TopicAlreadyExistsError) -> list[str]:
+        message = str(error)
+        response_message = re.search(r"topic_errors=\[.*?]", message)
+        topic_messages = re.findall(
+            r"topic='[^']*', error_code=" + str(self.__PYTHON_KAFKA_DUPLICATED_TOPIC_ERROR_CODE), response_message[0]  # type: ignore
+        )
+
+        return [re.search("'.*'", topic_message)[0].strip("'") for topic_message in topic_messages]  # type: ignore
+
+    def is_topic_created(
+        self,
+        topic: str,
+    ) -> bool:
+        return topic in self.get_topics()
+
+    def get_topics(
+        self,
+    ) -> set[str]:
+        return self.__resolve_cached_property(
+            TOPIC_CACHE_KEY, lambda: set(self._get_kafka_admin().list_topics()) - INTERNAL_KAFKA_TOPICS
+        )
+
+    def __resolve_cached_property(self, property_key: str, callback: Callable) -> Any:
+        value = self.__ttl_cache.get(property_key)
+        if value is not None:
+            return value
+        value = callback()
+        self.__ttl_cache[property_key] = value
+        return value
+
+    def delete_topics(
+        self,
+        *,
+        topics: set[str],
+    ) -> None:
+        self._get_kafka_admin().delete_topics(
+            topics=topics,
+        )
+        self.__remove_cache_property(TOPIC_CACHE_KEY)
+
+    def __remove_cache_property(self, property_key: str) -> None:
+        self.__ttl_cache.pop(property_key, None)
+
+    def delete_subscription_groups(
+        self,
+        *,
+        subscription_groups: set[str],
+    ) -> None:
+        self._get_kafka_admin().delete_consumer_groups(
+            group_ids=subscription_groups,
+        )
+
+    def get_cluster_consumer_groups(
+        self,
+    ) -> set[str]:
+        return set([consumer_group_tuple[0] for consumer_group_tuple in self._get_kafka_admin().list_consumer_groups()])
+
+    def _wait_for_cluster_update(self) -> None:
+        future = self._get_kafka_client().cluster.request_update()
+        self._get_kafka_client().poll(future=future)
+
+    def move_offsets_to_datetime(
+        self,
+        *,
+        consumer_group: str,
+        topic: str,
+        target_datetime: datetime,
+    ) -> None:
+        (consumer, topic_partitions) = self.__get_consumer_with_all_partitions_assigned(
+            consumer_group=consumer_group,
+            topic=topic,
+        )
+
+        offsets_for_date = self.__get_first_offset_after_date(
+            consumer=consumer,
+            topic_partitions=topic_partitions,
+            target_datetime=target_datetime,
+        )
+
+        try:
+            end_offsets = consumer.end_offsets(topic_partitions)
+
+            if end_offsets is None or len(end_offsets.keys()) != len(topic_partitions):
+                raise Exception(f'There was an error extracting the end offsets of the topic "{topic}"')
+
+            for topic_partition in topic_partitions:
+                offset_and_timestamp = offsets_for_date.get(topic_partition)
+                if offset_and_timestamp:
+                    self._logger.info(f'moving "{topic_partition}" to the offset "{offset_and_timestamp.offset}"')
+                    consumer.seek(topic_partition, offset_and_timestamp.offset)
+                else:
+                    self._logger.info(
+                        f'moving "{topic_partition}" to the end of the topic because there are no messages later than "{target_datetime}"'
+                    )
+                    consumer.seek(topic_partition, end_offsets[topic_partition])
+
+            consumer.commit()
+        except Exception as exception:
+            consumer.close()
+            raise exception
+
+        consumer.close()
+
+    def __get_consumer_with_all_partitions_assigned(
+        self,
+        consumer_group: str,
+        topic: str,
+    ) -> tuple[KafkaConsumer, Sequence[TopicPartition]]:
+        consumer = KafkaConsumer(
+            group_id=consumer_group,
+            enable_auto_commit=False,
+            auto_offset_reset=KafkaPythonConsumerInitialOffsetPositionTranslator.to_kafka_supported_format(
+                ConsumerInitialOffsetPosition.BEGINNING
+            ),
+            **self._config_in_library_format,
+        )
+
+        try:
+            partitions = self.get_number_of_partitions(topic)
+
+            topic_partitions = [TopicPartition(topic=topic, partition=partition) for partition in range(partitions)]
+
+            consumer.subscribe(topic)
+
+            self.__force_partition_assignment(consumer)
+
+            # We need all the partitions in order to update the offsets
+            if len(consumer.assignment()) != len(topic_partitions):
+                raise NotAllPartitionAssignedException(topic)
+
+            # This could produce a race condition, but it is a limitation of kafka admin (we are not able to check if all the partition are assigned using the manual assignment)
+            # https://github.com/dpkp/kafka-python/blob/master/kafka/consumer/group.py#L430
+            consumer.unsubscribe()
+            consumer.assign(topic_partitions)
+            self.__force_partition_assignment(consumer)
+
+            return (consumer, topic_partitions)
+        except Exception as exception:
+            consumer.close()
+            raise exception
+
+    def __get_first_offset_after_date(
+        self,
+        *,
+        consumer: KafkaConsumer,
+        topic_partitions: Sequence[TopicPartition],
+        target_datetime: datetime,
+    ) -> dict[TopicPartition, Optional[OffsetAndTimestamp]]:
+        offset_for_times: dict[TopicPartition, Optional[int]] = {}
+        timestamp_ms = int(target_datetime.timestamp() * 1000)
+
+        for topic_partition in topic_partitions:
+            offset_for_times[topic_partition] = timestamp_ms
+
+        return cast(
+            dict[TopicPartition, Optional[OffsetAndTimestamp]],
+            consumer.offsets_for_times(offset_for_times),
+        )
+
+    # We are not to commit the new offset, but we need to execute a polling in order to start the partition assignment
+    def __force_partition_assignment(self, consumer: KafkaConsumer) -> None:
+        consumer.poll(max_records=1, timeout_ms=0)
+
+    def increase_topic_partitions_and_set_offset_of_related_consumer_groups_to_the_beginning_of_the_new_ones(
+        self,
+        *,
+        topic: str,
+        new_number_of_partitions: int,
+    ) -> None:
+        self._logger.info(
+            f'Increasing topic "{topic}" partitions: Verifying the new number of partitions "{new_number_of_partitions}"'
+        )
+
+        previous_partitions_number = self.get_number_of_partitions(topic)
+        topic_partitions = [
+            TopicPartition(topic=topic, partition=partition) for partition in range(previous_partitions_number)
+        ]
+
+        if previous_partitions_number >= new_number_of_partitions:
+            raise NotValidPartitionNumberException(
+                partition_number=new_number_of_partitions,
+                min_partition_number=len(topic_partitions),
+            )
+
+        self._logger.info(f'Increasing topic "{topic}" partitions: Extracting related consumer groups')
+        related_consumer_groups = self.__get_consumer_groups_related_to_a_topic(topic_partitions)
+
+        self._logger.info(
+            f'Increasing topic "{topic}" partitions: The following consumer groups will be updated:"{related_consumer_groups}"'
+        )
+
+        consumers_to_update: list[KafkaConsumer] = []
+        new_partitions_consumer: Optional[KafkaConsumer] = None
+
+        try:
+            for consumer_group in related_consumer_groups:
+                self._logger.info(
+                    f'Increasing topic "{topic}" partitions: Requesting the assignment of the partitions of the group "{consumer_group}"'
+                )
+                (consumer_with_all_partitions, _) = self.__get_consumer_with_all_partitions_assigned(
+                    consumer_group=consumer_group,
+                    topic=topic,
+                )
+                consumers_to_update.append(consumer_with_all_partitions)
+
+            self._logger.info(
+                f'Increasing topic "{topic}" partitions: Incrementing the partition to "{new_number_of_partitions}"'
+            )
+
+            self._get_kafka_admin().create_partitions(
+                {
+                    topic: NewPartitions(total_count=new_number_of_partitions),
+                }
+            )
+
+            new_partitions = [
+                TopicPartition(
+                    topic=topic,
+                    partition=partition_index,
+                )
+                for partition_index in range(previous_partitions_number, new_number_of_partitions)
+            ]
+
+            for consumer_group in related_consumer_groups:
+                self._logger.info(
+                    f'Increasing topic "{topic}" partitions: Moving the offset of the consumer group "{consumer_group}" to the beginning of the new partitions'
+                )
+                # We need to create a new consumer because kafka-python has a limitation that does not allow to assign specific partitions to a consumer subscribed to an entire topic
+                new_partitions_consumer = KafkaConsumer(
+                    group_id=consumer_group,
+                    enable_auto_commit=False,
+                    auto_offset_reset=KafkaPythonConsumerInitialOffsetPositionTranslator.to_kafka_supported_format(
+                        ConsumerInitialOffsetPosition.BEGINNING
+                    ),
+                    **self._config_in_library_format,
+                )
+                new_partitions_consumer.assign(new_partitions)
+                for new_partition in new_partitions:
+                    new_partitions_consumer.seek(new_partition, 0)
+                new_partitions_consumer.commit()
+                new_partitions_consumer.close()
+
+            self._logger.info(f'Increasing topic "{topic}" partitions: Process complete')
+
+        except Exception as exception:
+            for consumer_with_all_partitions in consumers_to_update:
+                consumer_with_all_partitions.close()
+
+            if new_partitions_consumer is not None:
+                new_partitions_consumer.close()
+
+            self._logger.error(f'Increasing topic "{topic}" partitions: unexpected error {exception}')
+            raise exception
+
+        return
+
+    def get_number_of_partitions(self, topic: str) -> int:
+        consumer = KafkaConsumer(
+            enable_auto_commit=False,
+            auto_offset_reset=KafkaPythonConsumerInitialOffsetPositionTranslator.to_kafka_supported_format(
+                ConsumerInitialOffsetPosition.BEGINNING
+            ),
+            **self._config_in_library_format,
+        )
+
+        try:
+            partitions = consumer.partitions_for_topic(topic)
+            if partitions is None:
+                raise TopicNotFoundException(topic_name=topic)
+
+            return len(partitions)
+        except Exception as exception:
+            consumer.close()
+            raise exception
+
+    # The purpose of this function is to get all the consumer groups that are consuming from the topic
+    # It is a heavy tasks because we need to get the offset of all the partitions of the topic
+    def __get_consumer_groups_related_to_a_topic(self, topic_partitions: Sequence[TopicPartition]) -> set[str]:
+        cluster_consumer_groups = self.get_cluster_consumer_groups()
+
+        related_consumer_groups: set[str] = set()
+
+        for consumer_group in cluster_consumer_groups:
+            partitions_offsets = list(
+                self._get_kafka_admin()
+                .list_consumer_group_offsets(consumer_group, partitions=topic_partitions)
+                .values()
+            )
+
+            partitions_with_valid_offsets = [partition for partition in partitions_offsets if partition.offset != -1]
+
+            if len(partitions_with_valid_offsets) == 0:
+                continue
+
+            related_consumer_groups.add(consumer_group)
+
+        return related_consumer_groups

{buz-2.14.0 → buz-2.14.1}/src/buz/kafka/infrastructure/kafka_python/kafka_python_admin_test_client.py

@@ -88,5 +88,5 @@ class KafkaPythonAdminTestClient(KafkaPythonAdminClient, KafkaAdminTestClient):
         self,
     ) -> None:
         self.delete_topics(topics=self.get_topics())
-        self.delete_subscription_groups(subscription_groups=self.get_subscription_groups())
+        self.delete_subscription_groups(subscription_groups=self.get_cluster_consumer_groups())
         self._wait_for_cluster_update()