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