This directory contains clients implementing the official Kafka tests/kafkatest

Verifiable client test semantics.

FIXME: Instructions on how to use this.

import argparse, sys

from confluent_kafka import Consumer, KafkaError, KafkaException

from verifiable_client import VerifiableClient

class VerifiableConsumer (VerifiableClient):

"""

def __init__ (self, conf):

"""

super(VerifiableConsumer, self).__init__(conf)

self.conf['on_commit'] = self.on_commit

self.consumer = Consumer(**conf)

self.consumed_msgs = 0

self.consumed_msgs_last_reported = 0

self.consumed_msgs_at_last_commit = 0

self.use_auto_commit = False

self.use_async_commit = False

self.max_msgs = -1

self.assignment = []

self.assignment_dict = dict()

def find_assignment (self, topic, partition):

""" Find and return existing assignment based on \p topic and \p partition,

skey = '%s %d' % (topic, partition)

return self.assignment_dict.get(skey)

def send_records_consumed (self, immediate=False):

""" Send records_consumed, every 100 messages, on timeout,

if (self.consumed_msgs <= self.consumed_msgs_last_reported +

(0 if immediate else 100)):

return

if len(self.assignment) == 0:

return

d = {'name': 'records_consumed',

'count': self.consumed_msgs - self.consumed_msgs_last_reported,

'partitions': []}

for a in self.assignment:

if a.min_offset == -1:

# Skip partitions that havent had any messages since last time.

# This is to circumvent some minOffset checks in kafkatest.

continue

d['partitions'].append(a.to_dict())

a.min_offset = -1

self.send(d)

self.consumed_msgs_last_reported = self.consumed_msgs

def send_assignment (self, evtype, partitions):

""" Send assignment update, \p evtype is either 'assigned' or 'revoked' """

d = { 'name': 'partitions_' + evtype,

'partitions': [{'topic': x.topic, 'partition': x.partition} for x in partitions]}

self.send(d)

def on_assign (self, consumer, partitions):

""" Rebalance on_assign callback """

old_assignment = self.assignment

self.assignment = [AssignedPartition(p.topic, p.partition) for p in partitions]

# Move over our last seen offsets so that we can report a proper

# minOffset even after a rebalance loop.

for a in old_assignment:

b = self.find_assignment(a.topic, a.partition)

b.min_offset = a.min_offset

self.assignment_dict = {a.skey: a for a in self.assignment}

self.send_assignment('assigned', partitions)

def on_revoke (self, consumer, partitions):

""" Rebalance on_revoke callback """

# Send final consumed records prior to rebalancing to make sure

# latest consumed is in par with what is going to be committed.

self.send_records_consumed(immediate=True)

self.assignment = list()

self.assignment_dict = dict()

self.send_assignment('revoked', partitions)

self.do_commit(immediate=True)

def on_commit (self, consumer, err, partitions):

""" Offsets Committed callback """

if err is not None and err.code() == KafkaError._NO_OFFSET:

self.dbg('on_commit(): no offsets to commit')

return

# Report consumed messages to make sure consumed position >= committed position

self.send_records_consumed(immediate=True)

d = {'name': 'offsets_committed',

'offsets': []}

if err is not None:

d['success'] = False

d['error'] = str(err)

else:

d['success'] = True

d['error'] = ''

for p in partitions:

pd = {'topic': p.topic, 'partition': p.partition,

'offset': p.offset, 'error': str(p.error)}

d['offsets'].append(pd)

self.send(d)

def do_commit (self, immediate=False, async=None):

