Add new classes that describe parameters for kafka producer and consumer cuz AIOKafkaProducer and AIOKafkaConsumer must be initialize in async method or with passed loop parameter.

Author: chandr-andr

poetry.lock

@@ -27,6 +27,7 @@ packages = [{ include = "taskiq_aio_kafka" }]
 python = "^3.7"
 taskiq = "^0"
 aiokafka = "^0.8.0"
+pydantic = "^1.10.7"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.1.2"

pyproject.toml

@@ -1,14 +1,15 @@
 import asyncio
 import pickle  # noqa: S403
 from logging import getLogger
-from typing import AsyncGenerator, Callable, List, Optional, Set, TypeVar, Union
+from typing import Any, AsyncGenerator, Callable, List, Optional, Set, TypeVar, Union
 
 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
 from kafka.admin import KafkaAdminClient, NewTopic
 from taskiq import AsyncResultBackend, BrokerMessage
 from taskiq.abc.broker import AsyncBroker
 
 from taskiq_aio_kafka.exceptions import WrongAioKafkaBrokerParametersError
+from taskiq_aio_kafka.models import KafkaConsumerParameters, KafkaProducerParameters
 
 _T = TypeVar("_T")  # noqa: WPS111
 
@@ -86,9 +87,13 @@ def __init__(  # noqa: WPS211
             replication_factor=1,
         )
 
-        self._aiokafka_producer: Optional[AIOKafkaProducer] = aiokafka_producer
+        self._aiokafka_producer_params: KafkaProducerParameters = (
+            KafkaProducerParameters()
+        )
 
