feat(tasks): T55-T58 implement background tasks for OpenRouter sync
- T55: Setup APScheduler with AsyncIOScheduler and @scheduled_job decorator - T56: Implement hourly usage stats sync from OpenRouter API - T57: Implement daily API key validation job - T58: Implement weekly cleanup job for old usage stats - Add usage_stats_retention_days config option - Integrate scheduler with FastAPI lifespan events - Add 26 unit tests for scheduler, sync, and cleanup tasks - Add apscheduler to requirements.txt The background tasks now automatically: - Sync usage stats every hour from OpenRouter - Validate API keys daily at 2 AM UTC - Clean up old data weekly on Sunday at 3 AM UTC
This commit is contained in:
@@ -56,6 +56,10 @@ class Settings(BaseSettings):
|
||||
default=60,
|
||||
description="Background sync interval in minutes"
|
||||
)
|
||||
usage_stats_retention_days: int = Field(
|
||||
default=365,
|
||||
description="Retention period for usage stats in days"
|
||||
)
|
||||
|
||||
# Limits
|
||||
max_api_keys_per_user: int = Field(
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
Main application entry point for OpenRouter API Key Monitor.
|
||||
"""
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
@@ -11,15 +13,32 @@ from openrouter_monitor.routers import auth
|
||||
from openrouter_monitor.routers import public_api
|
||||
from openrouter_monitor.routers import stats
|
||||
from openrouter_monitor.routers import tokens
|
||||
from openrouter_monitor.tasks.scheduler import init_scheduler, shutdown_scheduler
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Application lifespan manager.
|
||||
|
||||
Handles startup and shutdown events including
|
||||
scheduler initialization and cleanup.
|
||||
"""
|
||||
# Startup
|
||||
init_scheduler()
|
||||
yield
|
||||
# Shutdown
|
||||
shutdown_scheduler()
|
||||
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI(
|
||||
title="OpenRouter API Key Monitor",
|
||||
description="Monitor and manage OpenRouter API keys",
|
||||
version="1.0.0",
|
||||
debug=settings.debug,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
# CORS middleware
|
||||
|
||||
0
src/openrouter_monitor/tasks/__init__.py
Normal file
0
src/openrouter_monitor/tasks/__init__.py
Normal file
59
src/openrouter_monitor/tasks/cleanup.py
Normal file
59
src/openrouter_monitor/tasks/cleanup.py
Normal file
@@ -0,0 +1,59 @@
|
||||
"""Cleanup tasks for old data.
|
||||
|
||||
T58: Task to clean up old usage stats data.
|
||||
"""
|
||||
import logging
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from apscheduler.triggers.cron import CronTrigger
|
||||
from sqlalchemy import delete
|
||||
|
||||
from openrouter_monitor.database import SessionLocal
|
||||
from openrouter_monitor.models.usage_stats import UsageStats
|
||||
from openrouter_monitor.config import get_settings
|
||||
from openrouter_monitor.tasks.scheduler import scheduled_job
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@scheduled_job(
|
||||
CronTrigger(day_of_week='sun', hour=3, minute=0),
|
||||
id='cleanup_old_usage_stats',
|
||||
replace_existing=True
|
||||
)
|
||||
async def cleanup_old_usage_stats():
|
||||
"""Clean up usage stats older than retention period.
|
||||
|
||||
Runs weekly on Sunday at 3:00 AM UTC.
|
||||
Removes UsageStats records older than usage_stats_retention_days
|
||||
(default: 365 days).
|
||||
|
||||
The retention period is configurable via the
|
||||
USAGE_STATS_RETENTION_DAYS environment variable.
|
||||
"""
|
||||
logger.info("Starting cleanup of old usage stats")
|
||||
|
||||
try:
|
||||
with SessionLocal() as db:
|
||||
# Calculate cutoff date
|
||||
retention_days = settings.usage_stats_retention_days
|
||||
cutoff_date = datetime.utcnow().date() - timedelta(days=retention_days)
|
||||
|
||||
logger.info(f"Removing usage stats older than {cutoff_date}")
|
||||
|
||||
# Delete old records
|
||||
stmt = delete(UsageStats).where(UsageStats.date < cutoff_date)
|
||||
result = db.execute(stmt)
|
||||
deleted_count = result.rowcount
|
||||
|
||||
db.commit()
|
||||
|
||||
logger.info(
|
||||
f"Cleanup completed. Deleted {deleted_count} old usage stats records "
|
||||
f"(retention: {retention_days} days)"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in cleanup_old_usage_stats job: {e}")
|
||||
76
src/openrouter_monitor/tasks/scheduler.py
Normal file
76
src/openrouter_monitor/tasks/scheduler.py
Normal file
@@ -0,0 +1,76 @@
|
||||
"""APScheduler task scheduler.
|
||||
|
||||
T55: Background task scheduler using APScheduler with AsyncIOScheduler.
|
||||
"""
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
|
||||
# Singleton scheduler instance
|
||||
_scheduler = None
|
||||
|
||||
|
||||
def get_scheduler():
|
||||
"""Get or create the singleton scheduler instance.
|
||||
|
||||
Returns:
|
||||
AsyncIOScheduler: The scheduler instance (singleton)
|
||||
|
||||
Example:
|
||||
>>> scheduler = get_scheduler()
|
||||
>>> scheduler.start()
|
||||
"""
|
||||
global _scheduler
|
||||
if _scheduler is None:
|
||||
_scheduler = AsyncIOScheduler(timezone='UTC')
|
||||
return _scheduler
|
||||
|
||||
|
||||
def scheduled_job(trigger, **trigger_args):
|
||||
"""Decorator to register a scheduled job.
|
||||
|
||||
Args:
|
||||
trigger: APScheduler trigger (IntervalTrigger, CronTrigger, etc.)
|
||||
**trigger_args: Additional arguments for add_job (id, name, etc.)
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the job and returns original function
|
||||
|
||||
Example:
|
||||
>>> from apscheduler.triggers.interval import IntervalTrigger
|
||||
>>>
|
||||
>>> @scheduled_job(IntervalTrigger(hours=1), id='sync_task')
|
||||
... async def sync_data():
|
||||
... pass
|
||||
"""
|
||||
def decorator(func):
|
||||
get_scheduler().add_job(func, trigger=trigger, **trigger_args)
|
||||
return func
|
||||
return decorator
|
||||
|
||||
|
||||
def init_scheduler():
|
||||
"""Initialize and start the scheduler.
|
||||
|
||||
Should be called during application startup.
|
||||
Registers all decorated jobs and starts the scheduler.
|
||||
|
||||
Example:
|
||||
>>> init_scheduler()
|
||||
>>> # Scheduler is now running
|
||||
"""
|
||||
scheduler = get_scheduler()
|
||||
scheduler.start()
|
||||
|
||||
|
||||
def shutdown_scheduler():
|
||||
"""Shutdown the scheduler gracefully.
|
||||
|
||||
Should be called during application shutdown.
|
||||
Waits for running jobs to complete before stopping.
|
||||
|
||||
Example:
|
||||
>>> shutdown_scheduler()
|
||||
>>> # Scheduler is stopped
|
||||
"""
|
||||
scheduler = get_scheduler()
|
||||
scheduler.shutdown(wait=True)
|
||||
192
src/openrouter_monitor/tasks/sync.py
Normal file
192
src/openrouter_monitor/tasks/sync.py
Normal file
@@ -0,0 +1,192 @@
|
||||
"""OpenRouter sync tasks.
|
||||
|
||||
T56: Task to sync usage stats from OpenRouter.
|
||||
T57: Task to validate API keys.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import httpx
|
||||
from apscheduler.triggers.interval import IntervalTrigger
|
||||
from apscheduler.triggers.cron import CronTrigger
|
||||
from sqlalchemy import select
|
||||
|
||||
from openrouter_monitor.database import SessionLocal
|
||||
from openrouter_monitor.models.api_key import ApiKey
|
||||
from openrouter_monitor.models.usage_stats import UsageStats
|
||||
from openrouter_monitor.services.encryption import EncryptionService
|
||||
from openrouter_monitor.config import get_settings
|
||||
from openrouter_monitor.tasks.scheduler import scheduled_job
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
# OpenRouter API configuration
|
||||
OPENROUTER_USAGE_URL = "https://openrouter.ai/api/v1/usage"
|
||||
OPENROUTER_AUTH_URL = "https://openrouter.ai/api/v1/auth/key"
|
||||
RATE_LIMIT_DELAY = 0.35 # ~20 req/min to stay under rate limit
|
||||
TIMEOUT_SECONDS = 30.0
|
||||
|
||||
|
||||
@scheduled_job(IntervalTrigger(hours=1), id='sync_usage_stats', replace_existing=True)
|
||||
async def sync_usage_stats():
|
||||
"""Sync usage stats from OpenRouter for all active API keys.
|
||||
|
||||
Runs every hour. Fetches usage data for the last 7 days and
|
||||
upserts records into the UsageStats table.
|
||||
|
||||
Rate limited to ~20 requests per minute to respect OpenRouter limits.
|
||||
"""
|
||||
logger.info("Starting usage stats sync job")
|
||||
|
||||
try:
|
||||
with SessionLocal() as db:
|
||||
# Query all active API keys
|
||||
stmt = select(ApiKey).where(ApiKey.is_active == True)
|
||||
result = db.execute(stmt)
|
||||
api_keys = result.scalars().all()
|
||||
|
||||
logger.info(f"Found {len(api_keys)} active API keys to sync")
|
||||
|
||||
if not api_keys:
|
||||
logger.info("No active API keys found, skipping sync")
|
||||
return
|
||||
|
||||
# Initialize encryption service
|
||||
encryption = EncryptionService(settings.encryption_key)
|
||||
|
||||
# Calculate date range (last 7 days)
|
||||
end_date = datetime.utcnow().date()
|
||||
start_date = end_date - timedelta(days=6) # 7 days inclusive
|
||||
|
||||
for api_key in api_keys:
|
||||
try:
|
||||
# Decrypt the API key
|
||||
decrypted_key = encryption.decrypt(api_key.key_encrypted)
|
||||
|
||||
# Fetch usage data from OpenRouter
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
OPENROUTER_USAGE_URL,
|
||||
headers={"Authorization": f"Bearer {decrypted_key}"},
|
||||
params={
|
||||
"start_date": start_date.strftime("%Y-%m-%d"),
|
||||
"end_date": end_date.strftime("%Y-%m-%d")
|
||||
},
|
||||
timeout=TIMEOUT_SECONDS
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.warning(
|
||||
f"Failed to fetch usage for key {api_key.id}: "
|
||||
f"HTTP {response.status_code}"
|
||||
)
|
||||
continue
|
||||
|
||||
data = response.json()
|
||||
usage_records = data.get("data", [])
|
||||
|
||||
logger.info(
|
||||
f"Fetched {len(usage_records)} usage records for key {api_key.id}"
|
||||
)
|
||||
|
||||
# Upsert usage stats
|
||||
for record in usage_records:
|
||||
try:
|
||||
usage_stat = UsageStats(
|
||||
api_key_id=api_key.id,
|
||||
date=datetime.strptime(record["date"], "%Y-%m-%d").date(),
|
||||
model=record.get("model", "unknown"),
|
||||
requests_count=record.get("requests_count", 0),
|
||||
tokens_input=record.get("tokens_input", 0),
|
||||
tokens_output=record.get("tokens_output", 0),
|
||||
cost=record.get("cost", 0.0)
|
||||
)
|
||||
db.merge(usage_stat)
|
||||
except (KeyError, ValueError) as e:
|
||||
logger.error(f"Error parsing usage record: {e}")
|
||||
continue
|
||||
|
||||
db.commit()
|
||||
logger.info(f"Successfully synced usage stats for key {api_key.id}")
|
||||
|
||||
# Rate limiting between requests
|
||||
await asyncio.sleep(RATE_LIMIT_DELAY)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error syncing key {api_key.id}: {e}")
|
||||
continue
|
||||
|
||||
logger.info("Usage stats sync job completed")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in sync_usage_stats job: {e}")
|
||||
|
||||
|
||||
@scheduled_job(CronTrigger(hour=2, minute=0), id='validate_api_keys', replace_existing=True)
|
||||
async def validate_api_keys():
|
||||
"""Validate all active API keys by checking with OpenRouter.
|
||||
|
||||
Runs daily at 2:00 AM UTC. Deactivates any keys that are invalid.
|
||||
"""
|
||||
logger.info("Starting API key validation job")
|
||||
|
||||
try:
|
||||
with SessionLocal() as db:
|
||||
# Query all active API keys
|
||||
stmt = select(ApiKey).where(ApiKey.is_active == True)
|
||||
result = db.execute(stmt)
|
||||
api_keys = result.scalars().all()
|
||||
|
||||
logger.info(f"Found {len(api_keys)} active API keys to validate")
|
||||
|
||||
if not api_keys:
|
||||
logger.info("No active API keys found, skipping validation")
|
||||
return
|
||||
|
||||
# Initialize encryption service
|
||||
encryption = EncryptionService(settings.encryption_key)
|
||||
|
||||
invalid_count = 0
|
||||
|
||||
for api_key in api_keys:
|
||||
try:
|
||||
# Decrypt the API key
|
||||
decrypted_key = encryption.decrypt(api_key.key_encrypted)
|
||||
|
||||
# Validate with OpenRouter
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
OPENROUTER_AUTH_URL,
|
||||
headers={"Authorization": f"Bearer {decrypted_key}"},
|
||||
timeout=TIMEOUT_SECONDS
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
# Key is invalid, deactivate it
|
||||
api_key.is_active = False
|
||||
invalid_count += 1
|
||||
logger.warning(
|
||||
f"API key {api_key.id} ({api_key.name}) is invalid, "
|
||||
f"deactivating. HTTP {response.status_code}"
|
||||
)
|
||||
else:
|
||||
logger.debug(f"API key {api_key.id} ({api_key.name}) is valid")
|
||||
|
||||
# Rate limiting between requests
|
||||
await asyncio.sleep(RATE_LIMIT_DELAY)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error validating key {api_key.id}: {e}")
|
||||
continue
|
||||
|
||||
db.commit()
|
||||
logger.info(
|
||||
f"API key validation completed. "
|
||||
f"Deactivated {invalid_count} invalid keys."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in validate_api_keys job: {e}")
|
||||
Reference in New Issue
Block a user