1 # Copyright (c) 2019, CMCC Technologies. Co., Ltd.
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
7 # http://www.apache.org/licenses/LICENSE-2.0
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 implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
21 from hashlib import sha1
24 from apscheduler.schedulers.background import BackgroundScheduler
25 from catalog.pub.Dmaap_lib.pub.exceptions import DmaapClientException
27 logger = logging.getLogger(__name__)
30 class BatchPublisherClient:
31 def __init__(self, base_url, topic, partition="", contenttype="text/plain", max_batch_size=100,
32 max_batch_age_ms=1000):
33 self.base_url = base_url
35 self.partition = partition
36 self.contenttype = contenttype
37 self.max_batch_size = max_batch_size
38 self.max_batch_age_ms = max_batch_age_ms
41 self.dont_send_until_ms = 0
42 self.scheduler = BackgroundScheduler()
44 # @self.scheduler.interval_schedule(second=1)
45 @self.scheduler.scheduled_job(second=1)
47 self.send_message(False)
49 self.scheduler.start()
51 def set_api_credentials(self, api_key, api_secret):
52 self.api_key = api_key
53 self.api_secret = api_secret
55 def send(self, partition, msg):
58 raise DmaapClientException("The publisher was closed.")
59 message = Message(partition, msg)
60 self.pending.append(message)
61 return len(self.pending)
62 except Exception as e:
63 raise DmaapClientException("append message failed: " + str(e))
65 def send_message(self, force):
66 if force or self.should_send_now():
67 if not self.send_batch():
68 logger.error("Send failed, %s message to send.", len(self.pending))
70 def should_send_now(self):
72 if len(self.pending) > 0:
73 now_ms = int(time.time() * 1000)
74 should_send = len(self.pending) >= self.max_batch_size
76 send_at_ms = self.pending[0].timestamp_ms
77 should_send = send_at_ms <= now_ms
79 should_send = should_send and now_ms >= self.dont_send_until_ms
84 if len(self.pending) < 1:
86 now_ms = int(time.time() * 1000)
87 url = self.create_url()
88 logger.info("sending %s msgs to %s . Oldest: %s ms", len(self.pending), url,
89 str(now_ms - self.pending[0].timestamp_ms))
92 if self.contenttype == "application/json":
93 str_msg = self.parse_json()
94 elif self.contenttype == "text/plain":
95 for m in self.pending:
98 elif self.contenttype == "application/cambria":
99 for m in self.pending:
100 str_msg += str(len(m.partition))
102 str_msg += str(len(m.msg))
104 str_msg += m.partition
108 for m in self.pending:
110 msg = bytearray(str_msg)
112 start_ms = int(time.time() * 1000)
114 headers = self.create_headers()
116 headers = {'content-type': self.contenttype}
117 ret = requests.post(url=url, data=msg, headers=headers)
118 if ret.status_code < 200 or ret.status_code > 299:
120 logger.info("MR reply ok (%s ms): %s", start_ms - int(time.time() * 1000), ret.json())
124 except Exception as e:
128 def create_headers(self):
129 data = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S') + '-04:00'
130 hmac_code = hmac.new(self.api_secret.encode(), data.encode(), sha1).digest()
131 signature = base64.b64encode(hmac_code).decode()
132 auth = self.api_key + ':' + signature
134 'X-CambriaDate': data,
135 'X-CambriaAuth': auth,
136 'content-type': self.contenttype
140 def create_url(self):
141 url = self.base_url + "/events/%s" % self.topic
143 url = url + "?partitionKey=" + self.partition
146 def parse_json(self):
148 for message in self.pending:
149 msg = json.loads(message.msg)
152 return json.dumps(data)
154 def close(self, timeout):
157 self.scheduler.shutdown()
158 now_ms = int(time.time() * 1000)
159 wait_in_ms = now_ms + timeout * 1000
161 while int(time.time() * 1000) < wait_in_ms and len(self.pending) > 0:
162 self.send_message(True)
165 except Exception as e:
166 raise DmaapClientException("send message failed: " + str(e))
170 def __init__(self, partition, msg):
174 self.partition = partition
176 self.timestamp_ms = int(time.time() * 1000)