ClickHouse/utils/kafka/status.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-

# `pip install …`
import kafka  # … kafka-python

import argparse


def main():
    parser = argparse.ArgumentParser(description='Kafka client to get groups and topics status')
    parser.add_argument('--server', type=str, metavar='HOST', default='localhost',
        help='Kafka bootstrap-server address')
    parser.add_argument('--port', type=int, metavar='PORT', default=9092,
        help='Kafka bootstrap-server port')
    parser.add_argument('--client', type=str, default='ch-kafka-python',
        help='custom client id for this producer')

    args = parser.parse_args()
    config = {
        'bootstrap_servers': f'{args.server}:{args.port}',
        'client_id': args.client,
    }

    client = kafka.KafkaAdminClient(**config)
    consumer = kafka.KafkaConsumer(**config)
    cluster = client._client.cluster

    topics = cluster.topics()
    for topic in topics:
        print(f'Topic "{topic}":', end='')
        for partition in cluster.partitions_for_topic(topic):
            tp = kafka.TopicPartition(topic, partition)
            print(f' {partition} (begin: {consumer.beginning_offsets([tp])[tp]}, end: {consumer.end_offsets([tp])[tp]})', end='')
        print()

    groups = client.list_consumer_groups()
    for group in groups:
        print(f'Group "{group[0]}" ({group[1]}):')

        consumer = kafka.KafkaConsumer(**config, group_id=group[0])
        offsets = client.list_consumer_group_offsets(group[0])
        for topic, offset in offsets.items():
            print(f'\t{topic.topic}[{topic.partition}]: {consumer.beginning_offsets([topic])[topic]}, {offset.offset}, {consumer.end_offsets([topic])[topic]}')
        consumer.close()

    client.close()
    return 0


if __name__ == "__main__":
    exit(main())