"""APScheduler background scheduler for upstream checks and daily tasks.""" from __future__ import annotations import json import logging from datetime import datetime, timezone from typing import Any from apscheduler.executors.pool import ThreadPoolExecutor from apscheduler.schedulers.background import BackgroundScheduler from apscheduler.triggers.cron import CronTrigger from sqlalchemy.orm import Session from app.database import SessionLocal from app.models.upstream import Upstream from app.models.snapshot import UpstreamRateSnapshot from app.services.auth_config import normalize_auth_config from app.services.upstream_client import UpstreamClient, build_snapshot from app.services.snapshot_service import diff_snapshots, prune_snapshots from app.services import webhook_service from app.services import website_sync from app.config import get_settings from app.routers.external_api_logs import clean_expired_logs logger = logging.getLogger(__name__) _scheduler = BackgroundScheduler(timezone="UTC", executors={"default": ThreadPoolExecutor(max_workers=1)}) def get_scheduler() -> BackgroundScheduler: return _scheduler def _persist_auth_config_update(upstream: Upstream, updated_config: dict[str, Any]) -> None: config_json = json.dumps( normalize_auth_config(upstream.auth_type, updated_config), ensure_ascii=False, ) updated_at = datetime.now(timezone.utc) db = SessionLocal() try: row = db.query(Upstream).filter(Upstream.id == upstream.id).first() if row: row.auth_config_json = config_json row.updated_at = updated_at db.commit() upstream.auth_config_json = row.auth_config_json upstream.updated_at = updated_at finally: db.close() def _check_upstream(upstream_id: int) -> None: """Full upstream check executed by scheduler (runs in thread). Phase 1 — upstream API call + snapshot write (single transaction). Phase 2 — webhook/website sync (separate sessions, so a notification failure never rolls back the snapshot). """ settings = get_settings() # ── Phase 1: upstream check + DB write ────────────────────────── db: Session = SessionLocal() try: upstream = db.query(Upstream).filter(Upstream.id == upstream_id).first() if not upstream or not upstream.enabled: _remove_job(upstream_id) return auth_config = json.loads(upstream.auth_config_json or "{}") was_unhealthy = upstream.last_status == "unhealthy" balance_alert_triggered = False snapshot = None changes = None with UpstreamClient( base_url=upstream.base_url, api_prefix=upstream.api_prefix, auth_type=upstream.auth_type, auth_config=auth_config, timeout=float(upstream.timeout_seconds), on_auth_config_update=lambda updated: _persist_auth_config_update(upstream, updated), target_id=upstream.id, target_name=upstream.name, ) as client: try: client.ensure_authenticated() groups = client.get_available_groups(upstream.groups_endpoint) raw_rates = client.get_group_rates(upstream.rate_endpoint) snapshot = build_snapshot( upstream.id, upstream.base_url, upstream.api_prefix, groups, raw_rates ) # ── Balance fetch (inside with block, client still open) ── balance: Optional[float] = None if upstream.balance_endpoint and upstream.balance_response_path: try: raw_balance = client.get_balance(upstream.balance_endpoint, upstream.balance_response_path) if raw_balance is not None: divisor = upstream.balance_divisor or 1.0 balance = raw_balance / divisor except Exception as exc: logger.warning("upstream %s balance fetch failed: %s", upstream.name, exc) if balance is not None: upstream.balance = balance upstream.balance_updated_at = datetime.now(timezone.utc) # ── 余额告警阈值检查 ── threshold = upstream.balance_alert_threshold if threshold is not None and threshold > 0: if balance < threshold and not upstream.balance_alert_notified: upstream.balance_alert_notified = True balance_alert_triggered = True elif balance >= threshold and upstream.balance_alert_notified: upstream.balance_alert_notified = False except Exception as exc: # failure path upstream.consecutive_failures = (upstream.consecutive_failures or 0) + 1 upstream.last_error = str(exc) upstream.last_checked_at = datetime.now(timezone.utc) threshold = settings.unhealthy_threshold became_unhealthy = ( upstream.consecutive_failures >= threshold and upstream.last_status != "unhealthy" ) if became_unhealthy: upstream.last_status = "unhealthy" db.commit() logger.warning("upstream %s check failed: %s", upstream.name, exc) # Phase 2: notify unhealthy in a fresh session if became_unhealthy: _notify_status(upstream.id, upstream.name, upstream.base_url, "upstream_unhealthy", str(exc)) return # success path (client auto-closed by `with`) prev_snapshot_row = ( db.query(UpstreamRateSnapshot) .filter(UpstreamRateSnapshot.upstream_id == upstream_id) .order_by(UpstreamRateSnapshot.captured_at.desc()) .first() ) previous = json.loads(prev_snapshot_row.snapshot_json) if prev_snapshot_row else None changes = diff_snapshots(previous, snapshot) # save new snapshot new_row = UpstreamRateSnapshot( upstream_id=upstream_id, snapshot_json=json.dumps(snapshot, ensure_ascii=False), captured_at=datetime.now(timezone.utc), ) db.add(new_row) prune_snapshots(db, upstream_id, settings.snapshot_retention_count) # update upstream status upstream.last_status = "healthy" upstream.last_checked_at = datetime.now(timezone.utc) upstream.last_error = None upstream.consecutive_failures = 0 db.commit() logger.info( "upstream %s: %d rate change(s)" if changes else "upstream %s: no changes", upstream.name, len(changes) if changes else 0, ) finally: db.close() # ── Phase 2: key sync (independent session) ─────────────────── if snapshot: captured_at = snapshot.get("captured_at") if isinstance(captured_at, str): from datetime import datetime as dt try: captured_at = dt.fromisoformat(captured_at) except Exception: captured_at = datetime.now(timezone.utc) elif captured_at is None: captured_at = datetime.now(timezone.utc) _sync_upstream_keys(upstream_id, snapshot, captured_at) # ── Phase 3: notifications (independent sessions) ────────────── if was_unhealthy: _notify_status(upstream_id, upstream.name, upstream.base_url, "upstream_recovered") if changes: _notify_rate_changed(upstream_id, upstream.name, upstream.base_url, changes) _sync_website_bindings(upstream_id, changes) _sync_account_priorities(upstream_id) if balance_alert_triggered: _notify_balance_low( upstream_id, upstream.name, upstream.base_url, upstream.balance, upstream.balance_alert_threshold, ) def _notify_status( upstream_id: int, upstream_name: str, base_url: str, event: str, error: str = "", ) -> None: db = SessionLocal() try: webhook_service.send_status_event(db, upstream_id, upstream_name, base_url, event, error) except Exception: logger.exception("status webhook failed for upstream %s", upstream_name) finally: db.close() def _notify_rate_changed( upstream_id: int, upstream_name: str, base_url: str, changes: list[dict[str, Any]], ) -> None: db = SessionLocal() try: webhook_service.send_rate_changed(db, upstream_id, upstream_name, base_url, changes) except Exception: logger.exception("rate webhook failed for upstream %s", upstream_name) finally: db.close() def _notify_balance_low( upstream_id: int, upstream_name: str, base_url: str, balance: float, threshold: float, ) -> None: db = SessionLocal() try: webhook_service.send_balance_low(db, upstream_id, upstream_name, base_url, balance, threshold) except Exception: logger.exception("balance low webhook failed for upstream %s", upstream_name) finally: db.close() def _sync_upstream_keys(upstream_id: int, snapshot: dict[str, Any], captured_at: datetime) -> None: """上游检测成功后同步 SmartUp Key 状态(远端删除/分组删除)。 委托给 website_sync.reconcile_upstream_keys 实现核心逻辑。 """ db = SessionLocal() try: active_group_ids = set(snapshot.get("groups", {}).keys()) # 获取倒数第二个快照作为 previous snapshot prev_active_group_ids: set[str] | None = None from app.models.snapshot import UpstreamRateSnapshot snapshots = ( db.query(UpstreamRateSnapshot) .filter(UpstreamRateSnapshot.upstream_id == upstream_id) .order_by(UpstreamRateSnapshot.captured_at.desc()) .limit(2) .all() ) if len(snapshots) >= 2: prev_snapshot = snapshots[1] try: prev_data = json.loads(prev_snapshot.snapshot_json or "{}") prev_groups = prev_data.get("groups") or {} if isinstance(prev_groups, dict): prev_active_group_ids = set(prev_groups.keys()) except Exception: pass # 用 UpstreamClient 查询远端活跃 Key ID 集合 remote_key_ids: set[str] | None = None try: upstream = db.query(Upstream).filter(Upstream.id == upstream_id).first() if upstream: auth_config = json.loads(upstream.auth_config_json or "{}") with UpstreamClient( base_url=upstream.base_url, api_prefix=upstream.api_prefix, auth_type=upstream.auth_type, auth_config=auth_config, timeout=float(upstream.timeout_seconds), on_auth_config_update=lambda updated: _persist_auth_config_update(upstream, updated), target_id=upstream.id, target_name=upstream.name, ) as client: client.ensure_authenticated() remote_key_ids = website_sync._fetch_remote_managed_key_ids(db, client, upstream_id) except Exception as exc: logger.warning("sync upstream keys list failed for %s: %s", upstream_id, exc) website_sync.reconcile_upstream_keys(db, upstream_id, active_group_ids, remote_key_ids, captured_at, prev_active_group_ids=prev_active_group_ids) db.commit() except Exception: logger.exception("key sync failed for upstream %s", upstream_id) finally: db.close() def _sync_account_priorities(upstream_id: int) -> None: """倍率变更后自动更新已导入下游账号的 priority。""" db = SessionLocal() try: website_sync.sync_account_priorities_for_upstream(db, upstream_id) except Exception: logger.exception("account priority sync failed for upstream %s", upstream_id) finally: db.close() def _sync_website_bindings(upstream_id: int, changes: list[dict[str, Any]]) -> None: db = SessionLocal() try: website_sync.sync_affected_bindings(db, upstream_id, changes) except Exception: logger.exception("website sync failed for upstream %s", upstream_id) finally: db.close() def _remove_job(upstream_id: int) -> None: job_id = f"upstream_{upstream_id}" if _scheduler.get_job(job_id): _scheduler.remove_job(job_id) def refresh_upstream(upstream_id: int, interval_seconds: int = 0, enabled: bool = True) -> None: """Add/update/remove a scheduler job for the given upstream.""" job_id = f"upstream_{upstream_id}" if not enabled or interval_seconds <= 0: _remove_job(upstream_id) return _scheduler.add_job( _check_upstream, "interval", seconds=interval_seconds, id=job_id, args=[upstream_id], replace_existing=True, coalesce=True, max_instances=1, misfire_grace_time=60, jitter=30, ) logger.info("scheduler job %s set to %ds interval", job_id, interval_seconds) def start_scheduler() -> None: """Start scheduler and load all enabled upstreams.""" _scheduler.start() db: Session = SessionLocal() try: upstreams = db.query(Upstream).filter(Upstream.enabled == True).all() for u in upstreams: refresh_upstream(u.id, u.check_interval_seconds, u.enabled) # Daily cleanup of expired external API logs _scheduler.add_job( _cleanup_external_logs, "interval", hours=24, id="cleanup_external_api_logs", replace_existing=True, coalesce=True, max_instances=1, misfire_grace_time=3600, ) # Daily finance summary at 00:05 Asia/Shanghai _scheduler.add_job( _compute_finance_daily_summary, trigger=CronTrigger(timezone="Asia/Shanghai", hour=0, minute=5), id="finance_daily_summary", replace_existing=True, coalesce=True, max_instances=1, misfire_grace_time=3600, ) logger.info("scheduler started with %d upstream job(s)", len(upstreams)) finally: db.close() def _cleanup_external_logs() -> None: """Delete external API logs older than the configured retention period.""" db = SessionLocal() try: deleted = clean_expired_logs(db) if deleted: logger.info("cleaned %d expired external API log(s)", deleted) except Exception: logger.exception("failed to clean expired external API logs") finally: db.close() def _compute_finance_daily_summary() -> None: """Compute and save finance daily summary for yesterday (Asia/Shanghai). Runs daily at 00:05 Asia/Shanghai. Delegates to save_summary_if_absent which handles the skip-if-exists logic and concurrent-save race. """ from app.services.finance_service import yesterday_shanghai, save_summary_if_absent target = yesterday_shanghai() db = SessionLocal() try: saved = save_summary_if_absent(db, target) if saved: logger.info("finance daily summary for %s saved", target) else: logger.info("finance daily summary for %s already exists, skipping", target) except Exception: logger.exception("failed to compute finance daily summary for %s", target) finally: db.close() def stop_scheduler() -> None: if _scheduler.running: _scheduler.shutdown(wait=True)