Merge branch 'main' into develop

Author: chandr-andr

README.md

@@ -13,6 +13,24 @@ async def test() -> None:
     print("The best task ever!")
 ```
 
+## Non-obvious things
+
+You can configure kafka producer and consumer with special methods `configure_producer` and `configure_consumer`.  
+Example:
+```python
+from taskiq_aio_kafka import AioKafkaBroker
+
+broker = AioKafkaBroker(bootstrap_servers="localhost")
+
+# configure producer, you can set any parameter from
+# base AIOKafkaProducer, except `loop` and `bootstrap_servers`
+broker.configure_producer(request_timeout_ms=100000)
+
+# configure consumer, you can set any parameter from
+# base AIOKafkaConsumer, except `loop` and `bootstrap_servers`
+broker.configure_consumer(group_id="the best group ever.")
+```
+
 ## Configuration
 
 AioKafkaBroker parameters:

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,21 +87,12 @@ def __init__(  # noqa: WPS211
             replication_factor=1,
         )
 
-        self._aiokafka_producer: AIOKafkaProducer = (
-            aiokafka_producer
-            or AIOKafkaProducer(
-                bootstrap_servers=self._bootstrap_servers,
-                loop=self._loop,
-            )
+        self._aiokafka_producer_params: KafkaProducerParameters = (
+            KafkaProducerParameters()
         )
 
-        self._aiokafka_consumer: AIOKafkaConsumer = (
-            aiokafka_consumer
-            or AIOKafkaConsumer(
-                self._kafka_topic.name,
-                bootstrap_servers=self._bootstrap_servers,
-                loop=self._loop,
-            )
+        self._aiokafka_consumer_params: KafkaConsumerParameters = (
+            KafkaConsumerParameters()
         )
 
         self._kafka_admin_client: KafkaAdminClient = (
@@ -118,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.
 
@@ -127,18 +143,29 @@ 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 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_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(
@@ -183,12 +210,11 @@ async def kick(self, message: BrokerMessage) -> None:
         if not self._is_producer_started:
             raise ValueError("Please run startup before kicking.")
 
-        kafka_message: bytes = pickle.dumps(message)
         topic_name: str = self._kafka_topic.name
 
-        await self._aiokafka_producer.send(
+        await self._aiokafka_producer.send(  # type: ignore
             topic=topic_name,
-            value=kafka_message,
+            value=message.message,
         )
 
     async def listen(
@@ -205,5 +231,5 @@ async def listen(
         if not self._is_consumer_started:
             raise ValueError("Please run startup before listening.")
 
-        async for raw_kafka_message in self._aiokafka_consumer:
+        async for raw_kafka_message in self._aiokafka_consumer:  # type: ignore
             yield raw_kafka_message.value

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

@@ -1,5 +1,6 @@
 import os
 from typing import AsyncGenerator
+from uuid import uuid4
 
 import pytest
 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
@@ -35,7 +36,7 @@ def base_topic_name() -> str:
 
     :returns: topic name.
     """
-    return "taskiq_topic"
+    return uuid4().hex
 
 
 @pytest.fixture()
@@ -115,6 +116,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,13 +127,15 @@ 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.
     """
     broker = AioKafkaBroker(
         bootstrap_servers=kafka_url,
         aiokafka_producer=test_kafka_producer,
         aiokafka_consumer=test_kafka_consumer,
+        kafka_topic=base_topic,
         delete_topic_on_shutdown=True,
     )
     broker.is_worker_process = True

tests/conftest.py

@@ -48,12 +48,7 @@ async def test_kick_success(broker: AioKafkaBroker) -> None:
         get_first_task(broker),
         timeout=1,
     )
-    assert pickle.dumps(message_to_send) == received_message_bytes
-
-    received_message: BrokerMessage = pickle.loads(
-        received_message_bytes,
-    )
-    assert message_to_send == received_message
+    assert received_message_bytes == message_to_send.message
 
 
 @pytest.mark.anyio
@@ -78,14 +73,13 @@ async def test_startup(
         str
     ] = broker_without_arguments._kafka_admin_client.list_topics()  # noqa: WPS437
 
-    assert base_topic_name in all_kafka_topics
+    assert broker_without_arguments._kafka_topic.name in all_kafka_topics
 
 
 @pytest.mark.anyio
 async def test_listen(
     broker: AioKafkaBroker,
     test_kafka_producer: AIOKafkaProducer,
-    base_topic_name: str,
 ) -> None:
     """Test that message are read correctly.
 
@@ -94,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)
@@ -109,22 +103,13 @@ async def test_listen(
     )
 
     await test_kafka_producer.send(
-        topic=base_topic_name,
-        value=pickle.dumps(message_to_send),
+        topic=broker._kafka_topic.name,
+        value=message_to_send.message,
     )
 
     received_message_bytes: bytes = await asyncio.wait_for(
         get_first_task(broker),
         timeout=1,
     )
 
-    assert pickle.dumps(message_to_send) == received_message_bytes
-
-    received_message: BrokerMessage = pickle.loads(
-        received_message_bytes,
-    )
-
-    assert received_message.message == message
-    assert received_message.labels == labels
-    assert received_message.task_id == task_id
-    assert received_message.task_name == task_name
+    assert received_message_bytes == message_to_send.message