""" Usage Tracking Scheduler Tracks daily usage snapshots for all active tenants """ import asyncio import structlog from datetime import datetime, timedelta, timezone from typing import Optional from sqlalchemy import select, func logger = structlog.get_logger() class UsageTrackingScheduler: """Scheduler for daily usage tracking""" def __init__(self, db_manager, redis_client, config): self.db_manager = db_manager self.redis = redis_client self.config = config self._running = False self._task: Optional[asyncio.Task] = None def seconds_until_target_time(self) -> float: """Calculate seconds until next target time (default 2am UTC)""" now = datetime.now(timezone.utc) target = now.replace( hour=self.config.USAGE_TRACKING_HOUR, minute=self.config.USAGE_TRACKING_MINUTE, second=0, microsecond=0 ) if target <= now: target += timedelta(days=1) return (target - now).total_seconds() async def _get_tenant_usage(self, session, tenant_id: str) -> dict: """Get current usage counts for a tenant""" usage = {} try: # Import models here to avoid circular imports from app.models.tenants import TenantMember # Users count result = await session.execute( select(func.count()).select_from(TenantMember).where(TenantMember.tenant_id == tenant_id) ) usage['users'] = result.scalar() or 0 # Get counts from other services via their databases # For now, we'll track basic metrics. More metrics can be added by querying other service databases # Training jobs today (from Redis quota tracking) today_key = f"quota:training_jobs:{tenant_id}:{datetime.now(timezone.utc).strftime('%Y-%m-%d')}" training_count = await self.redis.get(today_key) usage['training_jobs'] = int(training_count) if training_count else 0 # Forecasts today (from Redis quota tracking) forecast_key = f"quota:forecasts:{tenant_id}:{datetime.now(timezone.utc).strftime('%Y-%m-%d')}" forecast_count = await self.redis.get(forecast_key) usage['forecasts'] = int(forecast_count) if forecast_count else 0 # API calls this hour (from Redis quota tracking) hour_key = f"quota:api_calls:{tenant_id}:{datetime.now(timezone.utc).strftime('%Y-%m-%d-%H')}" api_count = await self.redis.get(hour_key) usage['api_calls'] = int(api_count) if api_count else 0 # Storage (placeholder - implement based on file storage system) usage['storage'] = 0.0 except Exception as e: logger.error("Error getting usage for tenant", tenant_id=tenant_id, error=str(e), exc_info=True) return {} return usage async def _track_metrics(self, tenant_id: str, usage: dict): """Track metrics to Redis""" from app.api.usage_forecast import track_usage_snapshot for metric_name, value in usage.items(): try: await track_usage_snapshot(tenant_id, metric_name, value) except Exception as e: logger.error( "Failed to track metric", tenant_id=tenant_id, metric=metric_name, error=str(e) ) async def _run_cycle(self): """Execute one tracking cycle""" start_time = datetime.now(timezone.utc) logger.info("Starting daily usage tracking cycle") try: async with self.db_manager.get_session() as session: # Import models here to avoid circular imports from app.models.tenants import Tenant, Subscription from sqlalchemy import select # Get all active tenants result = await session.execute( select(Tenant, Subscription) .join(Subscription, Tenant.id == Subscription.tenant_id) .where(Tenant.is_active == True) .where(Subscription.status.in_(['active', 'trialing', 'cancelled'])) ) tenants_data = result.all() total_tenants = len(tenants_data) success_count = 0 error_count = 0 logger.info(f"Found {total_tenants} active tenants to track") # Process each tenant for tenant, subscription in tenants_data: try: usage = await self._get_tenant_usage(session, tenant.id) if usage: await self._track_metrics(tenant.id, usage) success_count += 1 else: logger.warning( "No usage data available for tenant", tenant_id=tenant.id ) error_count += 1 except Exception as e: logger.error( "Error tracking tenant usage", tenant_id=tenant.id, error=str(e), exc_info=True ) error_count += 1 end_time = datetime.now(timezone.utc) duration = (end_time - start_time).total_seconds() logger.info( "Daily usage tracking completed", total_tenants=total_tenants, success=success_count, errors=error_count, duration_seconds=duration ) except Exception as e: logger.error("Usage tracking cycle failed", error=str(e), exc_info=True) async def _run_scheduler(self): """Main scheduler loop""" logger.info( "Usage tracking scheduler loop started", target_hour=self.config.USAGE_TRACKING_HOUR, target_minute=self.config.USAGE_TRACKING_MINUTE ) # Initial delay to target time delay = self.seconds_until_target_time() logger.info(f"Waiting {delay/3600:.2f} hours until next run at {self.config.USAGE_TRACKING_HOUR:02d}:{self.config.USAGE_TRACKING_MINUTE:02d} UTC") try: await asyncio.sleep(delay) except asyncio.CancelledError: logger.info("Scheduler cancelled during initial delay") return while self._running: try: await self._run_cycle() except Exception as e: logger.error("Scheduler cycle error", error=str(e), exc_info=True) # Wait 24 hours until next run try: await asyncio.sleep(86400) except asyncio.CancelledError: logger.info("Scheduler cancelled during sleep") break def start(self): """Start the scheduler""" if not self.config.USAGE_TRACKING_ENABLED: logger.info("Usage tracking scheduler disabled by configuration") return if self._running: logger.warning("Usage tracking scheduler already running") return self._running = True self._task = asyncio.create_task(self._run_scheduler()) logger.info("Usage tracking scheduler started successfully") async def stop(self): """Stop the scheduler gracefully""" if not self._running: logger.debug("Scheduler not running, nothing to stop") return logger.info("Stopping usage tracking scheduler") self._running = False if self._task: self._task.cancel() try: await self._task except asyncio.CancelledError: logger.info("Scheduler task cancelled successfully") logger.info("Usage tracking scheduler stopped") # Global instance _scheduler: Optional[UsageTrackingScheduler] = None async def start_scheduler(db_manager, redis_client, config): """Start the usage tracking scheduler""" global _scheduler try: _scheduler = UsageTrackingScheduler(db_manager, redis_client, config) _scheduler.start() logger.info("Usage tracking scheduler module initialized") except Exception as e: logger.error("Failed to start usage tracking scheduler", error=str(e), exc_info=True) raise async def stop_scheduler(): """Stop the usage tracking scheduler""" global _scheduler if _scheduler: await _scheduler.stop() _scheduler = None logger.info("Usage tracking scheduler module stopped")