Add deployment for ADML agent with http interface
[anna.git] / example / diameter / launcher / deployments / aots-setup / agents / KAFKA / Consumer.py
1 #!/usr/bin/env python
2
3 import os, sys, re, datetime, logging, json, signal
4
5 from kafka import KafkaConsumer
6 from json import loads
7 from argparse import ArgumentParser
8
9
10 def _exit(message = None, rc=0):
11   if (message): printMsg(message)
12   printMsg("Exiting [rc={}]".format(rc))
13   sys.exit(rc)
14
15 def timeFormat():
16   return '%b.%d.%Y-%H.%M.%S.%f'
17
18
19 def printMsg(message):
20   print ("[{}] [{}] {}".format(datetime.datetime.now().strftime(timeFormat()), GNode, message))
21
22
23 class MyKafkaConsumer:
24     kafka_service = os.environ.get('KAFKA_SERVER')
25     kafka_port = os.environ.get('KAFKA_PORT')
26
27     if not kafka_service: kafka_service = "localhost"
28     if not kafka_port: kafka_port = "9092"
29
30     server = kafka_service + ':' + kafka_port
31
32     def __init__(self, topic, mregexp, aor, showReadWhenMatchingEnabled):
33
34         # Auto offset reset:
35         printMsg("Auto offset reset: '{}'".format(aor))
36
37         # Topic:
38         printMsg("Working topic: '{}'".format(topic))
39
40         # Show fulfilled:
41         self.showReadWhenMatchingEnabled = showReadWhenMatchingEnabled
42
43         # Json content:
44         self.json = None
45         if GJfile:
46           with open(GJfile) as json_file:
47             self.json = json.load(json_file)
48           printMsg("Established break condition for expected json ('{}'): {}".format(GJfile, self.json))
49         else:
50           printMsg("No expected json ('any' selected)")
51
52         # Interpreted as regexp indicator:
53         self.mregexp = mregexp
54         if GJfile:
55           if (self.mregexp): printMsg("Json file will be matched as regular expression")
56           else: printMsg("Json file will be matched as exact comparison")
57
58
59         self.consumer = KafkaConsumer(topic,
60                         bootstrap_servers=[self.server],
61                         auto_offset_reset=aor,
62                         enable_auto_commit=True,
63                         value_deserializer=lambda x:
64                         json.loads(x.decode('utf-8')))
65
66     def match(self, json):
67        if self.mregexp:
68          for key, value in  sorted(json.items()):
69            match = re.match(self.json[key], value)
70            if not match: return False
71          return True
72
73        else:
74          return (sorted(self.json.items()) == sorted(json.items()))
75
76     def receive_data(self):
77         printMsg('Reading data from kafka ...')
78
79         for message in self.consumer:
80             # Show read
81             if (self.showReadWhenMatchingEnabled if self.json else True):
82               printMsg('Data read: {}'.format(message))
83
84             # Break condition (in case that a expected json is blocking consumer):
85             if (self.json):
86               #if (sorted(self.json.items()) == sorted(message.value.items())):
87               if (self.match(message.value)):
88                 printMsg('Message retrieved: condition fulfilled !')
89                 return message
90               else:
91                 global GSequence
92                 GSequence += 1
93                 #saved='{}.read.{}.json'.format(os.path.basename(GJfile), GSequence)
94                 saved='{}.read.{}.json'.format(GJfile, GSequence)
95                 with open(saved, 'w') as outfile:
96                   json.dump(message.value, outfile)
97
98                 printMsg('Message retrieved: condition not fulfilled (saved {}) !'.format(saved))
99
100
101     def close (self):
102       self.consumer.close()
103
104
105 def handler_function(signum, frame):
106   rc=0
107   if (GJfile): rc=1
108   _exit('Timeout !', rc)
109
110
111 def parse_arguments():
112
113   parser = ArgumentParser(description='Kafka Consumer Simulator')
114   parser.add_argument('-n', '--node', help='Node name (spaces are removed)', required=True)
115   parser.add_argument('-t', '--topic', help='Kafka topic', required=True)
116   parser.add_argument('-o', '--auto-offset-reset', help='Auto offset reset (earliest, latest, etc.). By default: latest', required=False)
117   parser.add_argument('-f', '--file', help='Json file to consume (consumer will break execution with rc=0)', required=False)
118   parser.add_argument('-m', '--match-regexp', help='Match file as regexp (exact match by default)', required=False, action='store_true')
119   parser.add_argument('-w', '--timeout', help='Timeout (in seconds) for completion. Timeout will break execution with rc=1', required=False)
120   parser.add_argument('-d', '--debug', help='Debug kafka', required=False, action='store_true')
121   parser.add_argument('-r', '--show-read-when-matching-enabled', help='Show read messages. False by default when json file for matching is provided.', required=False, action='store_true')
122
123   arguments = parser.parse_args()
124
125   return arguments
126
127
128 #####################
129 #      M A I N      #
130 #####################
131
132 if __name__ == "__main__":
133
134   arguments = parse_arguments()
135
136   GNode = arguments.node.replace(" ", "")
137   topic = arguments.topic
138   auto_offset_reset = arguments.auto_offset_reset
139   GJfile = arguments.file
140   mregexp = arguments.match_regexp
141   timeout = arguments.timeout
142   debug = arguments.debug
143   show_read_when_matching_enabled = arguments.show_read_when_matching_enabled
144
145   GSequence = 0
146
147   printMsg("Starting")
148
149   if (debug):
150     now = datetime.datetime.now().strftime(timeFormat())
151     logfile='{}-Consumer-{}.log'.format(GNode, now)
152     printMsg("Logging kafka in: '{}'".format(logfile))
153     logging.basicConfig(filename=logfile, format='%(asctime)s - %(message)s', datefmt=timeFormat(), level=logging.DEBUG)
154
155   if (timeout):
156     printMsg("Established timeout for consumer: '{} seconds'".format(timeout))
157     # Signal to terminate after timeout:
158     signal.signal(signal.SIGALRM,handler_function)
159     signal.alarm(int(timeout))
160
161
162   if not GJfile:
163     if (mregexp): printMsg("To match json file as regexp, you must provide the file. See help (-h) !")
164     printMsg("Show messages read (matching is enabled): {}".format(str(show_read_when_matching_enabled)));
165
166   if not auto_offset_reset: auto_offset_reset='latest'
167
168   consumer = None
169   try:
170     consumer = MyKafkaConsumer(topic, mregexp, auto_offset_reset, show_read_when_matching_enabled)
171     consumer.receive_data()
172     consumer.close()
173     _exit()
174
175   except Exception as e:
176     consumer.close()
177     _exit(e, 1)