class QueueKafka(QueueBase.QueueBase):
@QueueBase.catch
def __init__(self, name, host='web14', port=51092, **kwargs):
QueueBase.QueueBase.__init__(self, name, host, port)
self.__queue = []
self.__kafka = KafkaClient('%s:%d' % (host, port))
self.__producer = SimpleProducer(self.__kafka, async=kwargs.get('async', False))
self.__producer.client.ensure_topic_exists(self.name)
self.__consumer = SimpleConsumer(self.__kafka, self.name + '_consumer', self.name, auto_commit_every_n=1)
def __del__(self):
if self.__kafka:
[self.put(x.message.value) for x in self.__queue]
self.__kafka.close()
@QueueBase.catch
def put(self, value, *args, **kwargs):
if isinstance(value, dict) or isinstance(value, list):
self.__producer.send_messages(self.name, json.dumps(value))
else:
self.__producer.send_messages(self.name, value.encode('utf-8') if isinstance(value, unicode) else value)
@QueueBase.catch
def get(self, *args, **kwargs):
if not self.__queue:
self.__consumer._fetch()
kq = self.__consumer.queue
while not kq.empty():
partition, result = kq.get_nowait()
self.__queue.append(result)
self.__consumer.offsets[partition] += 1
self.__consumer.count_since_commit += 1
self.__consumer.queue = Queue()
self.__consumer.commit()
return self.__queue.pop().message.value if self.__queue else None
@QueueBase.catch
def size(self, *args, **kwargs):
count = 0
for k, v in self.__consumer.offsets.items():
reqs = [common.OffsetRequest(self.name, k, -1, 1)]
(resp, ) = self.__consumer.client.send_offset_request(reqs)
count += (resp.offsets[0] - v)
return count + len(self.__queue)
class kafka:
def __init__(self, host, port, table_name, **args):
"""
:param host
:param port
:param table_name
:return: kafka
"""
self.queue = []
self.queue_name = table_name.replace(":", "_")
# self.kafka = KafkaClient('%s:%d' % (host, port))
self.kafka = KafkaClient(hosts=host, client_id=self.queue_name)
self.producer = SimpleProducer(
self.kafka,
async=args['async'] if args.has_key('async') else False)
self.producer.client.ensure_topic_exists(self.queue_name)
self.consumer = SimpleConsumer(self.kafka,
self.queue_name + "_consumer",
self.queue_name,
auto_commit_every_n=1,
max_buffer_size=None)
print 'success init kafka connection'
def __del__(self):
if self.kafka:
[self.save(x.message.value) for x in self.queue]
self.kafka.close()
def save(self, data, **args):
try:
if isinstance(data, dict) or isinstance(data, list):
self.producer.send_messages(self.queue_name, json.dumps(data))
elif isinstance(data, unicode):
self.producer.send_messages(self.queue_name,
data.encode('utf-8'))
else:
self.producer.send_messages(self.queue_name, data)
except Exception as e:
print e
time.sleep(60)
def get(self, **args):
# self.consumer.seek(369600, 0)
if not self.queue:
try:
self.consumer._fetch()
except Exception as e:
print e
kq = self.consumer.queue
while not kq.empty():
partition, result = kq.get_nowait()
self.queue.append(result)
self.consumer.offsets[partition] += 1
self.consumer.count_since_commit += 1
self.consumer.queue = Queue()
self.consumer.commit()
if self.queue:
return self.queue.pop().message.value
else:
return None
def size(self, **args):
count = 0
for k, v in self.consumer.offsets.items():
reqs = [common.OffsetRequest(self.queue_name, k, -1, 1)]
(resp, ) = self.consumer.client.send_offset_request(reqs)
count += (resp.offsets[0] - v)
return count + len(self.queue)
# 切换队列
def select_queue(self, name):
self.queue_name = name.replace(":", "_")
self.consumer = SimpleConsumer(self.kafka,
self.queue_name + "_consumer",
self.queue_name,
max_buffer_size=None)
