zrb 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

zrb/builtin/generator/fastapp/template/src/kebab-zrb-app-name/src/core/messagebus/kafka/publisher.py

@@ -1,27 +1,27 @@
+import asyncio
+import logging
 from typing import Any, Optional
+
+from aiokafka import AIOKafkaProducer
+from aiokafka.producer.producer import DefaultPartitioner, _missing
+from core.messagebus.kafka.admin import KafkaAdmin, must_get_kafka_admin
 from core.messagebus.messagebus import (
-    Publisher, MessageSerializer, must_get_message_serializer
-)
-from core.messagebus.kafka.admin import (
-    KafkaAdmin, must_get_kafka_admin
+    MessageSerializer,
+    Publisher,
+    must_get_message_serializer,
 )
-from aiokafka import AIOKafkaProducer
-from aiokafka.producer.producer import _missing, DefaultPartitioner
 from pydantic import BaseModel
 
-import logging
-import asyncio
-
 
 class KafkaPublisher(Publisher):
     def __init__(
         self,
         logger: logging.Logger,
-        bootstrap_servers: str = 'localhost',
+        bootstrap_servers: str = "localhost",
         client_id: Optional[Any] = None,
         metadata_max_age_ms=300000,
         request_timeout_ms=40000,
-        api_version='auto',
+        api_version="auto",
         acks=_missing,
         key_serializer=None,
         value_serializer=None,
@@ -41,14 +41,14 @@ class KafkaPublisher(Publisher):
         sasl_mechanism="PLAIN",
         sasl_plain_password=None,
         sasl_plain_username=None,
-        sasl_kerberos_service_name='kafka',
+        sasl_kerberos_service_name="kafka",
         sasl_kerberos_domain_name=None,
         sasl_oauth_token_provider=None,
         serializer: Optional[MessageSerializer] = None,
         kafka_admin: Optional[KafkaAdmin] = None,
         retry: int = 3,
-        retry_interval: int = 3,
-        identifier='kafka-publisher'
+        retry_interval: int = 10,
+        identifier="kafka-publisher",
     ):
         self.logger = logger
         self.serializer = must_get_message_serializer(serializer)
@@ -100,33 +100,30 @@ class KafkaPublisher(Publisher):
         await self.kafka_admin.create_events([event_name])
         topic_name = self.kafka_admin.get_topic_name(event_name)
         if isinstance(message, BaseModel):
-            message = message.dict()
+            message = message.model_dump()
         for attempt in range(self.retry):
             try:
                 await self._connect()
                 encoded_value = self.serializer.encode(event_name, message)
                 self.logger.info(
-                    f'🐼 [{self.identifier}] Publish to "{topic_name}": ' +
-                    f'{message}'
-                )
-                return await self.producer.send_and_wait(
-                    topic_name, encoded_value
+                    f'🐼 [{self.identifier}] Publish to "{topic_name}": ' + f"{message}"
                 )
+                return await self.producer.send_and_wait(topic_name, encoded_value)
             except (asyncio.CancelledError, GeneratorExit, Exception) as e:
                 self.logger.error(
-                    f'🐼 [{self.identifier}] Failed to publish message: {e}'
+                    f"🐼 [{self.identifier}] Failed to publish message: {e}"
                 )
                 await self._disconnect()
                 await asyncio.sleep(self.retry_interval)
                 continue
         self.logger.error(
-            f'🐼 [{self.identifier}] Failed to publish message after ' +
-            f'{self.retry} attempts'
+            f"🐼 [{self.identifier}] Failed to publish message after "
+            + f"{self.retry} attempts"
         )
-        raise RuntimeError('Failed to publish message after retrying')
+        raise RuntimeError("Failed to publish message after retrying")
 
     async def _connect(self):
-        self.logger.info(f'🐼 [{self.identifier}] Create kafka producer')
+        self.logger.info(f"🐼 [{self.identifier}] Create kafka producer")
         self.producer = AIOKafkaProducer(
             bootstrap_servers=self.bootstrap_servers,
             client_id=self.client_id,
@@ -156,14 +153,14 @@ class KafkaPublisher(Publisher):
             sasl_kerberos_domain_name=self.sasl_kerberos_domain_name,
             sasl_oauth_token_provider=self.sasl_oauth_token_provider,
         )
