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