gitscope
← marzneshin / v0.6.5 / app/tasks/data_usage_percent_reached.py
https://dev.kronoz.win/r/marzneshin.git
app/tasks/data_usage_percent_reached.py · Python
Blame Raw
 1import asyncio
 2
 3from sqlalchemy import case
 4from sqlalchemy.orm import Session
 5
 6from app.config import NOTIFY_REACHED_USAGE_PERCENT
 7from app.db.models import User
 8from app.models.notification import UserNotification
 9from app.models.user import UserResponse
10from app.notification.notifiers import notify
11
12
13async def data_usage_percent_reached(db: Session, users_usage: list) -> None:
14    """
15    Monitors data usage of active users and sends a notification if usage exceeds NOTIFY_REACHED_USAGE_PERCENT.
16
17    This function performs the following steps:
18    - Receives a list of user data usage, where each user has an ID and the amount of data they have used from all nodes.
19    - If the new total data usage exceeds NOTIFY_REACHED_USAGE_PERCENT of the user's data limit, it sends a notification.
20
21    This function is called when user data usage updates occur.
22
23    """
24
25    users_usage_dict = {user["id"]: user["value"] for user in users_usage}
26
27    case_stmt = case(
28        *[
29            (User.id == user_id, User.used_traffic + increase_amount)
30            for user_id, increase_amount in users_usage_dict.items()
31        ],
32        else_=User.used_traffic,
33    )
34
35    exceeding_users = (
36        db.query(User)
37        .filter(
38            User.id.in_(users_usage_dict.keys()),
39            User.data_limit.isnot(None),
40            User.data_limit > 0,
41            (User.used_traffic / User.data_limit) * 100
42            < NOTIFY_REACHED_USAGE_PERCENT,
43            (case_stmt / User.data_limit) * 100 > NOTIFY_REACHED_USAGE_PERCENT,
44        )
45        .all()
46    )
47
48    for user in exceeding_users:
49        added_traffic = users_usage_dict[user.id]
50        user.used_traffic += added_traffic
51        asyncio.ensure_future(
52            notify(
53                action=UserNotification.Action.reached_usage_percent,
54                user=UserResponse.model_validate(user),
55            )
56        )
57
58    db.expunge_all()