diff --git a/app/db/crud.py b/app/db/crud.py index 0992032..dfda48b 100644 --- a/app/db/crud.py +++ b/app/db/crud.py @@ -6,7 +6,7 @@ from enum import Enum from types import NoneType from typing import List, Optional, Tuple, Union -from sqlalchemy import and_, delete, update, select, func +from sqlalchemy import and_, delete, update, select, func, cast, Date from sqlalchemy.orm import Session from app.db.models import ( @@ -310,6 +310,59 @@ def get_users( return query.all() +def get_user_total_usage( + db: Session, user: User, start: datetime, end: datetime, per_day=False +): + usages = defaultdict(int) + + query = db.query( + ( + cast(NodeUserUsage.created_at, Date).label("day") + if per_day + else NodeUserUsage.created_at + ), + func.sum(NodeUserUsage.used_traffic), + ).filter( + and_( + NodeUserUsage.user_id == user.id, + NodeUserUsage.created_at >= start, + NodeUserUsage.created_at <= end, + ) + ) + if per_day: + query = query.group_by(cast(NodeUserUsage.created_at, Date)) + else: + query = query.group_by(NodeUserUsage.created_at) + + for date, used_traffic in query: + if per_day: + timestamp = datetime( + date.year, date.month, date.day, tzinfo=timezone.utc + ).timestamp() + usages[timestamp] += int(used_traffic) + else: + usages[date.replace(tzinfo=timezone.utc).timestamp()] += int( + used_traffic + ) + + result = TrafficUsageSeries(usages=[]) + current = start.astimezone(timezone.utc).replace( + minute=0, second=0, microsecond=0 + ) + if per_day: + current = current.replace(hour=0) + + step = timedelta(days=1) if per_day else timedelta(hours=1) + while current <= end.replace(tzinfo=timezone.utc): + current_usage = usages.get(current.timestamp()) or 0 + result.usages.append((int(current.timestamp()), current_usage)) + result.total += current_usage + current += step + + result.step = int(step.total_seconds()) + return result + + def get_total_usages( db: Session, admin: Admin, start: datetime, end: datetime ) -> TrafficUsageSeries: diff --git a/app/models/system.py b/app/models/system.py index be885cf..9a6bca0 100644 --- a/app/models/system.py +++ b/app/models/system.py @@ -40,5 +40,7 @@ class NodesStats(BaseModel): class TrafficUsageSeries(BaseModel): + step: int = 3600 + total: int = 0 usages: list[tuple[int, int]] total: int diff --git a/app/routes/subscription.py b/app/routes/subscription.py index f889d0b..572e47d 100644 --- a/app/routes/subscription.py +++ b/app/routes/subscription.py @@ -9,6 +9,7 @@ from app.db import crud from app.db.models import Settings from app.dependencies import DBDep, SubUserDep, StartDateDep, EndDateDep from app.models.settings import SubscriptionSettings +from app.models.system import TrafficUsageSeries from app.models.user import UserResponse from app.utils.share import ( encode_title, @@ -124,16 +125,17 @@ def user_subscription_info(db_user: SubUserDep): return db_user -@router.get("/{username}/{key}/usage") +@router.get("/{username}/{key}/usage", response_model=TrafficUsageSeries) def user_get_usage( db_user: SubUserDep, db: DBDep, start_date: StartDateDep, end_date: EndDateDep, ): - usages = crud.get_user_usages(db, db_user, start_date, end_date) - - return {"usages": usages, "username": db_user.username} + per_day = (end_date - start_date).total_seconds() > 3 * 86400 + return crud.get_user_total_usage( + db, db_user, start_date, end_date, per_day=per_day + ) client_type_mime_type = {