-        self._aiokafka_consumer: Optional[AIOKafkaConsumer] = aiokafka_consumer
+        self._aiokafka_consumer_params: KafkaConsumerParameters = (
+            KafkaConsumerParameters()
+        )
 
         self._kafka_admin_client: KafkaAdminClient = (
             kafka_admin_client
@@ -105,6 +110,30 @@ def __init__(  # noqa: WPS211
         self._is_producer_started = False
         self._is_consumer_started = False
 
+    def configure_producer(self, **producer_parameters: Any) -> None:
+        """Configure kafka producer.
+
+        You can pass here any configuration parameters
+        accepted by the kafka producer.
+
+        :param producer_parameters: producer parameters kwargs.
+        """
+        self._aiokafka_producer_params = KafkaProducerParameters(
+            **producer_parameters,
+        )
+
+    def configure_consumer(self, **consumer_parameters: Any) -> None:
+        """Configure kafka consumer.
+
+        You can pass here any configuration parameters
+        accepted by the kafka consumer.
+
+        :param consumer_parameters: consumer parameters kwargs.
+        """
+        self._aiokafka_consumer_params = KafkaConsumerParameters(
+            **consumer_parameters,
+        )
+
     async def startup(self) -> None:
         """Setup AIOKafkaProducer, AIOKafkaConsumer and kafka topics.
 
@@ -114,30 +143,28 @@ async def startup(self) -> None:
         if there are no producer and consumer passed.
         """
         await super().startup()
-
-        is_topic_available: bool = bool(
-            self._kafka_admin_client.describe_topics([self._kafka_topic.name]),
+        available_condition: bool = (
+            self._kafka_topic.name not in self._kafka_admin_client.list_topics()
         )
-        if not is_topic_available:
+        if available_condition:
             self._kafka_admin_client.create_topics(
                 new_topics=[self._kafka_topic],
                 validate_only=False,
             )
+        self._aiokafka_producer = AIOKafkaProducer(
+            bootstrap_servers=self._bootstrap_servers,
+            loop=self._loop,
+            **self._aiokafka_producer_params.dict(),
+        )
+        await self._aiokafka_producer.start()
 
-        if not self._aiokafka_producer:
-            self._aiokafka_producer = AIOKafkaProducer(
+        if self.is_worker_process:
+            self._aiokafka_consumer = AIOKafkaConsumer(
+                self._kafka_topic.name,
                 bootstrap_servers=self._bootstrap_servers,
                 loop=self._loop,
+                **self._aiokafka_consumer_params.dict(),
             )
-        await self._aiokafka_producer.start()
-
-        if self.is_worker_process:
-            if not self._aiokafka_consumer:
-                self._aiokafka_consumer = AIOKafkaConsumer(
-                    self._kafka_topic.name,
-                    bootstrap_servers=self._bootstrap_servers,
-                    loop=self._loop,
-                )
 
             await self._aiokafka_consumer.start()
             self._is_consumer_started = True
@@ -148,10 +175,10 @@ async def shutdown(self) -> None:
         """Close all connections on shutdown."""
         await super().shutdown()
 
-        if self._aiokafka_producer:
+        if self._is_producer_started:
             await self._aiokafka_producer.stop()
 
-        if self._aiokafka_consumer:
+        if self._is_consumer_started:
             await self._aiokafka_consumer.stop()
 
         topic_delete_condition: bool = all(

taskiq_aio_kafka/broker.py

@@ -0,0 +1,76 @@
+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
+
+
+class KafkaProducerParameters(BaseModel):
+    """Parameters to kafka producer."""
+
+    client_id: Optional[str] = None
+    metadata_max_age_ms: int = 300000
+    request_timeout_ms: int = 40000
+    api_version: str = "auto"
+    acks: Optional[Union[str, int]] = 1
+    key_serializer: Optional[Callable[..., bytes]] = None
+    value_serializer: Optional[Callable[..., bytes]] = None
+    compression_type: Optional[str] = None
+    max_batch_size: int = 16384
+    partitioner: Callable[..., Any] = DefaultPartitioner()
+    max_request_size: int = 1048576
+    linger_ms: int = 0
+    send_backoff_ms: int = 100
+    retry_backoff_ms: int = 100
+    security_protocol: str = "PLAINTEXT"
+    ssl_context: Optional[Any] = None
+    connections_max_idle_ms: int = 540000
+    enable_idempotence: bool = False
+    transactional_id: Optional[Any] = None
+    transaction_timeout_ms: int = 60000
+    sasl_mechanism: str = "PLAIN"
+    sasl_plain_password: Optional[str] = None
+    sasl_plain_username: Optional[str] = None
+    sasl_kerberos_service_name: Any = "kafka"
+    sasl_kerberos_domain_name: Any = None
+    sasl_oauth_token_provider: Any = None
+
+
+class KafkaConsumerParameters(BaseModel):
+    """Parameters to kafka consumer."""
+
+    client_id: str = "aiokafka-" + __version__  # noqa: WPS336
+    group_id: Optional[str] = None
+    key_deserializer: Optional[Callable[..., Any]] = None
+    value_deserializer: Optional[Callable[..., Any]] = None
+    fetch_max_wait_ms: int = 500
+    fetch_max_bytes: int = 52428800
+    fetch_min_bytes: int = 1
+    max_partition_fetch_bytes: int = 1 * 1024 * 1024  # noqa: WPS345
+    request_timeout_ms: int = 40 * 1000  # noqa: WPS432
+    retry_backoff_ms: int = 100
+    auto_offset_reset: str = "latest"
+    enable_auto_commit: bool = True
+    auto_commit_interval_ms: int = 5000
+    check_crcs: bool = True
+    metadata_max_age_ms: int = 5 * 60 * 1000
+    partition_assignment_strategy: Any = (RoundRobinPartitionAssignor,)
+    max_poll_interval_ms: int = 300000
+    rebalance_timeout_ms: Optional[int] = None
+    session_timeout_ms: int = 10000
+    heartbeat_interval_ms: int = 3000
+    consumer_timeout_ms: int = 200
+    max_poll_records: Optional[int] = None
+    ssl_context: Optional[Any] = None
+    security_protocol: str = "PLAINTEXT"
+    api_version: str = "auto"
+    exclude_internal_topics: bool = True
+    connections_max_idle_ms: int = 540000
+    isolation_level: str = "read_uncommitted"
+    sasl_mechanism: str = "PLAIN"
+    sasl_plain_password: Optional[str] = None
+    sasl_plain_username: Optional[str] = None
+    sasl_kerberos_service_name: Optional[str] = "kafka"
+    sasl_kerberos_domain_name: Optional[str] = None
+    sasl_oauth_token_provider: Optional[str] = None

taskiq_aio_kafka/models.py

@@ -115,6 +115,7 @@ async def broker(
     kafka_url: str,
     test_kafka_producer: AIOKafkaProducer,
     test_kafka_consumer: AIOKafkaConsumer,
+    base_topic: NewTopic,
 ) -> AsyncGenerator[AioKafkaBroker, None]:
     """Yield new broker instance.
 
@@ -125,6 +126,7 @@ async def broker(
     :param kafka_url: url to kafka.
     :param test_kafka_producer: custom AIOKafkaProducer.
     :param test_kafka_consumer: custom AIOKafkaConsumer.
+    :param base_topic: base topic.
 
     :yields: broker.
     """
@@ -133,6 +135,7 @@ async def broker(
         aiokafka_producer=test_kafka_producer,
         aiokafka_consumer=test_kafka_consumer,
         delete_topic_on_shutdown=True,
+        kafka_topic=base_topic,
     )
     broker.is_worker_process = True
 

tests/conftest.py

@@ -80,7 +80,6 @@ async def test_startup(
 async def test_listen(
     broker: AioKafkaBroker,
     test_kafka_producer: AIOKafkaProducer,
-    base_topic_name: str,
 ) -> None:
     """Test that message are read correctly.
 
@@ -89,8 +88,8 @@ async def test_listen(
 
     :param broker: current broker.
     :param test_kafka_producer: AIOKafkaProducer.
-    :param base_topic_name: topic name.
     """
+    await test_kafka_producer.start()
     task_id: str = uuid4().hex
     task_name: str = uuid4().hex
     message: bytes = pickle.dumps(uuid4().hex)
@@ -104,7 +103,7 @@ async def test_listen(
     )
 
     await test_kafka_producer.send(
-        topic=base_topic_name,
+        topic=broker._kafka_topic.name,
         value=message_to_send.message,
     )