Rewrote the azure topic for better control

- Rewrote the azure topic for better control and ease of operation

Author: susrisha

src/example.py

@@ -13,16 +13,28 @@
 
 core = Core()
 print(f'Core version: {Core.__version__}')
-topic = 'temp-request'
+topic = 'temp-nar'
 subscription = 'temp'
 some_other_sub = 'temp'
+respond_topic = 'temp-response'
+# respond_topic_object = core.get_topic(topic_name=respond_topic)
+
+
+class MessageSender:
+    def __init__(self, topic_name):
+        self.topic = core.get_topic(topic_name=topic_name)
+
+    def send_message(self, message):
+        self.topic.publish(data=QueueMessage.data_from(message))
+
+
 
 
 def publish_messages(topic_name):
     topic_object = core.get_topic(topic_name=topic_name)
     queue_message = QueueMessage.data_from({
         'message': str(uuid.uuid4().hex),
-        'data': {'a': random.randint(60, 120)}
+        'data': {'a': random.randint(10, 25)}
     })
     topic_object.publish(data=queue_message)
     print('Message Published')
@@ -31,24 +43,46 @@ def publish_messages(topic_name):
 def long_running_task(sleep_time):
     # Simulate a long-running task
     time.sleep(sleep_time)
+    sender_obj = MessageSender(topic_name=respond_topic)
+    sender_obj.send_message({'message': 'Task Completed'})
+    print(f' > Task Completed: {sleep_time}')
 
 
 def subscribe(topic_name, subscription_name):
     def process(message):
         print(f'Message Received: {message.data}')
-        long_running_thread = threading.Thread(target=long_running_task, args=(message.data['a'],))
-        long_running_thread.start()
-        long_running_thread.join()
+        # long_running_thread = threading.Thread(target=long_running_task, args=(message.data['a'],))
+        # long_running_thread.start()
+        # long_running_thread.join()
+        long_running_task(message.data['a'])
         print(f' > Message Completed: {message.data}')
 
-    topic_object = core.get_topic(topic_name=topic_name)
+    topic_object = core.get_topic(topic_name=topic_name,max_concurrent_messages=2)
     try:
         topic_object.subscribe(subscription=subscription_name, callback=process)
     except Exception as e:
         print(e)
 
 
-subscribe(topic, subscription)
+
+import argparse
+
+parser = argparse.ArgumentParser(description='Process the core topic.')
+parser.add_argument('-m','--mode', type=str, help='Mode of operation (publish/subscribe)', required=True)
+parser.add_argument('-s','--size', type=int, help='Number of the messages to send', required=False)
+
+mode = parser.parse_args().mode
+size = parser.parse_args().size
+if mode == 'publish':
+    for x in range(size):
+        print(f'Publishing message {x+1}')
+        publish_messages(topic_name=topic)
+if mode == 'subscribe':
+    subscribe(topic, subscription)
+# subscribe(topic, subscription)
+# for x in range(10):
+#     publish_messages(topic_name=topic)
+    # subscribe(topic, subscription)
 # for x in range(10):
 #     publish_messages(topic_name=topic)
 

src/python_ms_core/__init__.py

@@ -3,6 +3,7 @@
 from .core.logger.logger import Logger
 from .core.logger.local_logger import LocalLogger
 from .core.topic.topic import Topic
+from .core.topic.azure_topic import AzureTopic
 from .core.topic.local_topic import LocalTopic
 from .core.storage.providers.azure.azure_storage_client import AzureStorageClient
 from .core.storage.providers.local.local_storage_client import LocalStorageClient
@@ -49,7 +50,7 @@ def get_topic(self, topic_name: str, max_concurrent_messages=os.cpu_count()):
         if topic_config.provider.upper() == LOCAL_ENV:
             return LocalTopic(config=topic_config, topic_name=topic_name)
         elif topic_config.provider.upper() == AZURE_ENV:
-            return Topic(config=topic_config, topic_name=topic_name, max_concurrent_messages=max_concurrent_messages)
+            return AzureTopic(config=topic_config, topic_name=topic_name, max_concurrent_messages=max_concurrent_messages)
         else:
             logging.error(f'Failed to initialize core.get_topic for provider: {topic_config.provider}')
 

src/python_ms_core/core/topic/azure_topic.py

