638 lines
20 KiB
Python
638 lines
20 KiB
Python
"""存储抽象层。
|
||
|
||
- RuleStore / EventStore / AlertStore：MySQL 语义（准），默认 InMemory 实现；
|
||
- Cache：Redis 语义（get/set NX/INCR/EXPIRE），默认 InMemory 实现；
|
||
- MessageBus：Kafka 语义（publish），默认 InMemory 实现。
|
||
|
||
真实中间件适配器放在同文件末尾，采用惰性导入（可选依赖），未安装驱动时给出
|
||
清晰错误；测试与本地运行均使用 InMemory 实现，保证零依赖可编译、可运行。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import itertools
|
||
import json
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
from .models import Alert, Event, MetricRule
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 接口定义
|
||
# --------------------------------------------------------------------------- #
|
||
class RuleStore(ABC):
|
||
"""规则仓库（Source of Truth：MySQL metric_rule + rule_scope）。"""
|
||
|
||
@abstractmethod
|
||
def list_rules(self) -> List[MetricRule]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def get_rule(self, rule_id: str) -> Optional[MetricRule]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def create_rule(self, rule: MetricRule) -> MetricRule:
|
||
...
|
||
|
||
@abstractmethod
|
||
def update_rule(self, rule: MetricRule) -> MetricRule:
|
||
...
|
||
|
||
@abstractmethod
|
||
def delete_rule(self, rule_id: str) -> bool:
|
||
...
|
||
|
||
@abstractmethod
|
||
def get_version(self) -> int:
|
||
...
|
||
|
||
|
||
class EventStore(ABC):
|
||
"""事件仓库（MySQL event 表）。"""
|
||
|
||
@abstractmethod
|
||
def insert(self, event: Event) -> None:
|
||
...
|
||
|
||
@abstractmethod
|
||
def list_events(
|
||
self,
|
||
host_id: Optional[str] = None,
|
||
rule_id: Optional[str] = None,
|
||
status: Optional[str] = None,
|
||
severity: Optional[str] = None,
|
||
start: Optional[float] = None,
|
||
end: Optional[float] = None,
|
||
page: int = 1,
|
||
page_size: int = 20,
|
||
) -> Tuple[int, List[Event]]:
|
||
...
|
||
|
||
|
||
class AlertStore(ABC):
|
||
"""告警仓库（MySQL alert 表）。"""
|
||
|
||
@abstractmethod
|
||
def upsert_firing(self, alert: Alert) -> Alert:
|
||
...
|
||
|
||
@abstractmethod
|
||
def get(self, alert_id: str) -> Optional[Alert]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def list_alerts(
|
||
self,
|
||
severity: Optional[str] = None,
|
||
status: Optional[str] = None,
|
||
ack_status: Optional[str] = None,
|
||
start: Optional[float] = None,
|
||
end: Optional[float] = None,
|
||
page: int = 1,
|
||
page_size: int = 20,
|
||
) -> Tuple[int, List[Alert]]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def ack(self, alert_id: str, ack_by: str = "system") -> Optional[Alert]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def close(self, alert_id: str) -> Optional[Alert]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def resolve_by_dedup_key(self, dedup_key: str, at: float) -> Optional[Alert]:
|
||
...
|
||
|
||
|
||
class Cache(ABC):
|
||
"""缓存抽象（Redis 语义子集）。"""
|
||
|
||
@abstractmethod
|
||
def get(self, key: str) -> Optional[str]:
|
||
...
|
||
|
||
@abstractmethod
|
||
def set(self, key: str, value: str, ttl: Optional[int] = None) -> bool:
|
||
...
|
||
|
||
@abstractmethod
|
||
def set_nx(self, key: str, value: str, ttl: Optional[int] = None) -> bool:
|
||
...
|
||
|
||
@abstractmethod
|
||
def incr(self, key: str, ttl: Optional[int] = None) -> int:
|
||
...
|
||
|
||
@abstractmethod
|
||
def delete(self, key: str) -> bool:
|
||
...
|
||
|
||
|
||
class MessageBus(ABC):
|
||
"""消息总线抽象（Kafka 语义子集）。"""
|
||
|
||
@abstractmethod
|
||
def publish(self, topic: str, key: str, value: Dict[str, Any]) -> None:
|
||
...
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# InMemory 实现（线程安全）
|
||
# --------------------------------------------------------------------------- #
|
||
class InMemoryRuleStore(RuleStore):
|
||
def __init__(self) -> None:
|
||
self._rules: Dict[str, MetricRule] = {}
|
||
self._version = 0
|
||
self._lock = threading.RLock()
|
||
|
||
def list_rules(self) -> List[MetricRule]:
|
||
with self._lock:
|
||
return list(self._rules.values())
|
||
|
||
def get_rule(self, rule_id: str) -> Optional[MetricRule]:
|
||
with self._lock:
|
||
return self._rules.get(rule_id)
|
||
|
||
def create_rule(self, rule: MetricRule) -> MetricRule:
|
||
with self._lock:
|
||
if rule.rule_id in self._rules:
|
||
raise ValueError(f"rule {rule.rule_id} already exists")
|
||
rule.version = rule.version + 1
|
||
self._rules[rule.rule_id] = rule
|
||
self._version += 1
|
||
return rule
|
||
|
||
def update_rule(self, rule: MetricRule) -> MetricRule:
|
||
with self._lock:
|
||
if rule.rule_id not in self._rules:
|
||
raise ValueError(f"rule {rule.rule_id} not found")
|
||
rule.version = self._rules[rule.rule_id].version + 1
|
||
self._rules[rule.rule_id] = rule
|
||
self._version += 1
|
||
return rule
|
||
|
||
def delete_rule(self, rule_id: str) -> bool:
|
||
with self._lock:
|
||
if rule_id not in self._rules:
|
||
return False
|
||
del self._rules[rule_id]
|
||
self._version += 1
|
||
return True
|
||
|
||
def get_version(self) -> int:
|
||
with self._lock:
|
||
return self._version
|
||
|
||
|
||
class InMemoryEventStore(EventStore):
|
||
def __init__(self) -> None:
|
||
self._events: List[Event] = []
|
||
self._lock = threading.RLock()
|
||
|
||
def insert(self, event: Event) -> None:
|
||
with self._lock:
|
||
self._events.append(event)
|
||
|
||
def list_events(
|
||
self,
|
||
host_id: Optional[str] = None,
|
||
rule_id: Optional[str] = None,
|
||
status: Optional[str] = None,
|
||
severity: Optional[str] = None,
|
||
start: Optional[float] = None,
|
||
end: Optional[float] = None,
|
||
page: int = 1,
|
||
page_size: int = 20,
|
||
) -> Tuple[int, List[Event]]:
|
||
with self._lock:
|
||
items = [e for e in self._events if self._match_event(e, host_id, rule_id, status, severity, start, end)]
|
||
items.sort(key=lambda e: e.fired_at, reverse=True)
|
||
total = len(items)
|
||
start_idx = (page - 1) * page_size
|
||
return total, items[start_idx : start_idx + page_size]
|
||
|
||
@staticmethod
|
||
def _match_event(
|
||
e: Event,
|
||
host_id: Optional[str],
|
||
rule_id: Optional[str],
|
||
status: Optional[str],
|
||
severity: Optional[str],
|
||
start: Optional[float],
|
||
end: Optional[float],
|
||
) -> bool:
|
||
if host_id and e.host_id != host_id:
|
||
return False
|
||
if rule_id and e.rule_id != rule_id:
|
||
return False
|
||
if status and e.status != status:
|
||
return False
|
||
if severity and e.severity != severity:
|
||
return False
|
||
if start is not None and e.fired_at < start:
|
||
return False
|
||
if end is not None and e.fired_at > end:
|
||
return False
|
||
return True
|
||
|
||
|
||
class InMemoryAlertStore(AlertStore):
|
||
def __init__(self) -> None:
|
||
self._alerts: Dict[str, Alert] = {}
|
||
self._lock = threading.RLock()
|
||
|
||
def upsert_firing(self, alert: Alert) -> Alert:
|
||
with self._lock:
|
||
# 同 aggregate_key 且仍在 firing 的告警做聚合：计数 + 1，更新 last_at 与明细。
|
||
for existing in self._alerts.values():
|
||
if existing.aggregate_key == alert.aggregate_key and existing.status == "firing":
|
||
existing.count += 1
|
||
existing.last_at = alert.last_at
|
||
existing.detail = self._merge_detail(existing.detail, alert.detail)
|
||
return existing
|
||
self._alerts[alert.alert_id] = alert
|
||
return alert
|
||
|
||
@staticmethod
|
||
def _merge_detail(a: Dict[str, Any], b: Dict[str, Any]) -> Dict[str, Any]:
|
||
merged = dict(a)
|
||
hosts = list(dict.fromkeys(list(merged.get("hosts", [])) + list(b.get("hosts", []))))
|
||
merged["hosts"] = hosts
|
||
for key in ("current", "threshold", "metric"):
|
||
if key in b:
|
||
merged[key] = b[key]
|
||
return merged
|
||
|
||
def get(self, alert_id: str) -> Optional[Alert]:
|
||
with self._lock:
|
||
return self._alerts.get(alert_id)
|
||
|
||
def list_alerts(
|
||
self,
|
||
severity: Optional[str] = None,
|
||
status: Optional[str] = None,
|
||
ack_status: Optional[str] = None,
|
||
start: Optional[float] = None,
|
||
end: Optional[float] = None,
|
||
page: int = 1,
|
||
page_size: int = 20,
|
||
) -> Tuple[int, List[Alert]]:
|
||
with self._lock:
|
||
items = [
|
||
a
|
||
for a in self._alerts.values()
|
||
if self._match_alert(a, severity, status, ack_status, start, end)
|
||
]
|
||
items.sort(key=lambda a: a.last_at, reverse=True)
|
||
total = len(items)
|
||
start_idx = (page - 1) * page_size
|
||
return total, items[start_idx : start_idx + page_size]
|
||
|
||
@staticmethod
|
||
def _match_alert(
|
||
a: Alert,
|
||
severity: Optional[str],
|
||
status: Optional[str],
|
||
ack_status: Optional[str],
|
||
start: Optional[float],
|
||
end: Optional[float],
|
||
) -> bool:
|
||
if severity and a.severity != severity:
|
||
return False
|
||
if status and a.status != status:
|
||
return False
|
||
if ack_status and a.ack_status != ack_status:
|
||
return False
|
||
if start is not None and a.last_at < start:
|
||
return False
|
||
if end is not None and a.last_at > end:
|
||
return False
|
||
return True
|
||
|
||
def ack(self, alert_id: str, ack_by: str = "system") -> Optional[Alert]:
|
||
with self._lock:
|
||
a = self._alerts.get(alert_id)
|
||
if a is None:
|
||
return None
|
||
a.ack_status = "acked"
|
||
a.ack_by = ack_by
|
||
a.ack_at = time.time()
|
||
return a
|
||
|
||
def close(self, alert_id: str) -> Optional[Alert]:
|
||
with self._lock:
|
||
a = self._alerts.get(alert_id)
|
||
if a is None:
|
||
return None
|
||
a.ack_status = "closed"
|
||
return a
|
||
|
||
def resolve_by_dedup_key(self, dedup_key: str, at: float) -> Optional[Alert]:
|
||
with self._lock:
|
||
for a in self._alerts.values():
|
||
if a.dedup_key == dedup_key and a.status == "firing":
|
||
a.status = "resolved"
|
||
a.last_at = at
|
||
return a
|
||
return None
|
||
|
||
|
||
class InMemoryCache(Cache):
|
||
def __init__(self) -> None:
|
||
self._data: Dict[str, Tuple[str, Optional[float]]] = {}
|
||
self._lock = threading.RLock()
|
||
|
||
def _now(self) -> float:
|
||
return time.monotonic()
|
||
|
||
def _purge(self, key: str) -> None:
|
||
if key in self._data:
|
||
value, exp = self._data[key]
|
||
if exp is not None and exp <= self._now():
|
||
del self._data[key]
|
||
|
||
def get(self, key: str) -> Optional[str]:
|
||
with self._lock:
|
||
self._purge(key)
|
||
item = self._data.get(key)
|
||
return item[0] if item else None
|
||
|
||
def set(self, key: str, value: str, ttl: Optional[int] = None) -> bool:
|
||
with self._lock:
|
||
exp = self._now() + ttl if ttl else None
|
||
self._data[key] = (value, exp)
|
||
return True
|
||
|
||
def set_nx(self, key: str, value: str, ttl: Optional[int] = None) -> bool:
|
||
with self._lock:
|
||
self._purge(key)
|
||
if key in self._data:
|
||
return False
|
||
exp = self._now() + ttl if ttl else None
|
||
self._data[key] = (value, exp)
|
||
return True
|
||
|
||
def incr(self, key: str, ttl: Optional[int] = None) -> int:
|
||
with self._lock:
|
||
self._purge(key)
|
||
item = self._data.get(key)
|
||
current = int(item[0]) + 1 if item else 1
|
||
exp = self._now() + ttl if ttl else None
|
||
self._data[key] = (str(current), exp)
|
||
return current
|
||
|
||
def delete(self, key: str) -> bool:
|
||
with self._lock:
|
||
if key in self._data:
|
||
del self._data[key]
|
||
return True
|
||
return False
|
||
|
||
|
||
class InMemoryMessageBus(MessageBus):
|
||
def __init__(self) -> None:
|
||
self.messages: List[Tuple[str, str, Dict[str, Any]]] = []
|
||
self._lock = threading.RLock()
|
||
|
||
def publish(self, topic: str, key: str, value: Dict[str, Any]) -> None:
|
||
with self._lock:
|
||
self.messages.append((topic, key, value))
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 工厂函数：根据配置选择驱动（未安装可选依赖时安全回退 memory）
|
||
# --------------------------------------------------------------------------- #
|
||
def _safe(loader, fallback):
|
||
try:
|
||
return loader()
|
||
except Exception:
|
||
return fallback()
|
||
|
||
|
||
def create_rule_store(config) -> RuleStore:
|
||
if config.rule_store_driver == "mysql":
|
||
return _safe(lambda: MySQLRuleStore(config.mysql_dsn), InMemoryRuleStore)
|
||
return InMemoryRuleStore()
|
||
|
||
|
||
def create_event_store(config) -> EventStore:
|
||
if config.event_store_driver == "mysql":
|
||
return _safe(lambda: MySQLEventStore(config.mysql_dsn), InMemoryEventStore)
|
||
return InMemoryEventStore()
|
||
|
||
|
||
def create_alert_store(config) -> AlertStore:
|
||
if config.alert_store_driver == "mysql":
|
||
return _safe(lambda: MySQLAlertStore(config.mysql_dsn), InMemoryAlertStore)
|
||
return InMemoryAlertStore()
|
||
|
||
|
||
def create_cache(config) -> Cache:
|
||
if config.cache_driver == "redis":
|
||
return _safe(lambda: RedisCache(config.redis_url), InMemoryCache)
|
||
return InMemoryCache()
|
||
|
||
|
||
def create_bus(config) -> MessageBus:
|
||
if config.bus_driver == "kafka":
|
||
return _safe(lambda: KafkaBus(config.kafka_bootstrap), InMemoryMessageBus)
|
||
return InMemoryMessageBus()
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 可选中间件适配器（惰性导入；未安装驱动时抛出 RuntimeError）
|
||
# --------------------------------------------------------------------------- #
|
||
class MySQLRuleStore(RuleStore):
|
||
"""基于 pymysql 的规则仓库。需安装：pip install pymysql"""
|
||
|
||
def __init__(self, dsn: str) -> None:
|
||
try:
|
||
import pymysql # noqa: F401
|
||
except ImportError as exc: # pragma: no cover
|
||
raise RuntimeError("MySQLRuleStore requires pymysql") from exc
|
||
self._dsn = dsn
|
||
self._conn = None
|
||
|
||
def _cursor(self):
|
||
import pymysql
|
||
|
||
if self._conn is None:
|
||
self._conn = pymysql.connect(**self._parse_dsn(self._dsn))
|
||
return self._conn.cursor()
|
||
|
||
@staticmethod
|
||
def _parse_dsn(dsn: str) -> Dict[str, Any]:
|
||
# 简化：mysql://user:pass@host:port/db
|
||
import urllib.parse
|
||
|
||
parsed = urllib.parse.urlparse(dsn)
|
||
return {
|
||
"host": parsed.hostname or "127.0.0.1",
|
||
"port": parsed.port or 3306,
|
||
"user": parsed.username or "root",
|
||
"password": parsed.password or "",
|
||
"database": parsed.path.lstrip("/") or "hms",
|
||
}
|
||
|
||
def list_rules(self) -> List[MetricRule]:
|
||
cur = self._cursor()
|
||
cur.execute(
|
||
"SELECT rule_id,name,metric,aggregation,operator,threshold,threshold2,"
|
||
"for_duration,severity,labels,notify_channels,enabled,version "
|
||
"FROM metric_rule"
|
||
)
|
||
rules = []
|
||
for row in cur.fetchall():
|
||
scopes = self._load_scopes(row[0])
|
||
rules.append(
|
||
MetricRule(
|
||
rule_id=row[0], name=row[1], metric=row[2], aggregation=row[3],
|
||
operator=row[4], threshold=float(row[5]),
|
||
threshold2=float(row[6]) if row[6] is not None else None,
|
||
for_duration=row[7], severity=row[8],
|
||
labels=json.loads(row[9]) if row[9] else {},
|
||
notify_channels=json.loads(row[10]) if row[10] else [],
|
||
enabled=bool(row[11]), version=int(row[12]),
|
||
scope=scopes[0] if scopes else RuleScope(),
|
||
)
|
||
)
|
||
return rules
|
||
|
||
def _load_scopes(self, rule_id: str) -> List[RuleScope]:
|
||
cur = self._cursor()
|
||
cur.execute(
|
||
"SELECT scope_type,host_id,host_group,service FROM rule_scope WHERE rule_id=%s",
|
||
(rule_id,),
|
||
)
|
||
scopes = []
|
||
for st, host_id, host_group, service in cur.fetchall():
|
||
scopes.append(
|
||
RuleScope(
|
||
scope_type=st,
|
||
host_ids=[host_id] if host_id else [],
|
||
host_group=host_group or "",
|
||
service=service or "",
|
||
)
|
||
)
|
||
return scopes
|
||
|
||
def get_rule(self, rule_id: str) -> Optional[MetricRule]:
|
||
for r in self.list_rules():
|
||
if r.rule_id == rule_id:
|
||
return r
|
||
return None
|
||
|
||
def create_rule(self, rule: MetricRule) -> MetricRule:
|
||
raise NotImplementedError("MySQL write path 由上层 API 通过事务实现")
|
||
|
||
def update_rule(self, rule: MetricRule) -> MetricRule:
|
||
raise NotImplementedError
|
||
|
||
def delete_rule(self, rule_id: str) -> bool:
|
||
raise NotImplementedError
|
||
|
||
def get_version(self) -> int:
|
||
cur = self._cursor()
|
||
cur.execute("SELECT COALESCE(MAX(version),0) FROM metric_rule")
|
||
return int(cur.fetchone()[0])
|
||
|
||
|
||
class MySQLEventStore(EventStore):
|
||
def __init__(self, dsn: str) -> None:
|
||
try:
|
||
import pymysql # noqa: F401
|
||
except ImportError as exc: # pragma: no cover
|
||
raise RuntimeError("MySQLEventStore requires pymysql") from exc
|
||
self._dsn = dsn
|
||
|
||
def insert(self, event: Event) -> None:
|
||
raise NotImplementedError("接入真实 MySQL 时实现 event 表写入")
|
||
|
||
def list_events(self, **kwargs) -> Tuple[int, List[Event]]:
|
||
raise NotImplementedError
|
||
|
||
|
||
class MySQLAlertStore(AlertStore):
|
||
def __init__(self, dsn: str) -> None:
|
||
try:
|
||
import pymysql # noqa: F401
|
||
except ImportError as exc: # pragma: no cover
|
||
raise RuntimeError("MySQLAlertStore requires pymysql") from exc
|
||
self._dsn = dsn
|
||
|
||
def upsert_firing(self, alert: Alert) -> Alert:
|
||
raise NotImplementedError("接入真实 MySQL 时实现 alert 表 upsert")
|
||
|
||
def get(self, alert_id: str) -> Optional[Alert]:
|
||
raise NotImplementedError
|
||
|
||
def list_alerts(self, **kwargs) -> Tuple[int, List[Alert]]:
|
||
raise NotImplementedError
|
||
|
||
def ack(self, alert_id: str, ack_by: str = "system") -> Optional[Alert]:
|
||
raise NotImplementedError
|
||
|
||
def close(self, alert_id: str) -> Optional[Alert]:
|
||
raise NotImplementedError
|
||
|
||
def resolve_by_dedup_key(self, dedup_key: str, at: float) -> Optional[Alert]:
|
||
raise NotImplementedError
|
||
|
||
|
||
class RedisCache(Cache):
|
||
"""基于 redis-py 的缓存。需安装：pip install redis"""
|
||
|
||
def __init__(self, url: str) -> None:
|
||
try:
|
||
import redis # noqa: F401
|
||
except ImportError as exc: # pragma: no cover
|
||
raise RuntimeError("RedisCache requires redis") from exc
|
||
import redis
|
||
|
||
self._client = redis.Redis.from_url(url, decode_responses=True)
|
||
|
||
def get(self, key: str) -> Optional[str]:
|
||
return self._client.get(key)
|
||
|
||
def set(self, key: str, value: str, ttl: Optional[int] = None) -> bool:
|
||
return bool(self._client.set(key, value, ex=ttl))
|
||
|
||
def set_nx(self, key: str, value: str, ttl: Optional[int] = None) -> bool:
|
||
return bool(self._client.set(key, value, ex=ttl, nx=True))
|
||
|
||
def incr(self, key: str, ttl: Optional[int] = None) -> int:
|
||
value = self._client.incr(key)
|
||
if ttl and value == 1:
|
||
self._client.expire(key, ttl)
|
||
return int(value)
|
||
|
||
def delete(self, key: str) -> bool:
|
||
return bool(self._client.delete(key))
|
||
|
||
|
||
class KafkaBus(MessageBus):
|
||
"""基于 kafka-python 的消息总线。需安装：pip install kafka-python"""
|
||
|
||
def __init__(self, bootstrap: str) -> None:
|
||
try:
|
||
from kafka import KafkaProducer # noqa: F401
|
||
except ImportError as exc: # pragma: no cover
|
||
raise RuntimeError("KafkaBus requires kafka-python") from exc
|
||
from kafka import KafkaProducer
|
||
|
||
self._producer = KafkaProducer(
|
||
bootstrap_servers=bootstrap,
|
||
value_serializer=lambda v: json.dumps(v).encode("utf-8"),
|
||
key_serializer=lambda k: k.encode("utf-8"),
|
||
)
|
||
|
||
def publish(self, topic: str, key: str, value: Dict[str, Any]) -> None:
|
||
self._producer.send(topic, key=key, value=value)
|