AirGate/backend/utils/scheduler.py

"""定时消费监控任务"""

import logging
from datetime import datetime
from decimal import Decimal

logger = logging.getLogger(__name__)

_scheduler_started = False


def check_spending():
    """定时检查所有子账号消费"""
    from apps.monitor.models import VolcAccount, IAMUser, GlobalConfig, AlertRecord
    from utils.crypto import decrypt
    from utils.billing_service import BillingService
    from utils.iam_service import IAMService
    from utils.feishu import send_feishu_alert

    bill_period = datetime.now().strftime("%Y-%m")
    config = GlobalConfig.get_solo()

    for volc_account in VolcAccount.objects.filter(is_active=True):
        ak = decrypt(volc_account.access_key_enc)
        sk = decrypt(volc_account.secret_key_enc)
        if not ak or not sk:
            logger.warning(f"主账号 {volc_account.name} 密钥为空，跳过")
            continue

        billing = BillingService(ak, sk)
        iam_svc = IAMService(ak, sk)

        users = IAMUser.objects.filter(
            volc_account=volc_account,
            monitor_enabled=True,
        ).exclude(status=IAMUser.Status.DISABLED)

        for user in users:
            try:
                spending = billing.get_spending_by_project(
                    bill_period, user.project_name or None
                )
                user.current_month_spending = spending
                user.spending_updated_at = datetime.now()
                user.save(update_fields=['current_month_spending', 'spending_updated_at'])

                disable_threshold = user.get_disable_threshold()
                alert_threshold = user.get_alert_threshold()

                # Check disable threshold
                if (user.auto_disable_enabled
                        and disable_threshold
                        and spending >= disable_threshold):

                    already_disabled = AlertRecord.objects.filter(
                        iam_user=user,
                        alert_type=AlertRecord.AlertType.DISABLE,
                        created_at__month=datetime.now().month,
                        created_at__year=datetime.now().year,
                    ).exists()

                    if not already_disabled:
                        try:
                            iam_svc.disable_user(user.username)
                            user.status = IAMUser.Status.DISABLED
                            user.save(update_fields=['status'])
                        except Exception as e:
                            logger.error(f"停用用户 {user.username} 失败: {e}")

                        alert = AlertRecord.objects.create(
                            iam_user=user,
                            alert_type=AlertRecord.AlertType.DISABLE,
                            title=f"子账号 {user.username} 已自动停用",
                            content=f"本月消费 ¥{spending:.2f}，达到停用阈值 ¥{disable_threshold:.2f}",
                            spending_amount=spending,
                            threshold_amount=disable_threshold,
                        )
                        webhook = config.feishu_webhook_url
                        send_feishu_alert(
                            webhook,
                            "🚨 子账号已自动停用",
                            f"**用户**: {user.username}\n"
                            f"**消费**: ¥{spending:.2f}\n"
                            f"**阈值**: ¥{disable_threshold:.2f}\n"
                            f"如需恢复，请在 AirGate 管理后台操作。",
                            template="red",
                        )
                        alert.notified = True
                        alert.save(update_fields=['notified'])

                # Check alert threshold
                elif alert_threshold and spending >= alert_threshold:
                    already_alerted = AlertRecord.objects.filter(
                        iam_user=user,
                        alert_type=AlertRecord.AlertType.WARNING,
                        created_at__month=datetime.now().month,
                        created_at__year=datetime.now().year,
                    ).exists()

                    if not already_alerted:
                        alert = AlertRecord.objects.create(
                            iam_user=user,
                            alert_type=AlertRecord.AlertType.WARNING,
                            title=f"子账号 {user.username} 消费告警",
                            content=f"本月消费 ¥{spending:.2f}，达到告警阈值 ¥{alert_threshold:.2f}",
                            spending_amount=spending,
                            threshold_amount=alert_threshold,
                        )
                        webhook = config.feishu_webhook_url
                        send_feishu_alert(
                            webhook,
                            "⚠️ 子账号消费告警",
                            f"**用户**: {user.username}\n"
                            f"**消费**: ¥{spending:.2f}\n"
                            f"**告警阈值**: ¥{alert_threshold:.2f}\n"
                            f"**停用阈值**: ¥{disable_threshold:.2f}",
                            template="orange",
                        )
                        alert.notified = True
                        alert.save(update_fields=['notified'])

            except Exception as e:
                logger.error(f"检查用户 {user.username} 消费失败: {e}")


def start_scheduler():
    """启动定时任务"""
    global _scheduler_started
    if _scheduler_started:
        return
    _scheduler_started = True

    try:
        from apscheduler.schedulers.background import BackgroundScheduler
        from django.conf import settings

        scheduler = BackgroundScheduler()
        interval = getattr(settings, 'MONITOR_INTERVAL', 3600)
        scheduler.add_job(check_spending, 'interval', seconds=interval,
                          id='check_spending', replace_existing=True)
        scheduler.start()
        logger.info(f"消费监控定时任务已启动，间隔 {interval} 秒")
    except Exception as e:
        logger.error(f"启动定时任务失败: {e}")