@@ -0,0 +1,122 @@
+import json
+
+import logging
+import time
+from ..config.config import TopicConfig
+from ..resource_errors import ExceptionHandler
+from concurrent.futures import ThreadPoolExecutor
+from .abstract.topic_abstract import TopicAbstract
+from ..queue.models.queue_message import QueueMessage
+from azure.servicebus import ServiceBusClient, ServiceBusMessage
+from azure.servicebus import AutoLockRenewer
+import concurrent.futures as cf
+
+logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
+logger = logging.getLogger('AzureTopic')
+logger.setLevel(logging.INFO)
+
+
+"""
+AzureTopic class represents a topic in Azure Service Bus.
+Attributes:
+    topic (str): The name of the topic.
+    client (ServiceBusClient): The ServiceBusClient object used to interact with the Service Bus.
+    max_concurrent_messages (int): The maximum number of concurrent messages to process.
+    topic_name (str): The name of the topic.
+    publisher (TopicSender): The TopicSender object used to send messages to the topic.
+    executor (ThreadPoolExecutor): The ThreadPoolExecutor object used to execute callback functions.
+    internal_count (int): The internal count of concurrent messages being processed.
+    lock_renewal (AutoLockRenewer): The AutoLockRenewer object used to renew message locks.
+    max_renewal_duration (int): The maximum duration in seconds to renew a message lock.
+    wait_time_for_message (int): The maximum wait time in seconds to receive messages.
+Methods:
+    publish(data: QueueMessage) -> None:
+        Publishes a message to the topic.
+    subscribe(subscription: str, callback) -> None:
+        Subscribes to a subscription of the topic and processes incoming messages.
+    internal_callback(message, callbackfn) -> ServiceBusMessage:
+        Internal callback function that processes a message and invokes the callback function.
+    settle_message(x: cf.Future) -> None:
+        Sets the message as completed and updates the internal count.
+"""
+class AzureTopic(TopicAbstract):
+    def __init__(self, config: TopicConfig=None, topic_name=None, max_concurrent_messages:int=1):
+        self.topic = topic_name
+        self.client = ServiceBusClient.from_connection_string(conn_str=config.connection_string, retry_total=10, retry_backoff_factor=1, retry_backoff_max=30)
+        self.max_concurrent_messages = max_concurrent_messages
+        self.topic_name = topic_name
+        self.publisher = self.client.get_topic_sender(topic_name=topic_name)
+        self.executor = ThreadPoolExecutor(max_workers=max_concurrent_messages)
+        self.internal_count = 0
+        self.lock_renewal = AutoLockRenewer(max_workers=4)
+        self.max_renewal_duration = 86400 # Renew the message upto 1 day
+        self.wait_time_for_message = 5
+    
+    
+    def publish(self, data: QueueMessage):
+        """
+        Publishes a message to the topic.
+        Args:
+            data (QueueMessage): The message to publish.
+        
+        """
+        message = QueueMessage.to_dict(data)
+        self.publisher.send_messages(ServiceBusMessage(json.dumps(message)))
+
+    def subscribe(self, subscription: str, callback):
+
+        """
+        Subscribes to a subscription of the topic and processes incoming messages.
+        Args:
+            subscription (str): The name of the subscription to subscribe to.
+            callback (function): The callback function to invoke for each message.
+        """
+        self.receiver = self.client.get_subscription_receiver(topic_name=self.topic_name, subscription_name=subscription)
+        while True:
+                try:
+                    to_receive = (self.max_concurrent_messages - self.internal_count)
+                    if to_receive > 0:
+                        messages = self.receiver.receive_messages(max_message_count=to_receive, max_wait_time=self.wait_time_for_message)
+                        if not messages or len(messages) == 0:
+                            continue
+                        for message in messages: 
+                            self.lock_renewal.register(self.receiver, message, max_lock_renewal_duration=self.max_renewal_duration)
+                            execution_task = self.executor.submit(self.internal_callback, message, callback)
+                            execution_task.add_done_callback(lambda x: self.settle_message(x))
+                    else:
+                        time.sleep(self.wait_time_for_message)
+                except Exception as e:
+                    logger.error(f'Error in receiving messages: {e}')
+
+    
+    def internal_callback(self, message, callbackfn):
+        """
+        Internal callback function that processes a message and invokes the callback function.
+        Args:
+            message (ServiceBusMessage): The message to process.
+            callbackfn (function): The callback function to invoke.
+        Returns:
+            ServiceBusMessage: The processed message.
+        """
+        try:
+            self.internal_count += 1
+            queue_message = QueueMessage.data_from(str(message))
+            callbackfn(queue_message)
+        except Exception as e:
+            logger.error(f'Error in processing message: {e}')
+        finally: 
+            return message                 
+
+    
+    def settle_message(self, x: cf.Future):
+        """
+        Sets the message as completed and updates the internal count.
+        Args:
+            x (cf.Future): The future object representing the message processing.
+        """
+        self.internal_count -= 1
+        incoming_message = x.result()
+        self.receiver.complete_message(incoming_message)
+        return  
+
+