-        self.logger.info(f'🐼 [{self.identifier}] Start kafka producer')
+        self.logger.info(f"🐼 [{self.identifier}] Start kafka producer")
         await self.producer.start()
-        self.logger.info(f'🐼 [{self.identifier}] Kafka producer started')
+        self.logger.info(f"🐼 [{self.identifier}] Kafka producer started")
         return self
 
     async def _disconnect(self):
-        self.logger.info(f'🐼 [{self.identifier}] Stop kafka producer')
+        self.logger.info(f"🐼 [{self.identifier}] Stop kafka producer")
         if self.producer is not None:
             await self.producer.stop()
-        self.logger.info(f'🐼 [{self.identifier}] Kafka producer stopped')
+        self.logger.info(f"🐼 [{self.identifier}] Kafka producer stopped")
         self.producer = None

zrb/builtin/generator/fastapp/template/src/kebab-zrb-app-name/src/core/messagebus/messagebus.py

@@ -1,6 +1,7 @@
-from typing import Any, Callable, List, Mapping, Optional
 from abc import ABC, abstractmethod
-from core.serializer.serializer import Serializer, JsonSerializer
+from typing import Any, Callable, List, Mapping, Optional
+
+from core.serializer.serializer import JsonSerializer, Serializer
 
 TEventHandler = Callable[[Any], Any]
 
@@ -23,9 +24,7 @@ class Publisher(ABC):
 
 class Consumer(ABC):
     @abstractmethod
-    def register(
-        self, event_name: str
-    ) -> Callable[[TEventHandler], Any]:
+    def register(self, event_name: str) -> Callable[[TEventHandler], Any]:
         pass
 
     @abstractmethod
@@ -37,11 +36,8 @@ class Consumer(ABC):
         pass
 
 
-class MessageSerializer():
-    def __init__(
-        self,
-        serializers: Optional[Mapping[str, Serializer]] = None
-    ):
+class MessageSerializer:
+    def __init__(self, serializers: Optional[Mapping[str, Serializer]] = None):
         serializers = serializers if serializers is not None else {}
         self.serializers: Mapping[str, Serializer] = serializers
         self.default_serializer = JsonSerializer()
@@ -55,13 +51,11 @@ class MessageSerializer():
         return serializer.decode(encoded_message)
 
     def _get_serializer(self, event_name: str) -> Serializer:
-        return self.serializers.get(
-            event_name, self.default_serializer
-        )
+        return self.serializers.get(event_name, self.default_serializer)
 
 
 def must_get_message_serializer(
-    serializer: Optional[MessageSerializer] = None
+    serializer: Optional[MessageSerializer] = None,
 ) -> MessageSerializer:
     if serializer is None:
         return MessageSerializer()

zrb/builtin/generator/fastapp/template/src/kebab-zrb-app-name/src/core/messagebus/mock.py

@@ -1,16 +1,19 @@
-from typing import Any, Callable, List, Mapping
-from core.messagebus.messagebus import (
-    Admin, Publisher, Consumer, MessageSerializer, TEventHandler
-)
 import asyncio
 import inspect
 import logging
+from typing import Any, Callable, List, Mapping
+
+from core.messagebus.messagebus import (
+    Admin,
+    Consumer,
+    MessageSerializer,
+    Publisher,
+    TEventHandler,
+)
 
 
 class MockConsumer(Consumer):
-    def __init__(
-        self, logger: logging.Logger, serializer: MessageSerializer
-    ):
+    def __init__(self, logger: logging.Logger, serializer: MessageSerializer):
         self.logger = logger
         self.serializer = serializer
         self._handlers: Mapping[str, TEventHandler] = {}
@@ -20,6 +23,7 @@ class MockConsumer(Consumer):
             self.logger.warning(f'🪵 Register handler for "{event_name}"')
             self._handlers[event_name] = handler
             return handler
+
         return wrapper
 
     async def handle(self, event_name: str, encoded_value: Any):
@@ -42,7 +46,7 @@ class MockPublisher(Publisher):
         self,
         logger: logging.Logger,
         consumer: MockConsumer,
-        serializer: MessageSerializer
+        serializer: MessageSerializer,
     ):
         self.logger = logger
         self.consumer = consumer

zrb/builtin/generator/fastapp/template/src/kebab-zrb-app-name/src/core/messagebus/rabbitmq/admin.py

@@ -1,15 +1,16 @@
-from typing import Any, Mapping, List, Optional
-from core.messagebus.messagebus import Admin
+import asyncio
 import logging
