osm_common/msgkafka.py

   1 # -*- coding: utf-8 -*-
   2
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #    http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
  12 # implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 import asyncio
  17 import logging
  18
  19 from aiokafka import AIOKafkaConsumer
  20 from aiokafka import AIOKafkaProducer
  21 from aiokafka.errors import KafkaError
  22 from osm_common.msgbase import MsgBase, MsgException
  23 import yaml
  24
  25 __author__ = (
  26     "Alfonso Tierno <alfonso.tiernosepulveda@telefonica.com>, "
  27     "Guillermo Calvino <guillermo.calvinosanchez@altran.com>"
  28 )
  29
  30
  31 class MsgKafka(MsgBase):
  32     def __init__(self, logger_name="msg", lock=False):
  33         super().__init__(logger_name, lock)
  34         self.host = None
  35         self.port = None
  36         self.consumer = None
  37         self.producer = None
  38         self.loop = None
  39         self.broker = None
  40         self.group_id = None
  41
  42     def connect(self, config):
  43         try:
  44             if "logger_name" in config:
  45                 self.logger = logging.getLogger(config["logger_name"])
  46             self.host = config["host"]
  47             self.port = config["port"]
  48             self.loop = config.get("loop") or asyncio.get_event_loop()
  49             self.broker = str(self.host) + ":" + str(self.port)
  50             self.group_id = config.get("group_id")
  51
  52         except Exception as e:  # TODO refine
  53             raise MsgException(str(e))
  54
  55     def disconnect(self):
  56         try:
  57             pass
  58             # self.loop.close()
  59         except Exception as e:  # TODO refine
  60             raise MsgException(str(e))
  61
  62     def write(self, topic, key, msg):
  63         """
  64         Write a message at kafka bus
  65         :param topic: message topic, must be string
  66         :param key: message key, must be string
  67         :param msg: message content, can be string or dictionary
  68         :return: None or raises MsgException on failing
  69         """
  70         retry = 2  # Try two times
  71         while retry:
  72             try:
  73                 self.loop.run_until_complete(
  74                     self.aiowrite(topic=topic, key=key, msg=msg)
  75                 )
  76                 break
  77             except Exception as e:
  78                 retry -= 1
  79                 if retry == 0:
  80                     raise MsgException(
  81                         "Error writing {} topic: {}".format(topic, str(e))
  82                     )
  83
  84     def read(self, topic):
  85         """
  86         Read from one or several topics.
  87         :param topic: can be str: single topic; or str list: several topics
  88         :return: topic, key, message; or None
  89         """
  90         try:
  91             return self.loop.run_until_complete(self.aioread(topic, self.loop))
  92         except MsgException:
  93             raise
  94         except Exception as e:
  95             raise MsgException("Error reading {} topic: {}".format(topic, str(e)))
  96
  97     async def aiowrite(self, topic, key, msg, loop=None):
  98         """
  99         Asyncio write
 100         :param topic: str kafka topic
 101         :param key: str kafka key
 102         :param msg: str or dictionary  kafka message
 103         :param loop: asyncio loop. To be DEPRECATED! in near future!!!  loop must be provided inside config at connect
 104         :return: None
 105         """
 106
 107         if not loop:
 108             loop = self.loop
 109         try:
 110             self.producer = AIOKafkaProducer(
 111                 loop=loop,
 112                 key_serializer=str.encode,
 113                 value_serializer=str.encode,
 114                 bootstrap_servers=self.broker,
 115             )
 116             await self.producer.start()
 117             await self.producer.send(
 118                 topic=topic, key=key, value=yaml.safe_dump(msg, default_flow_style=True)
 119             )
 120         except Exception as e:
 121             raise MsgException(
 122                 "Error publishing topic '{}', key '{}': {}".format(topic, key, e)
 123             )
 124         finally:
 125             await self.producer.stop()
 126
 127     async def aioread(
 128         self,
 129         topic,
 130         loop=None,
 131         callback=None,
 132         aiocallback=None,
 133         group_id=None,
 134         from_beginning=None,
 135         **kwargs
 136     ):
 137         """
 138         Asyncio read from one or several topics.
 139         :param topic: can be str: single topic; or str list: several topics
 140         :param loop: asyncio loop. To be DEPRECATED! in near future!!!  loop must be provided inside config at connect
 141         :param callback: synchronous callback function that will handle the message in kafka bus
 142         :param aiocallback: async callback function that will handle the message in kafka bus
 143         :param group_id: kafka group_id to use. Can be False (set group_id to None), None (use general group_id provided
 144                          at connect inside config), or a group_id string
 145         :param from_beginning: if True, messages will be obtained from beginning instead of only new ones.
 146                                If group_id is supplied, only the not processed messages by other worker are obtained.
 147                                If group_id is None, all messages stored at kafka are obtained.
 148         :param kwargs: optional keyword arguments for callback function
 149         :return: If no callback defined, it returns (topic, key, message)
 150         """
 151
 152         if not loop:
 153             loop = self.loop
 154         if group_id is False:
 155             group_id = None
 156         elif group_id is None:
 157             group_id = self.group_id
 158         try:
 159             if isinstance(topic, (list, tuple)):
 160                 topic_list = topic
 161             else:
 162                 topic_list = (topic,)
 163             self.consumer = AIOKafkaConsumer(
 164                 loop=loop,
 165                 bootstrap_servers=self.broker,
 166                 group_id=group_id,
 167                 auto_offset_reset="earliest" if from_beginning else "latest",
 168             )
 169             await self.consumer.start()
 170             self.consumer.subscribe(topic_list)
 171
 172             async for message in self.consumer:
 173                 if callback:
 174                     callback(
 175                         message.topic,
 176                         yaml.safe_load(message.key),
 177                         yaml.safe_load(message.value),
 178                         **kwargs
 179                     )
 180                 elif aiocallback:
 181                     await aiocallback(
 182                         message.topic,
 183                         yaml.safe_load(message.key),
 184                         yaml.safe_load(message.value),
 185                         **kwargs
 186                     )
 187                 else:
 188                     return (
 189                         message.topic,
 190                         yaml.safe_load(message.key),
 191                         yaml.safe_load(message.value),
 192                     )
 193         except KafkaError as e:
 194             raise MsgException(str(e))
 195         finally:
 196             await self.consumer.stop()