Added set_list at database.

[osm/common.git] / osm_common / msgkafka.py

diff --git a/osm_common/msgkafka.py b/osm_common/msgkafka.py
index 382bdef..aa756c4 100644 (file)
--- a/osm_common/msgkafka.py
+++ b/osm_common/msgkafka.py
@@ -1,3 +1,18 @@
+# -*- coding: utf-8 -*-
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import logging
 import asyncio
 import yaml
@@ -20,6 +35,7 @@ class MsgKafka(MsgBase):
         self.producer = None
         self.loop = None
         self.broker = None
+        self.group_id = None
 
     def connect(self, config):
         try:
@@ -29,6 +45,7 @@ class MsgKafka(MsgBase):
             self.port = config["port"]
             self.loop = asyncio.get_event_loop()
             self.broker = str(self.host) + ":" + str(self.port)
+            self.group_id = config.get("group_id")
 
         except Exception as e:  # TODO refine
             raise MsgException(str(e))
@@ -41,17 +58,22 @@ class MsgKafka(MsgBase):
             raise MsgException(str(e))
 
     def write(self, topic, key, msg):
+        """
+        Write a message at kafka bus
+        :param topic: message topic, must be string
+        :param key: message key, must be string
+        :param msg: message content, can be string or dictionary
+        :return: None or raises MsgException on failing
+        """
         try:
-            self.loop.run_until_complete(self.aiowrite(topic=topic, key=key,
-                                                       msg=yaml.safe_dump(msg, default_flow_style=True),
-                                                       loop=self.loop))
+            self.loop.run_until_complete(self.aiowrite(topic=topic, key=key, msg=msg))
 
         except Exception as e:
             raise MsgException("Error writing {} topic: {}".format(topic, str(e)))
 
     def read(self, topic):
         """
-        Read from one or several topics. it is non blocking returning None if nothing is available
+        Read from one or several topics.
         :param topic: can be str: single topic; or str list: several topics
         :return: topic, key, message; or None
         """
@@ -94,7 +116,7 @@ class MsgKafka(MsgBase):
             else:
                 topic_list = (topic,)
 
-            self.consumer = AIOKafkaConsumer(loop=loop, bootstrap_servers=self.broker)
+            self.consumer = AIOKafkaConsumer(loop=loop, bootstrap_servers=self.broker, group_id=self.group_id)
             await self.consumer.start()
             self.consumer.subscribe(topic_list)