+from typing import Any, List, Mapping, Optional
+
 import aiormq
-import asyncio
+from core.messagebus.messagebus import Admin
 
 
-class RMQEventConfig():
+class RMQEventConfig:
     def __init__(
         self,
         queue_name: str,
-        exchange_name: str = '',
+        exchange_name: str = "",
     ):
         self.queue_name = queue_name
         self.exchange_name = exchange_name
@@ -20,7 +21,7 @@ class RMQAdmin(Admin):
         self,
         logger: logging.Logger,
         configs: Mapping[str, RMQEventConfig],
-        connection_string: str
+        connection_string: str,
     ):
         self.logger = logger
         self.connection_string = connection_string
@@ -30,7 +31,8 @@ class RMQAdmin(Admin):
     async def create_events(self, event_names: List[str]):
         # Only handle non-existing events
         event_names = [
-            event_name for event_name in event_names
+            event_name
+            for event_name in event_names
             if event_name not in self._existing_events
         ]
         if len(event_names) == 0:
@@ -40,7 +42,7 @@ class RMQAdmin(Admin):
             channel = await connection.channel()
             for event_name in event_names:
                 config = self.get_config(event_name)
-                if config.exchange_name != '':
+                if config.exchange_name != "":
                     await self._declare_fanned_out_exchange(channel, config)
                     self._existing_events[event_name] = True
                     continue
@@ -49,17 +51,20 @@ class RMQAdmin(Admin):
             await self._clean_up(connection, channel)
         except (asyncio.CancelledError, GeneratorExit, Exception):
             self.logger.error(
-                ' '.join([
-                    '🐰 [rabbitmq-admin] Something wrong when ',
-                    f'creating events: {event_names}'
-                ]),
-                exc_info=True
+                " ".join(
+                    [
+                        "🐰 [rabbitmq-admin] Something wrong when ",
+                        f"creating events: {event_names}",
+                    ]
+                ),
+                exc_info=True,
             )
 
     async def delete_events(self, event_names: List[str]):
         # Only handle existing events
         event_names = [
-            event_name for event_name in event_names
+            event_name
+            for event_name in event_names
             if event_name in self._existing_events
         ]
         if len(event_names) == 0:
@@ -71,57 +76,51 @@ class RMQAdmin(Admin):
                 config = self.get_config(event_name)
                 # delete the queue
                 await channel.queue_delete(queue=config.queue_name)
-                if config.exchange_name != '':
+                if config.exchange_name != "":
                     # delete the exchange
-                    await channel.exchange_delete(
-                        exchange_name=config.exchange_name
-                    )
+                    await channel.exchange_delete(exchange_name=config.exchange_name)
                 del self._existing_events[event_name]
             await self._clean_up(connection, channel)
         except (asyncio.CancelledError, GeneratorExit, Exception):
             self.logger.error(
-                ' '.join([
-                    '🐰 [rabbitmq-admin] Something wrong when ',
-                    f'deleting events: {event_names}'
-                ]),
-                exc_info=True
+                " ".join(
+                    [
+                        "🐰 [rabbitmq-admin] Something wrong when ",
+                        f"deleting events: {event_names}",
+                    ]
+                ),
+                exc_info=True,
             )
 
-    async def _clean_up(
-        self, connection: aiormq.Connection, channel: aiormq.Channel
-    ):
+    async def _clean_up(self, connection: aiormq.Connection, channel: aiormq.Channel):
         try:
             await channel.close()
             await connection.close()
         except (asyncio.CancelledError, GeneratorExit, Exception):
-            self.logger.error('🐰 [rabbitmq-admin]', exc_info=True)
+            self.logger.error("🐰 [rabbitmq-admin]", exc_info=True)
 
     async def _declare_fanned_out_exchange(
         self, channel: aiormq.Channel, config: RMQEventConfig
     ):
         await channel.exchange_declare(
             exchange=config.exchange_name,
-            exchange_type='fanout',
+            exchange_type="fanout",
             durable=False,
-            auto_delete=False
+            auto_delete=False,
         )
         # declare a queue
         result = await self._declare_queue(channel, config)