""" Commit every 1000 messages or whenever there is a consume timeout

if (self.use_auto_commit or

self.consumed_msgs_at_last_commit + (0 if immediate else 1000) >

self.consumed_msgs):

return

# Make sure we report consumption before commit,

# otherwise tests may fail because of commit > consumed

if self.consumed_msgs_at_last_commit < self.consumed_msgs:

self.send_records_consumed(immediate=True)

if async is None:

async_mode = self.use_async_commit

else:

async_mode = async

self.dbg('Committing %d messages (Async=%s)' %

(self.consumed_msgs - self.consumed_msgs_at_last_commit,

async_mode))

try:

self.consumer.commit(async=async_mode)

except KafkaException as e:

if e.args[0].code() == KafkaError._WAIT_COORD:

self.dbg('Ignoring commit failure, still waiting for coordinator')

elif e.args[0].code() == KafkaError._NO_OFFSET:

self.dbg('No offsets to commit')

else:

raise

self.consumed_msgs_at_last_commit = self.consumed_msgs

def msg_consume (self, msg):

""" Handle consumed message (or error event) """

if msg.error():

if msg.error().code() == KafkaError._PARTITION_EOF:

# ignore EOF

pass

else:

self.err('Consume failed: %s' % msg.error(), term=True)

return

if False:

self.dbg('Read msg from %s [%d] @ %d' % \

(msg.topic(), msg.partition(), msg.offset()))

if self.max_msgs >= 0 and self.consumed_msgs >= self.max_msgs:

return # ignore extra messages

# Find assignment.

a = self.find_assignment(msg.topic(), msg.partition())

if a is None:

self.err('Received message on unassigned partition %s [%d] @ %d' %

(msg.topic(), msg.partition(), msg.offset()), term=True)

a.consumed_msgs += 1

if a.min_offset == -1:

a.min_offset = msg.offset()

if a.max_offset < msg.offset():

a.max_offset = msg.offset()

self.consumed_msgs += 1

self.send_records_consumed(immediate=False)

self.do_commit(immediate=False)

class AssignedPartition (object):

""" Local state container for assigned partition. """

def __init__ (self, topic, partition):

super(AssignedPartition, self).__init__()

self.topic = topic

self.partition = partition

self.skey = '%s %d' % (self.topic, self.partition)

self.consumed_msgs = 0

self.min_offset = -1

self.max_offset = 0

def to_dict (self):

""" Return a dict of this partition's state """

return {'topic': self.topic, 'partition': self.partition,

'minOffset': self.min_offset, 'maxOffset': self.max_offset}

if __name__ == '__main__':

parser = argparse.ArgumentParser(description='Verifiable Python Consumer')

parser.add_argument('--topic', action='append', type=str, required=True)

parser.add_argument('--group-id', dest='group.id', required=True)

parser.add_argument('--broker-list', dest='bootstrap.servers', required=True)

parser.add_argument('--session-timeout', type=int, dest='session.timeout.ms', default=6000)

parser.add_argument('--enable-autocommit', action='store_true', dest='enable.auto.commit', default=False)

parser.add_argument('--max-messages', type=int, dest='max_messages', default=-1)

parser.add_argument('--assignment-strategy', dest='partition.assignment.strategy')

parser.add_argument('--reset-policy', dest='topic.auto.offset.reset', default='earliest')

parser.add_argument('--consumer.config', dest='consumer_config')

args = vars(parser.parse_args())

conf = {'broker.version.fallback': '0.9.0',

'default.topic.config': dict()}

VerifiableClient.set_config(conf, args)

vc = VerifiableConsumer(conf)

vc.use_auto_commit = args['enable.auto.commit']

vc.max_msgs = args['max_messages']

vc.dbg('Using config: %s' % conf)

vc.dbg('Subscribing to %s' % args['topic'])

vc.consumer.subscribe(args['topic'],

on_assign=vc.on_assign, on_revoke=vc.on_revoke)

try:

while vc.run:

msg = vc.consumer.poll(timeout=1.0)

if msg is None:

# Timeout.

# Try reporting consumed messages

vc.send_records_consumed(immediate=True)

# Commit every poll() timeout instead of on every message.

# Also commit on every 1000 messages, whichever comes first.

vc.do_commit(immediate=True)

continue

# Handle message (or error event)

vc.msg_consume(msg)

except KeyboardInterrupt:

pass

vc.dbg('Closing consumer')

vc.send_records_consumed(immediate=True)

if not vc.use_auto_commit:

vc.do_commit(immediate=True, async=False)

vc.consumer.close()

vc.send({'name': 'shutdown_complete'})

vc.dbg('All done')