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 it 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 + bindparam("value")
 55            )
 56            .where(
 57                and_(
 58                    NodeUserUsage.user_id == bindparam("uid"),
 59                    NodeUserUsage.node_id == node_id,
 60                    NodeUserUsage.created_at == created_at,
 61                )
 62            )
 63        )
 64        db.connection().execute(
 65            stmt,
 66            [
 67                {**usage, "value": int(usage["value"] * consumption_factor)}
 68                for usage in params
 69            ],
 70            execution_options={"synchronize_session": None},
 71        )
 72        db.commit()
 73
 74
 75def record_node_stats(node_id: int, usage: int):
 76    if not usage:
 77        return
 78
 79    created_at = datetime.fromisoformat(
 80        datetime.utcnow().strftime("%Y-%m-%dT%H:00:00")
 81    )
 82
 83    with GetDB() as db:
 84        # make node usage row if doesn't exist
 85        select_stmt = select(NodeUsage.node_id).where(
 86            and_(
 87                NodeUsage.node_id == node_id,
 88                NodeUsage.created_at == created_at,
 89            )
 90        )
 91        notfound = db.execute(select_stmt).first() is None
 92        if notfound:
 93            stmt = insert(NodeUsage).values(
 94                created_at=created_at, node_id=node_id, uplink=0, downlink=0
 95            )
 96            db.execute(stmt)
 97
 98        # record
 99        stmt = (
100            update(NodeUsage)
101            .values(
102                downlink=NodeUsage.downlink + usage,
103            )
104            .where(
105                and_(
106                    NodeUsage.node_id == node_id,
107                    NodeUsage.created_at == created_at,
108                )
109            )
110        )
111
112        db.execute(stmt)
113        db.commit()
114
115
116async def get_users_stats(
117    node_id: int, node: MarzNodeBase
118) -> tuple[int, list[dict]]:
119    try:
120        params = list()
121        for stat in await asyncio.wait_for(node.fetch_users_stats(), 10):
122            if stat.usage:
123                params.append({"uid": stat.uid, "value": stat.usage})
124        return node_id, params
125    except:
126        return node_id, []
127
128
129async def record_user_usages():
130    # usage_coefficient = {None: 1}  # default usage coefficient for the main api instance
131
132    results = await asyncio.gather(
133        *[
134            get_users_stats(node_id, node)
135            for node_id, node in marznode.nodes.items()
136        ]
137    )
138    api_params = {node_id: params for node_id, params in list(results)}
139
140    users_usage = defaultdict(int)
141    for node_id, params in api_params.items():
142        coefficient = (
143            node.usage_coefficient
144            if (node := marznode.nodes.get(node_id))
145            else 1
146        )
147        node_usage = 0
148        for param in params:
149            users_usage[param["uid"]] += int(
150                param["value"] * coefficient
151            )  # apply the usage coefficient
152            node_usage += param["value"]
153        record_node_stats(node_id, node_usage)
154
155    users_usage = list(
156        {"id": uid, "value": value} for uid, value in users_usage.items()
157    )
158    if not users_usage:
159        return
160
161    # record users usage
162    with GetDB() as db:
163        await data_usage_percent_reached(db, users_usage)
164
165        stmt = update(User).values(
166            used_traffic=User.used_traffic + bindparam("value"),
167            lifetime_used_traffic=User.lifetime_used_traffic
168            + bindparam("value"),
169            online_at=datetime.utcnow(),
170        )
171
172        db.execute(
173            stmt, users_usage, execution_options={"synchronize_session": None}
174        )
175        db.commit()
176
177    for node_id, params in api_params.items():
178        record_user_usage_logs(
179            params,
180            node_id,
181            (
182                node.usage_coefficient
183                if (node := marznode.nodes.get(node_id))
184                else 1
185            ),
186        )