ems-core/backend/app/services/quota_checker.py

"""配额检测服务 - 计算配额使用率，超限时生成告警事件"""
import logging
from datetime import datetime, timezone, timedelta
from sqlalchemy import select, func, and_
from sqlalchemy.ext.asyncio import AsyncSession
from app.models.quota import EnergyQuota, QuotaUsage
from app.models.alarm import AlarmEvent
from app.models.energy import EnergyDailySummary
from app.hooks import get_hooks

logger = logging.getLogger("quota_checker")


def _get_period_range(period: str, now: datetime) -> tuple[datetime, datetime]:
    """根据配额周期计算当前统计区间"""
    if period == "monthly":
        start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        # 下月1号
        if now.month == 12:
            end = start.replace(year=now.year + 1, month=1)
        else:
            end = start.replace(month=now.month + 1)
    else:  # yearly
        start = now.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0)
        end = start.replace(year=now.year + 1)
    return start, end


async def check_quotas(session: AsyncSession):
    """主配额检测循环，计算每个活跃配额的使用率并更新记录"""
    now = datetime.now(timezone.utc)

    result = await session.execute(
        select(EnergyQuota).where(EnergyQuota.is_active == True)
    )
    quotas = result.scalars().all()

    hooks = get_hooks()
    for quota in quotas:
        period_start, period_end = _get_period_range(quota.period, now)

        # 从 EnergyDailySummary 汇总实际用量
        # target_id 对应 device_groups，这里按 device_id 关联
        # 简化处理：按 energy_type 汇总所有匹配设备的消耗
        usage_query = select(func.coalesce(func.sum(EnergyDailySummary.total_consumption), 0)).where(
            and_(
                EnergyDailySummary.energy_type == quota.energy_type,
                EnergyDailySummary.date >= period_start,
                EnergyDailySummary.date < period_end,
            )
        )
        actual_value = (await session.execute(usage_query)).scalar() or 0

        # 计算使用率
        usage_rate_pct = (actual_value / quota.quota_value * 100) if quota.quota_value > 0 else 0

        # 确定状态
        if usage_rate_pct >= quota.alert_threshold_pct:
            status = "exceeded"
        elif usage_rate_pct >= quota.warning_threshold_pct:
            status = "warning"
        else:
            status = "normal"

        # 更新或创建 QuotaUsage 记录
        existing_result = await session.execute(
            select(QuotaUsage).where(
                and_(
                    QuotaUsage.quota_id == quota.id,
                    QuotaUsage.period_start == period_start,
                    QuotaUsage.period_end == period_end,
                )
            )
        )
        usage_record = existing_result.scalar_one_or_none()

        if usage_record:
            usage_record.actual_value = actual_value
            usage_record.usage_rate_pct = usage_rate_pct
            usage_record.status = status
            usage_record.calculated_at = now
        else:
            usage_record = QuotaUsage(
                quota_id=quota.id,
                period_start=period_start,
                period_end=period_end,
                actual_value=actual_value,
                quota_value=quota.quota_value,
                usage_rate_pct=usage_rate_pct,
                status=status,
            )
            session.add(usage_record)

        # 超过预警阈值时生成告警事件
        if status in ("warning", "exceeded"):
            # 检查是否已存在未解决的同配额告警
            active_alarm = await session.execute(
                select(AlarmEvent).where(
                    and_(
                        AlarmEvent.title == f"配额预警: {quota.name}",
                        AlarmEvent.status.in_(["active", "acknowledged"]),
                    )
                )
            )
            if not active_alarm.scalar_one_or_none():
                severity = "critical" if status == "exceeded" else "warning"
                event = AlarmEvent(
                    rule_id=None,
                    device_id=quota.target_id,
                    severity=severity,
                    title=f"配额预警: {quota.name}",
                    description=f"当前使用 {actual_value:.1f}{quota.unit}，"
                                f"配额 {quota.quota_value:.1f}{quota.unit}，"
                                f"使用率 {usage_rate_pct:.1f}%",
                    value=actual_value,
                    threshold=quota.quota_value,
                    status="active",
                    triggered_at=now,
                )
                session.add(event)
                logger.info(
                    f"Quota alert: {quota.name} | usage={actual_value:.1f} "
                    f"quota={quota.quota_value:.1f} rate={usage_rate_pct:.1f}%"
                )

                # Customer hook: on_quota_exceeded
                try:
                    await hooks.on_quota_exceeded(quota, usage_record, session)
                except Exception as _he:
                    logger.error(f"Hook on_quota_exceeded error: {_he}")

    await session.flush()