lcm/osm_common/msgkafka.py

   1 from aiokafka import AIOKafkaConsumer
   2 from aiokafka import AIOKafkaProducer
   3 from aiokafka.errors import KafkaError
   4 from msgbase import MsgBase, MsgException
   5 import asyncio
   6 import yaml
   7 #import json
   8
   9 class MsgKafka(MsgBase):
  10     def __init__(self):
  11         self.host = None
  12         self.port = None
  13         self.consumer = None
  14         self.producer = None
  15         # create a different file for each topic
  16         #self.files = {}
  17
  18     def connect(self, config):
  19         try:
  20             self.host = config["host"]
  21             self.port = config["port"]
  22             self.topic_lst = []
  23             self.loop = asyncio.get_event_loop()
  24             self.broker = str(self.host) + ":" + str(self.port)
  25
  26         except Exception as e:  # TODO refine
  27             raise MsgException(str(e))
  28
  29     def write(self, topic, key, msg):
  30         try:
  31             self.loop.run_until_complete(self.aiowrite(topic=topic, key=key, msg=yaml.safe_dump(msg, default_flow_style=True)))
  32
  33         except Exception as e:
  34             raise MsgException("Error writing {} topic: {}".format(topic, str(e)))
  35
  36     def read(self, topic):
  37         #self.topic_lst.append(topic)
  38         try:
  39             return self.loop.run_until_complete(self.aioread(topic))
  40         except Exception as e:
  41             raise MsgException("Error reading {} topic: {}".format(topic, str(e)))
  42
  43     async def aiowrite(self, topic, key, msg, loop=None):
  44         try:
  45             if not loop:
  46                 loop = self.loop
  47             self.producer = AIOKafkaProducer(loop=loop, key_serializer=str.encode, value_serializer=str.encode,
  48                                              bootstrap_servers=self.broker)
  49             await self.producer.start()
  50             await self.producer.send(topic=topic, key=key, value=msg)
  51         except Exception as e:
  52             raise MsgException("Error publishing to {} topic: {}".format(topic, str(e)))
  53         finally:
  54             await self.producer.stop()
  55
  56     async def aioread(self, topic, loop=None):
  57         if not loop:
  58             loop = self.loop
  59         self.consumer = AIOKafkaConsumer(loop=loop, bootstrap_servers=self.broker)
  60         await self.consumer.start()
  61         self.consumer.subscribe([topic])
  62         try:
  63             async for message in self.consumer:
  64                 return yaml.load(message.key), yaml.load(message.value)
  65         except KafkaError as e:
  66             raise MsgException(str(e))
  67         finally:
  68             await self.consumer.stop()
  69
  70