python多线程消费kafka数据并记录offset

2019年2月14日 14:56 by wst

大数据

---

对于kafka里的数据，如果用普通的代码（比如python脚本）读取，避免不了要处理offset的问题。

关于offset的处理可使用以下两种方式：

1. 使用kafka本身的机制，kafka-python也实现了相应的功能，即自动提交kafka的offset。

2. 手动记录offset，下一次接着消费。

这里着重实现第二种方式。

思想：

1. 在partition不是太多的情况下，设置线程池里线程的个数和partition的个数相同。

2. 每次消费时，当满足固定条件（如读取2条数据，或者某条数据的内容包含某个字符串）时，停止读取，记录下最后一条数据的topic、partition、offset到本地数据库sqlite里。

3. 下次消费时，从sqlite里读取出上次消费的位置，接着消费。

完整的代码如下（所有代码均为python3版本）：

处理数据库的读写的模块：op_offset.py

import sqlite3

class OpOffset(object):
    def __init__(self):
        self.db = "kafka.db"
        self.table = "offset"

    def judge_table(self):
        conn = sqlite3.connect(self.db)
        cur = conn.cursor()
        cur.execute('''
            CREATE TABLE IF NOT EXISTS %s(
            topic          varchar(50),
            partition        integer ,
            offset  integer ,
            PRIMARY KEY( topic, partition )
            )
            '''%self.table)
        conn.commit()
        cur.close()
        conn.close()

    def update_offset(self, topic, partition, offset):
        conn = sqlite3.connect(self.db)
        cur = conn.cursor()
        sql = """replace into %s(topic, partition, offset)
        VALUES('%s', %s, %s)""" % (self.table, topic, partition, offset)
        cur.execute(sql)
        conn.commit()
        cur.close()
        conn.close()
    def read_data(self, topic, partition):
        conn = sqlite3.connect(self.db)
        cur = conn.cursor()
        sql = "select topic, partition, offset from %s where topic='%s' and partition='%s'" % (self.table, topic, partition)
        cur.execute(sql)
        values = cur.fetchall()
        cur.close()
        conn.close()
        return values

    def view_data(self):
        conn = sqlite3.connect(self.db)
        cur = conn.cursor()
        sql = "select * from %s" % self.table
        cur.execute(sql)
        values = cur.fetchall()
        cur.close()
        conn.close()
        return values

if __name__ == "__main__":
    data = ('ad', 1, 100)
    obj = OpOffset()
    obj.judge_table()
    obj.update_offset(data[0], data[1], data[2])
    print(obj.read_data(data[0], 101))
    print(obj.view_data())

主模块：kafka01.py

from kafka import KafkaConsumer
from kafka import TopicPartition
from concurrent.futures import ThreadPoolExecutor, as_completed
from op_offset import OpOffset

KAFKA_HOSTS = ['127.0.0.1:9092',]

def consume_data(topic, partition, db):
    print("[%s]线程被执行了。。。"%partition)
    consumer = KafkaConsumer(bootstrap_servers=KAFKA_HOSTS, group_id='wst01', auto_offset_reset='earliest')
    tp = TopicPartition(topic, partition)
    result = db.read_data(topic, partition)
    rs = []
    offset = 0
    if result:
        offset = result[0][2]

    print("offset:", offset)

    consumer.assign([tp])
    consumer.seek(tp, offset)
    i = 1
    for message in consumer:

        print("partition[%s]--record[%s]--%s:%d:%d: key=%s value=%s" % (partition, i,
            message.topic, message.partition, message.offset, message.key, message.value))
        rs.append(message.value)
        i += 1
        if i>2:
            db.update_offset(message.topic, message.partition, message.offset+1)
            break
    return rs

if __name__=='__main__':
    print("Start...")

    consumer_ = KafkaConsumer(bootstrap_servers=KAFKA_HOSTS, group_id='wst01', auto_offset_reset='earliest')
    topic_ = "test"
    partitions = consumer_.partitions_for_topic(topic_)
    print("Topic[%s] have %s partitions" % (topic_, partitions))
    db = OpOffset()
    db.judge_table()

    end_data = []


    with ThreadPoolExecutor(len(partitions)) as executor:
        # Start the load operations and mark each future with its URL
        future_to_read = {executor.submit(consume_data, topic_, partition_, db): partition_ for partition_ in partitions}
        for future in as_completed(future_to_read):
            url = future_to_read[future]
            try:
                data = future.result()
                end_data.extend(data)
                print("data:", data)
            except Exception as exc:
                print('%r generated an exception: %s' % (url, exc))
            else:
                print('%r partition have %d records' % (url, len(data)))

    print("Offset condition:", db.view_data())

    print("end_data:", end_data)

 

这两个模块已经包含了所有的源码， 如果你复制粘贴都懒得做，可以关注以下公众号，回复“读取kafka”获取所有源码，及requirement.txt

 

---