|
import uuid
|
|
from datetime import datetime, date, timezone
|
|
from decimal import Decimal
|
|
from typing import Optional
|
|
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)
|
|
}
|