1import asyncio
  2from collections import defaultdict
  3from datetime import datetime
  4
  5from sqlalchemy import and_, select, insert, update, bindparam
  6
  7from app import marznode
  8from app.db import GetDB
  9from app.db.models import NodeUsage, NodeUserUsage, User
 10from app.marznode import MarzNodeBase
 11from app.tasks.data_usage_percent_reached import data_usage_percent_reached
 12
 13
 14def record_user_usage_logs(
 15    params: list, node_id: int, consumption_factor: int = 1
 16):
 17    if not params:
 18        return
 19
 20    created_at = datetime.fromisoformat(
 21        datetime.utcnow().strftime("%Y-%m-%dT%H:00:00")
 22    )
 23
 24    with GetDB() as db:
 25        # make user usage row if doesn't exist
 26        select_stmt = select(NodeUserUsage.user_id).where(
 27            and_(
 28                NodeUserUsage.node_id == node_id,
 29                NodeUserUsage.created_at == created_at,
 30            )
 31        )
 32        existings = [r[0] for r in db.execute(select_stmt).fetchall()]
 33        uids_to_insert = set()
 34
 35        for p in params:
 36            uid = p["uid"]
 37            if uid in existings:
 38                continue
 39            uids_to_insert.add(uid)
 40
 41        if uids_to_insert:
 42            stmt = insert(NodeUserUsage).values(
 43                user_id=bindparam("uid"),
 44                created_at=created_at,
 45                node_id=node_id,
 46                used_traffic=0,
 47            )
 48            db.execute(stmt, [{"uid": uid} for uid in uids_to_insert])
 49
 50        # record
 51        stmt = (
 52            update(NodeUserUsage)
 53            .values(
 54                used_traffic=NodeUserUsage.used_traffic
 55                + bindparam("value") * consumption_factor
 56            )
 57            .where(
 58                and_(
 59                    NodeUserUsage.user_id == bindparam("uid"),
 60                    NodeUserUsage.node_id == node_id,
 61                    NodeUserUsage.created_at == created_at,
 62                )
 63            )
 64        )
 65        db.connection().execute(
 66            stmt, params, execution_options={"synchronize_session": None}
 67        )
 68        db.commit()
 69
 70
 71def record_node_stats(node_id: int, usage: int):
 72    if not usage:
 73        return
 74
 75    created_at = datetime.fromisoformat(
 76        datetime.utcnow().strftime("%Y-%m-%dT%H:00:00")
 77    )
 78
 79    with GetDB() as db:
 80        # make node usage row if doesn't exist
 81        select_stmt = select(NodeUsage.node_id).where(
 82            and_(
 83                NodeUsage.node_id == node_id,
 84                NodeUsage.created_at == created_at,
 85            )
 86        )
 87        notfound = db.execute(select_stmt).first() is None
 88        if notfound:
 89            stmt = insert(NodeUsage).values(
 90                created_at=created_at, node_id=node_id, uplink=0, downlink=0
 91            )
 92            db.execute(stmt)
 93
 94        # record
 95        stmt = (
 96            update(NodeUsage)
 97            .values(
 98                downlink=NodeUsage.downlink + usage,
 99            )
100            .where(
101                and_(
102                    NodeUsage.node_id == node_id,
103                    NodeUsage.created_at == created_at,
104                )
105            )
106        )
107
108        db.execute(stmt)
109        db.commit()
110
111
112async def get_users_stats(
113    node_id: int, node: MarzNodeBase
114) -> tuple[int, list[dict]]:
115    try:
116        params = list()
117        for stat in await asyncio.wait_for(node.fetch_users_stats(), 10):
118            if stat.usage:
119                params.append({"uid": stat.uid, "value": stat.usage})
120        return node_id, params
121    except:
122        return node_id, []
123
124
125async def record_user_usages():
126    # usage_coefficient = {None: 1}  # default usage coefficient for the main api instance
127
128    results = await asyncio.gather(
129        *[
130            get_users_stats(node_id, node)
131            for node_id, node in marznode.nodes.items()
132        ]
133    )
134    api_params = {node_id: params for node_id, params in list(results)}
135
136    users_usage = defaultdict(int)
137    for node_id, params in api_params.items():
138        coefficient = (
139            node.usage_coefficient
140            if (node := marznode.nodes.get(node_id))
141            else 1
142        )
143        node_usage = 0
144        for param in params:
145            users_usage[param["uid"]] += (
146                param["value"] * coefficient
147            )  # apply the usage coefficient
148            node_usage += param["value"]
149        record_node_stats(node_id, node_usage)
150
151    users_usage = list(
152        {"id": uid, "value": value} for uid, value in users_usage.items()
153    )
154    if not users_usage:
155        return
156
157    # record users usage
158    with GetDB() as db:
159        await data_usage_percent_reached(db, users_usage)
160
161        stmt = update(User).values(
162            used_traffic=User.used_traffic + bindparam("value"),
163            lifetime_used_traffic=User.lifetime_used_traffic
164            + bindparam("value"),
165            online_at=datetime.utcnow(),
166        )
167
168        db.execute(
169            stmt, users_usage, execution_options={"synchronize_session": None}
170        )
171        db.commit()
172
173    for node_id, params in api_params.items():
174        record_user_usage_logs(
175            params,
176            node_id,
177            (
178                node.usage_coefficient
179                if (node := marznode.nodes.get(node_id))
180                else 1
181            ),
182        )