import uuid from datetime import datetime, date, timezone from tortoise.transactions import in_transaction from .models import UsageRecord, UsageAggregate from ..models import User class UsageTracker: @staticmethod async def track_storage( user: User, amount_bytes: int, resource_type: str = None, resource_id: int = None, metadata: dict = None, ): idempotency_key = f"storage_{user.id}_{datetime.now(timezone.utc).timestamp()}_{uuid.uuid4().hex[:8]}" await UsageRecord.create( user=user, record_type="storage", amount_bytes=amount_bytes, resource_type=resource_type, resource_id=resource_id, idempotency_key=idempotency_key, metadata=metadata, ) @staticmethod async def track_bandwidth( user: User, amount_bytes: int, direction: str = "down", resource_type: str = None, resource_id: int = None, metadata: dict = None, ): record_type = f"bandwidth_{direction}" idempotency_key = f"{record_type}_{user.id}_{datetime.now(timezone.utc).timestamp()}_{uuid.uuid4().hex[:8]}" await UsageRecord.create( user=user, record_type=record_type, amount_bytes=amount_bytes, resource_type=resource_type, resource_id=resource_id, idempotency_key=idempotency_key, metadata=metadata, ) @staticmethod async def aggregate_daily_usage(user: User, target_date: date = None): if target_date is None: target_date = date.today() start_of_day = datetime.combine(target_date, datetime.min.time()) end_of_day = datetime.combine(target_date, datetime.max.time()) storage_records = await UsageRecord.filter( user=user, record_type="storage", timestamp__gte=start_of_day, timestamp__lte=end_of_day, ).all() storage_avg = sum(r.amount_bytes for r in storage_records) // max( len(storage_records), 1 ) storage_peak = max((r.amount_bytes for r in storage_records), default=0) bandwidth_up = await UsageRecord.filter( user=user, record_type="bandwidth_up", timestamp__gte=start_of_day, timestamp__lte=end_of_day, ).all() bandwidth_down = await UsageRecord.filter( user=user, record_type="bandwidth_down", timestamp__gte=start_of_day, timestamp__lte=end_of_day, ).all() total_up = sum(r.amount_bytes for r in bandwidth_up) total_down = sum(r.amount_bytes for r in bandwidth_down) async with in_transaction(): aggregate, created = await UsageAggregate.get_or_create( user=user, date=target_date, defaults={ "storage_bytes_avg": storage_avg, "storage_bytes_peak": storage_peak, "bandwidth_up_bytes": total_up, "bandwidth_down_bytes": total_down, }, ) if not created: aggregate.storage_bytes_avg = storage_avg aggregate.storage_bytes_peak = storage_peak aggregate.bandwidth_up_bytes = total_up aggregate.bandwidth_down_bytes = total_down await aggregate.save() return aggregate @staticmethod async def get_current_storage(user: User) -> int: from ..models import File files = await File.filter(owner=user, is_deleted=False).all() return sum(f.size for f in files) @staticmethod async def get_monthly_usage(user: User, year: int, month: int) -> dict: from datetime import date from calendar import monthrange start_date = date(year, month, 1) _, last_day = monthrange(year, month) end_date = date(year, month, last_day) aggregates = await UsageAggregate.filter( user=user, date__gte=start_date, date__lte=end_date ).all() if not aggregates: return { "storage_gb_avg": 0, "storage_gb_peak": 0, "bandwidth_up_gb": 0, "bandwidth_down_gb": 0, "total_bandwidth_gb": 0, } storage_avg = sum(a.storage_bytes_avg for a in aggregates) / len(aggregates) storage_peak = max(a.storage_bytes_peak for a in aggregates) bandwidth_up = sum(a.bandwidth_up_bytes for a in aggregates) bandwidth_down = sum(a.bandwidth_down_bytes for a in aggregates) return { "storage_gb_avg": round(storage_avg / (1024**3), 4), "storage_gb_peak": round(storage_peak / (1024**3), 4), "bandwidth_up_gb": round(bandwidth_up / (1024**3), 4), "bandwidth_down_gb": round(bandwidth_down / (1024**3), 4), "total_bandwidth_gb": round((bandwidth_up + bandwidth_down) / (1024**3), 4), }