-        queue_name = result['queue']
+        queue_name = result["queue"]
         # bind the queue to the exchange
         await channel.queue_bind(
-            queue=queue_name,
-            exchange=config.exchange_name,
-            routing_key=''
+            queue=queue_name, exchange=config.exchange_name, routing_key=""
         )
 
     async def _declare_queue(
         self, channel: aiormq.Channel, config: RMQEventConfig
     ) -> Any:
         return await channel.queue_declare(
-            queue=config.queue_name,
-            durable=True,
-            auto_delete=False
+            queue=config.queue_name, durable=True, auto_delete=False
         )
 
     def get_config(self, event_name: str) -> RMQEventConfig:
@@ -139,14 +138,8 @@ class RMQAdmin(Admin):
 
 
 def must_get_rmq_admin(
-    logger: logging.Logger,
-    rmq_admin: Optional[RMQAdmin],
-    connection_string: str
+    logger: logging.Logger, rmq_admin: Optional[RMQAdmin], connection_string: str
 ) -> RMQAdmin:
     if rmq_admin is None:
-        return RMQAdmin(
-            logger=logger,
-            configs={},
-            connection_string=connection_string
-        )
+        return RMQAdmin(logger=logger, configs={}, connection_string=connection_string)
     return rmq_admin

zrb/builtin/generator/fastapp/template/src/kebab-zrb-app-name/src/core/messagebus/rabbitmq/consumer.py

@@ -1,14 +1,16 @@
-from typing import Any, Callable, Mapping, Optional
-from core.messagebus.messagebus import (
-    Consumer, TEventHandler, MessageSerializer, must_get_message_serializer
-)
-from core.messagebus.rabbitmq.admin import (
-    RMQAdmin, must_get_rmq_admin
-)
 import asyncio
-import aiormq
 import inspect
 import logging
+from typing import Any, Callable, Mapping, Optional
+
+import aiormq
+from core.messagebus.messagebus import (
+    Consumer,
+    MessageSerializer,
+    TEventHandler,
+    must_get_message_serializer,
+)
+from core.messagebus.rabbitmq.admin import RMQAdmin, must_get_rmq_admin
 
 
 class RMQConsumer(Consumer):
@@ -19,15 +21,13 @@ class RMQConsumer(Consumer):
         serializer: Optional[MessageSerializer] = None,
         rmq_admin: Optional[RMQAdmin] = None,
         retry: int = 5,
-        retry_interval: int = 5,
+        retry_interval: int = 10,
         prefetch_count: int = 20,
-        identifier='rmq-consumer'
+        identifier="rmq-consumer",
     ):
         self.logger = logger
         self.rmq_admin = must_get_rmq_admin(
-            logger=logger,
-            rmq_admin=rmq_admin,
-            connection_string=connection_string
+            logger=logger, rmq_admin=rmq_admin, connection_string=connection_string
         )
         self.connection_string = connection_string
         self.connection: Optional[aiormq.Connection] = None
@@ -48,6 +48,7 @@ class RMQConsumer(Consumer):
             )
             self._handlers[event_name] = handler
             return handler
+
         return wrapper
 
     async def start(self):
@@ -70,32 +71,29 @@ class RMQConsumer(Consumer):
             f'🐰 [{self.identifier}] Listening from "{event_names}"'
             for event_name in event_names:
                 queue_name = self.rmq_admin.get_queue_name(event_name)
-                on_message = self._create_consumer_callback(
-                    self.channel, event_name
-                )
+                on_message = self._create_consumer_callback(self.channel, event_name)
                 await self.channel.basic_consume(
                     queue=queue_name, consumer_callback=on_message
                 )
             retry = self.retry
             while not self._is_stop_triggered:
                 await asyncio.sleep(0.01)
-                if (
-                    not self._is_stop_triggered and
-                    (self.connection is None or self.connection.is_closed)
+                if not self._is_stop_triggered and (
+                    self.connection is None or self.connection.is_closed
                 ):
-                    raise Exception('Rabbitmq connection is closed')
+                    raise Exception("Rabbitmq connection is closed")
         except (asyncio.CancelledError, GeneratorExit, Exception) as e:
             if retry > 0:
-                self.logger.error(f'🐰 [{self.identifier}]', exc_info=True)
+                self.logger.error(f"🐰 [{self.identifier}]", exc_info=True)
             if retry == 0:
                 self.logger.error(
-                    f'🐰 [{self.identifier}] Failed to consume message after ' +
-                    f'{self.retry} attempts'
+                    f"🐰 [{self.identifier}] Failed to consume message after "
+                    + f"{self.retry} attempts"
                 )
                 raise e
             await self._disconnect()
             await asyncio.sleep(self.retry_interval)
-            await self._start(retry-1)
+            await self._start(retry - 1)
         finally:
             await self._disconnect()
 
@@ -103,54 +101,34 @@ class RMQConsumer(Consumer):
         try:
             connection_created = False
             if self.connection is None or self.connection.is_closed:
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Create consumer connection'
-                )
+                self.logger.info(f"🐰 [{self.identifier}] Create consumer connection")
                 self.connection = await aiormq.connect(self.connection_string)
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Consumer connection created'
-                )
+                self.logger.info(f"🐰 [{self.identifier}] Consumer connection created")
                 connection_created = True
