bakery-ia/scripts/track_daily_usage.py

#!/usr/bin/env python3
"""
Daily Usage Tracker - Cron Job Script

Tracks daily usage snapshots for all active tenants to enable trend forecasting.
Stores data in Redis with 60-day retention for predictive analytics.

Schedule: Run daily at 2 AM
Crontab: 0 2 * * * /usr/bin/python3 /path/to/scripts/track_daily_usage.py >> /var/log/usage_tracking.log 2>&1

Or use Kubernetes CronJob (see deployment checklist).
"""

import asyncio
import sys
import os
from datetime import datetime, timezone
from pathlib import Path

# Add parent directory to path to import from services
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

# Import from tenant service
from services.tenant.app.core.database import database_manager
from services.tenant.app.models.tenants import Tenant, Subscription, TenantMember
from services.tenant.app.api.usage_forecast import track_usage_snapshot
from services.tenant.app.core.redis_client import get_redis_client

# Import models for counting (adjust these imports based on your actual model locations)
# You may need to update these imports based on your project structure
try:
    from services.inventory.app.models import Product
    from services.inventory.app.models import Location
    from services.inventory.app.models import Recipe
    from services.inventory.app.models import Supplier
except ImportError:
    # Fallback: If models are in different locations, you'll need to update these
    print("Warning: Could not import all models. Some usage metrics may not be tracked.")
    Product = None
    Location = None
    Recipe = None
    Supplier = None


async def get_tenant_current_usage(session: AsyncSession, tenant_id: str) -> dict:
    """
    Get current usage counts for a tenant across all metrics.

    This queries the actual database to get real-time counts.
    """
    usage = {}

    try:
        # Products count
        result = await session.execute(
            select(func.count()).select_from(Product).where(Product.tenant_id == tenant_id)
        )
        usage['products'] = result.scalar() or 0

        # Users count
        result = await session.execute(
            select(func.count()).select_from(TenantMember).where(TenantMember.tenant_id == tenant_id)
        )
        usage['users'] = result.scalar() or 0

        # Locations count
        result = await session.execute(
            select(func.count()).select_from(Location).where(Location.tenant_id == tenant_id)
        )
        usage['locations'] = result.scalar() or 0

        # Recipes count
        result = await session.execute(
            select(func.count()).select_from(Recipe).where(Recipe.tenant_id == tenant_id)
        )
        usage['recipes'] = result.scalar() or 0

        # Suppliers count
        result = await session.execute(
            select(func.count()).select_from(Supplier).where(Supplier.tenant_id == tenant_id)
        )
        usage['suppliers'] = result.scalar() or 0

        # Training jobs today (from Redis)
        redis = await get_redis_client()
        today_key = f"quota:training_jobs:{tenant_id}:{datetime.now(timezone.utc).strftime('%Y-%m-%d')}"
        training_count = await redis.get(today_key)
        usage['training_jobs'] = int(training_count) if training_count else 0

        # Forecasts today (from Redis)
        forecast_key = f"quota:forecasts:{tenant_id}:{datetime.now(timezone.utc).strftime('%Y-%m-%d')}"
        forecast_count = await redis.get(forecast_key)
        usage['forecasts'] = int(forecast_count) if forecast_count else 0

        # Storage (placeholder - implement based on your file storage system)
        # For now, set to 0. Replace with actual storage calculation.
        usage['storage'] = 0.0

        # API calls this hour (from Redis)
        hour_key = f"quota:api_calls:{tenant_id}:{datetime.now(timezone.utc).strftime('%Y-%m-%d-%H')}"
        api_count = await redis.get(hour_key)
        usage['api_calls'] = int(api_count) if api_count else 0

    except Exception as e:
        print(f"Error getting usage for tenant {tenant_id}: {e}")
        # Return empty dict on error
        return {}

    return usage


async def track_all_tenants():
    """
    Main function to track usage for all active tenants.
    """
    start_time = datetime.now(timezone.utc)
    print(f"[{start_time}] Starting daily usage tracking")

    try:
        # Get database session
        async with database_manager.get_session() as session:
            # Query 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)
            print(f"Found {total_tenants} active tenants to track")

            success_count = 0
            error_count = 0

            # Process each tenant
            for tenant, subscription in tenants_data:
                try:
                    # Get current usage for this tenant
                    usage = await get_tenant_current_usage(session, tenant.id)

                    if not usage:
                        print(f"  ⚠️  {tenant.id}: No usage data available")
                        error_count += 1
                        continue

                    # Track each metric
                    metrics_tracked = 0
                    for metric_name, value in usage.items():
                        try:
                            await track_usage_snapshot(
                                tenant_id=tenant.id,
                                metric=metric_name,
                                value=value
                            )
                            metrics_tracked += 1
                        except Exception as e:
                            print(f"  ❌ {tenant.id} - {metric_name}: Error tracking - {e}")

                    print(f"  ✅ {tenant.id}: Tracked {metrics_tracked} metrics")
                    success_count += 1

                except Exception as e:
                    print(f"  ❌ {tenant.id}: Error processing tenant - {e}")
                    error_count += 1
                    continue

            # Summary
            end_time = datetime.now(timezone.utc)
            duration = (end_time - start_time).total_seconds()

            print("\n" + "="*60)
            print(f"Daily Usage Tracking Complete")
            print(f"Started:  {start_time.strftime('%Y-%m-%d %H:%M:%S UTC')}")
            print(f"Finished: {end_time.strftime('%Y-%m-%d %H:%M:%S UTC')}")
            print(f"Duration: {duration:.2f}s")
            print(f"Tenants:  {total_tenants} total")
            print(f"Success:  {success_count} tenants tracked")
            print(f"Errors:   {error_count} tenants failed")
            print("="*60)

            # Exit with error code if any failures
            if error_count > 0:
                sys.exit(1)
            else:
                sys.exit(0)

    except Exception as e:
        print(f"FATAL ERROR: Failed to track usage - {e}")
        import traceback
        traceback.print_exc()
        sys.exit(2)


def main():
    """Entry point"""
    try:
        asyncio.run(track_all_tenants())
    except KeyboardInterrupt:
        print("\n⚠️  Interrupted by user")
        sys.exit(130)
    except Exception as e:
        print(f"FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(2)


if __name__ == '__main__':
    main()