"""定时消费监控任务""" 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}")