feat: added support for both versions

Author: David Romaschenko

poetry.lock

@@ -27,7 +27,6 @@ packages = [{ include = "taskiq_aio_kafka" }]
 python = "^3.9"
 taskiq = "^0"
 aiokafka = "^0.10.0"
-pydantic = "^2.7.4"
 kafka-python = "^2.0.2"
 
 [tool.poetry.group.dev.dependencies]

pyproject.toml

@@ -4,8 +4,11 @@
 
 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
 from kafka.admin import KafkaAdminClient, NewTopic
+from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
+from kafka.partitioner.default import DefaultPartitioner
 from taskiq import AsyncResultBackend, BrokerMessage
 from taskiq.abc.broker import AsyncBroker
+from taskiq.compat import model_dump
 
 from taskiq_aio_kafka.exceptions import WrongAioKafkaBrokerParametersError
 from taskiq_aio_kafka.models import KafkaConsumerParameters, KafkaProducerParameters
@@ -146,19 +149,31 @@ async def startup(self) -> None:
                 new_topics=[self._kafka_topic],
                 validate_only=False,
             )
+
+        partitioner = self._aiokafka_producer_params.partitioner or DefaultPartitioner()
+        producer_kwargs = model_dump(self._aiokafka_producer_params)
+        producer_kwargs["partitioner"] = partitioner
         self._aiokafka_producer = AIOKafkaProducer(
             bootstrap_servers=self._bootstrap_servers,
             loop=self._loop,
-            **self._aiokafka_producer_params.model_dump(),
+            **producer_kwargs,
         )
         await self._aiokafka_producer.start()
 
         if self.is_worker_process:
+            partition_assignment_strategy = (
+                self._aiokafka_consumer_params.partition_assignment_strategy
+                or (RoundRobinPartitionAssignor,)
+            )
+            consumer_kwargs = model_dump(self._aiokafka_consumer_params)
+            consumer_kwargs["partition_assignment_strategy"] = (
+                partition_assignment_strategy
+            )
             self._aiokafka_consumer = AIOKafkaConsumer(
                 self._kafka_topic.name,
                 bootstrap_servers=self._bootstrap_servers,
                 loop=self._loop,
-                **self._aiokafka_consumer_params.model_dump(),
+                **consumer_kwargs,
             )
 
             await self._aiokafka_consumer.start()

taskiq_aio_kafka/broker.py

@@ -1,8 +1,6 @@
 from typing import Any, Callable, Optional, Union
 
 from aiokafka import __version__
-from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
-from kafka.partitioner.default import DefaultPartitioner
 from pydantic import BaseModel
 
 
@@ -18,7 +16,7 @@ class KafkaProducerParameters(BaseModel):
     value_serializer: Optional[Callable[..., bytes]] = None
     compression_type: Optional[str] = None
     max_batch_size: int = 16384
-    partitioner: Callable[..., Any] = DefaultPartitioner()
+    partitioner: Optional[Callable[..., Any]] = None
     max_request_size: int = 1048576
     linger_ms: int = 0
     send_backoff_ms: int = 100
@@ -55,7 +53,7 @@ class KafkaConsumerParameters(BaseModel):
     auto_commit_interval_ms: int = 5000
     check_crcs: bool = True
     metadata_max_age_ms: int = 5 * 60 * 1000
-    partition_assignment_strategy: Any = (RoundRobinPartitionAssignor,)
+    partition_assignment_strategy: Any = None
     max_poll_interval_ms: int = 300000
     rebalance_timeout_ms: Optional[int] = None
     session_timeout_ms: int = 10000