-            if (
-                connection_created or
-                self.channel is None or
-                self.channel.is_closed
-            ):
-                self.logger.info(f'🐰 [{self.identifier}] Get consumer channel')
+            if connection_created or self.channel is None or self.channel.is_closed:
+                self.logger.info(f"🐰 [{self.identifier}] Get consumer channel")
                 self.channel = await self.connection.channel()
-                await self.channel.basic_qos(
-                    prefetch_count=self.prefetch_count
-                )
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Consumer channel created'
-                )
+                await self.channel.basic_qos(prefetch_count=self.prefetch_count)
+                self.logger.info(f"🐰 [{self.identifier}] Consumer channel created")
         except (asyncio.CancelledError, GeneratorExit, Exception):
-            self.logger.error(f'🐰 [{self.identifier}]', exc_info=True)
-            raise Exception('Cannot connect')
+            self.logger.error(f"🐰 [{self.identifier}]", exc_info=True)
+            raise Exception("Cannot connect")
 
     async def _disconnect(self):
         try:
             if self.channel is not None and not self.channel.is_closed:
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Close consumer channel'
-                )
+                self.logger.info(f"🐰 [{self.identifier}] Close consumer channel")
                 await self.channel.close()
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Consumer channel closed'
-                )
+                self.logger.info(f"🐰 [{self.identifier}] Consumer channel closed")
         except (asyncio.CancelledError, GeneratorExit, Exception):
-            self.logger.error(f'🐰 [{self.identifier}]', exc_info=True)
+            self.logger.error(f"🐰 [{self.identifier}]", exc_info=True)
         try:
             if self.connection is not None and not self.connection.is_closed:
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Close consumer connection'
-                )
+                self.logger.info(f"🐰 [{self.identifier}] Close consumer connection")
                 await self.connection.close()
-                self.logger.info(
-                    f'🐰 [{self.identifier}] Consumer connection closed'
-                )
+                self.logger.info(f"🐰 [{self.identifier}] Consumer connection closed")
         except (asyncio.CancelledError, GeneratorExit, Exception):
-            self.logger.error(f'🐰 [{self.identifier}]', exc_info=True)
+            self.logger.error(f"🐰 [{self.identifier}]", exc_info=True)
         self.connection = None
         self.channel = None
 
@@ -161,24 +139,21 @@ class RMQConsumer(Consumer):
     ) -> Callable[[Any], Any]:
         async def on_message(message):
             try:
-                decoded_value = self.serializer.decode(
-                    event_name, message.body
-                )
+                decoded_value = self.serializer.decode(event_name, message.body)
                 handler = self._handlers.get(event_name)
                 queue_name = self.rmq_admin.get_queue_name(event_name)
                 self.logger.info(
-                    f'🐰 [{self.identifier}] Consume from "{queue_name}": ' +
-                    f'{decoded_value}'
+                    f'🐰 [{self.identifier}] Consume from "{queue_name}": '
+                    + f"{decoded_value}"
                 )
                 await self._run_handler(handler, decoded_value)
                 await channel.basic_ack(message.delivery_tag)
             except (asyncio.CancelledError, GeneratorExit, Exception):
-                self.logger.error(f'🐰 [{self.identifier}]', exc_info=True)
+                self.logger.error(f"🐰 [{self.identifier}]", exc_info=True)
+
         return on_message
 
-    async def _run_handler(
-        self, message_handler: TEventHandler, decoded_value: Any
-    ):
+    async def _run_handler(self, message_handler: TEventHandler, decoded_value: Any):
         if inspect.iscoroutinefunction(message_handler):
             return asyncio.create_task(message_handler(decoded_value))
         return message_handler(decoded_value)