tuoheng_alg_move_20240715/util/KafkaUtils.py

from kafka import KafkaProducer, KafkaConsumer, TopicPartition, OffsetAndMetadata
from kafka.errors import kafka_errors
import json
from loguru import logger


# 生产者
class CustomerKafkaProducer():

    def __init__(self, content, requestId):
        self.content = content
        self.configs = self.content["kafka"][self.content["dsp"]["active"]]["producer"]
        self.customerProducer = None
        self.requestId = requestId
        self.get_producer()

    # 获取kafka生产者
    def get_producer(self):
        if self.customerProducer is None:
            logger.info("配置kafka生产者, requestId:{}", self.requestId)
            self.customerProducer = KafkaProducer(bootstrap_servers=self.content["kafka"][self.content["dsp"]["active"]]["bootstrap_servers"],
                                                  acks=self.configs["acks"],
                                                  retries=self.configs["retries"],
                                                  linger_ms=self.configs["linger_ms"],
                                                  retry_backoff_ms=self.configs["retry_backoff_ms"],
                                                  max_in_flight_requests_per_connection=self.configs[
                                                      "max_in_flight_requests_per_connection"],
                                                  key_serializer=lambda m: json.dumps(m).encode('utf-8'),
                                                  value_serializer=lambda m: json.dumps(m).encode('utf-8'))

    # mode 模式1：异步发送 2：同步发送
    # def on_send_success(record_metadata): 成功回调
    # def on_send_error(excp): 失败回调
    def sender(self, topic, key, message, mode, customer_send_success=None, customer_send_error=None):
        self.get_producer()
        logger.info("kafka发送信息，topic:{}|key:{}|message:{}|mode:{}|requestId:{}", topic, key, message, mode, self.requestId)
        if mode == 1:
            if not customer_send_success:
                customer_send_success = CustomerKafkaProducer.on_send_success
            if not customer_send_error:
                customer_send_error = CustomerKafkaProducer.on_send_error
            self.customerProducer.send(topic=topic, key=key, value=message).add_callback(
                customer_send_success, self.requestId).add_errback(customer_send_error, self.requestId)
        if mode == 2:
            try:
                self.customerProducer.send(topic=topic, key=key, value=message).get(timeout=30)
                logger.info("kafka同步发送信息成功, requestId:{}", self.requestId)
            except kafka_errors as e:
                logger.exception("kafka同步发送消息异常: {}, requestId:{}", e, self.requestId)
                raise e

    def close_producer(self):
        self.customerProducer.flush()
        self.customerProducer.close()
        logger.info("kafka生产者关闭完成, requestId:{}", self.requestId)

    def on_send_success(requestId, record_metadata):
        logger.info("kafka异步发送信息成功，topic:{}|partition:{}|offset:{}|requestId:{}", record_metadata.topic,
                    record_metadata.partition, record_metadata.offset, requestId)

    def on_send_error(requestId, excp):
        logger.exception("kafka异步发送消息异常: {}, requestId:{}", excp, requestId)


# 生产者
class CustomerKafkaConsumer():

    def __init__(self, content):
        logger.info("初始化消费者")
        self.content = content
        self.configs = self.content["kafka"][self.content["dsp"]["active"]]["consumer"]
        self.customerConsumer = None
        logger.info("初始化消费者完成")

    def get_consumer(self):
        if self.customerConsumer is None:
            logger.info("获取消费者！")
            self.customerConsumer = KafkaConsumer(bootstrap_servers=self.content["kafka"][self.content["dsp"]["active"]]["bootstrap_servers"],
                                                  client_id=self.configs["client_id"],
                                                  group_id=self.configs["group_id"],
                                                  auto_offset_reset=self.configs["auto_offset_reset"],
                                                  enable_auto_commit=self.configs["enable_auto_commit"],
                                                  max_poll_records=self.configs["max_poll_records"],
                                                  value_deserializer=lambda m: json.loads(m.decode('utf-8')))

    def subscribe(self, topics=()):
        logger.info("kafka生产者订阅topic:{}", topics)
        if topics is None or len(topics) == 0:
            logger.error("消费者订阅topic不能为空！")
            raise Exception("消费者订阅topic不能为空！")
        # 手动配置分区
        # customer_partition = []
        # for topic in topics:
        #     for p in self.content["kafka"]["topic"][topic]["partition"]:
        #         customer_partition.append(TopicPartition(topic, p))
        # self.customerConsumer.assign(customer_partition)
        # 自动配置
        self.customerConsumer.subscribe(topics=topics)
        logger.info("kafka生产者订阅topic完成")

    def commit_offset(self, message):
        logger.info("消费者开始提交offset,topic:{}|offset:{}|partition:{}", message.topic, message.offset + 1,
                    message.partition)
        tp = TopicPartition(topic=message.topic, partition=message.partition)
        self.customerConsumer.commit(offsets={tp: (OffsetAndMetadata(message.offset + 1, None))})
        logger.info("消费者提交offset完成,topic:{}|offset:{}|partition:{}", message.topic, message.offset + 1,
                    message.partition)

# if __name__=="__main__":
#     try:
#         1/0
#     except Exception as e:
#         logger.exception("aaaaa:{} {}", e, "11111")