139 lines
5.4 KiB
Python
139 lines
5.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Kafka 适配器
|
|
基于 aiokafka 异步客户端
|
|
支持：生产者/消费者、消费组、分区、偏移量管理
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
|
|
from aiokafka.errors import KafkaError
|
|
|
|
from uas.comm.base import CommBaseAdapter, ConnectionConfig, MessageHandler
|
|
from uas.comm.message import Message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class KafkaAdapterComm(CommBaseAdapter):
|
|
"""
|
|
Kafka 通信适配器
|
|
URL 格式: kafka://[user:pass@]host[:port][/group_id][?options]
|
|
示例: kafka://broker:9092/my-group?auto_offset_reset=earliest&acks=all
|
|
"""
|
|
|
|
def __init__(self, config: ConnectionConfig):
|
|
super().__init__(config)
|
|
self._producer: Optional[AIOKafkaProducer] = None
|
|
self._consumers: dict[str, AIOKafkaConsumer] = {}
|
|
self._consume_tasks: dict[str, asyncio.Task] = {}
|
|
|
|
@property
|
|
def _bootstrap_servers(self) -> str:
|
|
return f"{self.config.host}:{self.config.port}"
|
|
|
|
@property
|
|
def _group_id(self) -> str:
|
|
return self.config.database or "default-group"
|
|
|
|
# ── 生命周期 ──────────────────────────────────────────────────
|
|
|
|
async def connect(self) -> None:
|
|
opts = self.config.options
|
|
self._producer = AIOKafkaProducer(
|
|
bootstrap_servers = self._bootstrap_servers,
|
|
acks = opts.get("acks", "all"),
|
|
compression_type = opts.get("compression_type", None),
|
|
value_serializer = lambda v: json.dumps(v, ensure_ascii=False).encode(),
|
|
)
|
|
await self._producer.start()
|
|
self._running = True
|
|
self.logger.info(f"✅ Kafka Producer 已连接: {self._bootstrap_servers}")
|
|
|
|
async def disconnect(self) -> None:
|
|
self._running = False
|
|
# 停止所有消费任务
|
|
for topic, task in self._consume_tasks.items():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
# 停止所有消费者
|
|
for consumer in self._consumers.values():
|
|
await consumer.stop()
|
|
# 停止生产者
|
|
if self._producer:
|
|
await self._producer.stop()
|
|
self.logger.info("🔌 Kafka 连接已关闭")
|
|
|
|
# ── 核心功能 ──────────────────────────────────────────────────
|
|
|
|
async def publish(self, message: Message) -> bool:
|
|
if not self._producer:
|
|
raise RuntimeError("Kafka 未连接")
|
|
try:
|
|
await self._producer.send_and_wait(
|
|
message.topic,
|
|
value = message.to_dict(),
|
|
headers = [(k, v.encode()) for k, v in message.headers.items()],
|
|
)
|
|
self.logger.debug(f"📤 Kafka 发布 topic={message.topic!r}")
|
|
return True
|
|
except KafkaError as e:
|
|
self.logger.error(f"Kafka 发布失败: {e}")
|
|
return False
|
|
|
|
async def subscribe(self, topic: str, handler: MessageHandler) -> None:
|
|
opts = self.config.options
|
|
self._register_handler(topic, handler)
|
|
|
|
if topic not in self._consumers:
|
|
consumer = AIOKafkaConsumer(
|
|
topic,
|
|
bootstrap_servers = self._bootstrap_servers,
|
|
group_id = self._group_id,
|
|
auto_offset_reset = opts.get("auto_offset_reset", "latest"),
|
|
enable_auto_commit = opts.get("enable_auto_commit", "true") == "true",
|
|
value_deserializer = lambda v: json.loads(v.decode()),
|
|
)
|
|
await consumer.start()
|
|
self._consumers[topic] = consumer
|
|
self._consume_tasks[topic] = asyncio.create_task(
|
|
self._consume_loop(topic, consumer)
|
|
)
|
|
self.logger.info(f"📥 订阅 Kafka topic: {topic!r} group={self._group_id}")
|
|
|
|
async def unsubscribe(self, topic: str) -> None:
|
|
if topic in self._consume_tasks:
|
|
self._consume_tasks[topic].cancel()
|
|
del self._consume_tasks[topic]
|
|
if topic in self._consumers:
|
|
await self._consumers[topic].stop()
|
|
del self._consumers[topic]
|
|
self._handlers.pop(topic, None)
|
|
self.logger.info(f"🚫 取消订阅 Kafka topic: {topic!r}")
|
|
|
|
# ── 内部消费循环 ──────────────────────────────────────────────
|
|
|
|
async def _consume_loop(self, topic: str, consumer: AIOKafkaConsumer):
|
|
self.logger.info(f"🔄 Kafka 消费循环启动: {topic!r}")
|
|
try:
|
|
async for record in consumer:
|
|
try:
|
|
message = Message.from_dict(record.value)
|
|
await self._dispatch(message)
|
|
except Exception as e:
|
|
self.logger.error(f"消息处理异常: {e}", exc_info=True)
|
|
except asyncio.CancelledError:
|
|
self.logger.info(f"🛑 Kafka 消费循环停止: {topic!r}")
|
|
|
|
async def health_check(self) -> bool:
|
|
try:
|
|
return self._producer is not None and self._running
|
|
except Exception:
|
|
return False |