kafkatest: add --message-create-time and --producer.config support

edenhill · edenhill · commit 99f8eac1c59d · 2018-01-13T20:03:12.000+01:00
diff --git a/confluent_kafka/kafkatest/verifiable_client.py b/confluent_kafka/kafkatest/verifiable_client.py
@@ -70,17 +70,44 @@ def set_config(conf, args):
         for n, v in args.iteritems():
             if v is None:
                 continue
-            # Things to ignore
-            if '.' not in n:
+
+            if not n.startswith('conf_'):
                 # App config, skip
                 continue
-            if n.startswith('topic.'):
-                # Set "topic.<...>" properties on default topic conf dict
-                conf['default.topic.config'][n[6:]] = v
-            elif n == 'partition.assignment.strategy':
+
+            n = n[5:] # Remove conf_ prefix
+
+            if n == 'partition.assignment.strategy':
                 # Convert Java class name to config value.
                 # "org.apache.kafka.clients.consumer.RangeAssignor" -> "range"
                 conf[n] = re.sub(r'org.apache.kafka.clients.consumer.(\w+)Assignor',
                                  lambda x: x.group(1).lower(), v)
+                                # Handle known Java properties to librdkafka properties.
+            elif n == 'enable.idempotence':
+                # Ignore idempotence for now, best-effortly.
+                sys.stderr.write('%% WARN: Ignoring unsupported %s=%s\n' % (n, v))
             else:
                 conf[n] = v
+
+    @staticmethod
+    def read_config_file (path):
+        """Read (java client) config file and return dict with properties"""
+        conf = {}
+
+        with open(path, 'r') as f:
+            for line in f:
+                line = line.strip()
+
+                if line.startswith('#') or len(line) == 0:
+                    continue
+
+                fi = line.find('=')
+                if fi < 1:
+                    raise Exception('%s: invalid line, no key=value pair: %s' % (path, line))
+
+                k = line[:fi]
+                v = line[fi+1:]
+
+                conf[k] = v
+
+        return conf
diff --git a/confluent_kafka/kafkatest/verifiable_consumer.py b/confluent_kafka/kafkatest/verifiable_consumer.py
@@ -242,30 +242,32 @@ def to_dict(self):
 
     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('--group-id', dest='conf_group.id', required=True)
+    parser.add_argument('--broker-list', dest='conf_bootstrap.servers', required=True)
+    parser.add_argument('--session-timeout', type=int, dest='conf_session.timeout.ms', default=6000)
+    parser.add_argument('--enable-autocommit', action='store_true', dest='conf_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('--assignment-strategy', dest='conf_partition.assignment.strategy')
+    parser.add_argument('--reset-policy', dest='conf_auto.offset.reset', default='earliest')
     parser.add_argument('--consumer.config', dest='consumer_config')
     parser.add_argument('-X', nargs=1, dest='extra_conf', action='append', help='Configuration property', default=[])
     args = vars(parser.parse_args())
 
     conf = {'broker.version.fallback': '0.9.0',
-            'default.topic.config': dict(),
             # Do explicit manual offset stores to avoid race conditions
             # where a message is consumed from librdkafka but not yet handled
             # by the Python code that keeps track of last consumed offset.
             'enable.auto.offset.store': False}
 
-    VerifiableClient.set_config(conf, args)
+    if args.get('consumer_config', None) is not None:
+        args.update(VerifiableClient.read_config_file(args['consumer_config']))
+
+    args.update([x[0].split('=') for x in args.get('extra_conf', [])])
 
-    conf.update([x[0].split('=') for x in args.get('extra_conf', [])])
+    VerifiableClient.set_config(conf, args)
 
     vc = VerifiableConsumer(conf)
-    vc.use_auto_commit = args['enable.auto.commit']
+    vc.use_auto_commit = args['conf_enable.auto.commit']
     vc.max_msgs = args['max_messages']
 
     vc.dbg('Pid %d' % os.getpid())
diff --git a/confluent_kafka/kafkatest/verifiable_producer.py b/confluent_kafka/kafkatest/verifiable_producer.py
@@ -32,7 +32,6 @@ def __init__(self, conf):
         """
         super(VerifiableProducer, self).__init__(conf)
         self.conf['on_delivery'] = self.dr_cb
-        self.conf['default.topic.config']['produce.offset.report'] = True
         self.producer = Producer(**self.conf)
         self.num_acked = 0
         self.num_sent = 0
@@ -64,20 +63,24 @@ def dr_cb(self, err, msg):
     parser = argparse.ArgumentParser(description='Verifiable Python Producer')
     parser.add_argument('--topic', type=str, required=True)
     parser.add_argument('--throughput', type=int, default=0)
-    parser.add_argument('--broker-list', dest='bootstrap.servers', required=True)
+    parser.add_argument('--broker-list', dest='conf_bootstrap.servers', required=True)
     parser.add_argument('--max-messages', type=int, dest='max_msgs', default=1000000)  # avoid infinite
     parser.add_argument('--value-prefix', dest='value_prefix', type=str, default=None)
-    parser.add_argument('--acks', type=int, dest='topic.request.required.acks', default=-1)
+    parser.add_argument('--acks', type=int, dest='conf_request.required.acks', default=-1)
+    parser.add_argument('--message-create-time', type=int, dest='create_time', default=0)
     parser.add_argument('--producer.config', dest='producer_config')
     parser.add_argument('-X', nargs=1, dest='extra_conf', action='append', help='Configuration property', default=[])
     args = vars(parser.parse_args())
 
     conf = {'broker.version.fallback': '0.9.0',
-            'default.topic.config': dict()}
+            'produce.offset.report': True}
 
-    VerifiableClient.set_config(conf, args)
+    if args.get('producer_config', None) is not None:
+        args.update(VerifiableClient.read_config_file(args['producer_config']))
+
+    args.update([x[0].split('=') for x in args.get('extra_conf', [])])
 
-    conf.update([x[0].split('=') for x in args.get('extra_conf', [])])
+    VerifiableClient.set_config(conf, args)
 
     vp = VerifiableProducer(conf)
 
@@ -104,7 +107,8 @@ def dr_cb(self, err, msg):
             t_end = time.time() + delay
             while vp.run:
                 try:
-                    vp.producer.produce(topic, value=(value_fmt % i))
+                    vp.producer.produce(topic, value=(value_fmt % i),
+                                        timestamp=args.get('create_time', 0))
                     vp.num_sent += 1
                 except KafkaException as e:
                     vp.err('produce() #%d/%d failed: %s' % (i, vp.max_msgs, str(e)))
