buz 2.14.0__py3-none-any.whl → 2.14.1__py3-none-any.whl

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/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/kafka/infrastructure/kafka_python/kafka_python_admin_client.py

@@ -8,10 +8,12 @@ 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
@@ -150,10 +152,10 @@ class KafkaPythonAdminClient(KafkaAdminClient):
             group_ids=subscription_groups,
         )
 
-    def get_subscription_groups(
+    def get_cluster_consumer_groups(
         self,
     ) -> set[str]:
-        return set(self._get_kafka_admin().list_consumer_groups())
+        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()
@@ -166,6 +168,46 @@ class KafkaPythonAdminClient(KafkaAdminClient):
         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,
@@ -175,44 +217,29 @@ class KafkaPythonAdminClient(KafkaAdminClient):
             **self._config_in_library_format,
         )
 
-        partitions = consumer.partitions_for_topic(topic)
+        try:
+            partitions = self.get_number_of_partitions(topic)
 
-        if partitions is None:
-            raise TopicNotFoundException(topic)
+            topic_partitions = [TopicPartition(topic=topic, partition=partition) for partition in range(partitions)]
 
-        topic_partitions = [TopicPartition(topic, p) for p in partitions]
-        consumer.subscribe(topics=[topic])
+            consumer.subscribe(topic)
 
-        self.__force_partition_assignment(consumer)
+            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)
+            # We need all the partitions in order to update the offsets
+            if len(consumer.assignment()) != len(topic_partitions):
+                raise NotAllPartitionAssignedException(topic)
 
-        offsets_for_date = self.__get_first_offset_after_date(
-            consumer=consumer,
-            topic_partitions=topic_partitions,
-            target_datetime=target_datetime,
-        )
-
-        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])
+            # 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)
 
-        consumer.commit()
-        consumer.close()
+            return (consumer, topic_partitions)
+        except Exception as exception:
+            consumer.close()
+            raise exception
 
     def __get_first_offset_after_date(
         self,
@@ -235,3 +262,138 @@ class KafkaPythonAdminClient(KafkaAdminClient):
     # 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/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()

{buz-2.14.0.dist-info → buz-2.14.1.dist-info}/METADATA

@@ -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.dist-info → buz-2.14.1.dist-info}/RECORD

@@ -138,6 +138,7 @@ buz/kafka/__init__.py,sha256=R3fcyET-SNEAvk_XlBQbHIbQVb63Qiz6lVrif3nDhNU,3435
 buz/kafka/domain/exceptions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 buz/kafka/domain/exceptions/not_all_partition_assigned_exception.py,sha256=9zDWoh0SbHLRuCvpfIGcvrmcscKsXpbAPIxr5-z-GYg,296
 buz/kafka/domain/exceptions/not_valid_kafka_message_exception.py,sha256=Dn6I_-eGQnOuu5WW24oKGOdKOu4EdM8ByH3DLAbz5SY,57
+buz/kafka/domain/exceptions/not_valid_partition_number_exception.py,sha256=YZyGbblHk6ON9sBtjRQTDa-nC88i4oe14_VSO8vSTm0,337
 buz/kafka/domain/exceptions/topic_already_created_exception.py,sha256=UrisdveZGa2BB0ko4mS7-5fwy8eGsIu409_grtq1r9k,333
 buz/kafka/domain/exceptions/topic_not_found_exception.py,sha256=kLuqGqfsb6YTCe5UCKpMwBm_QAnU9Udfb8bWajPoA8k,201
 buz/kafka/domain/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -155,7 +156,7 @@ buz/kafka/domain/models/kafka_supported_sasl_mechanisms.py,sha256=ASyDaFgseQRcUJ
 buz/kafka/domain/models/kafka_supported_security_protocols.py,sha256=ffY2-9sOj4XIkJTSQVkqeOb4KnuqEYXISDarfDN8r9Q,161
 buz/kafka/domain/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 buz/kafka/domain/services/async_kafka_producer.py,sha256=gSq3WwEVux_gp3EKDAMN1WsM027uklB58E-WnKpyhPs,533
