|
|
|
|
@ -0,0 +1,276 @@
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
|
|
import json
|
|
|
|
|
import logging
|
|
|
|
|
import os
|
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
from typing import Any, Callable, Optional, Union
|
|
|
|
|
from iti.applications.service.iot.alert import add_endpoint_alert_log, add_node_alert_log
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
from rocketmq.client import ConsumeStatus, Message, Producer, PushConsumer
|
|
|
|
|
except Exception:
|
|
|
|
|
ConsumeStatus = None
|
|
|
|
|
Message = None
|
|
|
|
|
Producer = None
|
|
|
|
|
PushConsumer = None
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
RocketMQBody = Union[str, bytes, bytearray, dict, list]
|
|
|
|
|
ConsumeCallback = Callable[[Any], Any]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
|
|
|
class RocketMQConfig:
|
|
|
|
|
namesrv_addr: str
|
|
|
|
|
producer_group: str
|
|
|
|
|
consumer_group: str
|
|
|
|
|
default_topic: str
|
|
|
|
|
default_tags: str
|
|
|
|
|
charset: str
|
|
|
|
|
access_key: str
|
|
|
|
|
secret_key: str
|
|
|
|
|
security_token: str
|
|
|
|
|
instance_name: str
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def from_env() -> "RocketMQConfig":
|
|
|
|
|
return RocketMQConfig(
|
|
|
|
|
namesrv_addr=os.getenv("ROCKETMQ_NAMESRV_ADDR", "127.0.0.1:9876").strip(),
|
|
|
|
|
producer_group=os.getenv("ROCKETMQ_PRODUCER_GROUP", "iot-collect-group").strip(),
|
|
|
|
|
consumer_group=os.getenv("ROCKETMQ_CONSUMER_GROUP", "iot-collect-group").strip(),
|
|
|
|
|
default_topic=os.getenv("ROCKETMQ_TOPIC", "iot-collect-topic").strip(),
|
|
|
|
|
default_tags=os.getenv("ROCKETMQ_TAGS", "*").strip() or "*",
|
|
|
|
|
charset=os.getenv("ROCKETMQ_CHARSET", "utf-8").strip() or "utf-8",
|
|
|
|
|
access_key=os.getenv("ROCKETMQ_ACCESS_KEY", "").strip(),
|
|
|
|
|
secret_key=os.getenv("ROCKETMQ_SECRET_KEY", "").strip(),
|
|
|
|
|
security_token=os.getenv("ROCKETMQ_SECURITY_TOKEN", "").strip(),
|
|
|
|
|
instance_name=os.getenv("ROCKETMQ_INSTANCE_NAME", "").strip(),
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RocketMQMgr:
|
|
|
|
|
_instance = None
|
|
|
|
|
_initialized = False
|
|
|
|
|
|
|
|
|
|
def __new__(cls):
|
|
|
|
|
if cls._instance is None:
|
|
|
|
|
cls._instance = super().__new__(cls)
|
|
|
|
|
return cls._instance
|
|
|
|
|
|
|
|
|
|
def init_app(self, app) -> None:
|
|
|
|
|
if self._initialized:
|
|
|
|
|
logger.warning("rocketmq连接已初始化，跳过重复初始化")
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
if Producer is None or PushConsumer is None or Message is None:
|
|
|
|
|
logger.error(
|
|
|
|
|
"rocketmq客户端库不可用，请安装 rocketmq-client-python（并确保本机可加载 librocketmq）"
|
|
|
|
|
)
|
|
|
|
|
self._initialized = False
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
cfg = RocketMQConfig.from_env()
|
|
|
|
|
|
|
|
|
|
if not cfg.namesrv_addr:
|
|
|
|
|
logger.error("rocketmq配置不完整，缺少: ROCKETMQ_NAMESRV_ADDR")
|
|
|
|
|
self._initialized = False
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
self.cfg = cfg
|
|
|
|
|
self._app = app
|
|
|
|
|
self.producer: Optional[Any] = None
|
|
|
|
|
self.consumer: Optional[Any] = None
|
|
|
|
|
self._producer_started = False
|
|
|
|
|
self._consumer_started = False
|
|
|
|
|
|
|
|
|
|
self._initialized = True
|
|
|
|
|
logger.info("rocketmq初始化成功")
|
|
|
|
|
|
|
|
|
|
def _ensure_initialized(self) -> bool:
|
|
|
|
|
if not self._initialized:
|
|
|
|
|
logger.error("rocketmq未初始化，请先调用 iot_rocketmq.init_app(app)")
|
|
|
|
|
return False
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
def _apply_credentials_if_supported(self, client: Any) -> None:
|
|
|
|
|
if not self.cfg.access_key or not self.cfg.secret_key:
|
|
|
|
|
return
|
|
|
|
|
if hasattr(client, "set_session_credentials"):
|
|
|
|
|
try:
|
|
|
|
|
client.set_session_credentials(
|
|
|
|
|
self.cfg.access_key, self.cfg.secret_key, self.cfg.security_token
|
|
|
|
|
)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logger.warning(f"rocketmq设置ACL凭证失败: {e}")
|
|
|
|
|
|
|
|
|
|
def start_producer(self) -> bool:
|
|
|
|
|
if not self._ensure_initialized():
|
|
|
|
|
return False
|
|
|
|
|
if self._producer_started and self.producer:
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
producer = Producer(self.cfg.producer_group)
|
|
|
|
|
producer.set_name_server_address(self.cfg.namesrv_addr)
|
|
|
|
|
if self.cfg.instance_name and hasattr(producer, "set_instance_name"):
|
|
|
|
|
producer.set_instance_name(self.cfg.instance_name)
|
|
|
|
|
self._apply_credentials_if_supported(producer)
|
|
|
|
|
producer.start()
|
|
|
|
|
self.producer = producer
|
|
|
|
|
self._producer_started = True
|
|
|
|
|
logger.info("rocketmq producer启动成功")
|
|
|
|
|
return True
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logger.error(f"rocketmq producer启动失败: {e}", exc_info=True)
|
|
|
|
|
self.producer = None
|
|
|
|
|
self._producer_started = False
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
def shutdown_producer(self) -> None:
|
|
|
|
|
if not self.producer:
|
|
|
|
|
return
|
|
|
|
|
try:
|
|
|
|
|
self.producer.shutdown()
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logger.warning(f"rocketmq producer关闭失败: {e}")
|
|
|
|
|
finally:
|
|
|
|
|
self.producer = None
|
|
|
|
|
self._producer_started = False
|
|
|
|
|
|
|
|
|
|
def _encode_body(self, body: RocketMQBody) -> bytes:
|
|
|
|
|
if isinstance(body, bytes):
|
|
|
|
|
return body
|
|
|
|
|
if isinstance(body, bytearray):
|
|
|
|
|
return bytes(body)
|
|
|
|
|
if isinstance(body, str):
|
|
|
|
|
return body.encode(self.cfg.charset, errors="replace")
|
|
|
|
|
return json.dumps(body, ensure_ascii=False).encode(self.cfg.charset, errors="replace")
|
|
|
|
|
|
|
|
|
|
def send_sync(
|
|
|
|
|
self,
|
|
|
|
|
body: RocketMQBody,
|
|
|
|
|
topic: Optional[str] = None,
|
|
|
|
|
tags: Optional[str] = None,
|
|
|
|
|
keys: Optional[str] = None,
|
|
|
|
|
) -> Any:
|
|
|
|
|
if not self._ensure_initialized():
|
|
|
|
|
raise RuntimeError("rocketmq未初始化")
|
|
|
|
|
if not self.start_producer() or not self.producer:
|
|
|
|
|
raise RuntimeError("rocketmq producer未启动")
|
|
|
|
|
|
|
|
|
|
real_topic = (topic or self.cfg.default_topic).strip()
|
|
|
|
|
if not real_topic:
|
|
|
|
|
logger.error("topic不能为空（可通过参数传入或设置ROCKETMQ_TOPIC）")
|
|
|
|
|
raise ValueError("topic不能为空（可通过参数传入或设置ROCKETMQ_TOPIC）")
|
|
|
|
|
|
|
|
|
|
msg = Message(real_topic)
|
|
|
|
|
if tags:
|
|
|
|
|
if hasattr(msg, "set_tags"):
|
|
|
|
|
msg.set_tags(tags)
|
|
|
|
|
if keys:
|
|
|
|
|
if hasattr(msg, "set_keys"):
|
|
|
|
|
msg.set_keys(keys)
|
|
|
|
|
msg.set_body(self._encode_body(body))
|
|
|
|
|
return self.producer.send_sync(msg)
|
|
|
|
|
|
|
|
|
|
def start_consumer(
|
|
|
|
|
self,
|
|
|
|
|
topic: Optional[str] = None,
|
|
|
|
|
tags: Optional[str] = None,
|
|
|
|
|
callback: Optional[ConsumeCallback] = None,
|
|
|
|
|
) -> bool:
|
|
|
|
|
if not self._ensure_initialized():
|
|
|
|
|
return False
|
|
|
|
|
if self._consumer_started and self.consumer:
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
real_topic = (topic or self.cfg.default_topic).strip()
|
|
|
|
|
if not real_topic:
|
|
|
|
|
logger.error("consumer topic不能为空（可通过参数传入或设置ROCKETMQ_TOPIC）")
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
selector_expression = (tags or self.cfg.default_tags).strip() or "*"
|
|
|
|
|
consume_cb = callback or self._default_consume_callback
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
consumer = PushConsumer(self.cfg.consumer_group)
|
|
|
|
|
consumer.set_name_server_address(self.cfg.namesrv_addr)
|
|
|
|
|
if self.cfg.instance_name and hasattr(consumer, "set_instance_name"):
|
|
|
|
|
consumer.set_instance_name(self.cfg.instance_name)
|
|
|
|
|
self._apply_credentials_if_supported(consumer)
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
consumer.subscribe(real_topic, consume_cb, selector_expression)
|
|
|
|
|
except TypeError:
|
|
|
|
|
consumer.subscribe(real_topic, selector_expression, consume_cb)
|
|
|
|
|
|
|
|
|
|
consumer.start()
|
|
|
|
|
self.consumer = consumer
|
|
|
|
|
self._consumer_started = True
|
|
|
|
|
logger.info("rocketmq consumer启动成功")
|
|
|
|
|
return True
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logger.error(f"rocketmq consumer启动失败: {e}", exc_info=True)
|
|
|
|
|
self.consumer = None
|
|
|
|
|
self._consumer_started = False
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
def shutdown_consumer(self) -> None:
|
|
|
|
|
if not self.consumer:
|
|
|
|
|
return
|
|
|
|
|
try:
|
|
|
|
|
self.consumer.shutdown()
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logger.warning(f"rocketmq consumer关闭失败: {e}")
|
|
|
|
|
finally:
|
|
|
|
|
self.consumer = None
|
|
|
|
|
self._consumer_started = False
|
|
|
|
|
|
|
|
|
|
def _default_consume_callback(self, msg: Any) -> Any:
|
|
|
|
|
topic = getattr(msg, "topic", "")
|
|
|
|
|
tags = getattr(msg, "tags", "")
|
|
|
|
|
keys = getattr(msg, "keys", "")
|
|
|
|
|
body_text = getattr(msg, "body", b"").decode(self.cfg.charset, errors="replace")
|
|
|
|
|
|
|
|
|
|
# logger.info(
|
|
|
|
|
# f"[RocketMQ] topic={topic} tags={tags} keys={keys} body={body_text}"
|
|
|
|
|
# )
|
|
|
|
|
try:
|
|
|
|
|
body_dict = json.loads(body_text)
|
|
|
|
|
except Exception:
|
|
|
|
|
logger.error(f"[RocketMQ] 解析body失败: {body_text}")
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
app = getattr(self, "_app", None)
|
|
|
|
|
if app is None:
|
|
|
|
|
logger.error("[RocketMQ] 未绑定Flask app，无法创建application context")
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
with app.app_context():
|
|
|
|
|
if tags == b"netData" or tags == "netData":
|
|
|
|
|
add_endpoint_alert_log(body_dict["endpointId"])
|
|
|
|
|
|
|
|
|
|
elif tags == b"collectData" or tags == "collectData":
|
|
|
|
|
for node_data in body_dict["nodeDataList"]:
|
|
|
|
|
node_id = int(node_data["nodeId"])
|
|
|
|
|
alert_value = node_data["alertValue"]
|
|
|
|
|
add_node_alert_log(node_id, alert_value)
|
|
|
|
|
|
|
|
|
|
if ConsumeStatus is not None and hasattr(ConsumeStatus, "CONSUME_SUCCESS"):
|
|
|
|
|
return ConsumeStatus.CONSUME_SUCCESS
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
def shutdown(self) -> None:
|
|
|
|
|
self.shutdown_consumer()
|
|
|
|
|
self.shutdown_producer()
|
|
|
|
|
|
|
|
|
|
def __del__(self):
|
|
|
|
|
try:
|
|
|
|
|
self.shutdown()
|
|
|
|
|
except Exception:
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
iot_rocketmq = RocketMQMgr()
|
DESKTOP-1JS6RSM\Admin
3 months ago