lcm/osm_common/msgkafka.py

   1 import logging
   2 import asyncio
   3 import yaml
   4 from aiokafka import AIOKafkaConsumer
   5 from aiokafka import AIOKafkaProducer
   6 from aiokafka.errors import KafkaError
   7 from msgbase import MsgBase, MsgException
   8 #import json
   9
  10
  11 class MsgKafka(MsgBase):
  12     def __init__(self, logger_name='msg'):
  13         self.logger = logging.getLogger(logger_name)
  14         self.host = None
  15         self.port = None
  16         self.consumer = None
  17         self.producer = None
  18
  19     def connect(self, config):
  20         try:
  21             if "logger_name" in config:
  22                 self.logger = logging.getLogger(config["logger_name"])
  23             self.host = config["host"]
  24             self.port = config["port"]
  25             self.loop = asyncio.get_event_loop()
  26             self.broker = str(self.host) + ":" + str(self.port)
  27
  28         except Exception as e:  # TODO refine
  29             raise MsgException(str(e))
  30
  31     def disconnect(self):
  32         try:
  33             self.loop.close()
  34         except Exception as e:  # TODO refine
  35             raise MsgException(str(e))
  36
  37     def write(self, topic, key, msg):
  38         try:
  39             self.loop.run_until_complete(self.aiowrite(topic=topic, key=key,
  40                                                        msg=yaml.safe_dump(msg, default_flow_style=True),
  41                                                        loop=self.loop))
  42
  43         except Exception as e:
  44             raise MsgException("Error writing {} topic: {}".format(topic, str(e)))
  45
  46     def read(self, topic):
  47         """
  48         Read from one or several topics. it is non blocking returning None if nothing is available
  49         :param topic: can be str: single topic; or str list: several topics
  50         :return: topic, key, message; or None
  51         """
  52         try:
  53             return self.loop.run_until_complete(self.aioread(topic, self.loop))
  54         except MsgException:
  55             raise
  56         except Exception as e:
  57             raise MsgException("Error reading {} topic: {}".format(topic, str(e)))
  58
  59     async def aiowrite(self, topic, key, msg, loop=None):
  60
  61         if not loop:
  62             loop = self.loop
  63         try:
  64             self.producer = AIOKafkaProducer(loop=loop, key_serializer=str.encode, value_serializer=str.encode,
  65                                              bootstrap_servers=self.broker)
  66             await self.producer.start()
  67             await self.producer.send(topic=topic, key=key, value=yaml.safe_dump(msg, default_flow_style=True))
  68         except Exception as e:
  69             raise MsgException("Error publishing topic '{}', key '{}': {}".format(topic, key, e))
  70         finally:
  71             await self.producer.stop()
  72
  73     async def aioread(self, topic, loop=None, callback=None, *args):
  74         """
  75         Asyncio read from one or several topics. It blocks
  76         :param topic: can be str: single topic; or str list: several topics
  77         :param loop: asyncio loop
  78         :callback: callback function that will handle the message in kafka bus
  79         :*args: optional arguments for callback function
  80         :return: topic, key, message
  81         """
  82
  83         if not loop:
  84             loop = self.loop
  85         try:
  86             if isinstance(topic, (list, tuple)):
  87                 topic_list = topic
  88             else:
  89                 topic_list = (topic,)
  90
  91             self.consumer = AIOKafkaConsumer(loop=loop, bootstrap_servers=self.broker)
  92             await self.consumer.start()
  93             self.consumer.subscribe(topic_list)
  94
  95             async for message in self.consumer:
  96                 if callback:
  97                     callback(message.topic, yaml.load(message.key), yaml.load(message.value), *args)
  98                 else:
  99                     return message.topic, yaml.load(message.key), yaml.load(message.value)
 100         except KafkaError as e:
 101             raise MsgException(str(e))
 102         finally:
 103             await self.consumer.stop()
 104