-buz/kafka/domain/services/kafka_admin_client.py,sha256=XE6H-4JWsjygKjdtLtGMX5ELhWkJcpR7ai9CO2kB98Y,1222
+buz/kafka/domain/services/kafka_admin_client.py,sha256=Kh_w-qWEY8rsrlYjnJT1FrJLVZrO3l2LeRAWFyc_nOg,1558
 buz/kafka/domain/services/kafka_admin_test_client.py,sha256=91l_vFIo1yhJLQQCC_OmeXZ5F429zP7Hx5g4FNllpfE,1625
 buz/kafka/domain/services/kafka_producer.py,sha256=8bLTV328orrPHcARzkc6no4vyJzrArVtCsjmSRXDjos,506
 buz/kafka/infrastructure/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -180,8 +181,8 @@ buz/kafka/infrastructure/interfaces/async_connection_manager.py,sha256=JbaLu5UVV
 buz/kafka/infrastructure/interfaces/connection_manager.py,sha256=EWnvShJHOg8QYe6a3ma0urjKjmVMDBi7q8T2cv_i_MQ,200
 buz/kafka/infrastructure/kafka_python/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 buz/kafka/infrastructure/kafka_python/exception/consumer_interrupted_exception.py,sha256=fqhgV7HILdVdv-p1CsOIaaESKY2ZXBtRGYbrVSdPLg0,164
-buz/kafka/infrastructure/kafka_python/kafka_python_admin_client.py,sha256=7HbwkaoedXbxZ4LcGWytU2q_GTdLm_c1ziyGTBoKgF0,9038
-buz/kafka/infrastructure/kafka_python/kafka_python_admin_test_client.py,sha256=wLhnrHzyFJ9ETWNUfd-dmwm_CwZyiImaAP97cEdRgzE,2982
+buz/kafka/infrastructure/kafka_python/kafka_python_admin_client.py,sha256=OwyDCXXMA3-yhBRVW4iRCr9M9DJQ65QBtEKx032z1rs,16140
+buz/kafka/infrastructure/kafka_python/kafka_python_admin_test_client.py,sha256=5xP23dQ7FDuy7dIWNw39C3bMVmaUj9ZQhEEJISRv9ec,2986
 buz/kafka/infrastructure/kafka_python/kafka_python_producer.py,sha256=DkqqLSSXHBf4SXXf-IZwwLhxWrGE95Jg4MO_3RDsikU,3594
 buz/kafka/infrastructure/kafka_python/translators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 buz/kafka/infrastructure/kafka_python/translators/consumer_initial_offset_position_translator.py,sha256=hJ48_eyMcnbFL_Y5TOiMbGXrQSryuKk9CvP59MdqNOY,620
@@ -243,7 +244,7 @@ buz/serializer/message_to_json_bytes_serializer.py,sha256=RGZJ64t4t4Pz2FCASZZCv-
 buz/wrapper/__init__.py,sha256=GnRdJFcncn-qp0hzDG9dBHLmTJSbHFVjE_yr-MdW_n4,77
 buz/wrapper/async_to_sync.py,sha256=OfK-vrVUhuN-LLLvekLdMbQYtH0ue5lfbvuasj6ovMI,698
 buz/wrapper/event_loop.py,sha256=pfBJ1g-8A2a3YgW8Gf9Fg0kkewoh3-wgTy2KIFDyfHk,266
-buz-2.14.0.dist-info/LICENSE,sha256=Jytu2S-2SPEgsB0y6BF-_LUxIWY7402fl0JSh36TLZE,1062
-buz-2.14.0.dist-info/METADATA,sha256=f8tigTI_Ee1HBkPHTPfsYHjkOyuBZwlCX6dTeDpjzCQ,1617
-buz-2.14.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-buz-2.14.0.dist-info/RECORD,,
+buz-2.14.1.dist-info/LICENSE,sha256=Jytu2S-2SPEgsB0y6BF-_LUxIWY7402fl0JSh36TLZE,1062
+buz-2.14.1.dist-info/METADATA,sha256=WL8Cp60N7_OJhpYuwyBLADYJ5cVohFC29oKmBRDri00,1617
+buz-2.14.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+buz-2.14.1.dist-info/RECORD,,