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
 11
 12
 13def record_user_usage_logs(
 14    params: list, node_id: int, consumption_factor: int = 1
 15):
 16    if not params:
 17        return
 18
 19    created_at = datetime.fromisoformat(
 20        datetime.utcnow().strftime("%Y-%m-%dT%H:00:00")
 21    )
 22
 23    with GetDB() as db:
 24        # make user usage row if doesn't exist
 25        select_stmt = select(NodeUserUsage.user_id).where(
 26            and_(
 27                NodeUserUsage.node_id == node_id,
 28                NodeUserUsage.created_at == created_at,
 29            )
 30        )
 31        existings = [r[0] for r in db.execute(select_stmt).fetchall()]
 32        uids_to_insert = set()
 33
 34        for p in params:
 35            uid = p["uid"]
 36            if uid in existings:
 37                continue
 38            uids_to_insert.add(uid)
 39
 40        if uids_to_insert:
 41            stmt = insert(NodeUserUsage).values(
 42                user_id=bindparam("uid"),
 43                created_at=created_at,
 44                node_id=node_id,
 45                used_traffic=0,
 46            )
 47            db.execute(stmt, [{"uid": uid} for uid in uids_to_insert])
 48
 49        # record
 50        stmt = (
 51            update(NodeUserUsage)
 52            .values(
 53                used_traffic=NodeUserUsage.used_traffic
 54                + bindparam("value") * consumption_factor
 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, params, execution_options={"synchronize_session": None}
 66        )
 67        db.commit()
 68
 69
 70def record_node_stats(node_id: int, usage: int):
 71    if not usage:
 72        return
 73
 74    created_at = datetime.fromisoformat(
 75        datetime.utcnow().strftime("%Y-%m-%dT%H:00:00")
 76    )
 77
 78    with GetDB() as db:
 79
 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        stmt = update(User).values(
160            used_traffic=User.used_traffic + bindparam("value"),
161            lifetime_used_traffic=User.lifetime_used_traffic
162            + bindparam("value"),
163            online_at=datetime.utcnow(),
164        )
165
166        db.execute(
167            stmt, users_usage, execution_options={"synchronize_session": None}
168        )
169        db.commit()
170
171    for node_id, params in api_params.items():
172        record_user_usage_logs(
173            params,
174            node_id,
175            (
176                node.usage_coefficient
177                if (node := marznode.nodes.get(node_id))
178                else 1
179            ),
180        )